RUM-7999: Prevent sampled out spans from sending

Datadog/IntegrationUnitTests/Trace/HeadBasedSamplingTests.swift

@@ -53,11 +53,11 @@ class HeadBasedSamplingTests: XCTestCase {
         parent.finish()
 
         let spans = core.waitAndReturnSpanEvents()
-        XCTAssertEqual(spans.count, 3, "It must send all spans")
 
-        let allKept = spans.filter({ $0.isKept }).count == 3
-        let allDropped = spans.filter({ !$0.isKept }).count == 3
-        XCTAssertTrue(allKept || allDropped, "All spans must be either kept or dropped")
+        guard spans.isEmpty else {
+            XCTAssertEqual(spans.filter({ $0.isKept }).count, 3, "All spans must be either kept or dropped")
+            return
+        }
     }
 
     func testSamplingLocalTraceWithImplicitParent() throws {
@@ -82,11 +82,11 @@ class HeadBasedSamplingTests: XCTestCase {
         parent.finish()
 
         let spans = core.waitAndReturnSpanEvents()
-        XCTAssertEqual(spans.count, 3, "It must send all spans")
 
-        let allKept = spans.filter({ $0.isKept }).count == 3
-        let allDropped = spans.filter({ !$0.isKept }).count == 3
-        XCTAssertTrue(allKept || allDropped, "All spans must be either kept or dropped")
+        guard spans.isEmpty else {
+            XCTAssertEqual(spans.filter({ $0.isKept }).count, 3, "All spans must be either kept or dropped")
+            return
+        }
     }
 
     // MARK: - Distributed Tracing (through network instrumentation API)
@@ -154,18 +154,10 @@ class HeadBasedSamplingTests: XCTestCase {
         let request = try sendURLSessionRequest(to: "https://foo.com/request", using: InstrumentedSessionDelegate())
 
         // Then
-        let span = try XCTUnwrap(core.waitAndReturnSpanEvents().first, "It should send span event")
-        XCTAssertEqual(span.operationName, "urlsession.request")
-        XCTAssertEqual(span.samplingRate, 0, "Span must use distributed trace sample rate")
-        XCTAssertFalse(span.isKept, "Span must be dropped")
-
-        // Then
-        let expectedTraceIDField = span.traceID.toString(representation: .decimal)
-        let expectedSpanIDField = span.spanID.toString(representation: .decimal)
-        let expectedTagsField = "_dd.p.tid=\(span.traceID.idHiHex)"
-        XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.traceIDField), expectedTraceIDField)
-        XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.parentSpanIDField), expectedSpanIDField)
-        XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.tagsField), expectedTagsField)
+        XCTAssertNil(core.waitAndReturnSpanEvents().first, "It should not send span event")
+        XCTAssertNotNil(request.value(forHTTPHeaderField: TracingHTTPHeaders.traceIDField))
+        XCTAssertNotNil(request.value(forHTTPHeaderField: TracingHTTPHeaders.parentSpanIDField))
+        XCTAssertNotNil(request.value(forHTTPHeaderField: TracingHTTPHeaders.tagsField))
         XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.samplingPriorityField), "0")
     }
 
