From 295abe34646165c7c0b4be6bcc3c50d1f41ef1fa Mon Sep 17 00:00:00 2001 From: Batuhan Saka Date: Wed, 20 Nov 2024 15:20:55 +0100 Subject: [PATCH 1/5] add default baggage provider to urlsessioninstrumentation --- ...LSessionInstrumentationConfiguration.swift | 14 ++++- .../URLSession/URLSessionLogger.swift | 25 +++++++-- .../Baggage/DefaultBaggageBuilder.swift | 1 + .../URLSessionInstrumentationTests.swift | 54 ++++++++++++++++--- 4 files changed, 83 insertions(+), 11 deletions(-) diff --git a/Sources/Instrumentation/URLSession/URLSessionInstrumentationConfiguration.swift b/Sources/Instrumentation/URLSession/URLSessionInstrumentationConfiguration.swift index 9566f206..2636244d 100644 --- a/Sources/Instrumentation/URLSession/URLSessionInstrumentationConfiguration.swift +++ b/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)?)? + } diff --git a/Sources/Instrumentation/URLSession/URLSessionLogger.swift b/Sources/Instrumentation/URLSession/URLSessionLogger.swift index 3c4e291e..8f856f1b 100644 --- a/Sources/Instrumentation/URLSession/URLSessionLogger.swift +++ b/Sources/Instrumentation/URLSession/URLSessionLogger.swift @@ -160,7 +160,14 @@ 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 } } @@ -168,7 +175,7 @@ class URLSessionLogger { 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 } } diff --git a/Sources/OpenTelemetryApi/Baggage/DefaultBaggageBuilder.swift b/Sources/OpenTelemetryApi/Baggage/DefaultBaggageBuilder.swift index 53acf878..b8714f5f 100644 --- a/Sources/OpenTelemetryApi/Baggage/DefaultBaggageBuilder.swift +++ b/Sources/OpenTelemetryApi/Baggage/DefaultBaggageBuilder.swift @@ -61,6 +61,7 @@ public class DefaultBaggageBuilder: BaggageBuilder { } } } + print("Entries count: \(entries.count)") return entries.isEmpty ? EmptyBaggage.instance : DefaultBaggage(entries: combined) } } diff --git a/Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift b/Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift index 9a21702b..5535155a 100644 --- a/Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift +++ b/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,13 +79,15 @@ 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) @@ -90,11 +95,15 @@ class URLSessionInstrumentationTests: XCTestCase { 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() { @@ -250,7 +260,7 @@ class URLSessionInstrumentationTests: XCTestCase { XCTAssertTrue(URLSessionInstrumentationTests.checker.receivedErrorCalled) } - public func testShouldInstrumentRequest() throws { + public func testShouldInstrumentRequest_PropagateCombinedActiveAndDefaultBaggages() throws { let request1 = URLRequest(url: URL(string: "http://defaultName.com")!) let request2 = URLRequest(url: URL(string: "http://dontinstrument.com")!) @@ -268,7 +278,39 @@ class URLSessionInstrumentationTests: XCTestCase { 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]) + + // 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 + 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) From 6f12d982198448281b9384e94efa5bba4aede484 Mon Sep 17 00:00:00 2001 From: Batuhan Saka Date: Wed, 20 Nov 2024 18:03:02 +0100 Subject: [PATCH 2/5] rm redundant print calls --- Sources/OpenTelemetryApi/Baggage/DefaultBaggageBuilder.swift | 1 - 1 file changed, 1 deletion(-) diff --git a/Sources/OpenTelemetryApi/Baggage/DefaultBaggageBuilder.swift b/Sources/OpenTelemetryApi/Baggage/DefaultBaggageBuilder.swift index b8714f5f..53acf878 100644 --- a/Sources/OpenTelemetryApi/Baggage/DefaultBaggageBuilder.swift +++ b/Sources/OpenTelemetryApi/Baggage/DefaultBaggageBuilder.swift @@ -61,7 +61,6 @@ public class DefaultBaggageBuilder: BaggageBuilder { } } } - print("Entries count: \(entries.count)") return entries.isEmpty ? EmptyBaggage.instance : DefaultBaggage(entries: combined) } } From 4036def8654e5b2cdfc5ae540999a9a8978a722e Mon Sep 17 00:00:00 2001 From: Batuhan Saka Date: Thu, 21 Nov 2024 11:02:45 +0100 Subject: [PATCH 3/5] add missing `DataCompression` dep to `OpenTelemetryProtocolExporterHttp` --- Package.swift | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/Package.swift b/Package.swift index 15a43e33..2d091fe5 100644 --- a/Package.swift +++ b/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", From f406c140869b0ea46fc0f09b7bf2b6a5c0310813 Mon Sep 17 00:00:00 2001 From: Batuhan Saka Date: Thu, 21 Nov 2024 12:02:28 +0100 Subject: [PATCH 4/5] add missing deps to all package files --- Package@swift-5.6.swift | 3 ++- Package@swift-5.9.swift | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/Package@swift-5.6.swift b/Package@swift-5.6.swift index 2d25e584..3340cc29 100644 --- a/Package@swift-5.6.swift +++ b/Package@swift-5.6.swift @@ -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", diff --git a/Package@swift-5.9.swift b/Package@swift-5.9.swift index 49ff7747..3f4f5ddb 100644 --- a/Package@swift-5.9.swift +++ b/Package@swift-5.9.swift @@ -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", From 88e649994f2cd8e6a7723af0f3f195bcbbd193c0 Mon Sep 17 00:00:00 2001 From: Batuhan Saka Date: Thu, 21 Nov 2024 15:54:55 +0100 Subject: [PATCH 5/5] extract baggage evaluations to separate tests for `testShouldInstrumentRequest` --- .../URLSessionInstrumentationTests.swift | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift b/Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift index 5535155a..d941df53 100644 --- a/Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift +++ b/Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift @@ -260,6 +260,24 @@ class URLSessionInstrumentationTests: XCTestCase { XCTAssertTrue(URLSessionInstrumentationTests.checker.receivedErrorCalled) } + public func testShouldInstrumentRequest() 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) + + 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")!)