Fix the main path of the reported leak

initial work on passing promises around

more promise propagation

make sure the `runPromise` is fulfilled or precondition

refine preconditions

fix test failing on nightlies

revert a change which also fixes a warning

increase the required accuracy for the failing nightly test

add one missing promise succession call

format

revert typo fixes

eliminate most of the preconditions

some more promises taken care of

remove the rest of the preconditions

better fulfillment of a promise

better semantics in handling of another promise

try another way

Author: MahdiBM

Sources/PostgresNIO/New/Connection State Machine/AuthenticationStateMachine.swift

@@ -16,10 +16,10 @@ struct AuthenticationStateMachine {
     }
 
     enum Action {
-        case sendStartupMessage(AuthContext)
-        case sendPassword(PasswordAuthencationMode, AuthContext)
-        case sendSaslInitialResponse(name: String, initialResponse: [UInt8])
-        case sendSaslResponse([UInt8])
+        case sendStartupMessage(AuthContext, promise: EventLoopPromise<Void>?)
+        case sendPassword(PasswordAuthencationMode, AuthContext, promise: EventLoopPromise<Void>?)
+        case sendSaslInitialResponse(name: String, initialResponse: [UInt8], promise: EventLoopPromise<Void>?)
+        case sendSaslResponse([UInt8], promise: EventLoopPromise<Void>?)
         case wait
         case authenticated
 
@@ -34,12 +34,12 @@ struct AuthenticationStateMachine {
         self.state = .initialized
     }
 
-    mutating func start() -> Action {
+    mutating func start(_ promise: EventLoopPromise<Void>?) -> Action {
         guard case .initialized = self.state else {
             preconditionFailure("Unexpected state")
         }
         self.state = .startupMessageSent
-        return .sendStartupMessage(self.authContext)
+        return .sendStartupMessage(self.authContext, promise: promise)
     }
 
     mutating func authenticationMessageReceived(_ message: PostgresBackendMessage.Authentication) -> Action {
@@ -54,10 +54,10 @@ struct AuthenticationStateMachine {
                     return self.setAndFireError(PSQLError(code: .authMechanismRequiresPassword))
                 }
                 self.state = .passwordAuthenticationSent
-                return .sendPassword(.md5(salt: salt), self.authContext)
+                return .sendPassword(.md5(salt: salt), self.authContext, promise: nil)
             case .plaintext:
                 self.state = .passwordAuthenticationSent
-                return .sendPassword(.cleartext, authContext)
+                return .sendPassword(.cleartext, authContext, promise: nil)
             case .kerberosV5:
                 return self.setAndFireError(.unsupportedAuthMechanism(.kerberosV5))
             case .scmCredential:
@@ -89,7 +89,7 @@ struct AuthenticationStateMachine {
                     }
 
                     self.state = .saslInitialResponseSent(saslManager)
-                    return .sendSaslInitialResponse(name: SASLMechanism.SCRAM.SHA256.name, initialResponse: output)
+                    return .sendSaslInitialResponse(name: SASLMechanism.SCRAM.SHA256.name, initialResponse: output, promise: nil)
                 } catch {
                     return self.setAndFireError(.sasl(underlying: error))
                 }
@@ -122,7 +122,7 @@ struct AuthenticationStateMachine {
                 }
 
                 self.state = .saslChallengeResponseSent(saslManager)
-                return .sendSaslResponse(output)
+                return .sendSaslResponse(output, promise: nil)
             } catch {
                 return self.setAndFireError(.sasl(underlying: error))
             }

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

@@ -10,7 +10,7 @@ struct CloseStateMachine {
     }
 
     enum Action {
-        case sendCloseSync(CloseTarget)
+        case sendCloseSync(CloseTarget, promise: EventLoopPromise<Void>?)
         case succeedClose(CloseCommandContext)
         case failClose(CloseCommandContext, with: PSQLError)
 
@@ -24,14 +24,14 @@ struct CloseStateMachine {
         self.state = .initialized(closeContext)
     }
 
-    mutating func start() -> Action {
+    mutating func start(_ promise: EventLoopPromise<Void>?) -> Action {
         guard case .initialized(let closeContext) = self.state else {
             preconditionFailure("Start should only be called, if the query has been initialized")
         }
 
         self.state = .closeSyncSent(closeContext)
 
-        return .sendCloseSync(closeContext.target)
+        return .sendCloseSync(closeContext.target, promise: promise)
     }
 
     mutating func closeCompletedReceived() -> Action {

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

@@ -64,8 +64,8 @@ struct ConnectionStateMachine {
 
         case read
         case wait
-        case sendSSLRequest
-        case establishSSLConnection
+        case sendSSLRequest(EventLoopPromise<Void>?)
+        case establishSSLConnection(EventLoopPromise<Void>?)
         case provideAuthenticationContext
         case forwardNotificationToListeners(PostgresBackendMessage.NotificationResponse)
         case fireEventReadyForQuery
@@ -77,16 +77,16 @@ struct ConnectionStateMachine {
         case closeConnectionAndCleanup(CleanUpContext)
 
         // Auth Actions
-        case sendStartupMessage(AuthContext)
-        case sendPasswordMessage(PasswordAuthencationMode, AuthContext)
-        case sendSaslInitialResponse(name: String, initialResponse: [UInt8])
-        case sendSaslResponse([UInt8])
+        case sendStartupMessage(AuthContext, promise: EventLoopPromise<Void>?)
+        case sendPasswordMessage(PasswordAuthencationMode, AuthContext, promise: EventLoopPromise<Void>?)
+        case sendSaslInitialResponse(name: String, initialResponse: [UInt8], promise: EventLoopPromise<Void>?)
+        case sendSaslResponse([UInt8], promise: EventLoopPromise<Void>?)
 
         // Connection Actions
 
         // --- general actions
-        case sendParseDescribeBindExecuteSync(PostgresQuery)
-        case sendBindExecuteSync(PSQLExecuteStatement)
+        case sendParseDescribeBindExecuteSync(PostgresQuery, promise: EventLoopPromise<Void>?)
+        case sendBindExecuteSync(PSQLExecuteStatement, promise: EventLoopPromise<Void>?)
         case failQuery(EventLoopPromise<PSQLRowStream>, with: PSQLError, cleanupContext: CleanUpContext?)
         case succeedQuery(EventLoopPromise<PSQLRowStream>, with: QueryResult)
 
@@ -97,12 +97,12 @@ struct ConnectionStateMachine {
         case forwardStreamError(PSQLError, read: Bool, cleanupContext: CleanUpContext?)
 
         // Prepare statement actions
-        case sendParseDescribeSync(name: String, query: String, bindingDataTypes: [PostgresDataType])
+        case sendParseDescribeSync(name: String, query: String, bindingDataTypes: [PostgresDataType], promise: EventLoopPromise<Void>?)
         case succeedPreparedStatementCreation(EventLoopPromise<RowDescription?>, with: RowDescription?)
         case failPreparedStatementCreation(EventLoopPromise<RowDescription?>, with: PSQLError, cleanupContext: CleanUpContext?)
 
         // Close actions
-        case sendCloseSync(CloseTarget)
+        case sendCloseSync(CloseTarget, promise: EventLoopPromise<Void>?)
         case succeedClose(CloseCommandContext)
         case failClose(CloseCommandContext, with: PSQLError, cleanupContext: CleanUpContext?)
     }
@@ -131,7 +131,7 @@ struct ConnectionStateMachine {
         case require
     }
 
-    mutating func connected(tls: TLSConfiguration) -> ConnectionAction {
+    mutating func connected(tls: TLSConfiguration, promise: EventLoopPromise<Void>?) -> ConnectionAction {
         switch self.state {
         case .initialized:
             switch tls {
@@ -141,11 +141,11 @@ struct ConnectionStateMachine {
 
             case .prefer:
                 self.state = .sslRequestSent(.prefer)
-                return .sendSSLRequest
+                return .sendSSLRequest(promise)
 
             case .require:
                 self.state = .sslRequestSent(.require)
-                return .sendSSLRequest
+                return .sendSSLRequest(promise)
             }
 
         case .sslRequestSent,
@@ -164,8 +164,11 @@ struct ConnectionStateMachine {
         }
     }
 
-    mutating func provideAuthenticationContext(_ authContext: AuthContext) -> ConnectionAction {
-        self.startAuthentication(authContext)
+    mutating func provideAuthenticationContext(
+        _ authContext: AuthContext,
+        promise: EventLoopPromise<Void>?
+    ) -> ConnectionAction {
+        self.startAuthentication(authContext, promise: promise)
     }
 
     mutating func gracefulClose(_ promise: EventLoopPromise<Void>?) -> ConnectionAction {
@@ -233,8 +236,8 @@ struct ConnectionStateMachine {
                 return self.closeConnectionAndCleanup(.receivedUnencryptedDataAfterSSLRequest)
             }
             self.state = .sslNegotiated
-            return .establishSSLConnection
-            
+            return .establishSSLConnection(nil)
+
         case .initialized,
              .sslNegotiated,
              .sslHandlerAdded,
@@ -583,14 +586,16 @@ struct ConnectionStateMachine {
         }
 
         switch task {
-        case .extendedQuery(let queryContext):
+        case .extendedQuery(let queryContext, let writePromise):
+            writePromise?.fail(psqlErrror) /// Use `cleanupContext` or not?
             switch queryContext.query {
             case .executeStatement(_, let promise), .unnamed(_, let promise):
                 return .failQuery(promise, with: psqlErrror, cleanupContext: nil)
             case .prepareStatement(_, _, _, let promise):
                 return .failPreparedStatementCreation(promise, with: psqlErrror, cleanupContext: nil)
             }
-        case .closeCommand(let closeContext):
+        case .closeCommand(let closeContext, let writePromise):
+            writePromise?.fail(psqlErrror) /// Use `cleanupContext` or not?
             return .failClose(closeContext, with: psqlErrror, cleanupContext: nil)
         }
     }
@@ -800,14 +805,17 @@ struct ConnectionStateMachine {
 
     // MARK: - Private Methods -
 
-    private mutating func startAuthentication(_ authContext: AuthContext) -> ConnectionAction {
+    private mutating func startAuthentication(
+        _ authContext: AuthContext,
+        promise: EventLoopPromise<Void>?
+    ) -> ConnectionAction {
         guard case .waitingToStartAuthentication = self.state else {
             preconditionFailure("Can only start authentication after connect or ssl establish")
         }
 
         self.state = .modifying // avoid CoW
         var authState = AuthenticationStateMachine(authContext: authContext)
-        let action = authState.start()
+        let action = authState.start(promise)
         self.state = .authenticating(authState)
         return self.modify(with: action)
     }
@@ -934,17 +942,17 @@ struct ConnectionStateMachine {
         }
 
         switch task {
-        case .extendedQuery(let queryContext):
+        case .extendedQuery(let queryContext, let promise):
             self.state = .modifying // avoid CoW
             var extendedQuery = ExtendedQueryStateMachine(queryContext: queryContext)
-            let action = extendedQuery.start()
+            let action = extendedQuery.start(promise)
             self.state = .extendedQuery(extendedQuery, connectionContext)
             return self.modify(with: action)
 
-        case .closeCommand(let closeContext):
+        case .closeCommand(let closeContext, let promise):
             self.state = .modifying // avoid CoW
             var closeStateMachine = CloseStateMachine(closeContext: closeContext)
-            let action = closeStateMachine.start()
+            let action = closeStateMachine.start(promise)
             self.state = .closeCommand(closeStateMachine, connectionContext)
             return self.modify(with: action)
         }
@@ -1031,10 +1039,10 @@ extension ConnectionStateMachine {
 extension ConnectionStateMachine {
     mutating func modify(with action: ExtendedQueryStateMachine.Action) -> ConnectionStateMachine.ConnectionAction {
         switch action {
-        case .sendParseDescribeBindExecuteSync(let query):
-            return .sendParseDescribeBindExecuteSync(query)
-        case .sendBindExecuteSync(let executeStatement):
-            return .sendBindExecuteSync(executeStatement)
+        case .sendParseDescribeBindExecuteSync(let query, let promise):
+            return .sendParseDescribeBindExecuteSync(query, promise: promise)
+        case .sendBindExecuteSync(let executeStatement, let promise):
+            return .sendBindExecuteSync(executeStatement, promise: promise)
         case .failQuery(let requestContext, with: let error):
             let cleanupContext = self.setErrorAndCreateCleanupContextIfNeeded(error)
             return .failQuery(requestContext, with: error, cleanupContext: cleanupContext)
@@ -1057,8 +1065,8 @@ extension ConnectionStateMachine {
             return .read
         case .wait:
             return .wait
-        case .sendParseDescribeSync(name: let name, query: let query, bindingDataTypes: let bindingDataTypes):
-            return .sendParseDescribeSync(name: name, query: query, bindingDataTypes: bindingDataTypes)
+        case .sendParseDescribeSync(name: let name, query: let query, bindingDataTypes: let bindingDataTypes, let promise):
+            return .sendParseDescribeSync(name: name, query: query, bindingDataTypes: bindingDataTypes, promise: promise)
         case .succeedPreparedStatementCreation(let promise, with: let rowDescription):
             return .succeedPreparedStatementCreation(promise, with: rowDescription)
         case .failPreparedStatementCreation(let promise, with: let error):
@@ -1071,14 +1079,14 @@ extension ConnectionStateMachine {
 extension ConnectionStateMachine {
     mutating func modify(with action: AuthenticationStateMachine.Action) -> ConnectionStateMachine.ConnectionAction {
         switch action {
-        case .sendStartupMessage(let authContext):
-            return .sendStartupMessage(authContext)
-        case .sendPassword(let mode, let authContext):
-            return .sendPasswordMessage(mode, authContext)
-        case .sendSaslInitialResponse(let name, let initialResponse):
-            return .sendSaslInitialResponse(name: name, initialResponse: initialResponse)
-        case .sendSaslResponse(let bytes):
-            return .sendSaslResponse(bytes)
+        case .sendStartupMessage(let authContext, let promise):
+            return .sendStartupMessage(authContext, promise: promise)
+        case .sendPassword(let mode, let authContext, let promise):
+            return .sendPasswordMessage(mode, authContext, promise: promise)
+        case .sendSaslInitialResponse(let name, let initialResponse, let promise):
+            return .sendSaslInitialResponse(name: name, initialResponse: initialResponse, promise: promise)
+        case .sendSaslResponse(let bytes, let promise):
+            return .sendSaslResponse(bytes, promise: promise)
         case .authenticated:
             self.state = .authenticated(nil, [:])
             return .wait
@@ -1094,8 +1102,8 @@ extension ConnectionStateMachine {
 extension ConnectionStateMachine {
     mutating func modify(with action: CloseStateMachine.Action) -> ConnectionStateMachine.ConnectionAction {
         switch action {
-        case .sendCloseSync(let sendClose):
-            return .sendCloseSync(sendClose)
+        case .sendCloseSync(let sendClose, let promise):
+            return .sendCloseSync(sendClose, promise: promise)
         case .succeedClose(let closeContext):
             return .succeedClose(closeContext)
         case .failClose(let closeContext, with: let error):

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

@@ -25,10 +25,10 @@ struct ExtendedQueryStateMachine {
     }
 
     enum Action {
-        case sendParseDescribeBindExecuteSync(PostgresQuery)
-        case sendParseDescribeSync(name: String, query: String, bindingDataTypes: [PostgresDataType])
-        case sendBindExecuteSync(PSQLExecuteStatement)
-        
+        case sendParseDescribeBindExecuteSync(PostgresQuery, promise: EventLoopPromise<Void>?)
+        case sendParseDescribeSync(name: String, query: String, bindingDataTypes: [PostgresDataType], promise: EventLoopPromise<Void>?)
+        case sendBindExecuteSync(PSQLExecuteStatement, promise: EventLoopPromise<Void>?)
+
         // --- general actions
         case failQuery(EventLoopPromise<PSQLRowStream>, with: PSQLError)
         case succeedQuery(EventLoopPromise<PSQLRowStream>, with: QueryResult)
@@ -56,7 +56,7 @@ struct ExtendedQueryStateMachine {
         self.state = .initialized(queryContext)
     }
 
-    mutating func start() -> Action {
+    mutating func start(_ promise: EventLoopPromise<Void>?) -> Action {
         guard case .initialized(let queryContext) = self.state else {
             preconditionFailure("Start should only be called, if the query has been initialized")
         }
@@ -65,7 +65,7 @@ struct ExtendedQueryStateMachine {
         case .unnamed(let query, _):
             return self.avoidingStateMachineCoW { state -> Action in
                 state = .messagesSent(queryContext)
-                return .sendParseDescribeBindExecuteSync(query)
+                return .sendParseDescribeBindExecuteSync(query, promise: promise)
             }
 
         case .executeStatement(let prepared, _):
@@ -76,13 +76,14 @@ struct ExtendedQueryStateMachine {
                 case .none:
                     state = .noDataMessageReceived(queryContext)
                 }
-                return .sendBindExecuteSync(prepared)
+                return .sendBindExecuteSync(prepared, promise: promise)
             }
 
+            /// Not my code, but this is ignoring the last argument which is a promise? is that fine?
         case .prepareStatement(let name, let query, let bindingDataTypes, _):
             return self.avoidingStateMachineCoW { state -> Action in
                 state = .messagesSent(queryContext)
-                return .sendParseDescribeSync(name: name, query: query, bindingDataTypes: bindingDataTypes)
+                return .sendParseDescribeSync(name: name, query: query, bindingDataTypes: bindingDataTypes, promise: promise)
             }
         }
     }

Sources/PostgresNIO/New/NotificationListener.swift

@@ -42,7 +42,10 @@ final class NotificationListener: @unchecked Sendable {
         self.state = .closure(context, closure)
     }
 
-    func startListeningSucceeded(handler: PostgresChannelHandler) {
+    func startListeningSucceeded(
+        handler: PostgresChannelHandler,
+        writePromise: EventLoopPromise<Void>?
+    ) {
         self.eventLoop.preconditionInEventLoop()
         let handlerLoopBound = NIOLoopBound(handler, eventLoop: self.eventLoop)
 
@@ -56,26 +59,28 @@ final class NotificationListener: @unchecked Sendable {
                 switch reason {
                 case .cancelled:
                     eventLoop.execute {
-                        handlerLoopBound.value.cancelNotificationListener(channel: channel, id: listenerID)
+                        handlerLoopBound.value.cancelNotificationListener(channel: channel, id: listenerID, writePromise: nil)
                     }
 
                 case .finished:
-                    break
+                    writePromise?.succeed()
 
                 @unknown default:
-                    break
+                    writePromise?.succeed()
                 }
             }
             self.state = .streamListening(continuation)
 
             let notificationSequence = PostgresNotificationSequence(base: stream)
             checkedContinuation.resume(returning: notificationSequence)
+            writePromise?.succeed(())
 
         case .streamListening, .done:
             fatalError("Invalid state: \(self.state)")
 
         case .closure:
-            break // ignore
+            writePromise?.succeed(())
+            // ignore
         }
     }