@@ -245,23 +237,13 @@ class HeadBasedSamplingTests: XCTestCase {
 
         // Then
         let spanEvents = core.waitAndReturnSpanEvents()
-        let activeSpan = try XCTUnwrap(spanEvents.first(where: { $0.operationName == "active.span" }))
-        let urlsessionSpan = try XCTUnwrap(spanEvents.first(where: { $0.operationName == "urlsession.request" }))
-
-        XCTAssertEqual(activeSpan.samplingRate, 0, "Span must use local trace sample rate")
-        XCTAssertFalse(activeSpan.isKept, "Span must not be sampled")
-        XCTAssertEqual(urlsessionSpan.samplingRate, 0, "Span must use local trace sample rate")
-        XCTAssertFalse(urlsessionSpan.isKept, "Span must not be sampled")
-        XCTAssertEqual(urlsessionSpan.traceID, activeSpan.traceID)
-        XCTAssertEqual(urlsessionSpan.parentID, activeSpan.spanID)
+        XCTAssertNil(spanEvents.first(where: { $0.operationName == "active.span" }))
+        XCTAssertNil(spanEvents.first(where: { $0.operationName == "urlsession.request" }))
 
         // Then
-        let expectedTraceIDField = activeSpan.traceID.toString(representation: .decimal)
-        let expectedSpanIDField = urlsessionSpan.spanID.toString(representation: .decimal)
-        let expectedTagsField = "_dd.p.tid=\(activeSpan.traceID.idHiHex)"
-        XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.traceIDField), expectedTraceIDField)
-        XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.parentSpanIDField), expectedSpanIDField)
-        XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.tagsField), expectedTagsField)
+        XCTAssertNotNil(request.value(forHTTPHeaderField: TracingHTTPHeaders.traceIDField))
+        XCTAssertNotNil(request.value(forHTTPHeaderField: TracingHTTPHeaders.parentSpanIDField))
+        XCTAssertNotNil(request.value(forHTTPHeaderField: TracingHTTPHeaders.tagsField))
         XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.samplingPriorityField), "0")
     }
 
@@ -329,18 +311,12 @@ class HeadBasedSamplingTests: XCTestCase {
         span.finish()
 
         // Then
-        let networkSpan = try XCTUnwrap(core.waitAndReturnSpanEvents().first, "It should send span event")
-        XCTAssertEqual(networkSpan.operationName, "network.span")
-        XCTAssertEqual(networkSpan.samplingRate, 0, "Span must use local trace sample rate")
-        XCTAssertFalse(networkSpan.isKept, "Span must be dropped")
+        XCTAssertNil(core.waitAndReturnSpanEvents().first, "It should not send span event")
 
         // Then
-        let expectedTraceIDField = networkSpan.traceID.toString(representation: .decimal)
-        let expectedSpanIDField = networkSpan.spanID.toString(representation: .decimal)
-        let expectedTagsField = "_dd.p.tid=\(networkSpan.traceID.idHiHex)"
-        XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.traceIDField), expectedTraceIDField)
-        XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.parentSpanIDField), expectedSpanIDField)
-        XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.tagsField), expectedTagsField)
+        XCTAssertNotNil(request.value(forHTTPHeaderField: TracingHTTPHeaders.traceIDField))
+        XCTAssertNotNil(request.value(forHTTPHeaderField: TracingHTTPHeaders.parentSpanIDField))
+        XCTAssertNotNil(request.value(forHTTPHeaderField: TracingHTTPHeaders.tagsField))
         XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.samplingPriorityField), "0")
     }
 
@@ -420,23 +396,13 @@ class HeadBasedSamplingTests: XCTestCase {
 
         // Then
         let spanEvents = core.waitAndReturnSpanEvents()
-        let activeSpan = try XCTUnwrap(spanEvents.first(where: { $0.operationName == "active.span" }))
-        let networkSpan = try XCTUnwrap(spanEvents.first(where: { $0.operationName == "network.span" }))
-
-        XCTAssertEqual(activeSpan.samplingRate, 0, "Span must use local trace sample rate")
-        XCTAssertFalse(activeSpan.isKept, "Span must be dropped")
-        XCTAssertEqual(networkSpan.samplingRate, 0, "Span must use local trace sample rate")
-        XCTAssertFalse(networkSpan.isKept, "Span must be dropped")
-        XCTAssertEqual(networkSpan.traceID, activeSpan.traceID)
-        XCTAssertEqual(networkSpan.parentID, activeSpan.spanID)
+        XCTAssertNil(spanEvents.first(where: { $0.operationName == "active.span" }))
+        XCTAssertNil(spanEvents.first(where: { $0.operationName == "network.span" }))
 
         // Then
-        let expectedTraceIDField = activeSpan.traceID.toString(representation: .decimal)
-        let expectedSpanIDField = networkSpan.spanID.toString(representation: .decimal)
-        let expectedTagsField = "_dd.p.tid=\(activeSpan.traceID.idHiHex)"
-        XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.traceIDField), expectedTraceIDField)
-        XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.parentSpanIDField), expectedSpanIDField)
-        XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.tagsField), expectedTagsField)
+        XCTAssertNotNil(request.value(forHTTPHeaderField: TracingHTTPHeaders.traceIDField))
+        XCTAssertNotNil(request.value(forHTTPHeaderField: TracingHTTPHeaders.parentSpanIDField))
+        XCTAssertNotNil(request.value(forHTTPHeaderField: TracingHTTPHeaders.tagsField))
         XCTAssertEqual(request.value(forHTTPHeaderField: TracingHTTPHeaders.samplingPriorityField), "0")
     }
 

