Merge pull request #1259 from DataDog/xgouchet/RUMM-3194/tracer_sampl…

…ing_rate

RUMM-3194 add sampling rate on Tracer

Sources/Datadog/Tracer.swift

@@ -65,6 +65,8 @@ public class Tracer: OTTracer {
 
     internal let activeSpansPool = ActiveSpansPool()
 
+    internal let sampler: Sampler
+
     /// Tracer Attributes shared with other Feature registered in core.
     internal struct Attributes {
         internal static let traceID = "dd.trace_id"
@@ -146,6 +148,7 @@ public class Tracer: OTTracer {
         self.dateProvider = dateProvider
         self.rumIntegration = rumIntegration
         self.loggingIntegration = loggingIntegration
+        self.sampler = Sampler(samplingRate: configuration.samplingRate)
     }
 
     // MARK: - Open Tracing interface

Sources/Datadog/TracerConfiguration.swift

@@ -27,6 +27,9 @@ extension Tracer {
         /// - Parameter enabled: `true` by default
         public var bundleWithRUM: Bool
 
+        /// TRhe sampling rate for Traces, as a Float between 0 and 100.
+        public var samplingRate: Float
+
         /// Initializes the Datadog Tracer configuration.
         /// - Parameter serviceName: the service name that will appear in traces (if not provided or `nil`, the SDK default `serviceName` will be used).
         /// - Parameter sendNetworkInfo: adds network connection info to every span and span logs (`false` by default).
@@ -37,11 +40,13 @@ extension Tracer {
             serviceName: String? = nil,
             sendNetworkInfo: Bool = false,
             bundleWithRUM: Bool = true,
+            samplingRate: Float = 100,
             globalTags: [String: Encodable]? = nil
         ) {
             self.serviceName = serviceName
             self.sendNetworkInfo = sendNetworkInfo
             self.bundleWithRUM = bundleWithRUM
+            self.samplingRate = samplingRate
             self.globalTags = globalTags
         }
     }

Sources/Datadog/Tracing/DDSpan.swift

@@ -123,14 +123,14 @@ internal class DDSpan: OTSpan {
             if let activity = activityReference {
                 ddTracer.removeSpan(activityReference: activity)
             }
-            sendSpan(finishTime: time)
+            sendSpan(finishTime: time, sampler: ddTracer.sampler)
         }
     }
 
     // MARK: - Writing SpanEvent
 
     /// Sends span event for given `DDSpan`.
-    private func sendSpan(finishTime: Date) {
+    private func sendSpan(finishTime: Date, sampler: Sampler) {
         guard let tracing = ddTracer.core.v1.scope(for: TracingFeature.self) else {
             return
         }
@@ -160,6 +160,8 @@ internal class DDSpan: OTSpan {
                     operationName: self.unsafeOperationName,
                     startTime: self.startTime,
                     finishTime: finishTime,
+                    samplingRate: sampler.samplingRate / 100.0,
+                    isKept: sampler.sample(),
                     tags: self.unsafeTags,
                     baggageItems: baggageItems,
                     logFields: self.unsafeLogFields

Sources/Datadog/Tracing/Span/SpanEventBuilder.swift

@@ -25,6 +25,8 @@ internal struct SpanEventBuilder {
         operationName: String,
         startTime: Date,
         finishTime: Date,
+        samplingRate: Float,
+        isKept: Bool,
         tags: [String: Encodable],
         baggageItems: [String: String],
         logFields: [[String: Encodable]]
@@ -60,6 +62,8 @@ internal struct SpanEventBuilder {
             isError: tagsReducer.extractedIsError ?? false,
             source: context.source,
             origin: context.ciAppOrigin,
+            samplingRate: samplingRate,
+            isKept: isKept,
             tracerVersion: context.sdkVersion,
             applicationVersion: context.version,
             networkConnectionInfo: sendNetworkInfo ? context.networkConnectionInfo : nil,

Sources/Datadog/Tracing/Span/SpanEventEncoder.swift

@@ -55,6 +55,10 @@ public struct SpanEvent: Encodable {
     internal let source: String
     /// The origin for the Span, it is used to label the spans used created under testing
     internal let origin: String?
+    /// The sampling rate for the span (between 0 and 100)
+    internal let samplingRate: Float
+    /// If the span is kept according to sampling rules
+    internal let isKept: Bool
 
     // MARK: - Meta
 
@@ -111,6 +115,7 @@ internal struct SpanEventEncoder {
 
         case isRootSpan = "metrics._top_level"
         case samplingPriority = "metrics._sampling_priority_v1"
+        case samplingRate = "_dd.agent_psr"
 
         // MARK: - Meta
 
@@ -177,8 +182,9 @@ internal struct SpanEventEncoder {
         // NOTE: RUMM-299 only numeric values are supported for `metrics.*` attributes
         if span.parentID == nil {
             try container.encode(1, forKey: .isRootSpan)
+            try container.encode(span.samplingRate, forKey: .samplingRate)
+            try container.encode((span.isKept ? 1 : 0), forKey: .samplingPriority)
         }
-        try container.encode(1, forKey: .samplingPriority)
     }
 
     /// Encodes default `meta.*` attributes

Tests/DatadogTests/Datadog/Mocks/TracingFeatureMocks.swift

@@ -146,6 +146,8 @@ extension SpanEvent: AnyMockable, RandomMockable {
         isError: Bool = .mockAny(),
         source: String = .mockAny(),
         origin: String? = nil,
+        samplingRate: Float = 100,
+        isKept: Bool = true,
         tracerVersion: String = .mockAny(),
         applicationVersion: String = .mockAny(),
         networkConnectionInfo: NetworkConnectionInfo? = .mockAny(),
@@ -165,6 +167,8 @@ extension SpanEvent: AnyMockable, RandomMockable {
             isError: isError,
             source: source,
             origin: origin,
+            samplingRate: samplingRate,
+            isKept: isKept,
             tracerVersion: tracerVersion,
             applicationVersion: applicationVersion,
             networkConnectionInfo: networkConnectionInfo,
@@ -189,6 +193,8 @@ extension SpanEvent: AnyMockable, RandomMockable {
             isError: .random(),
             source: .mockRandom(),
             origin: .mockRandom(),
+            samplingRate: .mockRandom(),
+            isKept: .mockRandom(),
             tracerVersion: .mockRandom(),
             applicationVersion: .mockRandom(),
             networkConnectionInfo: .mockRandom(),

Tests/DatadogTests/Datadog/TracerTests.swift

@@ -54,6 +54,7 @@ class TracerTests: XCTestCase {
         {
           "spans": [
             {
+              "_dd.agent_psr": 1,
               "trace_id": "1",
               "span_id": "2",
               "parent_id": "0",
@@ -132,6 +133,27 @@ class TracerTests: XCTestCase {
         XCTAssertEqual(try spanMatcher.meta.custom(keyPath: "meta.globaltag2"), "overwrittenValue")
     }
 
+    // MARK: - Tracer with sampling rate
+
+    func testUsingSamplingRate() throws {
+        let feature: TracingFeature = .mockAny()
+        core.register(feature: feature)
+
+        let tracer = Tracer.initialize(configuration: .init(samplingRate: 42), in: core).dd
+
+        let span = tracer.startSpan(
+            operationName: "operation",
+            startTime: .mockDecember15th2019At10AMUTC()
+        )
+        span.finish(at: .mockDecember15th2019At10AMUTC(addingTimeInterval: 0.5))
+
+        let spanMatcher = try core.waitAndReturnSpanMatchers()[0]
+        XCTAssertEqual(try spanMatcher.operationName(), "operation")
+        XCTAssertEqual(try spanMatcher.startTime(), 1_576_404_000_000_000_000)
+        XCTAssertEqual(try spanMatcher.duration(), 500_000_000)
+        XCTAssertEqual(try spanMatcher.dd.samplingRate(), 0.42)
+    }
+
     // MARK: - Sending Customized Spans
 
     func testSendingCustomizedSpan() throws {

Tests/DatadogTests/Datadog/Tracing/Span/SpanEventBuilderTests.swift

@@ -21,6 +21,8 @@ class SpanEventBuilderTests: XCTestCase {
             operationName: "operation-name",
             startTime: .mockDecember15th2019At10AMUTC(),
             finishTime: .mockDecember15th2019At10AMUTC(addingTimeInterval: 0.5),
+            samplingRate: .mockAny(),
+            isKept: .mockAny(),
             tags: [
                 "foo": "bar",
                 "bizz": 123
@@ -60,6 +62,8 @@ class SpanEventBuilderTests: XCTestCase {
             operationName: "original operation name",
             startTime: .mockAny(),
             finishTime: .mockAny(),
+            samplingRate: .mockAny(),
+            isKept: .mockAny(),
             tags: [:],
             baggageItems: [:],
             logFields: []
@@ -81,6 +85,8 @@ class SpanEventBuilderTests: XCTestCase {
             operationName: .mockAny(),
             startTime: .mockAny(),
             finishTime: .mockAny(),
+            samplingRate: .mockAny(),
+            isKept: .mockAny(),
             tags: [OTTags.error: true],
             baggageItems: [:],
             logFields: []
@@ -99,6 +105,8 @@ class SpanEventBuilderTests: XCTestCase {
             operationName: .mockAny(),
             startTime: .mockAny(),
             finishTime: .mockAny(),
+            samplingRate: .mockAny(),
+            isKept: .mockAny(),
             tags: [OTTags.error: false],
             baggageItems: [:],
             logFields: []
@@ -121,6 +129,8 @@ class SpanEventBuilderTests: XCTestCase {
             operationName: .mockAny(),
             startTime: .mockAny(),
             finishTime: .mockAny(),
+            samplingRate: .mockAny(),
+            isKept: .mockAny(),
             tags: [:],
             baggageItems: [:],
             logFields: [
@@ -141,6 +151,8 @@ class SpanEventBuilderTests: XCTestCase {
             operationName: .mockAny(),
             startTime: .mockAny(),
             finishTime: .mockAny(),
+            samplingRate: .mockAny(),
+            isKept: .mockAny(),
             tags: [:],
             baggageItems: [:],
             logFields: [
@@ -173,6 +185,8 @@ class SpanEventBuilderTests: XCTestCase {
             operationName: .mockAny(),
             startTime: .mockAny(),
             finishTime: .mockAny(),
+            samplingRate: .mockAny(),
+            isKept: .mockAny(),
             tags: [:],
             baggageItems: [:],
             logFields: [
@@ -199,6 +213,8 @@ class SpanEventBuilderTests: XCTestCase {
             operationName: .mockAny(),
             startTime: .mockAny(),
             finishTime: .mockAny(),
+            samplingRate: .mockAny(),
+            isKept: .mockAny(),
             tags: [
                 "error": true
             ],
@@ -220,6 +236,8 @@ class SpanEventBuilderTests: XCTestCase {
             operationName: .mockAny(),
             startTime: .mockAny(),
             finishTime: .mockAny(),
+            samplingRate: .mockAny(),
+            isKept: .mockAny(),
             tags: [
                 "error": false
             ],
@@ -245,6 +263,8 @@ class SpanEventBuilderTests: XCTestCase {
             operationName: .mockAny(),
             startTime: .mockAny(),
             finishTime: .mockAny(),
+            samplingRate: .mockAny(),
+            isKept: .mockAny(),
             tags: [
                 DDTags.resource: "custom resource name"
             ],
@@ -269,6 +289,8 @@ class SpanEventBuilderTests: XCTestCase {
             operationName: .mockAny(),
             startTime: .mockAny(),
             finishTime: .mockAny(),
+            samplingRate: .mockAny(),
+            isKept: .mockAny(),
             tags: [
                 "tag-name": "tag value"
             ],
@@ -287,6 +309,66 @@ class SpanEventBuilderTests: XCTestCase {
         XCTAssertEqual(span.tags["tag-name"], "tag value", "It should prefer tags over baggage items if their names duplicate")
     }
 
+    func testBuildingKeptRootSpanWithSamplingRateSet() {
+        let builder: SpanEventBuilder = .mockAny()
+
+        // When
+        let span = builder.createSpanEvent(
+            context: .mockAny(),
+            traceID: .mockAny(),
+            spanID: .mockAny(),
+            parentSpanID: nil,
+            operationName: .mockAny(),
+            startTime: .mockAny(),
+            finishTime: .mockAny(),
+            samplingRate: 0.42,
+            isKept: true,
+            tags: [
+                "tag-name": "tag value"
+            ],
+            baggageItems: [
+                "item-1": "value 1",
+                "item-2": "value 2",
+                "tag-name": "baggage item value"
+            ],
+            logFields: []
+        )
+
+        // Then
+        XCTAssertEqual(span.samplingRate, 0.42)
+        XCTAssertTrue(span.isKept)
+    }
+
+    func testBuildingDiscardedRootSpanWithSamplingRateSet() {
+        let builder: SpanEventBuilder = .mockAny()
+
+        // When
+        let span = builder.createSpanEvent(
+            context: .mockAny(),
+            traceID: .mockAny(),
+            spanID: .mockAny(),
+            parentSpanID: nil,
+            operationName: .mockAny(),
+            startTime: .mockAny(),
+            finishTime: .mockAny(),
+            samplingRate: 0.13,
+            isKept: false,
+            tags: [
+                "tag-name": "tag value"
+            ],
+            baggageItems: [
+                "item-1": "value 1",
+                "item-2": "value 2",
+                "tag-name": "baggage item value"
+            ],
+            logFields: []
+        )
+
+        // Then
+        XCTAssertEqual(span.samplingRate, 0.13)
+        XCTAssertFalse(span.isKept)
+    }
+
     // MARK: - Attributes Conversion
 
     private struct Foo: Encodable {
@@ -336,6 +418,8 @@ class SpanEventBuilderTests: XCTestCase {
             operationName: .mockAny(),
             startTime: .mockAny(),
             finishTime: .mockAny(),
+            samplingRate: .mockAny(),
+            isKept: .mockAny(),
             tags: createMockAttributes(),
             baggageItems: [:],
             logFields: []
@@ -366,6 +450,8 @@ class SpanEventBuilderTests: XCTestCase {
             operationName: .mockAny(),
             startTime: .mockAny(),
             finishTime: .mockAny(),
+            samplingRate: .mockAny(),
+            isKept: .mockAny(),
             tags: [:],
             baggageItems: [:],
             logFields: []
@@ -392,6 +478,8 @@ class SpanEventBuilderTests: XCTestCase {
             operationName: .mockAny(),
             startTime: .mockAny(),
             finishTime: .mockAny(),
+            samplingRate: .mockAny(),
+            isKept: .mockAny(),
             tags: [
                 "failing-tag": FailingEncodableMock(errorMessage: "Value cannot be encoded.")
             ],