Send terminate message on close (vapor#86)

* send postgres terminate message when closing

* fix logging

Author: tanner0101

Sources/PostgresNIO/Connection/PostgresConnection+Connect.swift

@@ -13,8 +13,8 @@ extension PostgresConnection {
             .channelOption(ChannelOptions.socket(SocketOptionLevel(SOL_SOCKET), SO_REUSEADDR), value: 1)
         return bootstrap.connect(to: socketAddress).flatMap { channel in
             return channel.pipeline.addHandlers([
-                ByteToMessageHandler(PostgresMessageDecoder()),
-                MessageToByteHandler(PostgresMessageEncoder()),
+                ByteToMessageHandler(PostgresMessageDecoder(logger: logger)),
+                MessageToByteHandler(PostgresMessageEncoder(logger: logger)),
                 PostgresRequestHandler(logger: logger),
                 PostgresErrorHandler(logger: logger)
             ]).map {

Sources/PostgresNIO/Connection/PostgresRequest.swift renamed to Sources/PostgresNIO/Connection/PostgresConnection+Database.swift

@@ -12,24 +12,17 @@ extension PostgresConnection: PostgresDatabase {
         self.channel.flush()
         return promise.futureResult
     }
-    
+
     public func withConnection<T>(_ closure: (PostgresConnection) -> EventLoopFuture<T>) -> EventLoopFuture<T> {
         closure(self)
     }
 }
 
-public protocol PostgresRequest {
-    // return nil to end request
-    func respond(to message: PostgresMessage) throws -> [PostgresMessage]?
-    func start() throws -> [PostgresMessage]
-    func log(to logger: Logger)
-}
-
 final class PostgresRequestContext {
     let delegate: PostgresRequest
     let promise: EventLoopPromise<Void>
     var lastError: Error?
-    
+
     init(delegate: PostgresRequest, promise: EventLoopPromise<Void>) {
         self.delegate = delegate
         self.promise = promise
@@ -40,23 +33,23 @@ final class PostgresRequestHandler: ChannelDuplexHandler {
     typealias InboundIn = PostgresMessage
     typealias OutboundIn = PostgresRequestContext
     typealias OutboundOut = PostgresMessage
-    
+
     private var queue: [PostgresRequestContext]
     let logger: Logger
-    
+
     public init(logger: Logger) {
         self.queue = []
         self.logger = logger
     }
-    
+
     private func _channelRead(context: ChannelHandlerContext, data: NIOAny) throws {
         let message = self.unwrapInboundIn(data)
         guard self.queue.count > 0 else {
             // discard packet
             return
         }
         let request = self.queue[0]
-        
+
         switch message.identifier {
         case .error:
             let error = try PostgresMessage.Error(message: message)
@@ -67,7 +60,7 @@ final class PostgresRequestHandler: ChannelDuplexHandler {
             self.logger.notice("\(notice)")
         default: break
         }
-        
+
         if let responses = try request.delegate.respond(to: message) {
             for response in responses {
                 context.write(self.wrapOutboundOut(response), promise: nil)
@@ -82,7 +75,7 @@ final class PostgresRequestHandler: ChannelDuplexHandler {
             }
         }
     }
-    
+
     func channelRead(context: ChannelHandlerContext, data: NIOAny) {
         do {
             try self._channelRead(context: context, data: data)
@@ -92,7 +85,7 @@ final class PostgresRequestHandler: ChannelDuplexHandler {
         // Regardless of error, also pass the message downstream; this is necessary for PostgresNotificationHandler (which is appended at the end) to receive notifications
         context.fireChannelRead(data)
     }
-    
+
     func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
         let request = self.unwrapOutboundIn(data)
         self.queue.append(request)
@@ -105,13 +98,16 @@ final class PostgresRequestHandler: ChannelDuplexHandler {
             self.errorCaught(context: context, error: error)
         }
     }
-    
+
     func close(context: ChannelHandlerContext, mode: CloseMode, promise: EventLoopPromise<Void>?) {
+        let terminate = try! PostgresMessage.Terminate().message()
+        context.write(self.wrapOutboundOut(terminate), promise: nil)
+        context.close(mode: mode, promise: promise)
+
         for current in self.queue {
             current.promise.fail(PostgresError.connectionClosed)
         }
         self.queue = []
-        context.close(mode: mode, promise: promise)
     }
 }
 

Sources/PostgresNIO/Connection/PostgresConnection.swift

@@ -31,17 +31,11 @@ public final class PostgresConnection {
             return self.eventLoop.makeSucceededFuture(())
         }
         self.didClose = true
-        
-        let promise = self.eventLoop.makePromise(of: Void.self)
-        self.eventLoop.submit {
-            switch self.channel.isActive {
-            case true:
-                promise.succeed(())
-            case false:
-                self.channel.close(mode: .all, promise: promise)
-            }
-        }.cascadeFailure(to: promise)
-        return promise.futureResult
+        if !self.isClosed {
+            return self.channel.close(mode: .all)
+        } else {
+            return self.eventLoop.makeSucceededFuture(())
+        }
     }
 
     deinit {

Sources/PostgresNIO/Message/PostgresMessage+BackendKeyData.swift

@@ -3,7 +3,11 @@ import NIO
 extension PostgresMessage {
     /// Identifies the message as cancellation key data.
     /// The frontend must save these values if it wishes to be able to issue CancelRequest messages later.
-    public struct BackendKeyData: CustomStringConvertible {
+    public struct BackendKeyData: PostgresMessageType {
+        public static var identifier: PostgresMessage.Identifier {
+            .backendKeyData
+        }
+        
         /// Parses an instance of this message type from a byte buffer.
         public static func parse(from buffer: inout ByteBuffer) throws -> BackendKeyData {
             guard let processID = buffer.readInteger(as: Int32.self) else {
@@ -20,10 +24,5 @@ extension PostgresMessage {
 
         /// The secret key of this backend.
         public var secretKey: Int32
-        
-        /// See `CustomStringConvertible`.
-        public var description: String {
-            return "processID: \(processID), secretKey: \(secretKey)"
-        }
     }
 }

Sources/PostgresNIO/Message/PostgresMessage+Terminate.swift

@@ -0,0 +1,9 @@
+extension PostgresMessage {
+    public struct Terminate: PostgresMessageType {
+        public static var identifier: PostgresMessage.Identifier {
+            .terminate
+        }
+
+        public func serialize(into buffer: inout ByteBuffer) { }
+    }
+}

Sources/PostgresNIO/Message/PostgresMessageDecoder.swift

@@ -9,10 +9,14 @@ public final class PostgresMessageDecoder: ByteToMessageDecoder {
 
     /// If `true`, the server has asked for authentication.
     public var hasSeenFirstMessage: Bool
+
+    /// Logger to send debug messages to.
+    let logger: Logger?
 
     /// Creates a new `PostgresMessageDecoder`.
-    public init() {
+    public init(logger: Logger? = nil) {
         self.hasSeenFirstMessage = false
+        self.logger = logger
     }
 
     /// See `ByteToMessageDecoder`.
@@ -48,6 +52,7 @@ public final class PostgresMessageDecoder: ByteToMessageDecoder {
 
         // there is sufficient data, use this buffer
         buffer = peekBuffer
+        self.logger?.trace("Decoded: PostgresMessage (\(message.identifier))")
         context.fireChannelRead(wrapInboundOut(message))
         return .continue
     }

Sources/PostgresNIO/Message/PostgresMessageEncoder.swift

@@ -3,11 +3,17 @@ import NIO
 public final class PostgresMessageEncoder: MessageToByteEncoder {
     /// See `MessageToByteEncoder`.
     public typealias OutboundIn = PostgresMessage
+
+    /// Logger to send debug messages to.
+    let logger: Logger?
+
+    /// Creates a new `PostgresMessageEncoder`.
+    public init(logger: Logger? = nil) {
+        self.logger = logger
+    }
 
     /// See `MessageToByteEncoder`.
     public func encode(data message: PostgresMessage, out: inout ByteBuffer) throws {
-        // print("PostgresMessage.ChannelEncoder.encode(\(data))")
-        
         // serialize identifier
         var message = message
         switch message.identifier {
@@ -25,6 +31,7 @@ public final class PostgresMessageEncoder: MessageToByteEncoder {
 
         // set message size
         out.setInteger(Int32(out.writerIndex - messageSizeIndex), at: messageSizeIndex)
+        self.logger?.trace("Encoded: PostgresMessage (\(message.identifier))")
     }
 }
 

Sources/PostgresNIO/Connection/PostgresClient+Query.swift renamed to Sources/PostgresNIO/PostgresDatabase+Query.swift

@@ -0,0 +1,8 @@
+import Logging
+
+public protocol PostgresRequest {
+    // return nil to end request
+    func respond(to message: PostgresMessage) throws -> [PostgresMessage]?
+    func start() throws -> [PostgresMessage]
+    func log(to logger: Logger)
+}

Sources/PostgresNIO/Connection/PostgresClient+SimpleQuery.swift renamed to Sources/PostgresNIO/PostgresDatabase+SimpleQuery.swift

@@ -10,7 +10,7 @@ final class PostgresNIOTests: XCTestCase {
     }
 
     override func setUp() {
-        testLogLevel = .info
+        XCTAssertTrue(isLoggingConfigured)
         self.group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
     }
 
@@ -1164,3 +1164,13 @@ private func prepareTableToMeasureSelectPerformance(
         _ = try conn.query(insertQuery, batchedFixtureData).wait()
     }
 }
+
+
+let isLoggingConfigured: Bool = {
+    LoggingSystem.bootstrap { label in
+        var handler = StreamLogHandler.standardOutput(label: label)
+        handler.logLevel = .debug
+        return handler
+    }
+    return true
+}()

Sources/PostgresNIO/Connection/PostgresClient.swift renamed to Sources/PostgresNIO/PostgresDatabase.swift

@@ -1,8 +1,6 @@
 import Logging
 import PostgresNIO
 
-var testLogLevel: Logger.Level = .info
-
 extension PostgresConnection {
     static func address() throws -> SocketAddress {
         #if os(Linux)
@@ -27,7 +25,6 @@ extension PostgresConnection {
                 database: "vapor_database",
                 password: "vapor_password"
             ).map {
-                conn.logger.logLevel = testLogLevel
                 return conn
             }
         }