DatadogCore/Tests/Objc/ObjcAPITests/DDTrace+apiTests.m

@@ -58,6 +58,8 @@ - (void)testDDTraceConfigurationAPI {
 }
 
 - (void)testDDTracerAPI {
+    id<OTSpan> rootSpan = [[DDTracer shared] startRootSpan:@"" tags:NULL startTime:NULL customSamplingRate:NULL];
+    [rootSpan setActive];
     [[DDTracer shared] startSpan:@""];
     [[DDTracer shared] startSpan:@"" tags:@{}];
     [[DDTracer shared] startSpan:@"" childOf:NULL];

DatadogTrace/Sources/DDNoOps.swift

@@ -46,6 +46,11 @@ internal class DDNoopTracer: OTTracer, OpenTelemetryApi.Tracer {
         return DDNoopGlobals.span
     }
 
+    func startRootSpan(operationName: String, tags: [String: any Encodable]?, startTime: Date?, customSamplingRate: SampleRate?) -> any OTSpan {
+        warn()
+        return DDNoopGlobals.span
+    }
+
     // MARK: - Open Telemetry
 
     func spanBuilder(spanName: String) -> OpenTelemetryApi.SpanBuilder {

DatadogTrace/Sources/DDSpan.swift

@@ -126,7 +126,9 @@ internal final class DDSpan: OTSpan {
             ddTracer.removeSpan(span: self)
             activity.leave()
         }
-        sendSpan(finishTime: time, sampler: ddTracer.localTraceSampler)
+        if self.ddContext.isKept {
+            sendSpan(finishTime: time, sampler: ddTracer.localTraceSampler)
+        }
     }
 
     // MARK: - Writing SpanEvent

DatadogTrace/Sources/DatadogTracer.swift

@@ -87,9 +87,14 @@ internal final class DatadogTracer: OTTracer, OpenTelemetryApi.Tracer {
         )
     }
 
