Add base support for async requests (#53)

Author: rockbruno

Sources/SourceKitBazelBSP/RequestHandlers/InitializeHandler.swift

@@ -58,7 +58,7 @@ final class InitializeHandler {
 
     func initializeBuild(
         _ request: InitializeBuildRequest,
-        _ id: RequestID
+        _ id: RequestID,
     ) throws -> (InitializeBuildResponse, InitializedServerConfig) {
         let taskId = TaskId(id: "initializeBuild-\(id.description)")
         connection?.startWorkTask(id: taskId, title: "Indexing: Initializing sourcekit-bazel-bsp")

Sources/SourceKitBazelBSP/Server/MessageHandler/BSPMessageHandler.swift

@@ -47,6 +47,18 @@ final class BSPMessageHandler: MessageHandler {
         state.notificationHandlers[Notification.method] = AnyNotificationHandler(handler: notificationHandler)
     }
 
+    /// Simple abstraction for registering requests that don't need async responses.
+    func register<Request: RequestType>(syncRequestHandler: @escaping BSPSyncRequestHandler<Request>) {
+        register(requestHandler: { (request: Request, id, completion) in
+            do {
+                let response = try syncRequestHandler(request, id)
+                completion(.success(response))
+            } catch {
+                completion(.failure(error))
+            }
+        })
+    }
+
     func handle<Notification: NotificationType>(_ notification: Notification) {
         logger.info("Received notification: \(Notification.method)")
         do {
@@ -63,16 +75,19 @@ final class BSPMessageHandler: MessageHandler {
         logger.info("Received request: \(Request.method)")
         do {
             let handler = try getHandler(for: request, id, reply, state: state)
-            let response = try handler(request, id)
-            logger.info("Replying to \(Request.method)")
-            reply(.success(response))
+            handler(request, id) { [buildLSPError] result in
+                do {
+                    let response = try result.get()
+                    logger.info("Replying to \(Request.method)")
+                    reply(.success(response))
+                } catch {
+                    logger.error("Error while replying to \(Request.method): \(error.localizedDescription)")
+                    reply(.failure(buildLSPError(error)))
+                }
+            }
         } catch {
             logger.error("Error while handling BSP request: \(error.localizedDescription)")
-            if let responseError = error as? ResponseError {
-                reply(.failure(responseError))
-            } else {
-                reply(.failure(ResponseError.internalError(error.localizedDescription)))
-            }
+            reply(.failure(buildLSPError(from: error)))
         }
     }
 
@@ -105,4 +120,11 @@ final class BSPMessageHandler: MessageHandler {
         }
         return handler
     }
+
+    private func buildLSPError(from error: Error) -> ResponseError {
+        guard let responseError = error as? ResponseError else {
+            return ResponseError.internalError(error.localizedDescription)
+        }
+        return responseError
+    }
 }

Sources/SourceKitBazelBSP/Server/MessageHandler/BSPRequestHandler.swift

@@ -20,7 +20,11 @@
 import Foundation
 import LanguageServerProtocol
 
-typealias BSPRequestHandler<Request: RequestType> = ((Request, RequestID) throws -> Request.Response)
+typealias BSPRequestHandlerCompletion<Request: RequestType> = ((Result<Request.Response, Error>) -> Void)
+typealias BSPRequestHandler<Request: RequestType> = (
+    (Request, RequestID, @escaping BSPRequestHandlerCompletion<Request>) -> Void
+)
+typealias BSPSyncRequestHandler<Request: RequestType> = ((Request, RequestID) throws -> Request.Response)
 
 /// A type-erased request handler wrapper to allow for dynamic registration of handlers.
 final class AnyRequestHandler {

Sources/SourceKitBazelBSP/Server/SourceKitBazelBSPServer.swift

@@ -38,13 +38,13 @@ package final class SourceKitBazelBSPServer {
         // Everything else will be registered post-init.
         let initHandler = InitializeHandler(baseConfig: baseConfig, connection: connection)
         let shutdownHandler = ShutdownHandler()
-        registry.register(requestHandler: { (request: InitializeBuildRequest, id: RequestID) in
+        registry.register(syncRequestHandler: { (request: InitializeBuildRequest, id: RequestID) in
             let result = try initHandler.initializeBuild(request, id)
             Self.registerPostInitHandlers(registry: registry, initializedConfig: result.1, connection: connection)
             return result.0
         })
         registry.register(notificationHandler: shutdownHandler.onBuildExit)
-        registry.register(requestHandler: shutdownHandler.buildShutdown)
+        registry.register(syncRequestHandler: shutdownHandler.buildShutdown)
         return registry
     }
 
@@ -57,7 +57,7 @@ package final class SourceKitBazelBSPServer {
         registry.register(notificationHandler: { (_: OnBuildInitializedNotification) in
             // no-op
         })
-        registry.register(requestHandler: { (_: WorkspaceWaitForBuildSystemUpdatesRequest, _: RequestID) in
+        registry.register(syncRequestHandler: { (_: WorkspaceWaitForBuildSystemUpdatesRequest, _: RequestID) in
             // FIXME: no-op, no special handling since the code today is not async, but I might be wrong here.
             VoidResponse()
         })
@@ -66,27 +66,27 @@ package final class SourceKitBazelBSPServer {
         // workspace/buildTargets
         let targetStore = BazelTargetStoreImpl(initializedConfig: initializedConfig)
         let buildTargetsHandler = BuildTargetsHandler(targetStore: targetStore, connection: connection)
-        registry.register(requestHandler: buildTargetsHandler.workspaceBuildTargets)
+        registry.register(syncRequestHandler: buildTargetsHandler.workspaceBuildTargets)
 
         // buildTarget/sources
         let targetSourcesHandler = TargetSourcesHandler(initializedConfig: initializedConfig, targetStore: targetStore)
-        registry.register(requestHandler: targetSourcesHandler.buildTargetSources)
+        registry.register(syncRequestHandler: targetSourcesHandler.buildTargetSources)
 
         // textDocument/sourceKitOptions
         let skOptionsHandler = SKOptionsHandler(
             initializedConfig: initializedConfig,
             targetStore: targetStore,
             connection: connection
         )
-        registry.register(requestHandler: skOptionsHandler.textDocumentSourceKitOptions)
+        registry.register(syncRequestHandler: skOptionsHandler.textDocumentSourceKitOptions)
 
         // buildTarget/prepare
         let prepareHandler = PrepareHandler(
             initializedConfig: initializedConfig,
             targetStore: targetStore,
             connection: connection
         )
-        registry.register(requestHandler: prepareHandler.prepareTarget)
+        registry.register(syncRequestHandler: prepareHandler.prepareTarget)
 
         // OnWatchedFilesDidChangeNotification
         let watchedFileChangeHandler = WatchedFileChangeHandler(

Tests/SourceKitBazelBSPTests/BSPMessageHandlerTests.swift

@@ -18,6 +18,7 @@
 // under the License.
 
 import BuildServerProtocol
+import Foundation
 import LanguageServerProtocol
 import LanguageServerProtocolJSONRPC
 import Testing
@@ -50,7 +51,7 @@ struct BSPMessageHandlerTests {
         let mockResponse = BuildTargetSourcesResponse(items: [])
 
         let handler = BSPMessageHandler()
-        handler.register(requestHandler: { (_: BuildTargetSourcesRequest, _) in mockResponse })
+        handler.register(syncRequestHandler: { (_: BuildTargetSourcesRequest, _) in mockResponse })
 
         let request = BuildTargetSourcesRequest(targets: [BuildTargetIdentifier(uri: try URI(string: "file:///test"))],
         )
@@ -69,7 +70,7 @@ struct BSPMessageHandlerTests {
     func failedRequest() throws {
         let mockError = ResponseError.internalError("Test error")
         let handler = BSPMessageHandler()
-        handler.register(requestHandler: { (_: BuildTargetSourcesRequest, _) in throw mockError })
+        handler.register(syncRequestHandler: { (_: BuildTargetSourcesRequest, _) in throw mockError })
 
         let request = BuildTargetSourcesRequest(targets: [BuildTargetIdentifier(uri: try URI(string: "file:///test"))],
         )
@@ -84,6 +85,36 @@ struct BSPMessageHandlerTests {
         }
     }
 
+    @Test
+    func asyncRequestHandlerCanReplyAsynchronously() throws {
+        let queue = DispatchQueue(label: "test", qos: .userInteractive)
+        let semaphore = DispatchSemaphore(value: 0)
+        let handler = BSPMessageHandler()
+        handler.register(requestHandler: { (request: BuildTargetSourcesRequest, id, completion) in
+            queue.asyncAfter(deadline: .now() + 1) {
+                completion(.success(BuildTargetSourcesResponse(items: [])))
+            }
+        })
+
+        let request = BuildTargetSourcesRequest(
+            targets: [BuildTargetIdentifier(uri: try URI(string: "file:///test"))]
+        )
+
+        var receivedResponse: LSPResult<BuildTargetSourcesResponse>?
+        handler.handle(request, id: RequestID.number(1)) { result in
+            receivedResponse = result
+            semaphore.signal()
+        }
+
+        semaphore.wait()
+
+        let result = try #require(receivedResponse)
+        switch result {
+        case .success(let response): #expect(response == BuildTargetSourcesResponse(items: []))
+        case .failure(let error): Issue.record("Expected success but got error: \(error)")
+        }
+    }
+
     @Test
     func unknownNotification() throws {
         var initialized = false