Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for Default Baggage Injection in URLSessionInstrumentation #639

Open
wants to merge 7 commits into
base: main
Choose a base branch
from
3 changes: 2 additions & 1 deletion Package.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
3 changes: 2 additions & 1 deletion [email protected]
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
3 changes: 2 additions & 1 deletion [email protected]
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -36,6 +37,7 @@ public struct URLSessionInstrumentationConfiguration {
self.receivedResponse = receivedResponse
self.receivedError = receivedError
self.delegateClassesToInstrument = delegateClassesToInstrument
self.defaultBaggageProvider = defaultBaggageProvider
}

// Instrumentation Callbacks
Expand Down Expand Up @@ -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)?)?

}
25 changes: 21 additions & 4 deletions Sources/Instrumentation/URLSession/URLSessionLogger.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand All @@ -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() {
Expand Down Expand Up @@ -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 {
nachoBonafonte marked this conversation as resolved.
Show resolved Hide resolved
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)

Expand Down
Loading