Merge pull request #575 from DataDog/ncreated/RUMM-1579-inject-sampli…

…ng-headers-when-propagating-trace

RUMM-1579 Inject sampling headers to instrumented requests

Sources/Datadog/Tracing/Propagation/HTTPHeadersWriter.swift

@@ -25,8 +25,6 @@ import Foundation
 ///
 ///
 public class HTTPHeadersWriter: OTHTTPHeadersWriter {
-    public init() {}
-
     /// A dictionary with HTTP Headers required to propagate the trace started in the mobile app
     /// to the backend instrumented with Datadog APM.
     ///
@@ -38,14 +36,23 @@ public class HTTPHeadersWriter: OTHTTPHeadersWriter {
     ///
     public private(set) var tracePropagationHTTPHeaders: [String: String] = [:]
 
+    /// Pre-computed headers with constant values.
+    private let constantHTTPHeaders: [String: String]
+
+    public init() {
+        constantHTTPHeaders = [
+            TracingHTTPHeaders.ddSamplingPriority.field: TracingHTTPHeaders.ddSamplingPriority.value,
+            TracingHTTPHeaders.ddSampled.field: TracingHTTPHeaders.ddSampled.value,
+        ]
+    }
+
     public func inject(spanContext: OTSpanContext) {
         guard let spanContext = spanContext.dd else {
             return
         }
 
-        tracePropagationHTTPHeaders = [
-            TracingHTTPHeaders.traceIDField: String(spanContext.traceID.rawValue),
-            TracingHTTPHeaders.parentSpanIDField: String(spanContext.spanID.rawValue)
-        ]
+        tracePropagationHTTPHeaders = constantHTTPHeaders
+        tracePropagationHTTPHeaders[TracingHTTPHeaders.traceIDField] = String(spanContext.traceID.rawValue)
+        tracePropagationHTTPHeaders[TracingHTTPHeaders.parentSpanIDField] = String(spanContext.spanID.rawValue)
     }
 }

Sources/Datadog/Tracing/Propagation/TracingHTTPHeaders.swift

@@ -6,13 +6,27 @@
 
 import Foundation
 
+/// Trace propagation headers as explained in
+/// https://docs.datadoghq.com/real_user_monitoring/connect_rum_and_traces/?tab=browserrum#how-are-rum-resources-linked-to-traces
 internal struct TracingHTTPHeaders {
+    /// Trace propagation header.
+    /// It is used both in Tracing and RUM features.
     static let traceIDField = "x-datadog-trace-id"
+
+    /// Trace propagation header.
+    /// In RUM - it allows Datadog to generate the first span from the trace.
+    /// In Tracing - it injects the `spanID` of mobile span so downstream spans can be properly linked in distributed tracing.
     static let parentSpanIDField = "x-datadog-parent-id"
-    static let originField = "x-datadog-origin"
-    /// Value for `originField` header field, indicating that the request is tracked as RUM Resource by the client.
-    static let rumOriginValue = "rum"
 
-    // TODO: RUMM-338 support `x-datadog-sampling-priority`. `dd-trace-ot` reference:
-    // https://github.com/DataDog/dd-trace-java/blob/4ba0ca0f9da748d4018310d026b1a72b607947f1/dd-trace-ot/src/main/java/datadog/opentracing/propagation/DatadogHttpCodec.java#L23
+    /// To make sure that the Agent keeps the trace.
+    /// It is used both in Tracing and RUM features.
+    static let ddSamplingPriority = (field: "x-datadog-sampling-priority", value: "1")
+
+    /// Indicates this request is selected for sampling.
+    /// It is used both in Tracing and RUM features.
+    static let ddSampled = (field: "x-datadog-sampled", value: "1")
+
+    /// To make sure the generated traces from RUM don’t affect APM Index Spans counts.
+    /// **Note:** it is only added to requests that we create RUM Resource for (it is not injected when RUM feature is disabled and only Tracing is used).
+    static let ddOrigin = (field: "x-datadog-origin", value: "rum")
 }

Sources/Datadog/URLSessionAutoInstrumentation/Interception/URLSessionInterceptor.swift

@@ -79,9 +79,9 @@ 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.
+                // 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.originField: TracingHTTPHeaders.rumOriginValue
+                    TracingHTTPHeaders.ddOrigin.field: TracingHTTPHeaders.ddOrigin.value
                 ]
             } else {
                 self.additionalHeadersForFirstPartyRequests = nil

Tests/DatadogIntegrationTests/Scenarios/RUM/RUMResourcesScenarioTests.swift

@@ -98,10 +98,17 @@ class RUMResourcesScenarioTests: IntegrationTests, RUMCommonAsserts {
             getSpanID(from: firstPartyPOSTRequest),
             "Tracing information should be propagated to `firstPartyPOSTResourceURL`."
         )
-        XCTAssertEqual(
-            getDatadogOrigin(from: firstPartyPOSTRequest),
-            "rum",
-            "Additional `x-datadog-origin: rum` header should be propagated to `firstPartyPOSTResourceURL`"
+        XCTAssertTrue(
+            firstPartyPOSTRequest.httpHeaders.contains("x-datadog-sampling-priority: 1"),
+            "`x-datadog-sampling-priority: 1` header must be set for `firstPartyPOSTResourceURL`"
+        )
+        XCTAssertTrue(
+            firstPartyPOSTRequest.httpHeaders.contains("x-datadog-sampled: 1"),
+            "`x-datadog-sampled: 1` header must be set for `firstPartyPOSTResourceURL`"
+        )
+        XCTAssertTrue(
+            firstPartyPOSTRequest.httpHeaders.contains("x-datadog-origin: rum"),
+            "`x-datadog-origin: rum` header must be set for `firstPartyPOSTResourceURL`"
         )
 
         // Get RUM Sessions with expected number of View visits and Resources
@@ -232,11 +239,4 @@ class RUMResourcesScenarioTests: IntegrationTests, RUMCommonAsserts {
         header?.removeFirst(prefix.count)
         return header
     }
-
-    private func getDatadogOrigin(from request: Request) -> String? {
-        let prefix = "x-datadog-origin: "
-        var header = request.httpHeaders.first { $0.hasPrefix(prefix) }
-        header?.removeFirst(prefix.count)
-        return header
-    }
 }

Tests/DatadogIntegrationTests/Scenarios/Tracing/TracingURLSessionScenarioTests.swift

@@ -124,19 +124,27 @@ class TracingURLSessionScenarioTests: IntegrationTests, TracingCommonAsserts {
         XCTAssertEqual(firstPartyRequests.count, 1)
 
         let firstPartyRequest = firstPartyRequests[0]
-        let expectedTraceIDHeader = "x-datadog-trace-id: \(try taskWithRequest.traceID().hexadecimalNumberToDecimal)"
-        let expectedSpanIDHeader = "x-datadog-parent-id: \(try taskWithRequest.spanID().hexadecimalNumberToDecimal)"
-        XCTAssertTrue(
-            firstPartyRequest.httpHeaders.contains(expectedTraceIDHeader),
-            """
-            Request `\(firstPartyRequest.path)` does not contain `\(expectedTraceIDHeader)` header.
-            - request.headers: \(firstPartyRequest.httpHeaders)
-            """
-        )
-        XCTAssertTrue(
-            firstPartyRequest.httpHeaders.contains(expectedSpanIDHeader),
+        let expectedHeaders = [
+            "x-datadog-trace-id: \(try taskWithRequest.traceID().hexadecimalNumberToDecimal)",
+            "x-datadog-parent-id: \(try taskWithRequest.spanID().hexadecimalNumberToDecimal)",
+            "x-datadog-sampling-priority: 1",
+            "x-datadog-sampled: 1"
+        ]
+
+        expectedHeaders.forEach { expectedHeader in
+            XCTAssertTrue(
+                firstPartyRequest.httpHeaders.contains(expectedHeader),
+                """
+                Request `\(firstPartyRequest.path)` must contain `\(expectedHeader)` header.
+                - request.headers: \(firstPartyRequest.httpHeaders)
+                """
+            )
+        }
+
+        XCTAssertFalse(
+            firstPartyRequest.httpHeaders.contains("x-datadog-origin: rum"),
             """
-            Request `\(firstPartyRequest.path)` does not contain `\(expectedSpanIDHeader)` header.
+            Request `\(firstPartyRequest.path)` must not contain `x-datadog-origin: rum` header.
             - request.headers: \(firstPartyRequest.httpHeaders)
             """
         )

Tests/DatadogTests/Datadog/TracerTests.swift

@@ -823,18 +823,35 @@ class TracerTests: XCTestCase {
 
     func testItInjectsSpanContextWithHTTPHeadersWriter() {
         let tracer: Tracer = .mockAny()
-        let spanContext = DDSpanContext(traceID: 1, spanID: 2, parentSpanID: .mockAny(), baggageItems: .mockAny())
+        let spanContext1 = DDSpanContext(traceID: 1, spanID: 2, parentSpanID: .mockAny(), baggageItems: .mockAny())
+        let spanContext2 = DDSpanContext(traceID: 3, spanID: 4, parentSpanID: .mockAny(), baggageItems: .mockAny())
 
         let httpHeadersWriter = HTTPHeadersWriter()
         XCTAssertEqual(httpHeadersWriter.tracePropagationHTTPHeaders, [:])
 
-        tracer.inject(spanContext: spanContext, writer: httpHeadersWriter)
+        // When
+        tracer.inject(spanContext: spanContext1, writer: httpHeadersWriter)
 
-        let expectedHTTPHeaders = [
+        // Then
+        let expectedHTTPHeaders1 = [
             "x-datadog-trace-id": "1",
             "x-datadog-parent-id": "2",
+            "x-datadog-sampling-priority": "1",
+            "x-datadog-sampled": "1",
+        ]
+        XCTAssertEqual(httpHeadersWriter.tracePropagationHTTPHeaders, expectedHTTPHeaders1)
+
+        // When
+        tracer.inject(spanContext: spanContext2, writer: httpHeadersWriter)
+
+        // Then
+        let expectedHTTPHeaders2 = [
+            "x-datadog-trace-id": "3",
+            "x-datadog-parent-id": "4",
+            "x-datadog-sampling-priority": "1",
+            "x-datadog-sampled": "1",
         ]
-        XCTAssertEqual(httpHeadersWriter.tracePropagationHTTPHeaders, expectedHTTPHeaders)
+        XCTAssertEqual(httpHeadersWriter.tracePropagationHTTPHeaders, expectedHTTPHeaders2)
     }
 
     func testItExtractsSpanContextWithHTTPHeadersReader() {

Tests/DatadogTests/Datadog/URLSessionAutoInstrumentation/Interception/URLSessionInterceptorTests.swift

@@ -90,7 +90,7 @@ class URLSessionInterceptorTests: XCTestCase {
         )
         XCTAssertEqual(
             interceptor.additionalHeadersForFirstPartyRequests,
-            [TracingHTTPHeaders.originField: TracingHTTPHeaders.rumOriginValue],
+            [TracingHTTPHeaders.ddOrigin.field: TracingHTTPHeaders.ddOrigin.value],
             "Additional `x-datadog-origin: rum` header should be injected when both Tracing and RUM instrumentations are enabled."
         )
     }
@@ -133,26 +133,30 @@ class URLSessionInterceptorTests: XCTestCase {
         // Then
         XCTAssertNotNil(interceptedFirstPartyRequest.allHTTPHeaderFields?[TracingHTTPHeaders.traceIDField])
         XCTAssertNotNil(interceptedFirstPartyRequest.allHTTPHeaderFields?[TracingHTTPHeaders.parentSpanIDField])
-        XCTAssertEqual(interceptedFirstPartyRequest.allHTTPHeaderFields?[TracingHTTPHeaders.originField], TracingHTTPHeaders.rumOriginValue)
+        XCTAssertEqual(interceptedFirstPartyRequest.allHTTPHeaderFields?[TracingHTTPHeaders.ddOrigin.field], TracingHTTPHeaders.ddOrigin.value)
         assertRequestsEqual(
             interceptedFirstPartyRequest
                 .removing(httpHeaderField: TracingHTTPHeaders.traceIDField)
                 .removing(httpHeaderField: TracingHTTPHeaders.parentSpanIDField)
-                .removing(httpHeaderField: TracingHTTPHeaders.originField),
+                .removing(httpHeaderField: TracingHTTPHeaders.ddSamplingPriority.field)
+                .removing(httpHeaderField: TracingHTTPHeaders.ddSampled.field)
+                .removing(httpHeaderField: TracingHTTPHeaders.ddOrigin.field),
             firstPartyRequest,
-            "The only modification of the original requests should be the addition of 3 tracing headers."
+            "The only modification of the original requests should be the addition of 5 tracing headers."
         )
 
         XCTAssertNotNil(interceptedCustomFirstPartyRequest.allHTTPHeaderFields?[TracingHTTPHeaders.traceIDField])
         XCTAssertNotNil(interceptedCustomFirstPartyRequest.allHTTPHeaderFields?[TracingHTTPHeaders.parentSpanIDField])
-        XCTAssertEqual(interceptedCustomFirstPartyRequest.allHTTPHeaderFields?[TracingHTTPHeaders.originField], TracingHTTPHeaders.rumOriginValue)
+        XCTAssertEqual(interceptedCustomFirstPartyRequest.allHTTPHeaderFields?[TracingHTTPHeaders.ddOrigin.field], TracingHTTPHeaders.ddOrigin.value)
         assertRequestsEqual(
             interceptedCustomFirstPartyRequest
                 .removing(httpHeaderField: TracingHTTPHeaders.traceIDField)
                 .removing(httpHeaderField: TracingHTTPHeaders.parentSpanIDField)
-                .removing(httpHeaderField: TracingHTTPHeaders.originField),
+                .removing(httpHeaderField: TracingHTTPHeaders.ddSamplingPriority.field)
+                .removing(httpHeaderField: TracingHTTPHeaders.ddSampled.field)
+                .removing(httpHeaderField: TracingHTTPHeaders.ddOrigin.field),
             alternativeFirstPartyRequest,
-            "The only modification of the original requests should be the addition of 3 tracing headers."
+            "The only modification of the original requests should be the addition of 5 tracing headers."
         )
 
         assertRequestsEqual(thirdPartyRequest, interceptedThirdPartyRequest, "Intercepted 3rd party request should not be modified.")
@@ -176,13 +180,15 @@ class URLSessionInterceptorTests: XCTestCase {
         // Then
         XCTAssertNotNil(interceptedFirstPartyRequest.allHTTPHeaderFields?[TracingHTTPHeaders.traceIDField])
         XCTAssertNotNil(interceptedFirstPartyRequest.allHTTPHeaderFields?[TracingHTTPHeaders.parentSpanIDField])
-        XCTAssertNil(interceptedFirstPartyRequest.allHTTPHeaderFields?[TracingHTTPHeaders.originField], "Origin header should not be added if RUM is disabled.")
+        XCTAssertNil(interceptedFirstPartyRequest.allHTTPHeaderFields?[TracingHTTPHeaders.ddOrigin.field], "Origin header should not be added if RUM is disabled.")
         assertRequestsEqual(
             interceptedFirstPartyRequest
                 .removing(httpHeaderField: TracingHTTPHeaders.traceIDField)
-                .removing(httpHeaderField: TracingHTTPHeaders.parentSpanIDField),
+                .removing(httpHeaderField: TracingHTTPHeaders.parentSpanIDField)
+                .removing(httpHeaderField: TracingHTTPHeaders.ddSamplingPriority.field)
+                .removing(httpHeaderField: TracingHTTPHeaders.ddSampled.field),
             firstPartyRequest,
-            "The only modification of the original requests should be the addition of 2 tracing headers."
+            "The only modification of the original requests should be the addition of 4 tracing headers."
         )
         assertRequestsEqual(thirdPartyRequest, interceptedThirdPartyRequest, "Intercepted 3rd party request should not be modified.")
         assertRequestsEqual(internalRequest, interceptedInternalRequest, "Intercepted internal request should not be modified.")

Tests/DatadogTests/DatadogObjc/DDTracerTests.swift

@@ -235,6 +235,8 @@ class DDTracerTests: XCTestCase {
         let expectedHTTPHeaders = [
             "x-datadog-trace-id": "1",
             "x-datadog-parent-id": "2",
+            "x-datadog-sampling-priority": "1",
+            "x-datadog-sampled": "1",
         ]
         XCTAssertEqual(objcWriter.tracePropagationHTTPHeaders, expectedHTTPHeaders)
     }