improve thread safety for expectation helper, and remove some flakey tests (#23)

* improve thread safety for `expectation` helper, and remove flakey tests

* fix flakey tests

* lint

Author: gsabran

.github/workflows/swift.yml

@@ -14,10 +14,8 @@ jobs:
     runs-on: macos-latest
     steps:
     - uses: swift-actions/setup-swift@v2
-    with:
-      swift-version: "6.0.1"
-
-    steps:
+      with:
+        swift-version: "6.0.1"
     - name: Get swift version
       run: swift --version
     - uses: actions/checkout@v4

MCPClient/Sources/stdioTransport/DataChannel+StdioProcess.swift

@@ -44,7 +44,7 @@ extension JSONRPCSetupError: LocalizedError {
 
 extension Transport {
 
-#if os(macOS)
+  #if os(macOS)
   /// Creates a new `Transport` by launching the given executable with the specified arguments and attaching to its standard IO.
   /// This functionality is only available on macOS.
   public static func stdioProcess(
@@ -286,7 +286,7 @@ extension Transport {
 
     return String(data: stdoutData, encoding: .utf8)?.trimmingCharacters(in: .whitespacesAndNewlines)
   }
-#endif
+  #endif
 }
 
 // MARK: - Lifetime

Package.swift

@@ -7,7 +7,7 @@ let package = Package(
   name: "MCP",
   platforms: [
     .macOS(.v14),
-    .iOS(.v17)
+    .iOS(.v17),
   ],
   products: [
     .library(

SwiftTestingUtils/Sources/Atomic.swift

@@ -0,0 +1,20 @@
+import os
+
+final class Atomic<Value: Sendable>: Sendable {
+
+  init(_ value: Value) {
+    lock = .init(initialState: value)
+  }
+
+  var value: Value {
+    lock.withLock { $0 }
+  }
+
+  @discardableResult
+  func mutate<Result: Sendable>(_ mutation: @Sendable (inout Value) -> Result) -> Result {
+    lock.withLock { mutation(&$0) }
+  }
+
+  private let lock: OSAllocatedUnfairLock<Value>
+
+}

SwiftTestingUtils/Sources/SwiftTestingUtils.swift

@@ -1,5 +1,5 @@
-
 import Foundation
+import os
 import OSLog
 import Testing
 
@@ -9,7 +9,7 @@ private let logger = Logger(subsystem: "testing", category: "testing")
 
 public enum SwiftTestingUtils {
 
-  public class Expectation {
+  public final class Expectation: Sendable {
 
     init(
       description: String,
@@ -19,16 +19,24 @@ public enum SwiftTestingUtils {
       location = _sourceLocation
     }
 
+    public var isFulfilled: Bool {
+      lock.withLock { $0.isFulfilled }
+    }
+
     public func fulfill() {
-      if isFulfilled {
+      let (wasFulfilled, fullfillmentActions) = lock.withLock { state in
+        let wasFulfilled = state.isFulfilled
+        state.isFulfilled = true
+        return (wasFulfilled, state.onFulfill)
+      }
+
+      if wasFulfilled {
         Issue.record("Expectation from \(location) already fulfilled.")
+        return
       }
-      isFulfilled = true
-      onFulfill()
+      for fullfillmentAction in fullfillmentActions { fullfillmentAction() }
     }
 
-    private(set) var isFulfilled = false
-
     let description: String
     let location: SourceLocation
 
@@ -37,34 +45,47 @@ public enum SwiftTestingUtils {
         return
       }
 
-      var hasTimedOut = false
-      var hasCompleted = false
+      let hasTimedOut = Atomic(false)
       let description = description
 
       try await withCheckedThrowingContinuation { continuation in
-        onFulfill = {
-          if !hasTimedOut {
-            if hasCompleted {
-              Issue.record("Expectation \(description) fulfilled multiple times.")
-            } else {
-              hasCompleted = true
-              continuation.resume(returning: ())
-            }
+        let onFulfill: @Sendable () -> Void = {
+          if !hasTimedOut.value {
+            continuation.resume(returning: ())
           } else {
             logger.error("Expectation \(description) fulfilled after timeout.")
           }
         }
+        let wasFulfilled = lock.withLock { state in
+          if state.isFulfilled {
+            return true
+          }
+          state.onFulfill.append(onFulfill)
+          return false
+        }
+        if wasFulfilled {
+          continuation.resume(returning: ())
+        }
         Task {
-          try await Task.sleep(nanoseconds: UInt64(timeout * 1_000_000_000))
-          if !hasCompleted {
-            hasTimedOut = true
-            continuation.resume(throwing: SwiftTesting.expectationTimeout(description: description))
+          do {
+            try await Task.sleep(nanoseconds: UInt64(timeout * 1_000_000_000))
+            if !isFulfilled {
+              hasTimedOut.mutate { $0 = true }
+              continuation.resume(throwing: SwiftTesting.expectationTimeout(description: description))
+            }
+          } catch {
+            continuation.resume(throwing: error)
           }
         }
       }
     }
 
-    private var onFulfill: () -> Void = { }
+    private struct InternalState {
+      var isFulfilled = false
+      var onFulfill: [@Sendable () -> Void] = []
+    }
+
+    private let lock = OSAllocatedUnfairLock<InternalState>(initialState: InternalState())
   }
 
   enum SwiftTesting: Error {