-    func startRootSpan(operationName: String, tags: [String: Encodable]? = nil, startTime: Date? = nil) -> OTSpan {
-        startSpan(
-            spanContext: createSpanContext(parentSpanContext: nil, using: localTraceSampler),
+    func startRootSpan(operationName: String, tags: [String: Encodable]? = nil, startTime: Date? = nil, customSamplingRate: SampleRate? = nil) -> OTSpan {
+        let sampler: Sampling = if let customSamplingRate {
+            Sampler(samplingRate: customSamplingRate)
+        } else {
+            localTraceSampler
+        }
+        return startSpan(
+            spanContext: createSpanContext(parentSpanContext: nil, using: sampler),
             operationName: operationName,
             tags: tags,
             startTime: startTime

DatadogTrace/Sources/Objc/OpenTracing/OTTracer+objc.swift

@@ -20,6 +20,7 @@ public protocol objc_OTTracer {
     func startSpan(_ operationName: String, childOf parent: objc_OTSpanContext?) -> objc_OTSpan
     func startSpan(_ operationName: String, childOf parent: objc_OTSpanContext?, tags: NSDictionary?) -> objc_OTSpan
     func startSpan(_ operationName: String, childOf parent: objc_OTSpanContext?, tags: NSDictionary?, startTime: Date?) -> objc_OTSpan
+    func startRootSpan(_ operationName: String, tags: NSDictionary?, startTime: Date?, customSamplingRate: NSNumber?) -> objc_OTSpan
     func inject(_ spanContext: objc_OTSpanContext, format: String, carrier: Any) throws
     func extractWithFormat(_ format: String, carrier: Any) throws
 }

DatadogTrace/Sources/Objc/Tracing/Trace+objc.swift

@@ -184,6 +184,23 @@ public final class objc_Tracer: NSObject, objc_OTTracer {
         )
     }
 
+    public func startRootSpan(
+        _ operationName: String,
+        tags: NSDictionary?,
+        startTime: Date?,
+        customSamplingRate: NSNumber?
+    ) -> any objc_OTSpan {
+        return objc_SpanObjc(
+            objcTracer: self,
+            swiftSpan: swiftTracer.startRootSpan(
+                operationName: operationName,
+                tags: tags.flatMap { castTagsToSwift($0) },
+                startTime: startTime,
+                customSamplingRate: customSamplingRate?.floatValue
+            )
+        )
+    }
+
     public func inject(_ spanContext: objc_OTSpanContext, format: String, carrier: Any) throws {
         if let objcWriter = carrier as? objc_HTTPHeadersWriter, format == OT.formatTextMap {
             guard let ddspanContext = spanContext.dd else {

DatadogTrace/Sources/OpenTracing/OTTracer.swift

@@ -1,4 +1,5 @@
 import Foundation
+import DatadogInternal
 
 /// Tracer is the starting point for all OpenTracing instrumentation. Use it
 /// to create OTSpans, inject/extract them between processes, and so on.
@@ -34,7 +35,8 @@ public protocol OTTracer {
     func startRootSpan(
         operationName: String,
         tags: [String: Encodable]?,
-        startTime: Date?
+        startTime: Date?,
+        customSamplingRate: SampleRate?
     ) -> OTSpan
 
     /// Transfer the span information into the carrier of the given format.
@@ -117,12 +119,14 @@ public extension OTTracer {
     func startRootSpan(
         operationName: String,
         tags: [String: Encodable]? = nil,
-        startTime: Date? = nil
+        startTime: Date? = nil,
+        customSamplingRate: SampleRate? = nil
     ) -> OTSpan {
         return self.startRootSpan(
             operationName: operationName,
             tags: tags,
-            startTime: startTime
+            startTime: startTime,
+            customSamplingRate: customSamplingRate
         )
     }
 }

DatadogTrace/Tests/DatadogTracer+SamplingTests.swift

@@ -35,9 +35,55 @@ class DatadogTracer_SamplingTests: XCTestCase {
 
         // Then
         let events = try XCTUnwrap(featureScope.spanEventsWritten())
-        XCTAssertEqual(events.filter({ $0.samplingRate == 0.42 }).count, 100, "All spans must encode sample rate")
+        XCTAssertTrue(events.allSatisfy({ $0.samplingRate == 0.42 }), "All kept spans must encode sample rate")
         XCTAssertGreaterThan(events.filter({ $0.isKept }).count, 1, "Some spans should be kept")
-        XCTAssertGreaterThan(events.filter({ !$0.isKept }).count, 1, "Some spans should be dropped")
+        XCTAssertEqual(events.filter({ !$0.isKept }).count, 0, "Not kept spans should be dropped")
+    }
+
+    func testRecordingCustomSampleRateInRootSpanEvent() throws {
+        // When
+        let tracer = createTracer(sampleRate: 0)
+        (0..<10).forEach { _ in
+            let span = tracer.startRootSpan(operationName: .mockAny(), customSamplingRate: 100)
+            span.finish()
+        }
+
+        // Then
+        let events = try XCTUnwrap(featureScope.spanEventsWritten())
+        XCTAssertEqual(events.filter({ $0.samplingRate == 1 }).count, 10)
+        XCTAssertEqual(events.filter({ $0.isKept }).count, 10)
+    }
+
+    func testRootSampleInOverridesTracerAndPropagatesToChildSpans() throws {
+        // When
+        let tracer = createTracer(sampleRate: 0)
+        let root = tracer.startRootSpan(operationName: .mockAny(), customSamplingRate: 100)
+        let child = tracer.startSpan(operationName: .mockAny(), childOf: root.context)
+        let grandChild = tracer.startSpan(operationName: .mockAny(), childOf: child.context)
+        grandChild.finish()
+        child.finish()
+        root.finish()
+
+        // Then
+        let events = try XCTUnwrap(featureScope.spanEventsWritten())
+        XCTAssertEqual(events.count, 3)
+        XCTAssertEqual(events.filter({ $0.samplingRate == 1 }).count, 3)
+        XCTAssertEqual(events.filter({ $0.isKept }).count, 3)
+    }
+
+    func testRootSampleOutOverridesTracerAndPropagatesToChildSpans() throws {
+        // When
+        let tracer = createTracer(sampleRate: 100)
+        let root = tracer.startRootSpan(operationName: .mockAny(), customSamplingRate: 0)
+        let child = tracer.startSpan(operationName: .mockAny(), childOf: root.context)
+        let grandChild = tracer.startSpan(operationName: .mockAny(), childOf: child.context)
+        grandChild.finish()
+        child.finish()
+        root.finish()
+
+        // Then
+        let events = try XCTUnwrap(featureScope.spanEventsWritten())
+        XCTAssertEqual(events.count, 0)
     }
 
     func testRecordingSampledSpan() throws {
@@ -59,28 +105,12 @@ class DatadogTracer_SamplingTests: XCTestCase {
         span.finish()
 
         // Then
-        let event = try XCTUnwrap(featureScope.spanEventsWritten().first)
-        XCTAssertEqual(event.samplingRate, 0)
-        XCTAssertFalse(event.isKept)
+        let event = try featureScope.spanEventsWritten().first
+        XCTAssertNil(event)
     }
 
     // MARK: - Head-based Sampling
 
-    func testRecordingSampleRateInChildSpanEvents() throws {
-        // When
-        let tracer = createTracer(sampleRate: 42)
-        let root = tracer.startSpan(operationName: .mockAny())
-        let child = tracer.startSpan(operationName: .mockAny(), childOf: root.context)
-        let grandChild = tracer.startSpan(operationName: .mockAny(), childOf: child.context)
-        grandChild.finish()
-        child.finish()
-        root.finish()
-
-        // Then
-        let events = try XCTUnwrap(featureScope.spanEventsWritten())
-        XCTAssertEqual(events.filter({ $0.samplingRate == 0.42 }).count, 3, "All spans must encode the same sample rate")
-    }
-
     func testWhenRootSpanIsSampled_thenAllChildSpansMustBeSampledTheSameWay() throws {
         // When
         let tracer = createTracer(sampleRate: 50)
@@ -93,10 +123,13 @@ class DatadogTracer_SamplingTests: XCTestCase {
 
         // Then
         let events = try XCTUnwrap(featureScope.spanEventsWritten())
-        XCTAssertEqual(events.count, 3)
-        let allKept = events.filter({ $0.isKept }).count == 3
-        let allDropped = events.filter({ !$0.isKept }).count == 3
-        XCTAssertTrue(allKept || allDropped, "All spans must be either kept or dropped")
+
+        if events.isEmpty {
+            XCTAssert(true, "No spans were collected")
+        } else {
+            XCTAssertEqual(events.filter({ $0.samplingRate == 0.5 }).count, 3, "All spans must encode the same sample rate")
+            XCTAssertEqual(events.filter({ $0.isKept }).count, 3, "All spans must be kept")
+        }
     }
 }
 

DatadogTrace/Tests/TracingURLSessionHandlerTests.swift

@@ -241,7 +241,7 @@ class TracingURLSessionHandlerTests: XCTestCase {
         let expectation = expectation(description: "Send span")
         core.onEventWriteContext = { _ in expectation.fulfill() }
         let sampleRate: Float = .mockRandom(min: 1, max: 100)
-        let isKept: Bool = .mockRandom()
+        let isKept: Bool = true
 
         // Given
         let interception = URLSessionTaskInterception(