Show the only way we can implement withSpan correctly inside an actor

Author: ktoso

Sources/Tracing/Tracer.swift

@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 import Dispatch
+import Distributed
 @_exported import Instrumentation
 @_exported import ServiceContextModule
 
@@ -382,7 +383,7 @@ public func withSpan<T>(
     function: String = #function,
     file fileID: String = #fileID,
     line: UInt = #line,
-    _ operation: @Sendable (any Span) async throws -> T
+    _ operation: (any Span) async throws -> T
 ) async rethrows -> T {
     try await InstrumentationSystem.legacyTracer.withAnySpan(
         operationName,
@@ -477,7 +478,7 @@ public func withSpan<T>(
     function: String = #function,
     file fileID: String = #fileID,
     line: UInt = #line,
-    _ operation: @Sendable (any Span) async throws -> T
+    _ operation: (any Span) async throws -> T
 ) async rethrows -> T {
     try await InstrumentationSystem.legacyTracer.withAnySpan(
         operationName,
@@ -491,4 +492,114 @@ public func withSpan<T>(
         try await operation(anySpan)
     }
 }
+
+// FIXME: type checker does not understand that AnyActor shall suffice for checking isolation domains,
+//        thus we can't declare the method once, but have to duplicate it for every actor kind:
+//@available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
+//extension AnyActor {
+//    public func withSpanWorkaround<T>(
+//        _ operationName: String,
+//        context: @autoclosure () -> ServiceContext = .current ?? .topLevel,
+//        ofKind kind: SpanKind = .internal,
+//        at instant: @autoclosure () -> some TracerInstant = DefaultTracerClock.now,
+//        function: String = #function,
+//        file fileID: String = #fileID,
+//        line: UInt = #line,
+//        _ operation: (any Span) async throws -> T
+//    ) async rethrows -> T {
+//        fatalError("Can't use this, must extend both types separately")
+//    }
+//}
+
+@available(macOS 10.15, iOS 13, tvOS 13, watchOS 6, *)
+extension Actor {
+
+    public func withSpanWorkaround<T>(
+        _ operationName: String,
+        context: @autoclosure () -> ServiceContext = .current ?? .topLevel,
+        ofKind kind: SpanKind = .internal,
+        at instant: @autoclosure () -> some TracerInstant = DefaultTracerClock.now,
+        function: String = #function,
+        file fileID: String = #fileID,
+        line: UInt = #line,
+        _ operation: (any Span) async throws -> T
+    ) async rethrows -> T {
+        let span = InstrumentationSystem.legacyTracer.startAnySpan(
+            operationName,
+            context: context(),
+            ofKind: kind,
+            at: instant(),
+            function: function,
+            file: fileID,
+            line: line
+        )
+        defer { span.end() }
+        do {
+            return try await ServiceContext.$current.withValue(span.context) {
+                try await operation(span)
+            }
+        } catch {
+            span.recordError(error)
+            throw error // rethrow
+        }
+//        // FIXME: can't use withSpan to implement this... must use startSpan and duplicate much code
+//        try await InstrumentationSystem.legacyTracer.withAnySpan(
+//            operationName,
+//            at: instant(),
+//            context: context(),
+//            ofKind: kind,
+//            function: function,
+//            file: fileID,
+//            line: line
+//        ) { anySpan in
+//            try await operation(anySpan)
+//        }
+    }
+}
+
+@available(macOS 13.0, iOS 16.0, watchOS 9.0, tvOS 16.0, *)
+extension DistributedActor {
+
+    public func withSpanWorkaround<T>(
+        _ operationName: String,
+        context: @autoclosure () -> ServiceContext = .current ?? .topLevel,
+        ofKind kind: SpanKind = .internal,
+        at instant: @autoclosure () -> some TracerInstant = DefaultTracerClock.now,
+        function: String = #function,
+        file fileID: String = #fileID,
+        line: UInt = #line,
+        _ operation: (any Span) async throws -> T
+    ) async rethrows -> T {
+        let span = InstrumentationSystem.legacyTracer.startAnySpan(
+            operationName,
+            context: context(),
+            ofKind: kind,
+            at: instant(),
+            function: function,
+            file: fileID,
+            line: line
+        )
+        defer { span.end() }
+        do {
+            return try await ServiceContext.$current.withValue(span.context) {
+                try await operation(span)
+            }
+        } catch {
+            span.recordError(error)
+            throw error // rethrow
+        }
+//        // FIXME: can't use withSpan to implement this... must use startSpan and duplicate much code
+//        try await InstrumentationSystem.legacyTracer.withAnySpan(
+//            operationName,
+//            at: instant(),
+//            context: context(),
+//            ofKind: kind,
+//            function: function,
+//            file: fileID,
+//            line: line
+//        ) { anySpan in
+//            try await operation(anySpan)
+//        }
+    }
+}
 #endif

Tests/TracingTests/ActorTracingTests.swift

@@ -12,9 +12,10 @@
 //
 //===----------------------------------------------------------------------===//
 
-@testable @_spi(Locking) import Instrumentation
+@testable import Instrumentation
 import ServiceContextModule
 import Tracing
+import Distributed
 import XCTest
 
 final class ActorTracingTests: XCTestCase {
@@ -27,17 +28,30 @@ final class ActorTracingTests: XCTestCase {
 func work() async {}
 
 actor Foo {
-    let bar: LockedValueBox<Int> = .init(0)
+    var bar = 0
+
+    func foo() async {
+        var num = 0
+        await self.withSpanWorkaround(#function) { _ in
+            self.bar += 1
+            await work()
+            num += 1
+        }
+    }
+}
+
+@available(macOS 13.0, iOS 16.0, watchOS 9.0, tvOS 16.0, *)
+distributed actor DistributedFoo {
+    typealias ActorSystem = LocalTestingDistributedActorSystem
+
+    var bar = 0
+
     func foo() async {
-        let num: LockedValueBox<Int> = .init(0)
-        await withSpan(#function) { _ in
-            self.bar.withValue { bar in
-                bar += 1
-            }
+        var num = 0
+        await self.withSpanWorkaround(#function) { _ in
+            self.bar += 1
             await work()
-            num.withValue { num in
-                num += 1
-            }
+            num += 1
         }
     }
 }