RUM-3470 Support head-based sampling for manual distributed traces

DataDog · Apr 24, 2024 · d8c2f2b · d8c2f2b
1 parent c452a2d
commit d8c2f2b
Show file tree

Hide file tree

Showing 26 changed files with 622 additions and 477 deletions.
diff --git a/Datadog/Datadog.xcodeproj/project.pbxproj b/Datadog/Datadog.xcodeproj/project.pbxproj
@@ -300,6 +300,10 @@
 		614CADD72510BAC000B93D2D /* Environment.swift in Sources */ = {isa = PBXBuildFile; fileRef = 614CADD62510BAC000B93D2D /* Environment.swift */; };
 		614E9EB3244719FA007EE3E1 /* BundleType.swift in Sources */ = {isa = PBXBuildFile; fileRef = 614E9EB2244719FA007EE3E1 /* BundleType.swift */; };
 		614ED36C260352DC00C8C519 /* CrashReporter.xcframework in Frameworks */ = {isa = PBXBuildFile; fileRef = 614ED36B260352DC00C8C519 /* CrashReporter.xcframework */; };
+		615192CD2BD6948B0005A782 /* HTTPHeadersWriterTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 615192CC2BD6948B0005A782 /* HTTPHeadersWriterTests.swift */; };
+		615192CE2BD6948B0005A782 /* HTTPHeadersWriterTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 615192CC2BD6948B0005A782 /* HTTPHeadersWriterTests.swift */; };
+		615192D02BD6B7C90005A782 /* DatadogTracer+InjectAndExtract.swift in Sources */ = {isa = PBXBuildFile; fileRef = 615192CF2BD6B7C90005A782 /* DatadogTracer+InjectAndExtract.swift */; };
+		615192D12BD6B7C90005A782 /* DatadogTracer+InjectAndExtract.swift in Sources */ = {isa = PBXBuildFile; fileRef = 615192CF2BD6B7C90005A782 /* DatadogTracer+InjectAndExtract.swift */; };
 		61570005246AADFA00E96950 /* DatadogObjc.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 61133BF0242397DA00786299 /* DatadogObjc.framework */; };
 		615A4A8324A3431600233986 /* Trace+objc.swift in Sources */ = {isa = PBXBuildFile; fileRef = 615A4A8224A3431600233986 /* Trace+objc.swift */; };
 		615A4A8924A34FD700233986 /* DDTracerTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 615A4A8824A34FD700233986 /* DDTracerTests.swift */; };
@@ -2244,6 +2248,8 @@
 		614CADD62510BAC000B93D2D /* Environment.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Environment.swift; sourceTree = "<group>"; };
 		614E9EB2244719FA007EE3E1 /* BundleType.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = BundleType.swift; sourceTree = "<group>"; };
 		614ED36B260352DC00C8C519 /* CrashReporter.xcframework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.xcframework; name = CrashReporter.xcframework; path = ../Carthage/Build/CrashReporter.xcframework; sourceTree = "<group>"; };
+		615192CC2BD6948B0005A782 /* HTTPHeadersWriterTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = HTTPHeadersWriterTests.swift; sourceTree = "<group>"; };
+		615192CF2BD6B7C90005A782 /* DatadogTracer+InjectAndExtract.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "DatadogTracer+InjectAndExtract.swift"; sourceTree = "<group>"; };
 		6152C84224BE2165006A1679 /* MockServerAddress.local.xcconfig */ = {isa = PBXFileReference; lastKnownFileType = text.xcconfig; path = MockServerAddress.local.xcconfig; sourceTree = "<group>"; };
 		615519252461BCE7002A85CF /* Datadog.xcconfig */ = {isa = PBXFileReference; lastKnownFileType = text.xcconfig; path = Datadog.xcconfig; sourceTree = "<group>"; };
 		615519262461BCE7002A85CF /* Datadog.local.xcconfig */ = {isa = PBXFileReference; lastKnownFileType = text.xcconfig; path = Datadog.local.xcconfig; sourceTree = "<group>"; };
@@ -5735,6 +5741,7 @@
 				619CE75D2A458CE1005588CB /* TraceConfigurationTests.swift */,
 				61AD4E172451C7FF006E34EA /* TracingFeatureMocks.swift */,
 				61F3E3622BC5556D00C7881E /* DatadogTracer+SamplingTests.swift */,
