+span additional withSpan to ease entering async world from non-async world (#64)

* +span additional withSpan to ease entering async world from non-async world

* add test

* fix API, always require Baggage, people should pass .topLevel

Author: ktoso

Sources/Tracing/Tracer.swift

@@ -150,5 +150,36 @@ extension Tracer {
             throw error // rethrow
         }
     }
+
+    /// Execute the given async operation within a newly created `Span`,
+    /// started as a child of the passed in `Baggage` or as a root span if `nil`.
+    ///
+    /// DO NOT `end()` the passed in span manually. It will be ended automatically when the `operation` returns.
+    ///
+    /// - Parameters:
+    ///   - operationName: The name of the operation being traced. This may be a handler function, database call, ...
+    ///   - baggage: The baggage to be used for the newly created span. It may be obtained by the user manually from the `Baggage.current`,
+    //               task local and modified before passing into this function. The baggage will be made the current task-local baggage for the duration of the `operation`.
+    ///   - kind: The `SpanKind` of the `Span` to be created. Defaults to `.internal`.
+    ///   - operation: operation to wrap in a span start/end and execute immediately
+    /// - Returns: the value returned by `operation`
+    /// - Throws: the error the `operation` has thrown (if any)
+    public func withSpan<T>(
+        _ operationName: String,
+        baggage: Baggage,
+        ofKind kind: SpanKind = .internal,
+        _ operation: (Span) async throws -> T
+    ) async rethrows -> T {
+        let span = self.startSpan(operationName, baggage: baggage, ofKind: kind)
+        defer { span.end() }
+        do {
+            return try await Baggage.$current.withValue(span.baggage) {
+                try await operation(span)
+            }
+        } catch {
+            span.recordError(error)
+            throw error // rethrow
+        }
+    }
 }
 #endif

Tests/TracingTests/TracerTests+XCTest.swift

@@ -33,6 +33,7 @@ extension TracerTests {
                 ("testWithSpan_automaticBaggagePropagation_sync", testWithSpan_automaticBaggagePropagation_sync),
                 ("testWithSpan_automaticBaggagePropagation_sync_throws", testWithSpan_automaticBaggagePropagation_sync_throws),
                 ("testWithSpan_automaticBaggagePropagation_async", testWithSpan_automaticBaggagePropagation_async),
+                ("testWithSpan_enterFromNonAsyncCode_passBaggage_asyncOperation", testWithSpan_enterFromNonAsyncCode_passBaggage_asyncOperation),
                 ("testWithSpan_automaticBaggagePropagation_async_throws", testWithSpan_automaticBaggagePropagation_async_throws),
            ]
    }

Tests/TracingTests/TracerTests.swift

@@ -187,6 +187,40 @@ final class TracerTests: XCTestCase {
         #endif
     }
 
+    func testWithSpan_enterFromNonAsyncCode_passBaggage_asyncOperation() throws {
+        #if swift(>=5.5) && canImport(_Concurrency)
+        guard #available(macOS 10.15, iOS 13.0, tvOS 13.0, watchOS 6.0, *) else {
+            throw XCTSkip("Task locals are not supported on this platform.")
+        }
+
+        let tracer = TestTracer()
+        InstrumentationSystem.bootstrapInternal(tracer)
+        defer {
+            InstrumentationSystem.bootstrapInternal(nil)
+        }
+
+        var spanEnded = false
+        tracer.onEndSpan = { _ in spanEnded = true }
+
+        func operation(span: Span) async -> String {
+            "world"
+        }
+
+        self.testAsync {
+            var fromNonAsyncWorld = Baggage.topLevel
+            fromNonAsyncWorld.traceID = "1234-5678"
+            let value = await tracer.withSpan("hello", baggage: fromNonAsyncWorld) { (span: Span) -> String in
+                XCTAssertEqual(span.baggage.traceID, Baggage.current?.traceID)
+                XCTAssertEqual(span.baggage.traceID, fromNonAsyncWorld.traceID)
+                return await operation(span: span)
+            }
+
+            XCTAssertEqual(value, "world")
+            XCTAssertTrue(spanEnded)
+        }
+        #endif
+    }
+
     func testWithSpan_automaticBaggagePropagation_async_throws() throws {
         #if swift(>=5.5) && canImport(_Concurrency)
         guard #available(macOS 10.15, iOS 13.0, tvOS 13.0, watchOS 6.0, *) else {