Merge pull request #196 from benlangmuir/shutdown-flush

Attempt to save index on shutdown

Author: benlangmuir

Sources/LanguageServerProtocolJSONRPC/JSONRPCConnection.swift

@@ -52,17 +52,15 @@ public final class JSONRPCConection {
   /// The set of currently outstanding outgoing requests along with information about how to decode and handle their responses.
   var outstandingRequests: [RequestID: OutstandingRequest] = [:]
 
-  var closeHandler: () -> Void
+  var closeHandler: (() -> Void)! = nil
 
   public init(
     protocol messageRegistry: MessageRegistry,
     inFD: Int32,
     outFD: Int32,
-    syncRequests: Bool = false,
-    closeHandler: @escaping () -> Void = {})
+    syncRequests: Bool = false)
   {
     state = .created
-    self.closeHandler = closeHandler
     self.messageRegistry = messageRegistry
     self.syncRequests = syncRequests
 
@@ -93,10 +91,11 @@ public final class JSONRPCConection {
   /// Start processing `inFD` and send messages to `receiveHandler`.
   ///
   /// - parameter receiveHandler: The message handler to invoke for requests received on the `inFD`.
-  public func start(receiveHandler: MessageHandler) {
+  public func start(receiveHandler: MessageHandler, closeHandler: @escaping () -> Void = {}) {
     precondition(state == .created)
     state = .running
     self.receiveHandler = receiveHandler
+    self.closeHandler = closeHandler
 
     receiveIO.read(offset: 0, length: Int.max, queue: queue) { done, data, errorCode in
       guard errorCode == 0 else {

Sources/SKTestSupport/SKTibsTestWorkspace.swift

@@ -38,13 +38,15 @@ public final class SKTibsTestWorkspace {
     immutableProjectDir: URL,
     persistentBuildDir: URL,
     tmpDir: URL,
+    removeTmpDir: Bool,
     toolchain: Toolchain,
     clientCapabilities: ClientCapabilities) throws
   {
     self.tibsWorkspace = try TibsTestWorkspace(
       immutableProjectDir: immutableProjectDir,
       persistentBuildDir: persistentBuildDir,
       tmpDir: tmpDir,
+      removeTmpDir: removeTmpDir,
       toolchain: TibsToolchain(toolchain))
 
     initWorkspace(clientCapabilities: clientCapabilities)
@@ -91,15 +93,22 @@ extension SKTibsTestWorkspace {
 
 extension XCTestCase {
 
-  public func staticSourceKitTibsWorkspace(name: String, clientCapabilities: ClientCapabilities = .init(), testFile: String = #file) throws -> SKTibsTestWorkspace? {
+  public func staticSourceKitTibsWorkspace(
+    name: String,
+    clientCapabilities: ClientCapabilities = .init(),
+    tmpDir: URL? = nil,
+    removeTmpDir: Bool = true,
+    testFile: String = #file
+  ) throws -> SKTibsTestWorkspace? {
     let testDirName = testDirectoryName
     let workspace = try SKTibsTestWorkspace(
       immutableProjectDir: inputsDirectory(testFile: testFile)
         .appendingPathComponent(name, isDirectory: true),
       persistentBuildDir: XCTestCase.productsDirectory
         .appendingPathComponent("sk-tests/\(testDirName)", isDirectory: true),
-      tmpDir: URL(fileURLWithPath: NSTemporaryDirectory(), isDirectory: true)
+      tmpDir: tmpDir ?? URL(fileURLWithPath: NSTemporaryDirectory(), isDirectory: true)
         .appendingPathComponent("sk-test-data/\(testDirName)", isDirectory: true),
+      removeTmpDir: removeTmpDir,
       toolchain: ToolchainRegistry.shared.default!,
       clientCapabilities: clientCapabilities)
 

Sources/SourceKit/SourceKitServer.swift

@@ -46,7 +46,7 @@ public final class SourceKitServer: LanguageServer {
 
   let fs: FileSystem
 
-  let onExit: () -> Void
+  var onExit: () -> Void
 
   /// Creates a language server for the given client.
   public init(client: Connection, fileSystem: FileSystem = localFileSystem, options: Options, onExit: @escaping () -> Void = {}) {
@@ -351,13 +351,47 @@ extension SourceKitServer {
     requestCancellation[key]?.cancel()
   }
 
+  /// The server is about to exit, and the server should flush any buffered state.
+  ///
+  /// The server shall not be used to handle more requests (other than possibly
+  /// `shutdown` and `exit`) and should attempt to flush any buffered state
+  /// immediately, such as sending index changes to disk.
+  public func prepareForExit() {
+    // Note: this method should be safe to call multiple times, since we want to
+    // be resilient against multiple possible shutdown sequences, including
+    // pipe failure.
+
+    // Close the index, which will flush to disk.
+    self.queue.sync {
+      self._prepareForExit()
+    }
+  }
+
+  func _prepareForExit() {
+    // Note: this method should be safe to call multiple times, since we want to
+    // be resilient against multiple possible shutdown sequences, including
+    // pipe failure.
+
+    // Close the index, which will flush to disk.
+    self.workspace?.index = nil
+  }
+
+
   func shutdown(_ request: Request<Shutdown>) {
-    // Nothing to do yet.
+    _prepareForExit()
     request.reply(VoidResponse())
   }
 
   func exit(_ notification: Notification<Exit>) {
-    onExit()
+    // Should have been called in shutdown, but allow misbehaving clients.
+    _prepareForExit()
+
+    // Call onExit only once, and hop off queue to allow the handler to call us back.
+    let onExit = self.onExit
+    self.onExit = {}
+    DispatchQueue.global().async {
+      onExit()
+    }
   }
 
   // MARK: - Text synchronization

Sources/sourcekit-lsp/main.swift

@@ -102,18 +102,18 @@ let clientConnection = JSONRPCConection(
   protocol: MessageRegistry.lspProtocol,
   inFD: STDIN_FILENO,
   outFD: STDOUT_FILENO,
-  syncRequests: options.syncRequests,
-  closeHandler: {
-  exit(0)
-})
+  syncRequests: options.syncRequests)
 
 let installPath = AbsolutePath(Bundle.main.bundlePath)
 ToolchainRegistry.shared = ToolchainRegistry(installPath: installPath, localFileSystem)
 
 let server = SourceKitServer(client: clientConnection, options: options.serverOptions, onExit: {
   clientConnection.close()
 })
-clientConnection.start(receiveHandler: server)
+clientConnection.start(receiveHandler: server, closeHandler: {
+  server.prepareForExit()
+  exit(0)
+})
 
 Logger.shared.addLogHandler { message, _ in
   clientConnection.send(LogMessage(type: .log, message: message))

Tests/LanguageServerProtocolJSONRPCTests/ConnectionTests.swift

@@ -213,18 +213,17 @@ class ConnectionTests: XCTestCase {
       let conn = JSONRPCConection(
         protocol: MessageRegistry(requests: [], notifications: []),
         inFD: to.fileHandleForReading.fileDescriptor,
-        outFD: from.fileHandleForWriting.fileDescriptor,
-        closeHandler: {
-        // We get an error from XCTest if this is fulfilled more than once.
-        expectation.fulfill()
-      })
+        outFD: from.fileHandleForWriting.fileDescriptor)
 
       final class DummyHandler: MessageHandler {
         func handle<N: NotificationType>(_: N, from: ObjectIdentifier) {}
         func handle<R: RequestType>(_: R, id: RequestID, from: ObjectIdentifier, reply: @escaping (LSPResult<R.Response>) -> Void) {}
       }
 
-      conn.start(receiveHandler: DummyHandler())
+      conn.start(receiveHandler: DummyHandler(), closeHandler: {
+        // We get an error from XCTest if this is fulfilled more than once.
+        expectation.fulfill()
+      })
 
 
       close(to.fileHandleForWriting.fileDescriptor)

Tests/SourceKitTests/SourceKitTests.swift

@@ -87,6 +87,61 @@ final class SKTests: XCTestCase {
     ])
   }
 
+  func testIndexShutdown() throws {
+
+    let tmpDir = URL(fileURLWithPath: NSTemporaryDirectory(), isDirectory: true)
+        .appendingPathComponent("sk-test-data/\(testDirectoryName)", isDirectory: true)
+
+    func listdir(_ url: URL) throws -> [URL] {
+      try FileManager.default.contentsOfDirectory(at: url, includingPropertiesForKeys: nil)
+    }
+
+    func checkRunningIndex(build: Bool) throws -> URL? {
+      guard let ws = try staticSourceKitTibsWorkspace(
+        name: "SwiftModules", tmpDir: tmpDir, removeTmpDir: false)
+      else {
+        return nil
+      }
+
+      if build {
+        try ws.buildAndIndex()
+      }
+
+      let locDef = ws.testLoc("aaa:def")
+      let locRef = ws.testLoc("aaa:call:c")
+      try ws.openDocument(locRef.url, language: .swift)
+      let jump = try ws.sk.sendSync(DefinitionRequest(
+        textDocument: locRef.docIdentifier,
+        position: locRef.position))
+      XCTAssertEqual(jump.count, 1)
+      XCTAssertEqual(jump.first?.uri, DocumentURI(locDef.url))
+      XCTAssertEqual(jump.first?.range.lowerBound, locDef.position)
+
+      let tmpContents = try listdir(tmpDir)
+      guard let versionedPath = tmpContents.filter({ $0.lastPathComponent.starts(with: "v") }).spm_only else {
+        XCTFail("expected one version path 'v[0-9]*', found \(tmpContents)")
+        return nil
+      }
+
+      let versionContentsBefore = try listdir(versionedPath)
+      XCTAssertEqual(versionContentsBefore.count, 1)
+      XCTAssert(versionContentsBefore.first?.lastPathComponent.starts(with: "p") ?? false)
+
+      _ = try ws.sk.sendSync(Shutdown())
+      return versionedPath
+    }
+
+    guard let versionedPath = try checkRunningIndex(build: true) else { return }
+    
+    let versionContentsAfter = try listdir(versionedPath)
+    XCTAssertEqual(versionContentsAfter.count, 1)
+    XCTAssertEqual(versionContentsAfter.first?.lastPathComponent, "saved")
+
+    _ = try checkRunningIndex(build: true)
+
+    try FileManager.default.removeItem(atPath: tmpDir.path)
+  }
+
   func testCodeCompleteSwiftTibs() throws {
     guard let ws = try staticSourceKitTibsWorkspace(name: "CodeCompleteSingleModule") else { return }
     let loc = ws.testLoc("cc:A")

Tests/SourceKitTests/XCTestManifests.swift

@@ -122,6 +122,7 @@ extension SKTests {
     // to regenerate.
     static let __allTests__SKTests = [
         ("testCodeCompleteSwiftTibs", testCodeCompleteSwiftTibs),
+        ("testIndexShutdown", testIndexShutdown),
         ("testIndexSwiftModules", testIndexSwiftModules),
         ("testInitJSON", testInitJSON),
         ("testInitLocal", testInitLocal),