+				615192CF2BD6B7C90005A782 /* DatadogTracer+InjectAndExtract.swift */,
 				61C5A89824509C1100DA608C /* DDSpanTests.swift */,
 				61F1A620249A45E400075390 /* DDSpanContextTests.swift */,
 				D2F1B81426D8E5FF009F3293 /* DDNoopTracerTests.swift */,
@@ -6049,6 +6056,7 @@
 				61B558D32469CDD8001460D3 /* TraceIDGeneratorTests.swift */,
 				3CCECDAE2BC688120013C125 /* SpanIDGeneratorTests.swift */,
 				61F3E3652BC595F600C7881E /* HTTPHeadersReaderTests.swift */,
+				615192CC2BD6948B0005A782 /* HTTPHeadersWriterTests.swift */,
 				A79B0F5A292B7C06008742B3 /* B3HTTPHeadersWriterTests.swift */,
 				A79B0F60292BB071008742B3 /* B3HTTPHeadersReaderTests.swift */,
 				A728ADA22934DB5000397996 /* W3CHTTPHeadersWriterTests.swift */,
@@ -8561,6 +8569,7 @@
 				61F3E3632BC5556D00C7881E /* DatadogTracer+SamplingTests.swift in Sources */,
 				D2C1A51C29C4C75700946C31 /* ContextMessageReceiverTests.swift in Sources */,
 				619CE7612A458D66005588CB /* TraceTests.swift in Sources */,
+				615192D02BD6B7C90005A782 /* DatadogTracer+InjectAndExtract.swift in Sources */,
 				D2C1A52029C4C75700946C31 /* DDSpanTests.swift in Sources */,
 				D2C1A51B29C4C75700946C31 /* DDSpanContextTests.swift in Sources */,
 				D2C1A52729C4C7D000946C31 /* TracingFeatureMocks.swift in Sources */,
@@ -8764,6 +8773,7 @@
 				61F3E3642BC5556D00C7881E /* DatadogTracer+SamplingTests.swift in Sources */,
 				D2C1A56529C4F2E800946C31 /* ContextMessageReceiverTests.swift in Sources */,
 				619CE7622A458D66005588CB /* TraceTests.swift in Sources */,
+				615192D12BD6B7C90005A782 /* DatadogTracer+InjectAndExtract.swift in Sources */,
 				D2C1A56629C4F2E800946C31 /* DDSpanTests.swift in Sources */,
 				D2C1A56729C4F2E800946C31 /* DDSpanContextTests.swift in Sources */,
 				D2C1A56829C4F2E800946C31 /* TracingFeatureMocks.swift in Sources */,
@@ -9193,6 +9203,7 @@
 				D2160CE929C0E00200FAA9A5 /* MethodSwizzlerTests.swift in Sources */,
 				D2216EC32A96649500ADAEC8 /* FeatureBaggageTests.swift in Sources */,
 				D2160CDC29C0DF6700FAA9A5 /* HostsSanitizerTests.swift in Sources */,
+				615192CD2BD6948B0005A782 /* HTTPHeadersWriterTests.swift in Sources */,
 				D2F44FB8299AA1DA0074B0D9 /* DataCompressionTests.swift in Sources */,
 				D2160CE029C0DF6700FAA9A5 /* URLSessionDelegateAsSuperclassTests.swift in Sources */,
 				D2EBEE3B29BA163E00B15732 /* B3HTTPHeadersReaderTests.swift in Sources */,
@@ -9239,6 +9250,7 @@
 				D2EBEE4229BA163F00B15732 /* W3CHTTPHeadersReaderTests.swift in Sources */,
 				D2216EC42A96649700ADAEC8 /* FeatureBaggageTests.swift in Sources */,
 				D2160CDD29C0DF6700FAA9A5 /* HostsSanitizerTests.swift in Sources */,
+				615192CE2BD6948B0005A782 /* HTTPHeadersWriterTests.swift in Sources */,
 				D2F44FB9299AA1DB0074B0D9 /* DataCompressionTests.swift in Sources */,
 				D2160CE129C0DF6700FAA9A5 /* URLSessionDelegateAsSuperclassTests.swift in Sources */,
 				D2EBEE3F29BA163F00B15732 /* B3HTTPHeadersReaderTests.swift in Sources */,

