From 3fdaa106445da293d84f2a5226d5613fda2d9d17 Mon Sep 17 00:00:00 2001
From: Alex Hoppen <ahoppen@apple.com>
Date: Wed, 5 Jun 2024 10:58:37 -0700
Subject: [PATCH] Fix a race condition in `LocalConnection`
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The comment on `LocalConnection` assumed that all mutable state was guarded by `queue` but `state` and `handler` actually weren’t. This was found by the thread sanitizer.
---
 Sources/InProcessClient/LocalConnection.swift | 48 ++++++++++++-------
 1 file changed, 32 insertions(+), 16 deletions(-)

diff --git a/Sources/InProcessClient/LocalConnection.swift b/Sources/InProcessClient/LocalConnection.swift
index 190e7f133..21bf01d10 100644
--- a/Sources/InProcessClient/LocalConnection.swift
+++ b/Sources/InProcessClient/LocalConnection.swift
@@ -26,10 +26,8 @@ import LanguageServerProtocol
 /// conn.send(...) // handled by server
 /// conn.close()
 /// ```
-///
-/// - Note: Unchecked sendable conformance because shared state is guarded by `queue`.
-public final class LocalConnection: Connection, @unchecked Sendable {
-  enum State {
+public final class LocalConnection: Connection, Sendable {
+  private enum State {
     case ready, started, closed
   }
 
@@ -37,36 +35,51 @@ public final class LocalConnection: Connection, @unchecked Sendable {
   private let name: String
 
   /// The queue guarding `_nextRequestID`.
-  let queue: DispatchQueue = DispatchQueue(label: "local-connection-queue")
+  private let queue: DispatchQueue = DispatchQueue(label: "local-connection-queue")
 
-  var _nextRequestID: Int = 0
+  /// - Important: Must only be accessed from `queue`
+  nonisolated(unsafe) private var _nextRequestID: Int = 0
 
-  var state: State = .ready
+  /// - Important: Must only be accessed from `queue`
+  nonisolated(unsafe) private var state: State = .ready
 
-  var handler: MessageHandler? = nil
+  /// - Important: Must only be accessed from `queue`
+  nonisolated(unsafe) private var handler: MessageHandler? = nil
 
   public init(name: String) {
     self.name = name
   }
 
   deinit {
-    if state != .closed {
-      close()
+    queue.sync {
+      if state != .closed {
+        closeAssumingOnQueue()
+      }
     }
   }
 
   public func start(handler: MessageHandler) {
-    precondition(state == .ready)
-    state = .started
-    self.handler = handler
+    queue.sync {
+      precondition(state == .ready)
+      state = .started
+      self.handler = handler
+    }
   }
 
-  public func close() {
+  /// - Important: Must only be called from `queue`
+  private func closeAssumingOnQueue() {
+    dispatchPrecondition(condition: .onQueue(queue))
     precondition(state != .closed)
     handler = nil
     state = .closed
   }
 
+  public func close() {
+    queue.sync {
+      closeAssumingOnQueue()
+    }
+  }
+
   func nextRequestID() -> RequestID {
     return queue.sync {
       _nextRequestID += 1
@@ -81,7 +94,10 @@ public final class LocalConnection: Connection, @unchecked Sendable {
       \(notification.forLogging)
       """
     )
-    self.handler?.handle(notification)
+    guard let handler = queue.sync(execute: { handler }) else {
+      return
+    }
+    handler.handle(notification)
   }
 
   public func send<Request: RequestType>(
@@ -97,7 +113,7 @@ public final class LocalConnection: Connection, @unchecked Sendable {
       """
     )
 
-    guard let handler = self.handler else {
+    guard let handler = queue.sync(execute: { handler }) else {
       logger.info(
         """
         Replying to request \(id, privacy: .public) with .serverCancelled because no handler is specified in \(self.name, privacy: .public)