Add LockIsolated For Recursive Registration

Package.swift

@@ -54,6 +54,10 @@ let package = Package(
             name: "BazelProtobufBindings",
             dependencies: [
                 .product(name: "SwiftProtobuf", package: "swift-protobuf")
+            ],
+            exclude: [
+                "README.md",
+                "analysis_v2.proto",
             ]
         ),
         .testTarget(

Sources/SourceKitBazelBSP/Server/MessageHandler/BSPMessageHandler.swift

@@ -37,30 +37,32 @@ final class BSPMessageHandler: MessageHandler {
 
     // We currently use a single-threaded setup for simplicity,
     // but we can eventually reply asynchronously if we find a need for it.
-    private let lock: OSAllocatedUnfairLock<Void> = .init()
-    // FIXME: Can't put state into the lock for now because of recursiveness in the registration.
-    nonisolated(unsafe) private var state: State = State()
+    private let state = LockIsolated<State>(.init())
 
     init() {}
 
     func register<Request: RequestType>(
         requestHandler: @escaping BSPRequestHandler<Request>
     ) {
-        state.requestHandlers[Request.method] = AnyRequestHandler(
-            handler: requestHandler
-        )
+        state.modify { state in
+            state.requestHandlers[Request.method] = AnyRequestHandler(
+                handler: requestHandler
+            )
+        }
     }
 
     func register<Notification: NotificationType>(
         notificationHandler: @escaping BSPNotificationHandler<Notification>
     ) {
-        state.notificationHandlers[Notification.method] = AnyNotificationHandler(
-            handler: notificationHandler
-        )
+        state.modify { state in
+            state.notificationHandlers[Notification.method] = AnyNotificationHandler(
+                handler: notificationHandler
+            )
+        }
     }
 
     func handle<Notification: NotificationType>(_ notification: Notification) {
-        lock.withLockUnchecked {
+        state.read { state in
             logger.info(
                 "Received notification: \(Notification.method)"
             )
@@ -78,7 +80,7 @@ final class BSPMessageHandler: MessageHandler {
         id: RequestID,
         reply: @escaping (LSPResult<Request.Response>) -> Void
     ) {
-        lock.withLockUnchecked {
+        state.read { state in
             logger.info(
                 "Received request: \(Request.method)"
             )
@@ -132,3 +134,24 @@ final class BSPMessageHandler: MessageHandler {
         return handler
     }
 }
+
+final class LockIsolated<Value>: @unchecked Sendable {
+    private var value: Value
+    private let lock = NSRecursiveLock()
+
+    init(_ value: Value) {
+        self.value = value
+    }
+
+    func read<T>(_ closure: (Value) throws -> T) rethrows -> T {
+        lock.lock()
+        defer { lock.unlock() }
+        return try closure(value)
+    }
+
+    func modify<T>(_ closure: (inout Value) throws -> T) rethrows -> T {
+        lock.lock()
+        defer { lock.unlock() }
+        return try closure(&value)
+    }
+}

Tests/SourceKitBazelBSPTests/BSPMessageHandlerTests.swift

@@ -127,4 +127,51 @@ import Testing
 
         #expect(initialized)
     }
+
+    @Test
+    func testRecursiveRegistration() async throws {
+        // setup mocks
+        let sourcesRequest = BuildTargetSourcesRequest(targets: [])
+        let sourcesResponse = BuildTargetSourcesResponse(items: [])
+        let initRequest = InitializeBuildRequest(
+            displayName: "mockInitRequest",
+            version: "1.0",
+            bspVersion: "2.0",
+            rootUri: URI(filePath: "/Users/BSP", isDirectory: false),
+            capabilities: .init(languageIds: [])
+        )
+        let initResponse = InitializeBuildResponse(
+            displayName: "mockInitResponse",
+            version: "1.0",
+            bspVersion: "2.0",
+            capabilities: .init()
+        )
+
+        // recursive registering
+        let registry = BSPMessageHandler()
+
+        registry.register { (request: InitializeBuildRequest, id: RequestID) in
+            registry.register { (request: BuildTargetSourcesRequest, id: RequestID) in
+                return sourcesResponse
+            }
+            return initResponse
+        }
+
+        // validating handler registered
+        registry.handle(initRequest, id: .number(1)) { result in
+            if case let .success(resp) = result {
+                #expect(resp.displayName == "mockInitResponse")
+            } else {
+                Issue.record("initRequest failed due to handler not found")
+            }
+        }
+
+        registry.handle(sourcesRequest, id: .number(2)) { result in
+            if case let .success(resp) = result {
+                #expect(resp.items.count == 0)
+            } else {
+                Issue.record("sourcesRequest failed due to handler not found")
+            }
+        }
+    }
 }