diff --git a/Datadog/Datadog.xcodeproj/xcshareddata/xcschemes/DatadogCore iOS.xcscheme b/Datadog/Datadog.xcodeproj/xcshareddata/xcschemes/DatadogCore iOS.xcscheme
@@ -209,12 +209,6 @@
                <Test
                   Identifier = "HeadBasedSamplingTests/testSendingDroppedDistributedTraceWithParent_throughURLSessionInstrumentationAPI()">
                </Test>
-               <Test
-                  Identifier = "HeadBasedSamplingTests/testSendingSampledDistributedTraceWithNoParent_throughTracerAPI()">
-               </Test>
-               <Test
-                  Identifier = "HeadBasedSamplingTests/testSendingSampledDistributedTraceWithParent_throughTracerAPI()">
-               </Test>
             </SkippedTests>
          </TestableReference>
          <TestableReference

diff --git a/Datadog/Example/Debugging/DebugManualTraceInjectionViewController.swift b/Datadog/Example/Debugging/DebugManualTraceInjectionViewController.swift
@@ -107,22 +107,19 @@ internal struct DebugManualTraceInjectionView: View {
 
         switch selectedTraceHeaderType {
         case .datadog:
-            let writer = HTTPHeadersWriter(sampleRate: sampleRate)
+            let writer = HTTPHeadersWriter(samplingStrategy: .custom(sampleRate: sampleRate))
             Tracer.shared().inject(spanContext: span.context, writer: writer)
             writer.traceHeaderFields.forEach { request.setValue($0.value, forHTTPHeaderField: $0.key) }
         case .w3c:
-            let writer = W3CHTTPHeadersWriter(
-                sampleRate: sampleRate,
-                tracestate: [:]
-            )
+            let writer = W3CHTTPHeadersWriter(samplingStrategy: .custom(sampleRate: sampleRate), tracestate: [:])
             Tracer.shared().inject(spanContext: span.context, writer: writer)
             writer.traceHeaderFields.forEach { request.setValue($0.value, forHTTPHeaderField: $0.key) }
         case .b3Single:
-            let writer = B3HTTPHeadersWriter(sampleRate: sampleRate, injectEncoding: .single)
+            let writer = B3HTTPHeadersWriter(samplingStrategy: .custom(sampleRate: sampleRate), injectEncoding: .single)
             Tracer.shared().inject(spanContext: span.context, writer: writer)
             writer.traceHeaderFields.forEach { request.setValue($0.value, forHTTPHeaderField: $0.key) }
         case .b3Multiple:
-            let writer = B3HTTPHeadersWriter(sampleRate: sampleRate, injectEncoding: .multiple)
+            let writer = B3HTTPHeadersWriter(samplingStrategy: .custom(sampleRate: sampleRate), injectEncoding: .multiple)
             Tracer.shared().inject(spanContext: span.context, writer: writer)
             writer.traceHeaderFields.forEach { request.setValue($0.value, forHTTPHeaderField: $0.key) }
         }

diff --git a/Datadog/IntegrationUnitTests/Trace/HeadBasedSamplingTests.swift b/Datadog/IntegrationUnitTests/Trace/HeadBasedSamplingTests.swift
@@ -269,7 +269,6 @@ class HeadBasedSamplingTests: XCTestCase {
 
     // MARK: - Distributed Tracing (through Tracer API)
 
-    // TODO: RUM-3470 Enable this test when head-based sampling is supported for distributed tracing through Tracer API
     func testSendingSampledDistributedTraceWithNoParent_throughTracerAPI() throws {
         /*
          This is the situation where distributed trace starts with the span created with Datadog tracer:
@@ -279,15 +278,14 @@ class HeadBasedSamplingTests: XCTestCase {
          */
 
         let localTraceSampling: Float = 100 // keep all
-        let distributedTraceSampling: Float = 0 // drop all
 
         // Given
         traceConfig.sampleRate = localTraceSampling
         Trace.enable(with: traceConfig, in: core)
 
         // When
         var request: URLRequest = .mockAny()
-        let writer = HTTPHeadersWriter(sampleRate: distributedTraceSampling)
+        let writer = HTTPHeadersWriter(samplingStrategy: .auto)
         let span = Tracer.shared(in: core).startSpan(operationName: "network.span")
         Tracer.shared(in: core).inject(spanContext: span.context, writer: writer)
         writer.traceHeaderFields.forEach { field, value in request.setValue(value, forHTTPHeaderField: field) }
@@ -309,7 +307,7 @@ class HeadBasedSamplingTests: XCTestCase {
         XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.samplingPriorityField), "1")
     }
 
-    // TODO: RUM-3470 Enable this test when head-based sampling is supported for distributed tracing through Tracer API
+    // TODO: RUM-3535 Enable this test when trace context injection control is implemented
     func testSendingDroppedDistributedTraceWithNoParent_throughTracerAPI() throws {
         /*
          This is the situation where distributed trace starts with the span created with Datadog tracer:
@@ -319,15 +317,14 @@ class HeadBasedSamplingTests: XCTestCase {
          */
 
         let localTraceSampling: Float = 0 // drop all
-        let distributedTraceSampling: Float = 100 // keep all
 
         // Given
         traceConfig.sampleRate = localTraceSampling
         Trace.enable(with: traceConfig, in: core)
 
         // When
         var request: URLRequest = .mockAny()
-        let writer = HTTPHeadersWriter(sampleRate: distributedTraceSampling)
+        let writer = HTTPHeadersWriter(samplingStrategy: .auto)
         let span = Tracer.shared(in: core).startSpan(operationName: "network.span")
         Tracer.shared(in: core).inject(spanContext: span.context, writer: writer)
         writer.traceHeaderFields.forEach { field, value in request.setValue(value, forHTTPHeaderField: field) }
@@ -349,7 +346,6 @@ class HeadBasedSamplingTests: XCTestCase {
         XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.samplingPriorityField), "0")
     }
 
