diff --git a/Sources/Instrumentation/URLSession/README.md b/Sources/Instrumentation/URLSession/README.md index 063ef1ae..7bca6998 100644 --- a/Sources/Instrumentation/URLSession/README.md +++ b/Sources/Instrumentation/URLSession/README.md @@ -28,4 +28,5 @@ This behaviour can be modified or augmented by using the optional callbacks defi `receivedError: ((Error, DataOrFile?, HTTPStatus, Span) -> Void)?` - Called after an error is received, it allows to add extra information to the Span +`baggageProvider: ((inout URLRequest, Span) -> (Baggage)?)?`: Provides baggage instance for instrumented requests that is merged with active baggage. The callback receives URLRequest and Span parameters to create dynamic baggage based on request context. The resulting baggage is injected into request headers using the configured propagator. diff --git a/Sources/Instrumentation/URLSession/URLSessionInstrumentationConfiguration.swift b/Sources/Instrumentation/URLSession/URLSessionInstrumentationConfiguration.swift index 2636244d..424191a1 100644 --- a/Sources/Instrumentation/URLSession/URLSessionInstrumentationConfiguration.swift +++ b/Sources/Instrumentation/URLSession/URLSessionInstrumentationConfiguration.swift @@ -25,7 +25,7 @@ public struct URLSessionInstrumentationConfiguration { receivedResponse: ((URLResponse, DataOrFile?, Span) -> Void)? = nil, receivedError: ((Error, DataOrFile?, HTTPStatus, Span) -> Void)? = nil, delegateClassesToInstrument: [AnyClass]? = nil, - defaultBaggageProvider: (() -> (Baggage)?)? = nil) + baggageProvider: ((inout URLRequest, Span?) -> (Baggage)?)? = nil) { self.shouldRecordPayload = shouldRecordPayload self.shouldInstrument = shouldInstrument @@ -37,7 +37,7 @@ public struct URLSessionInstrumentationConfiguration { self.receivedResponse = receivedResponse self.receivedError = receivedError self.delegateClassesToInstrument = delegateClassesToInstrument - self.defaultBaggageProvider = defaultBaggageProvider + self.baggageProvider = baggageProvider } // Instrumentation Callbacks @@ -76,13 +76,15 @@ public struct URLSessionInstrumentationConfiguration { /// The array of URLSession delegate classes that will be instrumented by the library, will autodetect if nil is passed. public var delegateClassesToInstrument: [AnyClass]? - /// Implement this callback to provide a default baggage instance for all instrumented requests. - /// The provided baggage is merged with active baggage (if any) to create a combined baggage. - /// The combined baggage is then injected into request headers using the configured `TextMapBaggagePropagator`. - /// This allows consistent propagation across requests, regardless of the active context. + /// Provides a baggage instance for instrumented requests that is merged with active baggage (if any). + /// The callback can be used to define static baggage for all requests or create dynamic baggage + /// based on the provided URLRequest and Span parameters. + /// + /// The resulting baggage is injected into request headers using the configured `TextMapBaggagePropagator`, + /// ensuring consistent propagation across requests, regardless of the active context. /// /// Note: The injected baggage depends on the propagator in use (e.g., W3C or custom). - /// Returns: A `Baggage` instance or `nil` if no default baggage is needed. - public let defaultBaggageProvider: (() -> (Baggage)?)? + /// Returns: A `Baggage` instance or `nil` if no baggage is needed. + public let baggageProvider: ((inout URLRequest, Span?) -> (Baggage)?)? } diff --git a/Sources/Instrumentation/URLSession/URLSessionLogger.swift b/Sources/Instrumentation/URLSession/URLSessionLogger.swift index c41fc34c..92a6b48c 100644 --- a/Sources/Instrumentation/URLSession/URLSessionLogger.swift +++ b/Sources/Instrumentation/URLSession/URLSessionLogger.swift @@ -157,14 +157,14 @@ class URLSessionLogger { return nil } instrumentation.configuration.injectCustomHeaders?(&request, span) + let customBaggage = instrumentation.configuration.baggageProvider?(&request, span) + var instrumentedRequest = request objc_setAssociatedObject(instrumentedRequest, URLSessionInstrumentation.instrumentedKey, true, .OBJC_ASSOCIATION_COPY_NONATOMIC) let propagators = OpenTelemetry.instance.propagators - let defaultBaggage = instrumentation.configuration.defaultBaggageProvider?() - var traceHeaders = tracePropagationHTTPHeaders(span: span, - defaultBaggage: defaultBaggage, + customBaggage: customBaggage, textMapPropagator: propagators.textMapPropagator, textMapBaggagePropagator: propagators.textMapBaggagePropagator) @@ -175,7 +175,7 @@ class URLSessionLogger { return instrumentedRequest } - private static func tracePropagationHTTPHeaders(span: Span?, defaultBaggage: Baggage?, textMapPropagator: TextMapPropagator, textMapBaggagePropagator: TextMapBaggagePropagator) -> [String: String] { + private static func tracePropagationHTTPHeaders(span: Span?, customBaggage: Baggage?, textMapPropagator: TextMapPropagator, textMapBaggagePropagator: TextMapBaggagePropagator) -> [String: String] { var headers = [String: String]() struct HeaderSetter: Setter { @@ -195,8 +195,8 @@ class URLSessionLogger { activeBaggage.getEntries().forEach { baggageBuilder.put(key: $0.key, value: $0.value, metadata: $0.metadata) } } - if let defaultBaggage { - defaultBaggage.getEntries().forEach { baggageBuilder.put(key: $0.key, value: $0.value, metadata: $0.metadata) } + if let customBaggage { + customBaggage.getEntries().forEach { baggageBuilder.put(key: $0.key, value: $0.value, metadata: $0.metadata) } } let combinedBaggage = baggageBuilder.build() diff --git a/Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift b/Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift index d941df53..33b52bbf 100644 --- a/Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift +++ b/Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift @@ -40,7 +40,7 @@ class URLSessionInstrumentationTests: XCTestCase { static var responseCopy: HTTPURLResponse! static var activeBaggage: Baggage! - static var defaultBaggage: Baggage! + static var customBaggage: Baggage! static var config = URLSessionInstrumentationConfiguration(shouldRecordPayload: nil, shouldInstrument: { req in @@ -80,8 +80,8 @@ class URLSessionInstrumentationTests: XCTestCase { receivedError: { _, _, _, _ in URLSessionInstrumentationTests.checker.receivedErrorCalled = true }, - defaultBaggageProvider: { - defaultBaggage + baggageProvider: { _, _ in + customBaggage }) static var checker = Check() @@ -95,7 +95,7 @@ class URLSessionInstrumentationTests: XCTestCase { OpenTelemetry.registerPropagators(textPropagators: [W3CTraceContextPropagator()], baggagePropagator: W3CBaggagePropagator()) OpenTelemetry.registerTracerProvider(tracerProvider: TracerProviderSdk()) - defaultBaggage = DefaultBaggageManager.instance.baggageBuilder() + customBaggage = DefaultBaggageManager.instance.baggageBuilder() .put(key: EntryKey(name: "bar")!, value: EntryValue(string: "baz")!, metadata: nil) .build() @@ -120,7 +120,7 @@ class URLSessionInstrumentationTests: XCTestCase { override class func tearDown() { server.stop() - defaultBaggage = nil + customBaggage = nil OpenTelemetry.instance.contextProvider.removeContextForBaggage(activeBaggage) } @@ -278,7 +278,7 @@ class URLSessionInstrumentationTests: XCTestCase { XCTAssertEqual("HTTP GET", span.name) } - public func testShouldInstrumentRequest_PropagateCombinedActiveAndDefaultBaggages() throws { + public func testShouldInstrumentRequest_PropagateCombinedActiveAndCustomBaggages() throws { let request1 = URLRequest(url: URL(string: "http://defaultName.com")!) let request2 = URLRequest(url: URL(string: "http://dontinstrument.com")!) @@ -301,7 +301,7 @@ class URLSessionInstrumentationTests: XCTestCase { // foo=bar propagated through active baggage defined in `setUp` XCTAssertTrue(baggageHeaderValue.contains("foo=bar")) - // bar=baz propagated through default baggage provided in `URLSessionInstrumentationConfiguration` + // bar=baz propagated through custom baggage provided in `URLSessionInstrumentationConfiguration` XCTAssertTrue(baggageHeaderValue.contains("bar=baz")) XCTAssertEqual(1, URLSessionLogger.runningSpans.count) @@ -311,7 +311,7 @@ class URLSessionInstrumentationTests: XCTestCase { } public func testShouldInstrumentRequest_PropagateOnlyActiveBaggage() throws { - Self.defaultBaggage = nil + Self.customBaggage = nil let request1 = URLRequest(url: URL(string: "http://defaultName.com")!)