Skip to content

Commit 8b944b2

Browse files
MahdiBMMahdiBM
committed
Properly fulfill write promises
1 parent d18b137 commit 8b944b2

12 files changed

+194
-134
lines changed

‎Sources/PostgresNIO/New/Connection State Machine/CloseStateMachine.swift‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ struct CloseStateMachine {
1010
}
1111

1212
enum Action {
13-
case sendCloseSync(CloseTarget)
13+
case sendCloseSync(CloseTarget, promise: EventLoopPromise<Void>?)
1414
case succeedClose(CloseCommandContext)
1515
case failClose(CloseCommandContext, with: PSQLError)
1616

@@ -24,14 +24,14 @@ struct CloseStateMachine {
2424
self.state = .initialized(closeContext)
2525
}
2626

27-
mutating func start() -> Action {
27+
mutating func start(_ promise: EventLoopPromise<Void>?) -> Action {
2828
guard case .initialized(let closeContext) = self.state else {
2929
preconditionFailure("Start should only be called, if the query has been initialized")
3030
}
3131

3232
self.state = .closeSyncSent(closeContext)
3333

34-
return .sendCloseSync(closeContext.target)
34+
return .sendCloseSync(closeContext.target, promise: promise)
3535
}
3636

3737
mutating func closeCompletedReceived() -> Action {

‎Sources/PostgresNIO/New/Connection State Machine/ConnectionStateMachine.swift‎

Lines changed: 22 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -85,8 +85,8 @@ struct ConnectionStateMachine {
8585
// Connection Actions
8686

8787
// --- general actions
88-
case sendParseDescribeBindExecuteSync(PostgresQuery)
89-
case sendBindExecuteSync(PSQLExecuteStatement)
88+
case sendParseDescribeBindExecuteSync(PostgresQuery, promise: EventLoopPromise<Void>?)
89+
case sendBindExecuteSync(PSQLExecuteStatement, promise: EventLoopPromise<Void>?)
9090
case failQuery(EventLoopPromise<PSQLRowStream>, with: PSQLError, cleanupContext: CleanUpContext?)
9191
case succeedQuery(EventLoopPromise<PSQLRowStream>, with: QueryResult)
9292

@@ -97,12 +97,12 @@ struct ConnectionStateMachine {
9797
case forwardStreamError(PSQLError, read: Bool, cleanupContext: CleanUpContext?)
9898

9999
// Prepare statement actions
100-
case sendParseDescribeSync(name: String, query: String, bindingDataTypes: [PostgresDataType])
100+
case sendParseDescribeSync(name: String, query: String, bindingDataTypes: [PostgresDataType], promise: EventLoopPromise<Void>?)
101101
case succeedPreparedStatementCreation(EventLoopPromise<RowDescription?>, with: RowDescription?)
102102
case failPreparedStatementCreation(EventLoopPromise<RowDescription?>, with: PSQLError, cleanupContext: CleanUpContext?)
103103

104104
// Close actions
105-
case sendCloseSync(CloseTarget)
105+
case sendCloseSync(CloseTarget, promise: EventLoopPromise<Void>?)
106106
case succeedClose(CloseCommandContext)
107107
case failClose(CloseCommandContext, with: PSQLError, cleanupContext: CleanUpContext?)
108108
}
@@ -234,7 +234,7 @@ struct ConnectionStateMachine {
234234
}
235235
self.state = .sslNegotiated
236236
return .establishSSLConnection
237-
237+
238238
case .initialized,
239239
.sslNegotiated,
240240
.sslHandlerAdded,
@@ -247,7 +247,7 @@ struct ConnectionStateMachine {
247247
.closing,
248248
.closed:
249249
return self.closeConnectionAndCleanup(.unexpectedBackendMessage(.sslSupported))
250-
250+
251251
case .modifying:
252252
preconditionFailure("Invalid state: \(self.state)")
253253
}
@@ -583,14 +583,16 @@ struct ConnectionStateMachine {
583583
}
584584

585585
switch task {
586-
case .extendedQuery(let queryContext):
586+
case .extendedQuery(let queryContext, let writePromise):
587+
writePromise?.fail(psqlErrror) /// Use `cleanupContext` or not?
587588
switch queryContext.query {
588589
case .executeStatement(_, let promise), .unnamed(_, let promise):
589590
return .failQuery(promise, with: psqlErrror, cleanupContext: nil)
590591
case .prepareStatement(_, _, _, let promise):
591592
return .failPreparedStatementCreation(promise, with: psqlErrror, cleanupContext: nil)
592593
}
593-
case .closeCommand(let closeContext):
594+
case .closeCommand(let closeContext, let writePromise):
595+
writePromise?.fail(psqlErrror) /// Use `cleanupContext` or not?
594596
return .failClose(closeContext, with: psqlErrror, cleanupContext: nil)
595597
}
596598
}
@@ -934,17 +936,17 @@ struct ConnectionStateMachine {
934936
}
935937

936938
switch task {
937-
case .extendedQuery(let queryContext):
939+
case .extendedQuery(let queryContext, let promise):
938940
self.state = .modifying // avoid CoW
939941
var extendedQuery = ExtendedQueryStateMachine(queryContext: queryContext)
940-
let action = extendedQuery.start()
942+
let action = extendedQuery.start(promise)
941943
self.state = .extendedQuery(extendedQuery, connectionContext)
942944
return self.modify(with: action)
943945

944-
case .closeCommand(let closeContext):
946+
case .closeCommand(let closeContext, let promise):
945947
self.state = .modifying // avoid CoW
946948
var closeStateMachine = CloseStateMachine(closeContext: closeContext)
947-
let action = closeStateMachine.start()
949+
let action = closeStateMachine.start(promise)
948950
self.state = .closeCommand(closeStateMachine, connectionContext)
949951
return self.modify(with: action)
950952
}
@@ -1031,10 +1033,10 @@ extension ConnectionStateMachine {
10311033
extension ConnectionStateMachine {
10321034
mutating func modify(with action: ExtendedQueryStateMachine.Action) -> ConnectionStateMachine.ConnectionAction {
10331035
switch action {
1034-
case .sendParseDescribeBindExecuteSync(let query):
1035-
return .sendParseDescribeBindExecuteSync(query)
1036-
case .sendBindExecuteSync(let executeStatement):
1037-
return .sendBindExecuteSync(executeStatement)
1036+
case .sendParseDescribeBindExecuteSync(let query, let promise):
1037+
return .sendParseDescribeBindExecuteSync(query, promise: promise)
1038+
case .sendBindExecuteSync(let executeStatement, let promise):
1039+
return .sendBindExecuteSync(executeStatement, promise: promise)
10381040
case .failQuery(let requestContext, with: let error):
10391041
let cleanupContext = self.setErrorAndCreateCleanupContextIfNeeded(error)
10401042
return .failQuery(requestContext, with: error, cleanupContext: cleanupContext)
@@ -1057,8 +1059,8 @@ extension ConnectionStateMachine {
10571059
return .read
10581060
case .wait:
10591061
return .wait
1060-
case .sendParseDescribeSync(name: let name, query: let query, bindingDataTypes: let bindingDataTypes):
1061-
return .sendParseDescribeSync(name: name, query: query, bindingDataTypes: bindingDataTypes)
1062+
case .sendParseDescribeSync(name: let name, query: let query, bindingDataTypes: let bindingDataTypes, let promise):
1063+
return .sendParseDescribeSync(name: name, query: query, bindingDataTypes: bindingDataTypes, promise: promise)
10621064
case .succeedPreparedStatementCreation(let promise, with: let rowDescription):
10631065
return .succeedPreparedStatementCreation(promise, with: rowDescription)
10641066
case .failPreparedStatementCreation(let promise, with: let error):
@@ -1094,8 +1096,8 @@ extension ConnectionStateMachine {
10941096
extension ConnectionStateMachine {
10951097
mutating func modify(with action: CloseStateMachine.Action) -> ConnectionStateMachine.ConnectionAction {
10961098
switch action {
1097-
case .sendCloseSync(let sendClose):
1098-
return .sendCloseSync(sendClose)
1099+
case .sendCloseSync(let sendClose, let promise):
1100+
return .sendCloseSync(sendClose, promise: promise)
10991101
case .succeedClose(let closeContext):
11001102
return .succeedClose(closeContext)
11011103
case .failClose(let closeContext, with: let error):

‎Sources/PostgresNIO/New/Connection State Machine/ExtendedQueryStateMachine.swift‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,10 @@ struct ExtendedQueryStateMachine {
2525
}
2626

2727
enum Action {
28-
case sendParseDescribeBindExecuteSync(PostgresQuery)
29-
case sendParseDescribeSync(name: String, query: String, bindingDataTypes: [PostgresDataType])
30-
case sendBindExecuteSync(PSQLExecuteStatement)
31-
28+
case sendParseDescribeBindExecuteSync(PostgresQuery, promise: EventLoopPromise<Void>?)
29+
case sendParseDescribeSync(name: String, query: String, bindingDataTypes: [PostgresDataType], promise: EventLoopPromise<Void>?)
30+
case sendBindExecuteSync(PSQLExecuteStatement, promise: EventLoopPromise<Void>?)
31+
3232
// --- general actions
3333
case failQuery(EventLoopPromise<PSQLRowStream>, with: PSQLError)
3434
case succeedQuery(EventLoopPromise<PSQLRowStream>, with: QueryResult)
@@ -56,7 +56,7 @@ struct ExtendedQueryStateMachine {
5656
self.state = .initialized(queryContext)
5757
}
5858

59-
mutating func start() -> Action {
59+
mutating func start(_ promise: EventLoopPromise<Void>?) -> Action {
6060
guard case .initialized(let queryContext) = self.state else {
6161
preconditionFailure("Start should only be called, if the query has been initialized")
6262
}
@@ -65,7 +65,7 @@ struct ExtendedQueryStateMachine {
6565
case .unnamed(let query, _):
6666
return self.avoidingStateMachineCoW { state -> Action in
6767
state = .messagesSent(queryContext)
68-
return .sendParseDescribeBindExecuteSync(query)
68+
return .sendParseDescribeBindExecuteSync(query, promise: promise)
6969
}
7070

7171
case .executeStatement(let prepared, _):
@@ -76,13 +76,14 @@ struct ExtendedQueryStateMachine {
7676
case .none:
7777
state = .noDataMessageReceived(queryContext)
7878
}
79-
return .sendBindExecuteSync(prepared)
79+
return .sendBindExecuteSync(prepared, promise: promise)
8080
}
8181

82+
/// Not my code, but this is ignoring the last argument which is a promise? is that fine?
8283
case .prepareStatement(let name, let query, let bindingDataTypes, _):
8384
return self.avoidingStateMachineCoW { state -> Action in
8485
state = .messagesSent(queryContext)
85-
return .sendParseDescribeSync(name: name, query: query, bindingDataTypes: bindingDataTypes)
86+
return .sendParseDescribeSync(name: name, query: query, bindingDataTypes: bindingDataTypes, promise: promise)
8687
}
8788
}
8889
}

‎Sources/PostgresNIO/New/NotificationListener.swift‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,10 @@ final class NotificationListener: @unchecked Sendable {
4242
self.state = .closure(context, closure)
4343
}
4444

45-
func startListeningSucceeded(handler: PostgresChannelHandler) {
45+
func startListeningSucceeded(
46+
handler: PostgresChannelHandler,
47+
writePromise: EventLoopPromise<Void>?
48+
) {
4649
self.eventLoop.preconditionInEventLoop()
4750
let handlerLoopBound = NIOLoopBound(handler, eventLoop: self.eventLoop)
4851

@@ -56,7 +59,7 @@ final class NotificationListener: @unchecked Sendable {
5659
switch reason {
5760
case .cancelled:
5861
eventLoop.execute {
59-
handlerLoopBound.value.cancelNotificationListener(channel: channel, id: listenerID)
62+
handlerLoopBound.value.cancelNotificationListener(channel: channel, id: listenerID, writePromise: nil)
6063
}
6164

6265
case .finished:
@@ -70,12 +73,14 @@ final class NotificationListener: @unchecked Sendable {
7073

7174
let notificationSequence = PostgresNotificationSequence(base: stream)
7275
checkedContinuation.resume(returning: notificationSequence)
76+
writePromise?.succeed(())
7377

7478
case .streamListening, .done:
7579
fatalError("Invalid state: \(self.state)")
7680

7781
case .closure:
78-
break // ignore
82+
writePromise?.succeed(())
83+
// ignore
7984
}
8085
}
8186

‎Sources/PostgresNIO/New/PSQLTask.swift‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,12 @@ enum HandlerTask {
1010
}
1111

1212
enum PSQLTask {
13-
case extendedQuery(ExtendedQueryContext)
14-
case closeCommand(CloseCommandContext)
13+
case extendedQuery(ExtendedQueryContext, writePromise: EventLoopPromise<Void>?)
14+
case closeCommand(CloseCommandContext, writePromise: EventLoopPromise<Void>?)
1515

1616
func failWithError(_ error: PSQLError) {
1717
switch self {
18-
case .extendedQuery(let extendedQueryContext):
18+
case .extendedQuery(let extendedQueryContext, let writePromise):
1919
switch extendedQueryContext.query {
2020
case .unnamed(_, let eventLoopPromise):
2121
eventLoopPromise.fail(error)
@@ -24,9 +24,11 @@ enum PSQLTask {
2424
case .prepareStatement(_, _, _, let eventLoopPromise):
2525
eventLoopPromise.fail(error)
2626
}
27+
writePromise?.fail(error)
2728

28-
case .closeCommand(let closeCommandContext):
29+
case .closeCommand(let closeCommandContext, let writePromise):
2930
closeCommandContext.promise.fail(error)
31+
writePromise?.fail(error)
3032
}
3133
}
3234
}

0 commit comments

Comments
 (0)