From ff57bcc23747bc41e2ab82e57e578e80639fc0a5 Mon Sep 17 00:00:00 2001 From: Batuhan Saka Date: Tue, 25 Jul 2023 19:08:36 +0200 Subject: [PATCH] Add support to inject baggage propagator headers to the instrumented request (#443) --- .../URLSession/URLSessionLogger.swift | 9 +++-- .../URLSessionInstrumentationTests.swift | 36 ++++++++++++++----- 2 files changed, 34 insertions(+), 11 deletions(-) diff --git a/Sources/Instrumentation/URLSession/URLSessionLogger.swift b/Sources/Instrumentation/URLSession/URLSessionLogger.swift index e731f96b..eae6a592 100644 --- a/Sources/Instrumentation/URLSession/URLSessionLogger.swift +++ b/Sources/Instrumentation/URLSession/URLSessionLogger.swift @@ -148,7 +148,8 @@ class URLSessionLogger { instrumentation.configuration.injectCustomHeaders?(&request, span) var instrumentedRequest = request objc_setAssociatedObject(instrumentedRequest, &URLSessionInstrumentation.instrumentedKey, true, .OBJC_ASSOCIATION_COPY_NONATOMIC) - var traceHeaders = tracePropagationHTTPHeaders(span: span, textMapPropagator: OpenTelemetry.instance.propagators.textMapPropagator) + let propagators = OpenTelemetry.instance.propagators + var traceHeaders = tracePropagationHTTPHeaders(span: span, textMapPropagator: propagators.textMapPropagator, textMapBaggagePropagator: propagators.textMapBaggagePropagator) if let originalHeaders = request.allHTTPHeaderFields { traceHeaders.merge(originalHeaders) { _, new in new } } @@ -156,7 +157,7 @@ class URLSessionLogger { return instrumentedRequest } - private static func tracePropagationHTTPHeaders(span: Span?, textMapPropagator: TextMapPropagator) -> [String: String] { + private static func tracePropagationHTTPHeaders(span: Span?, textMapPropagator: TextMapPropagator, textMapBaggagePropagator: TextMapBaggagePropagator) -> [String: String] { var headers = [String: String]() struct HeaderSetter: Setter { @@ -169,6 +170,10 @@ class URLSessionLogger { return headers } textMapPropagator.inject(spanContext: currentSpan.context, carrier: &headers, setter: HeaderSetter()) + + if let baggage = OpenTelemetry.instance.contextProvider.activeBaggage { + textMapBaggagePropagator.inject(baggage: baggage, carrier: &headers, setter: HeaderSetter()) + } return headers } } diff --git a/Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift b/Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift index 473eb1f6..ad945d3a 100644 --- a/Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift +++ b/Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift @@ -82,11 +82,20 @@ class URLSessionInstrumentationTests: XCTestCase { 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() + .put(key: EntryKey(name: "foo")!, value: EntryValue(string: "bar")!, metadata: nil) + .build() + + OpenTelemetry.instance.contextProvider.setActiveBaggage(baggage) + let sem = DispatchSemaphore(value: 0) DispatchQueue.global(qos: .default).async { do { @@ -102,6 +111,7 @@ class URLSessionInstrumentationTests: XCTestCase { override class func tearDown() { server.stop() + OpenTelemetry.instance.contextProvider.removeContextForBaggage(baggage) } override func setUp() { @@ -240,20 +250,30 @@ class URLSessionInstrumentationTests: XCTestCase { XCTAssertTrue(URLSessionInstrumentationTests.checker.receivedErrorCalled) } - public func testShouldInstrumentRequest() { + public func testShouldInstrumentRequest() throws { let request1 = URLRequest(url: URL(string: "http://defaultName.com")!) let request2 = URLRequest(url: URL(string: "http://dontinstrument.com")!) - URLSessionLogger.processAndLogRequest(request1, sessionTaskId: "111", instrumentation: URLSessionInstrumentationTests.instrumentation, shouldInjectHeaders: true) - URLSessionLogger.processAndLogRequest(request2, sessionTaskId: "222", instrumentation: URLSessionInstrumentationTests.instrumentation, shouldInjectHeaders: true) + 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 + XCTAssertTrue(processedHeaders1.contains(where: { $0.key == W3CBaggagePropagator.headerBaggage && $0.value == "foo=bar" })) + XCTAssertEqual(1, URLSessionLogger.runningSpans.count) - XCTAssertNotNil(URLSessionLogger.runningSpans["111"]) - if let span = URLSessionLogger.runningSpans["111"] { - XCTAssertEqual("HTTP GET", span.name) - } + + let span = try XCTUnwrap(URLSessionLogger.runningSpans["111"]) + XCTAssertEqual("HTTP GET", span.name) } public func testDataTaskWithRequestBlock() { @@ -388,8 +408,6 @@ class URLSessionInstrumentationTests: XCTestCase { XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled) XCTAssertNotNil(URLSessionInstrumentationTests.requestCopy?.allHTTPHeaderFields?[W3CTraceContextPropagator.traceparent]) } - - #if swift(>=5.5.2) @available(macOS 10.15, iOS 13.0, tvOS 13.0, *)