Merge pull request #1107 from DataDog/maciey/clean-up-tracing-origin

Add x-datadog-origin header only when instrumenting with datadog tracing
DataDog · Dec 30, 2022 · 218e2a2 · 218e2a2
2 parents 9f25eb2 + ec36928
commit 218e2a2
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 15 deletions.
diff --git a/Sources/Datadog/URLSessionAutoInstrumentation/Interception/URLSessionInterceptor.swift b/Sources/Datadog/URLSessionAutoInstrumentation/Interception/URLSessionInterceptor.swift
@@ -44,7 +44,7 @@ public class URLSessionInterceptor: URLSessionInterceptorType {
     internal let injectTracingHeadersToFirstPartyRequests: Bool
     /// Additional header injected to intercepted 1st party requests.
     /// Set to `x-datadog-origin: rum` if both RUM and Tracing instrumentations are enabled and `nil` in all other cases.
-    internal let additionalHeadersForFirstPartyRequests: [String: String]?
+    internal let datadogOriginHeader: (field: String, value: String)?
     /// Tracing sampler used to sample traces generated by the SDK or RUM BE.
     internal let tracingSampler: Sampler
 
@@ -88,15 +88,13 @@ public class URLSessionInterceptor: URLSessionInterceptorType {
             if configuration.instrumentRUM {
                 // If RUM instrumentation is enabled, additional `x-datadog-origin: rum` header is injected to the user request,
                 // so that user's backend instrumentation can further process it and count on RUM quota (w/o affecting APM Index Spans counts).
-                self.additionalHeadersForFirstPartyRequests = [
-                    TracingHTTPHeaders.ddOrigin.field: TracingHTTPHeaders.ddOrigin.value
-                ]
+                self.datadogOriginHeader = TracingHTTPHeaders.ddOrigin
             } else {
-                self.additionalHeadersForFirstPartyRequests = nil
+                self.datadogOriginHeader = nil
             }
         } else {
             self.injectTracingHeadersToFirstPartyRequests = false
-            self.additionalHeadersForFirstPartyRequests = nil
+            self.datadogOriginHeader = nil
         }
     }
 
@@ -250,6 +248,9 @@ public class URLSessionInterceptor: URLSessionInterceptorType {
             switch $0 {
             case .datadog:
                 writer = HTTPHeadersWriter(sampler: tracingSampler)
+                if let datadogOriginHeader = datadogOriginHeader {
+                    newRequest.setValue(datadogOriginHeader.value, forHTTPHeaderField: datadogOriginHeader.field)
+                }
             case .b3:
                 writer = OTelHTTPHeadersWriter(
                     sampler: tracingSampler,
@@ -268,10 +269,6 @@ public class URLSessionInterceptor: URLSessionInterceptorType {
             writer.tracePropagationHTTPHeaders.forEach { field, value in
                 newRequest.setValue(value, forHTTPHeaderField: field)
             }
-
-            additionalHeadersForFirstPartyRequests?.forEach { field, value in
-                newRequest.setValue(value, forHTTPHeaderField: field)
-            }
         }
         return newRequest
     }

diff --git a/...Tests/Datadog/URLSessionAutoInstrumentation/Interception/URLSessionInterceptorTests.swift b/...Tests/Datadog/URLSessionAutoInstrumentation/Interception/URLSessionInterceptorTests.swift
@@ -55,7 +55,7 @@ class URLSessionInterceptorTests: XCTestCase {
             "Tracing headers should be injected when only Tracing instrumentation is enabled."
         )
         XCTAssertNil(
-            interceptor.additionalHeadersForFirstPartyRequests,
+            interceptor.datadogOriginHeader,
             "Just the tracing headers should be injected when only Tracing instrumentation is enabled."
         )
     }
@@ -79,7 +79,7 @@ class URLSessionInterceptorTests: XCTestCase {
             "Tracing headers should not be injected when only RUM instrumentation is enabled."
         )
         XCTAssertNil(
-            interceptor.additionalHeadersForFirstPartyRequests,
+            interceptor.datadogOriginHeader,
             "No additional headers should be injected when only RUM instrumentation is enabled."
         )
     }
@@ -103,9 +103,14 @@ class URLSessionInterceptorTests: XCTestCase {
             "Tracing headers should be injected when both Tracing and RUM instrumentations are enabled."
         )
         XCTAssertEqual(
-            interceptor.additionalHeadersForFirstPartyRequests,
-            [TracingHTTPHeaders.ddOrigin.field: TracingHTTPHeaders.ddOrigin.value],
-            "Additional `x-datadog-origin: rum` header should be injected when both Tracing and RUM instrumentations are enabled."
+            interceptor.datadogOriginHeader?.field,
+            TracingHTTPHeaders.ddOrigin.field,
+            "Additional `x-datadog-origin` header field should be injected when both Tracing and RUM instrumentations are enabled."
+        )
+        XCTAssertEqual(
+            interceptor.datadogOriginHeader?.value,
+            TracingHTTPHeaders.ddOrigin.value,
+            "Additional `rum` header value should be injected when both Tracing and RUM instrumentations are enabled."
         )
     }