-    // TODO: RUM-3470 Enable this test when head-based sampling is supported for distributed tracing through Tracer API
     func testSendingSampledDistributedTraceWithParent_throughTracerAPI() throws {
         /*
          This is the situation where distributed trace starts with an active local span and is continued with the span
@@ -361,15 +357,14 @@ class HeadBasedSamplingTests: XCTestCase {
          */
 
         let localTraceSampling: Float = 100 // keep all
-        let distributedTraceSampling: Float = 0 // drop all
 
         // Given
         traceConfig.sampleRate = localTraceSampling
         Trace.enable(with: traceConfig, in: core)
 
         // When
         var request: URLRequest = .mockAny()
-        let writer = HTTPHeadersWriter(sampleRate: distributedTraceSampling)
+        let writer = HTTPHeadersWriter(samplingStrategy: .auto)
         let parentSpan = Tracer.shared(in: core).startSpan(operationName: "active.span").setActive()
         let span = Tracer.shared(in: core).startSpan(operationName: "network.span")
         Tracer.shared(in: core).inject(spanContext: span.context, writer: writer)
@@ -399,7 +394,7 @@ class HeadBasedSamplingTests: XCTestCase {
         XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.samplingPriorityField), "1")
     }
 
-    // TODO: RUM-3470 Enable this test when head-based sampling is supported for distributed tracing through Tracer API
+    // TODO: RUM-3535 Enable this test when trace context injection control is implemented
     func testSendingDroppedDistributedTraceWithParent_throughTracerAPI() throws {
         /*
          This is the situation where distributed trace starts with an active local span and is continued with the span
@@ -411,15 +406,14 @@ class HeadBasedSamplingTests: XCTestCase {
          */
 
         let localTraceSampling: Float = 0 // drop all
-        let distributedTraceSampling: Float = 100 // keep all
 
         // Given
         traceConfig.sampleRate = localTraceSampling
         Trace.enable(with: traceConfig, in: core)
 
         // When
         var request: URLRequest = .mockAny()
-        let writer = HTTPHeadersWriter(sampleRate: distributedTraceSampling)
+        let writer = HTTPHeadersWriter(samplingStrategy: .auto)
         let parentSpan = Tracer.shared(in: core).startSpan(operationName: "active.span").setActive()
         let span = Tracer.shared(in: core).startSpan(operationName: "network.span")
         Tracer.shared(in: core).inject(spanContext: span.context, writer: writer)