BaseStreamSocketChannel half-close allows outstanding writes to complete

Sources/NIOPosix/BaseSocketChannel.swift

@@ -595,8 +595,16 @@ class BaseSocketChannel<SocketType: BaseSocketProtocol>: SelectableChannel, Chan
             switch writeResult.writeResult {
             case .couldNotWriteEverything:
                 newWriteRegistrationState = .register
-            case .writtenCompletely:
+            case .writtenCompletely(let closeState):
                 newWriteRegistrationState = .unregister
+                switch closeState {
+                case .open, .pending:
+                    ()
+                case .readyForClose:
+                    self.close0(error: ChannelError.outputClosed, mode: .output, promise: nil)
+                case .closed:
+                    ()  // we can be flushed before becoming active
+                }
             }
 
             if !self.isOpen || !self.hasFlushedPendingWrites() {

Sources/NIOPosix/BaseStreamSocketChannel.swift

@@ -194,13 +194,28 @@ class BaseStreamSocketChannel<Socket: SocketProtocol>: BaseSocketChannel<Socket>
                     self.close0(error: error, mode: .all, promise: promise)
                     return
                 }
-                try self.shutdownSocket(mode: mode)
-                // Fail all pending writes and so ensure all pending promises are notified
-                self.pendingWrites.failAll(error: error, close: false)
                 self.unregisterForWritable()
-                promise?.succeed(())
 
