[JSONRPC] Do not exit until outstanding I/O has finished

Sources/LanguageServerProtocolJSONRPC/JSONRPCConnection.swift

@@ -64,16 +64,28 @@ public final class JSONRPCConnection {
     self.messageRegistry = messageRegistry
     self.syncRequests = syncRequests
 
+    let ioGroup = DispatchGroup()
+
+    ioGroup.enter()
     receiveIO = DispatchIO(type: .stream, fileDescriptor: inFD, queue: queue) { (error: Int32) in
       if error != 0 {
         log("IO error \(error)", level: .error)
       }
+      ioGroup.leave()
     }
 
+    ioGroup.enter()
     sendIO = DispatchIO(type: .stream, fileDescriptor: outFD, queue: sendQueue) { (error: Int32) in
       if error != 0 {
         log("IO error \(error)", level: .error)
       }
+      ioGroup.leave()
+    }
+
+    ioGroup.notify(queue: queue) { [weak self] in
+      guard let self = self else { return }
+      self.closeHandler()
+      self.receiveHandler = nil // break retain cycle
     }
 
     // We cannot assume the client will send us bytes in packets of any particular size, so set the lower limit to 1.
@@ -99,7 +111,9 @@ public final class JSONRPCConnection {
 
     receiveIO.read(offset: 0, length: Int.max, queue: queue) { done, data, errorCode in
       guard errorCode == 0 else {
-        log("IO error \(errorCode)", level: .error)
+        if errorCode != POSIXError.ECANCELED.rawValue {
+          log("IO error reading \(errorCode)", level: .error)
+        }
         if done { self._close() }
         return
       }
@@ -287,6 +301,9 @@ public final class JSONRPCConnection {
   }
 
   /// Close the connection.
+  ///
+  /// The user-provided close handler will be called *asynchronously* when all outstanding I/O
+  /// operations have completed. No new I/O will be accepted after `close` returns.
   public func close() {
     queue.sync { _close() }
   }
@@ -298,10 +315,10 @@ public final class JSONRPCConnection {
       state = .closed
 
       log("\(JSONRPCConnection.self): closing...")
+      // Attempt to close the reader immediately; we do not need to accept remaining inputs.
       receiveIO.close(flags: .stop)
-      sendIO.close(flags: .stop)
-      receiveHandler = nil // break retain cycle
-      closeHandler()
+      // Close the writer after it finishes outstanding work.
+      sendIO.close()
     }
   }
 

Tests/LanguageServerProtocolJSONRPCTests/ConnectionTests.swift

@@ -197,6 +197,20 @@ class ConnectionTests: XCTestCase {
     waitForExpectations(timeout: 10)
   }
 
+  func testSendBeforeClose() {
+    let client = connection.client
+    let server = connection.server
+
+    let expectation = self.expectation(description: "received notification")
+    client.handleNextNotification { (note: Notification<EchoNotification>) in
+      expectation.fulfill()
+    }
+
+    server.client.send(EchoNotification(string: "about to close!"))
+    connection.serverConnection.close()
+
+    waitForExpectations(timeout: 10)
+  }
 
   /// We can explicitly close a connection, but the connection also
   /// automatically closes itself if the pipe is closed (or has an error).

Tests/LanguageServerProtocolJSONRPCTests/XCTestManifests.swift

@@ -34,6 +34,7 @@ extension ConnectionTests {
         ("testMessageBuffer", testMessageBuffer),
         ("testRound", testRound),
         ("testSendAfterClose", testSendAfterClose),
+        ("testSendBeforeClose", testSendBeforeClose),
         ("testUnexpectedResponse", testUnexpectedResponse),
         ("testUnknownNotification", testUnknownNotification),
         ("testUnknownRequest", testUnknownRequest),