Skip to content

Commit

Permalink
Merge branch 'main' into session-instrumentation-AV-fix
Browse files Browse the repository at this point in the history
  • Loading branch information
nachoBonafonte authored Nov 22, 2024
2 parents 38e3837 + 7001892 commit fd59ade
Show file tree
Hide file tree
Showing 6 changed files with 105 additions and 13 deletions.
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 Package@swift-5.6.swift
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 Package@swift-5.9.swift
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 {
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

0 comments on commit fd59ade

Please sign in to comment.