Add Callback to customize Spans in URLSessionInstrumention (#340)

* Add ability to set a parent using a callback in URLSessionInstrumention, and associated tests
* Convert to generic customization of spans, instead of just customizing parent

Author: gjeltep

Sources/Instrumentation/URLSession/README.md

@@ -19,6 +19,8 @@ This behaviour can be modified or augmented by using the optional callbacks defi
 
 `nameSpan: ((URLRequest) -> (String)?)?` - Modifies the name for the given request instead of stantard Opentelemetry name
 
+`spanCustomization: ((URLRequest, SpanBuilder) -> Void)?` - Customizes the span while it's being built, such as by adding a parent, a link, attributes, etc.
+
 `createdRequest: ((URLRequest, Span) -> Void)?` - Called after request is created,  it allows to add extra information to the Span
 
 `receivedResponse: ((URLResponse, DataOrFile?, Span) -> Void)?`- Called after response is received,  it allows to add extra information to the Span

Sources/Instrumentation/URLSession/URLSessionInstrumentationConfiguration.swift

@@ -15,6 +15,7 @@ public struct URLSessionInstrumentationConfiguration {
     public init(shouldRecordPayload: ((URLSession) -> (Bool)?)? = nil,
                 shouldInstrument: ((URLRequest) -> (Bool)?)? = nil,
                 nameSpan: ((URLRequest) -> (String)?)? = nil,
+                spanCustomization: ((URLRequest, SpanBuilder) -> Void)? = nil,
                 shouldInjectTracingHeaders: ((URLRequest) -> (Bool)?)? = nil,
                 injectCustomHeaders: ((inout URLRequest, Span?) -> Void)? = nil,
                 createdRequest: ((URLRequest, Span) -> Void)? = nil,
@@ -26,6 +27,7 @@ public struct URLSessionInstrumentationConfiguration {
         self.shouldInjectTracingHeaders = shouldInjectTracingHeaders
         self.injectCustomHeaders = injectCustomHeaders
         self.nameSpan = nameSpan
+        self.spanCustomization = spanCustomization
         self.createdRequest = createdRequest
         self.receivedResponse = receivedResponse
         self.receivedError = receivedError
@@ -52,6 +54,9 @@ public struct URLSessionInstrumentationConfiguration {
     /// default name: `HTTP {method}` e.g. `HTTP PUT`
     public var nameSpan: ((URLRequest) -> (String)?)?
 
+    /// Implement this callback to customize the span, such as by adding a parent, a link, attributes, etc
+    public var spanCustomization: ((URLRequest, SpanBuilder) -> Void)?
+
     ///  Called before the span is created, it allows to add extra information to the Span
     public var createdRequest: ((URLRequest, Span) -> Void)?
 

Sources/Instrumentation/URLSession/URLSessionLogger.swift

@@ -71,6 +71,7 @@ class URLSessionLogger {
         attributes.forEach {
             spanBuilder.setAttribute(key: $0.key, value: $0.value)
         }
+        instrumentation.configuration.spanCustomization?(request, spanBuilder)
 
         let span = spanBuilder.startSpan()
         runningSpansQueue.sync {

Tests/InstrumentationTests/URLSessionTests/URLSessionInstrumentationTests.swift

@@ -13,6 +13,7 @@ class URLSessionInstrumentationTests: XCTestCase {
         public var shouldRecordPayloadCalled: Bool = false
         public var shouldInstrumentCalled: Bool = false
         public var nameSpanCalled: Bool = false
+        public var spanCustomizationCalled: Bool = false
         public var shouldInjectTracingHeadersCalled: Bool = false
         public var createdRequestCalled: Bool = false
         public var receivedResponseCalled: Bool = false
@@ -54,6 +55,12 @@ class URLSessionInstrumentationTests: XCTestCase {
                                                                    }
                                                                    return "new name"
                                                                },
+                                                               spanCustomization: { req, spanBuilder in
+                                                                   checker.spanCustomizationCalled = true
+                                                                   if !(req.url?.host?.contains("defaultName") ?? false) {
+                                                                       spanBuilder.setSpanKind(spanKind: .consumer)
+                                                                   }
+                                                               },
                                                                shouldInjectTracingHeaders: { _ in
                                                                    checker.shouldInjectTracingHeadersCalled = true
                                                                    return true
@@ -138,6 +145,32 @@ class URLSessionInstrumentationTests: XCTestCase {
         }
     }
 
+    public func testOverrideSpanCustomization() {
+            let request = URLRequest(url: URL(string: "http://google.com")!)
+
+            URLSessionLogger.processAndLogRequest(request, sessionTaskId: "id", instrumentation: URLSessionInstrumentationTests.instrumentation, shouldInjectHeaders: true)
+
+            XCTAssertTrue(URLSessionInstrumentationTests.checker.spanCustomizationCalled)
+
+            XCTAssertEqual(1, URLSessionLogger.runningSpans.count)
+            XCTAssertNotNil(URLSessionLogger.runningSpans["id"])
+            if let span = URLSessionLogger.runningSpans["id"] {
+                XCTAssertEqual(SpanKind.consumer, span.kind)
+            }
+        }
+
+    public func testDefaultSpanCustomization() {
+        let request = URLRequest(url: URL(string: "http://defaultName.com")!)
+
+        URLSessionLogger.processAndLogRequest(request, sessionTaskId: "id", instrumentation: URLSessionInstrumentationTests.instrumentation, shouldInjectHeaders: true)
+
+        XCTAssertEqual(1, URLSessionLogger.runningSpans.count)
+        XCTAssertNotNil(URLSessionLogger.runningSpans["id"])
+        if let span = URLSessionLogger.runningSpans["id"] {
+            XCTAssertEqual(SpanKind.client, span.kind)
+        }
+    }
+
     public func testConfigurationCallbacksCalledWhenSuccess() {
         let request = URLRequest(url: URL(string: "http://localhost:33333/success")!)
 
@@ -154,6 +187,7 @@ class URLSessionInstrumentationTests: XCTestCase {
 
         XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInstrumentCalled)
         XCTAssertTrue(URLSessionInstrumentationTests.checker.nameSpanCalled)
+        XCTAssertTrue(URLSessionInstrumentationTests.checker.spanCustomizationCalled)
         XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInjectTracingHeadersCalled)
         XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
         XCTAssertTrue(URLSessionInstrumentationTests.checker.receivedResponseCalled)
@@ -174,6 +208,7 @@ class URLSessionInstrumentationTests: XCTestCase {
 
         XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInstrumentCalled)
         XCTAssertTrue(URLSessionInstrumentationTests.checker.nameSpanCalled)
+        XCTAssertTrue(URLSessionInstrumentationTests.checker.spanCustomizationCalled)
         XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInjectTracingHeadersCalled)
         XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
         XCTAssertTrue(URLSessionInstrumentationTests.checker.receivedResponseCalled)
@@ -193,6 +228,7 @@ class URLSessionInstrumentationTests: XCTestCase {
 
         XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInstrumentCalled)
         XCTAssertTrue(URLSessionInstrumentationTests.checker.nameSpanCalled)
+        XCTAssertTrue(URLSessionInstrumentationTests.checker.spanCustomizationCalled)
         XCTAssertTrue(URLSessionInstrumentationTests.checker.shouldInjectTracingHeadersCalled)
         XCTAssertTrue(URLSessionInstrumentationTests.checker.createdRequestCalled)
         XCTAssertFalse(URLSessionInstrumentationTests.checker.receivedResponseCalled)