Add support for Default Baggage Injection in URLSessionInstrumentation (#639)

* add default baggage provider to urlsessioninstrumentation
* add missing `DataCompression` dep to `OpenTelemetryProtocolExporterHttp`
* add missing deps to all package files

Author: batuhansk

Package.swift

@@ -104,7 +104,8 @@ let package = Package(
                 path: "Sources/Exporters/OpenTelemetryProtocolCommon"),
         .target(name: "OpenTelemetryProtocolExporterHttp",
                 dependencies: ["OpenTelemetrySdk",
-                               "OpenTelemetryProtocolExporterCommon"],
+                               "OpenTelemetryProtocolExporterCommon",
+                               "DataCompression"],
                 path: "Sources/Exporters/OpenTelemetryProtocolHttp"),
         .target(name: "OpenTelemetryProtocolExporterGrpc",
                 dependencies: ["OpenTelemetrySdk",

[email protected]

@@ -109,7 +109,8 @@ let package = Package(
                 path: "Sources/Exporters/OpenTelemetryProtocolCommon"),
         .target(name: "OpenTelemetryProtocolExporterHttp",
                 dependencies: ["OpenTelemetrySdk",
-                               "OpenTelemetryProtocolExporterCommon"],
+                               "OpenTelemetryProtocolExporterCommon",
+                               "DataCompression"],
                 path: "Sources/Exporters/OpenTelemetryProtocolHttp"),
         .target(name: "OpenTelemetryProtocolExporterGrpc",
                 dependencies: ["OpenTelemetrySdk",

[email protected]

@@ -64,7 +64,8 @@ let package = Package(
                 path: "Sources/Exporters/OpenTelemetryProtocolCommon"),
         .target(name: "OpenTelemetryProtocolExporterHttp",
                 dependencies: ["OpenTelemetrySdk",
-                               "OpenTelemetryProtocolExporterCommon"],
+                               "OpenTelemetryProtocolExporterCommon",
+                               "DataCompression"],
                 path: "Sources/Exporters/OpenTelemetryProtocolHttp"),
         .target(name: "OpenTelemetryProtocolExporterGrpc",
                 dependencies: ["OpenTelemetrySdk",

Sources/Instrumentation/URLSession/URLSessionInstrumentationConfiguration.swift

@@ -24,7 +24,8 @@ public struct URLSessionInstrumentationConfiguration {
                 createdRequest: ((URLRequest, Span) -> Void)? = nil,
                 receivedResponse: ((URLResponse, DataOrFile?, Span) -> Void)? = nil,
                 receivedError: ((Error, DataOrFile?, HTTPStatus, Span) -> Void)? = nil,
-                delegateClassesToInstrument: [AnyClass]? = nil)
+                delegateClassesToInstrument: [AnyClass]? = nil,
+                defaultBaggageProvider: (() -> (Baggage)?)? = nil)
     {
         self.shouldRecordPayload = shouldRecordPayload
         self.shouldInstrument = shouldInstrument
@@ -36,6 +37,7 @@ public struct URLSessionInstrumentationConfiguration {
         self.receivedResponse = receivedResponse
         self.receivedError = receivedError
         self.delegateClassesToInstrument = delegateClassesToInstrument
+        self.defaultBaggageProvider = defaultBaggageProvider
     }
 
     // Instrumentation Callbacks
@@ -73,4 +75,14 @@ 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.
+    ///
+    /// 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)?)?
+
 }

Sources/Instrumentation/URLSession/URLSessionLogger.swift

@@ -160,15 +160,22 @@ class URLSessionLogger {
         var instrumentedRequest = request
         objc_setAssociatedObject(instrumentedRequest, URLSessionInstrumentation.instrumentedKey, true, .OBJC_ASSOCIATION_COPY_NONATOMIC)
         let propagators = OpenTelemetry.instance.propagators
-        var traceHeaders = tracePropagationHTTPHeaders(span: span, textMapPropagator: propagators.textMapPropagator, textMapBaggagePropagator: propagators.textMapBaggagePropagator)
+
+        let defaultBaggage = instrumentation.configuration.defaultBaggageProvider?()
+
+        var traceHeaders = tracePropagationHTTPHeaders(span: span,
+                                                       defaultBaggage: defaultBaggage,
+                                                       textMapPropagator: propagators.textMapPropagator,
+                                                       textMapBaggagePropagator: propagators.textMapBaggagePropagator)
+
         if let originalHeaders = request.allHTTPHeaderFields {
             traceHeaders.merge(originalHeaders) { _, new in new }
         }
         instrumentedRequest.allHTTPHeaderFields = traceHeaders
         return instrumentedRequest
     }
 
-    private static func tracePropagationHTTPHeaders(span: Span?, textMapPropagator: TextMapPropagator, textMapBaggagePropagator: TextMapBaggagePropagator) -> [String: String] {
+    private static func tracePropagationHTTPHeaders(span: Span?, defaultBaggage: Baggage?, textMapPropagator: TextMapPropagator, textMapBaggagePropagator: TextMapBaggagePropagator) -> [String: String] {
         var headers = [String: String]()
 
         struct HeaderSetter: Setter {
@@ -182,9 +189,19 @@ class URLSessionLogger {
         }
         textMapPropagator.inject(spanContext: currentSpan.context, carrier: &headers, setter: HeaderSetter())
 
-        if let baggage = OpenTelemetry.instance.contextProvider.activeBaggage {
-            textMapBaggagePropagator.inject(baggage: baggage, carrier: &headers, setter: HeaderSetter())
+        let baggageBuilder = OpenTelemetry.instance.baggageManager.baggageBuilder()
+
+        if let activeBaggage = OpenTelemetry.instance.contextProvider.activeBaggage {
+            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) }
+        }
+
+        let combinedBaggage = baggageBuilder.build()
+        textMapBaggagePropagator.inject(baggage: combinedBaggage, carrier: &headers, setter: HeaderSetter())
+
         return headers
     }
 }

Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift

@@ -39,6 +39,9 @@ class URLSessionInstrumentationTests: XCTestCase {
     static var requestCopy: URLRequest!
     static var responseCopy: HTTPURLResponse!
 
+    static var activeBaggage: Baggage!
+    static var defaultBaggage: Baggage!
+
     static var config = URLSessionInstrumentationConfiguration(shouldRecordPayload: nil,
                                                                shouldInstrument: { req in
                                                                    checker.shouldInstrumentCalled = true
@@ -76,25 +79,31 @@ class URLSessionInstrumentationTests: XCTestCase {
                                                                },
                                                                receivedError: { _, _, _, _ in
                                                                    URLSessionInstrumentationTests.checker.receivedErrorCalled = true
+                                                               },
+                                                               defaultBaggageProvider: {
+                                                                   defaultBaggage
                                                                })
 
     static var checker = Check()
     static var semaphore: DispatchSemaphore!
     var sessionDelegate: SessionDelegate!
     static var instrumentation: URLSessionInstrumentation!
-    static var baggage: Baggage!
 
     static let server = HttpTestServer(url: URL(string: "http://localhost:33333"), config: nil)
 
     override class func setUp() {
         OpenTelemetry.registerPropagators(textPropagators: [W3CTraceContextPropagator()], baggagePropagator: W3CBaggagePropagator())
         OpenTelemetry.registerTracerProvider(tracerProvider: TracerProviderSdk())
 
-        baggage = DefaultBaggageManager.instance.baggageBuilder()
+        defaultBaggage = DefaultBaggageManager.instance.baggageBuilder()
+            .put(key: EntryKey(name: "bar")!, value: EntryValue(string: "baz")!, metadata: nil)
+            .build()
+
+        activeBaggage = DefaultBaggageManager.instance.baggageBuilder()
             .put(key: EntryKey(name: "foo")!, value: EntryValue(string: "bar")!, metadata: nil)
             .build()
 
-        OpenTelemetry.instance.contextProvider.setActiveBaggage(baggage)
+        OpenTelemetry.instance.contextProvider.setActiveBaggage(activeBaggage)
 
         let sem = DispatchSemaphore(value: 0)
         DispatchQueue.global(qos: .default).async {
@@ -111,7 +120,8 @@ class URLSessionInstrumentationTests: XCTestCase {
 
     override class func tearDown() {
         server.stop()
-        OpenTelemetry.instance.contextProvider.removeContextForBaggage(baggage)
+        defaultBaggage = nil
+        OpenTelemetry.instance.contextProvider.removeContextForBaggage(activeBaggage)
     }
 
     override func setUp() {
@@ -262,13 +272,63 @@ class URLSessionInstrumentationTests: XCTestCase {
 
         XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInstrumentCalled)
 
+        XCTAssertEqual(1, URLSessionLogger.runningSpans.count)
+
+        let span = try XCTUnwrap(URLSessionLogger.runningSpans["111"])
+        XCTAssertEqual("HTTP GET", span.name)
+    }
+
+    public func testShouldInstrumentRequest_PropagateCombinedActiveAndDefaultBaggages() throws {
+        let request1 = URLRequest(url: URL(string: "http://defaultName.com")!)
+        let request2 = URLRequest(url: URL(string: "http://dontinstrument.com")!)
+
+        let processedRequest1 = try XCTUnwrap(URLSessionLogger.processAndLogRequest(request1, sessionTaskId: "111", instrumentation: URLSessionInstrumentationTests.instrumentation, shouldInjectHeaders: true))
+        let processedRequest2 = URLSessionLogger.processAndLogRequest(request2, sessionTaskId: "222", instrumentation: URLSessionInstrumentationTests.instrumentation, shouldInjectHeaders: true)
+
+        // `processedRequest2` is expected to be nil, because its URL was marked as not to be instrumented.
+        XCTAssertNil(processedRequest2)
+
+        XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInstrumentCalled)
+
+        let processedHeaders1 = try XCTUnwrap(processedRequest1.allHTTPHeaderFields)
+
+        // headers injected from `TextMapPropagator` implementation
+        XCTAssertTrue(processedHeaders1.contains(where: { $0.key == W3CTraceContextPropagator.traceparent }))
+
+        // headers injected from `TextMapBaggagePropagator` implementation
+        let baggageHeaderValue = try XCTUnwrap(processedHeaders1[W3CBaggagePropagator.headerBaggage])
+
+        // foo=bar propagated through active baggage defined in `setUp`
+        XCTAssertTrue(baggageHeaderValue.contains("foo=bar"))
+
+        // bar=baz propagated through default baggage provided in `URLSessionInstrumentationConfiguration`
+        XCTAssertTrue(baggageHeaderValue.contains("bar=baz"))
+
+        XCTAssertEqual(1, URLSessionLogger.runningSpans.count)
+
+        let span = try XCTUnwrap(URLSessionLogger.runningSpans["111"])
+        XCTAssertEqual("HTTP GET", span.name)
+    }
+
+    public func testShouldInstrumentRequest_PropagateOnlyActiveBaggage() throws {
+        Self.defaultBaggage = nil
+
+        let request1 = URLRequest(url: URL(string: "http://defaultName.com")!)
+
+        let processedRequest1 = try XCTUnwrap(URLSessionLogger.processAndLogRequest(request1, sessionTaskId: "111", instrumentation: URLSessionInstrumentationTests.instrumentation, shouldInjectHeaders: true))
+
+        XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInstrumentCalled)
+
         let processedHeaders1 = try XCTUnwrap(processedRequest1.allHTTPHeaderFields)
 
         // headers injected from `TextMapPropagator` implementation
         XCTAssertTrue(processedHeaders1.contains(where: { $0.key == W3CTraceContextPropagator.traceparent }))
 
         // headers injected from `TextMapBaggagePropagator` implementation
-        XCTAssertTrue(processedHeaders1.contains(where: { $0.key == W3CBaggagePropagator.headerBaggage && $0.value == "foo=bar" }))
+        let baggageHeaderValue = try XCTUnwrap(processedHeaders1[W3CBaggagePropagator.headerBaggage])
+
+        // bar=baz propagated through default baggage provided in `URLSessionInstrumentationConfiguration`
+        XCTAssertEqual(baggageHeaderValue, "foo=bar")
 
         XCTAssertEqual(1, URLSessionLogger.runningSpans.count)