diff --git a/Sources/Tracing/Tracer.swift b/Sources/Tracing/Tracer.swift index 241dda1..91857d7 100644 --- a/Sources/Tracing/Tracer.swift +++ b/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( + _ 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 diff --git a/Tests/TracingTests/TracerTests+XCTest.swift b/Tests/TracingTests/TracerTests+XCTest.swift index 1737ae5..360df37 100644 --- a/Tests/TracingTests/TracerTests+XCTest.swift +++ b/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), ] } diff --git a/Tests/TracingTests/TracerTests.swift b/Tests/TracingTests/TracerTests.swift index 02141cd..987f3e7 100644 --- a/Tests/TracingTests/TracerTests.swift +++ b/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 {