From 2afcc5c2a091fa46b0c480a7cd653952b5ce2b3e Mon Sep 17 00:00:00 2001
From: Martin Holman <martin@honeycomb.io>
Date: Wed, 16 Oct 2024 09:42:17 -0700
Subject: [PATCH] feat(samplers): Add deterministic sampler (#17)

## Short description of the changes

Adds a deterministic sampler modeled off the deterministic sampler in
honeycombs other distros.

## How to verify that this has the expected result

unit tests
---
 .../HoneycombDeterministicSampler.swift       | 60 +++++++++++++++++++
 .../HoneycombDeterministicSamplerTests.swift  | 53 ++++++++++++++++
 2 files changed, 113 insertions(+)
 create mode 100644 Sources/Honeycomb/HoneycombDeterministicSampler.swift
 create mode 100644 Tests/Honeycomb/HoneycombDeterministicSamplerTests.swift

diff --git a/Sources/Honeycomb/HoneycombDeterministicSampler.swift b/Sources/Honeycomb/HoneycombDeterministicSampler.swift
new file mode 100644
index 0000000..e5507aa
--- /dev/null
+++ b/Sources/Honeycomb/HoneycombDeterministicSampler.swift
@@ -0,0 +1,60 @@
+import Foundation
+import OpenTelemetryApi
+import OpenTelemetrySdk
+
+class HoneycombDeterministicSampler: Sampler {
+    private let inner: Sampler
+    private let rate: [String: AttributeValue]
+
+    init(sampleRate: Int) {
+        var inner: Sampler
+
+        switch sampleRate {
+        case Int.min..<1:
+            inner = Samplers.alwaysOff
+        case 1:
+            inner = Samplers.alwaysOn
+        default:
+            inner = Samplers.traceIdRatio(ratio: 1.0 / Double(sampleRate))
+        }
+
+        self.inner = inner
+        self.rate = ["SampleRate": AttributeValue(sampleRate)]
+    }
+
+    func shouldSample(
+        parentContext: SpanContext?,
+        traceId: TraceId,
+        name: String,
+        kind: SpanKind,
+        attributes: [String: AttributeValue],
+        parentLinks: [SpanData.Link]
+    ) -> any Decision {
+        var result = self.inner.shouldSample(
+            parentContext: parentContext,
+            traceId: traceId,
+            name: name,
+            kind: kind,
+            attributes: attributes,
+            parentLinks: parentLinks
+        )
+
+        if result.isSampled {
+            let attrs = result.attributes.merging(
+                rate,
+                uniquingKeysWith: { (_, new) in new }
+            )
+            result = HoneycombDecision(isSampled: result.isSampled, attributes: attrs)
+        }
+
+        return result
+    }
+
+    var description: String = "DeterministicSampler"
+}
+
+private struct HoneycombDecision: Decision {
+    let isSampled: Bool
+
+    let attributes: [String: AttributeValue]
+}
diff --git a/Tests/Honeycomb/HoneycombDeterministicSamplerTests.swift b/Tests/Honeycomb/HoneycombDeterministicSamplerTests.swift
new file mode 100644
index 0000000..d281ffc
--- /dev/null
+++ b/Tests/Honeycomb/HoneycombDeterministicSamplerTests.swift
@@ -0,0 +1,53 @@
+import OpenTelemetryApi
+import OpenTelemetrySdk
+import XCTest
+
+@testable import Honeycomb
+
+class HoneycombDeterministicSamplerTests: XCTestCase {
+    func testSampler() {
+        let testCases = [
+            (rate: 0, sampled: false),
+            (rate: 1, sampled: true),
+            (rate: 10, sampled: true),
+            (rate: 100, sampled: true),
+        ]
+
+        // static trace id to ensure the inner traceIdRatio
+        // sampler always samples.
+        let traceID = TraceId.init(idHi: 10, idLo: 10)
+        let spanID = SpanId.random()
+        let parentContext = SpanContext.create(
+            traceId: traceID,
+            spanId: spanID,
+            traceFlags: TraceFlags.init(),
+            traceState: TraceState.init()
+        )
+
+        for (rate, sampled) in testCases {
+            XCTContext.runActivity(
+                named: "",
+                block: { activity in
+                    let sampler = HoneycombDeterministicSampler(sampleRate: rate)
+                    let result = sampler.shouldSample(
+                        parentContext: parentContext,
+                        traceId: traceID,
+                        name: "test",
+                        kind: SpanKind.client,
+                        attributes: [:],
+                        parentLinks: []
+                    )
+                    XCTAssertEqual(result.isSampled, sampled)
+
+                    if sampled {
+                        guard let r = result.attributes["SampleRate"] else {
+                            XCTFail("sample rate attribute not found")
+                            return
+                        }
+                        XCTAssertEqual(AttributeValue.int(rate), r)
+                    }
+                }
+            )
+        }
+    }
+}