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

Specifically, we care that all outstanding **writes** are finished
before we call the close handler, because otherwise we may (a) send
corrupted output during shutdown, or (b) drop notifications and replies
sent during the shutdown process.  The former is a potential issue for
clients that are not robust about parse failures, and the latter is an
issue for reproducibility and robustness during testing/debugging - in
particular, we have some integration tests that send data without
waiting for individual replies and they need to finish outstanding
replies before exiting.

rdar://60159448

Author: benlangmuir

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),