-                self.pipeline.fireUserInboundEventTriggered(ChannelEvent.outputClosed)
+                let writesCloseResult = self.pendingWrites.close(promise)
+                switch writesCloseResult {
+                case .pending:
+                    ()  // promise is stored in `pendingWrites` state for completing on later call to `closeComplete`
+                case .readyForClose:
+                    // Shutdown the socket only when the pending writes are dealt with
+                    do {
+                        try self.shutdownSocket(mode: mode)
+                        self.pendingWrites.closeComplete()
+                    } catch let err {
+                        self.pendingWrites.closeComplete(err)
+                    }
+                    self.pipeline.fireUserInboundEventTriggered(ChannelEvent.outputClosed)
+                case .closed:
+                    promise?.succeed(())
+                case .open:
+                    promise?.fail(ChannelError.inappropriateOperationForState)
+                    assertionFailure("Close resulted in an open state, this should never happen")
+                }
+
             case .input:
                 if self.inputShutdown {
                     promise?.fail(ChannelError._inputClosed)
@@ -224,6 +239,7 @@ class BaseStreamSocketChannel<Socket: SocketProtocol>: BaseSocketChannel<Socket>
                 promise?.succeed(())
 
                 self.pipeline.fireUserInboundEventTriggered(ChannelEvent.inputClosed)
+
             case .all:
                 if let timeout = self.connectTimeoutScheduled {
                     self.connectTimeoutScheduled = nil

Sources/NIOPosix/PendingDatagramWritesManager.swift

@@ -419,6 +419,13 @@ final class PendingDatagramWritesManager: PendingWritesManager {
     internal var publishedWritability = true
     internal var writeSpinCount: UInt = 16
     private(set) var isOpen = true
+    var outboundCloseState: CloseState {
+        if self.isOpen {
+            .open
+        } else {
+            .closed
+        }
+    }
 
     /// Initialize with a pre-allocated array of message headers and storage references. We pass in these pre-allocated
     /// objects to save allocations. They can be safely be re-used for all `Channel`s on a given `EventLoop` as an

Sources/NIOPosix/PendingWritesManager.swift

@@ -97,9 +97,9 @@ internal enum OneWriteOperationResult {
 /// The result of trying to write all the outstanding flushed data. That naturally includes all `ByteBuffer`s and
 /// `FileRegions` and the individual writes have potentially been retried (see `WriteSpinOption`).
 internal struct OverallWriteResult {
-    enum WriteOutcome {
+    enum WriteOutcome: Equatable {
         /// Wrote all the data that was flushed. When receiving this result, we can unsubscribe from 'writable' notification.
-        case writtenCompletely
+        case writtenCompletely(CloseResult)
 
         /// Could not write everything. Before attempting further writes the eventing system should send a 'writable' notification.
         case couldNotWriteEverything
@@ -152,7 +152,7 @@ private struct PendingStreamWritesState {
         self.subtractOutstanding(bytes: bytes)
     }
 
-    /// Initialise a new, empty `PendingWritesState`.
+    /// Initialize a new, empty `PendingWritesState`.
     public init() {}
 
     /// Check if there are no outstanding writes.
@@ -310,6 +310,8 @@ final class PendingStreamWritesManager: PendingWritesManager {
 
     private(set) var isOpen = true
 
+    private(set) var outboundCloseState: CloseState = .open
+
     /// Mark the flush checkpoint.
     func markFlushCheckpoint() {
         self.state.markFlushCheckpoint()
@@ -337,7 +339,7 @@ final class PendingStreamWritesManager: PendingWritesManager {
     /// - result: If the `Channel` is still writable after adding the write of `data`.
     func add(data: IOData, promise: EventLoopPromise<Void>?) -> Bool {
         assert(self.isOpen)
-        self.state.append(.init(data: data, promise: promise))
+        self.state.append(PendingStreamWrite(data: data, promise: promise))
 
         if self.state.bytes > waterMark.high
             && channelWritabilityFlag.compareExchange(expected: true, desired: false, ordering: .relaxed).exchanged
@@ -367,7 +369,7 @@ final class PendingStreamWritesManager: PendingWritesManager {
         vectorBufferWriteOperation: (UnsafeBufferPointer<IOVector>) throws -> IOResult<Int>,
         scalarFileWriteOperation: (CInt, Int, Int) throws -> IOResult<Int>
     ) throws -> OverallWriteResult {
-        try self.triggerWriteOperations { writeMechanism in
+        var result = try self.triggerWriteOperations { writeMechanism in
             switch writeMechanism {
             case .scalarBufferWrite:
                 return try triggerScalarBufferWrite({ try scalarBufferWriteOperation($0) })
@@ -380,6 +382,28 @@ final class PendingStreamWritesManager: PendingWritesManager {
                 return .writtenCompletely
             }
         }
+
+        // If we have no more writes check if we have a pending close
+        if self.isEmpty {
+            switch result.writeResult {
+            case .writtenCompletely:
+                switch self.outboundCloseState {
+                case .open:
+                    ()
+                case .pending(let closePromise):
+                    self.outboundCloseState = .readyForClose(closePromise)
+                case .readyForClose:
+                    assertionFailure("Transitioned from readyForClose to readyForClose, shouldn't we have closed?")
+                case .closed:
+                    ()
+                }
+
+            case .couldNotWriteEverything:
+                assertionFailure("Write result is .couldNotWriteEverything but we have no more writes to perform.")
+            }
+            result.writeResult = .writtenCompletely(.init(self.outboundCloseState))
+        }
+        return result
     }
 
     /// To be called after a write operation (usually selected and run by `triggerAppropriateWriteOperation`) has
@@ -463,18 +487,85 @@ final class PendingStreamWritesManager: PendingWritesManager {
         return self.didWrite(itemCount: result.itemCount, result: result.writeResult)
     }
 
-    /// Fail all the outstanding writes. This is useful if for example the `Channel` is closed.
+    /// Fail all the outstanding writes.
     func failAll(error: Error, close: Bool) {
         if close {
             assert(self.isOpen)
             self.isOpen = false
+            switch self.outboundCloseState {
+            case .open, .closed:
+                self.outboundCloseState = .closed
+            case .pending(let closePromise), .readyForClose(let closePromise):
+                self.outboundCloseState = .closed
+                closePromise?.fail(error)
+            }
         }
 
         self.state.removeAll()?.fail(error)
 
         assert(self.state.isEmpty)
     }
 
+    /// Signal the intention to close. Takes a promise which MUST be completed via a call to `closeComplete`
+    ///
+    /// - Parameters:
+    ///   - promise: Optionally an `EventLoopPromise` that will be succeeded once all outstanding writes have been dealt with
+    func close(_ promise: EventLoopPromise<Void>?) -> CloseResult {
+        assert(self.isOpen)
+
+        if self.isEmpty {
+            switch self.outboundCloseState {
+            case .open:
+                self.outboundCloseState = .readyForClose(promise)
+            case .readyForClose(var closePromise), .pending(var closePromise):
+                closePromise.setOrCascade(to: promise)
+                self.outboundCloseState = .readyForClose(closePromise)
+            case .closed:
+                promise?.succeed(())
+            }
+
+        } else {
+            switch self.outboundCloseState {
+            case .open:
+                self.outboundCloseState = .pending(promise)
+            case .pending(var closePromise):
+                closePromise.setOrCascade(to: promise)
+                self.outboundCloseState = .pending(closePromise)
+            case .readyForClose:
+                preconditionFailure(
+                    "We are in .readyForClose state but we still have pending writes. This should never happen."
+                )
+            case .closed:
+                preconditionFailure(
+                    "We are in .closed state but we still have pending writes. This should never happen."
+                )
+            }
+
+        }
+
+        return .init(self.outboundCloseState)
+    }
+
+    /// Signal that the `Channel` is closed.
+    func closeComplete(_ error: Error? = nil) {
+        assert(self.isOpen)
+        assert(self.isEmpty)
+
+        switch self.outboundCloseState {
+        case .readyForClose(let closePromise):
+            self.outboundCloseState = .closed
+            if let error {
+                closePromise?.fail(error)
+            } else {
+                closePromise?.succeed(())
+            }
+        case .closed:
+            ()  // nothing to do
+        case .open, .pending:
+            preconditionFailure("close complete called on channel in unexpected state: \(self.outboundCloseState)")
+        }
+    }
+
     /// Initialize with a pre-allocated array of IO vectors and storage references. We pass in these pre-allocated
     /// objects to save allocations. They can be safely be re-used for all `Channel`s on a given `EventLoop` as an
     /// `EventLoop` always runs on one and the same thread. That means that there can't be any writes of more than
@@ -487,6 +578,33 @@ final class PendingStreamWritesManager: PendingWritesManager {
     }
 }
 
+internal enum CloseState {
+    case open
+    case pending(EventLoopPromise<Void>?)
+    case readyForClose(EventLoopPromise<Void>?)
+    case closed
+}
+
+internal enum CloseResult {
+    case open
+    case pending
+    case readyForClose
+    case closed
+
+    init(_ state: CloseState) {
+        switch state {
+        case .open:
+            self = .open
+        case .pending:
+            self = .pending
+        case .readyForClose:
+            self = .readyForClose
+        case .closed:
+            self = .closed
+        }
+    }
+}
+
 internal enum WriteMechanism {
     case scalarBufferWrite
     case vectorBufferWrite
@@ -496,6 +614,8 @@ internal enum WriteMechanism {
 
 internal protocol PendingWritesManager: AnyObject {
     var isOpen: Bool { get }
+    var isEmpty: Bool { get }
+    var outboundCloseState: CloseState { get }
     var isFlushPending: Bool { get }
     var writeSpinCount: UInt { get }
     var currentBestWriteMechanism: WriteMechanism { get }
@@ -522,7 +642,7 @@ extension PendingWritesManager {
             var oneResult: OneWriteOperationResult
             repeat {
                 guard self.isOpen && self.isFlushPending else {
-                    result.writeResult = .writtenCompletely
+                    result.writeResult = .writtenCompletely(.init(self.outboundCloseState))
                     break writeSpinLoop
                 }