Add Request Context to URLSessionInstrumentation Baggage Provider

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.
 

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)?)?
 
 }

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()

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")!)