diff --git a/ExampleApp/ExampleApp/AppDelegate.swift b/ExampleApp/ExampleApp/AppDelegate.swift index 65d7f41b..c41e9a9b 100644 --- a/ExampleApp/ExampleApp/AppDelegate.swift +++ b/ExampleApp/ExampleApp/AppDelegate.swift @@ -46,7 +46,7 @@ final class AppDelegate: NSObject, UIApplicationDelegate { _ application: UIApplication, didFinishLaunchingWithOptions launchOptions: [UIApplication.LaunchOptionsKey : Any]? = nil ) -> Bool { - print(once) + _ = once return true } } diff --git a/ExampleApp/ExampleApp/ContentView.swift b/ExampleApp/ExampleApp/ContentView.swift index 6255cf43..7a7ff947 100644 --- a/ExampleApp/ExampleApp/ContentView.swift +++ b/ExampleApp/ExampleApp/ContentView.swift @@ -8,6 +8,8 @@ import SwiftUI struct ContentView: View { + @Environment(Browser.self) var browser + var body: some View { VStack(spacing: 32) { Button { @@ -15,8 +17,23 @@ struct ContentView: View { } label: { Text("Crash") } - NetworkRequestView() - FeatureFlagView() + Button { + browser.navigate(to: .automaticInstrumentation) + } label: { + Text("Automatic Instrumentation") + } + Button { + browser.navigate(to: .evaluation) + } label: { + Text("Flag evaluation") + } + Button { + browser.navigate(to: .manualInstrumentation) + } label: { + Text("Manual Instrumentation") + } +// NetworkRequestView() +// FeatureFlagView() } .padding() } diff --git a/ExampleApp/ExampleApp/DMButtonStyle.swift b/ExampleApp/ExampleApp/DMButtonStyle.swift new file mode 100644 index 00000000..ad4bc21e --- /dev/null +++ b/ExampleApp/ExampleApp/DMButtonStyle.swift @@ -0,0 +1,18 @@ +import SwiftUI + +struct DMButtonStyle: ButtonStyle { + func makeBody(configuration: Configuration) -> some View { + configuration.label + .padding() + .background(Color.blue) + .foregroundColor(.white) + .cornerRadius(8) + .bold() + .frame(width: 100) + .opacity(configuration.isPressed ? 0.8 : 1) + } +} + +extension ButtonStyle where Self == DMButtonStyle { + static var ldStyle: Self { .init() } +} diff --git a/ExampleApp/ExampleApp/ExampleAppApp.swift b/ExampleApp/ExampleApp/ExampleAppApp.swift index f274671e..9071ddeb 100644 --- a/ExampleApp/ExampleApp/ExampleAppApp.swift +++ b/ExampleApp/ExampleApp/ExampleAppApp.swift @@ -3,10 +3,27 @@ import SwiftUI @main struct ExampleAppApp: App { @UIApplicationDelegateAdaptor(AppDelegate.self) private var appDelegate + @State private var browser = Browser() + var body: some Scene { WindowGroup { - ContentView() + NavigationStack(path: $browser.path) { + ContentView() + .navigationDestination(for: Path.self) { path in + switch path { + case .home: + ContentView() + case .manualInstrumentation: + InstrumentationView() + case .automaticInstrumentation: + NetworkRequestView() + case .evaluation: + FeatureFlagView() + } + } + } + .environment(browser) } } } diff --git a/ExampleApp/ExampleApp/NetworkRequestView.swift b/ExampleApp/ExampleApp/Instrumentation/Automatic/NetworkRequestView.swift similarity index 100% rename from ExampleApp/ExampleApp/NetworkRequestView.swift rename to ExampleApp/ExampleApp/Instrumentation/Automatic/NetworkRequestView.swift diff --git a/ExampleApp/ExampleApp/Instrumentation/Manual/InstrumentationView.swift b/ExampleApp/ExampleApp/Instrumentation/Manual/InstrumentationView.swift new file mode 100644 index 00000000..7b307b01 --- /dev/null +++ b/ExampleApp/ExampleApp/Instrumentation/Manual/InstrumentationView.swift @@ -0,0 +1,16 @@ +import SwiftUI +import LaunchDarklyObservability +import OpenTelemetryApi + +struct InstrumentationView: View { + var body: some View { + VStack { + TraceView() + } + .padding() + } +} + +#Preview { + InstrumentationView() +} diff --git a/ExampleApp/ExampleApp/Instrumentation/Manual/TraceView.swift b/ExampleApp/ExampleApp/Instrumentation/Manual/TraceView.swift new file mode 100644 index 00000000..8d65912c --- /dev/null +++ b/ExampleApp/ExampleApp/Instrumentation/Manual/TraceView.swift @@ -0,0 +1,39 @@ +import SwiftUI +import LaunchDarklyObservability +import OpenTelemetryApi + + +struct TraceView: View { + @State private var name: String = "" + @State private var started = false + @State private var span = Optional.none + + var body: some View { + VStack(alignment: .leading, spacing: 4) { + Text("Traces") + .bold() + HStack { + TextField(text: $name) { + Text("Span name:") + } + .textCase(.lowercase) + .textInputAutocapitalization(.never) + .textFieldStyle(.roundedBorder) + Spacer() + Text("is started") + Toggle(isOn: $started) { + Text("started") + } + .labelsHidden() + .disabled(name.isEmpty) + .task(id: started) { + guard started else { + span?.end() + return name = "" + } + span = LDObserve.shared.startSpan(name: name) + } + } + } + } +} diff --git a/ExampleApp/ExampleApp/Path.swift b/ExampleApp/ExampleApp/Path.swift new file mode 100644 index 00000000..8a3fc3f9 --- /dev/null +++ b/ExampleApp/ExampleApp/Path.swift @@ -0,0 +1,24 @@ +import SwiftUI + +enum Path: Hashable { + case home + case manualInstrumentation + case automaticInstrumentation + case evaluation +} + +@Observable final class Browser { + var path = NavigationPath() + + func navigate(to path: Path) { + self.path.append(path) + } + + func pop() { + self.path.removeLast() + } + + func reset() { + self.path = NavigationPath() + } +} diff --git a/Package.swift b/Package.swift index 1d0b4663..682492a7 100644 --- a/Package.swift +++ b/Package.swift @@ -34,6 +34,36 @@ let package = Package( .product(name: "Installations", package: "KSCrash") ] ), + .target( + name: "Sampling", + dependencies: [ + .product(name: "OpenTelemetryApi", package: "opentelemetry-swift"), + .product(name: "OpenTelemetrySdk", package: "opentelemetry-swift") + ] + ), + .target( + name: "SamplingLive", + dependencies: [ + "Sampling", + "Common", + .product(name: "OpenTelemetryApi", package: "opentelemetry-swift"), + .product(name: "OpenTelemetrySdk", package: "opentelemetry-swift") + ] + ), + .testTarget( + name: "SamplingLiveTests", + dependencies: [ + "Sampling", + "SamplingLive", + "Common", + .product(name: "OpenTelemetryApi", package: "opentelemetry-swift"), + .product(name: "OpenTelemetrySdk", package: "opentelemetry-swift") + ], + resources: [ + .copy("Resources/Stubs/Config.json"), + .copy("Resources/Stubs/MinConfig.json") + ] + ), .target( name: "Observability", dependencies: [ @@ -41,9 +71,15 @@ let package = Package( "API", "CrashReporter", "CrashReporterLive", + "Sampling", + "SamplingLive", .product(name: "OpenTelemetrySdk", package: "opentelemetry-swift"), + .product(name: "OpenTelemetryApi", package: "opentelemetry-swift"), .product(name: "OpenTelemetryProtocolExporterHTTP", package: "opentelemetry-swift"), .product(name: "URLSessionInstrumentation", package: "opentelemetry-swift"), + ], + resources: [ + .copy("Resources/Config.json"), ] ), .target( diff --git a/Sources/Common/SemanticConventionLD.swift b/Sources/Common/SemanticConventionLD.swift index 726acddb..c52f966f 100644 --- a/Sources/Common/SemanticConventionLD.swift +++ b/Sources/Common/SemanticConventionLD.swift @@ -1,3 +1,7 @@ -public enum SemanticConvention: String { - case highlightSessionId = "highlight.session_id" +public enum SemanticConvention { + public static let highlightSessionId = "highlight.session_id" +} + +public enum LDSemanticAttribute { + public static let ATTR_SAMPLING_RATIO = "launchdarkly.sampling.ratio" } diff --git a/Sources/LaunchDarklyObservability/LDObserve.swift b/Sources/LaunchDarklyObservability/LDObserve.swift index 03921913..99ba8250 100644 --- a/Sources/LaunchDarklyObservability/LDObserve.swift +++ b/Sources/LaunchDarklyObservability/LDObserve.swift @@ -65,19 +65,19 @@ public final class LDObserve: @unchecked Sendable, Observe { client.recordUpDownCounter(metric: metric) } - public func recordError(error: any Error, attributes: [String : AttributeValue]) { + public func recordError(error: any Error, attributes: [String : AttributeValue] = [:]) { lock.lock() defer { lock.unlock() } client.recordError(error: error, attributes: attributes) } - public func recordLog(message: String, severity: Severity, attributes: [String : AttributeValue]) { + public func recordLog(message: String, severity: Severity, attributes: [String : AttributeValue] = [:]) { lock.lock() defer { lock.unlock() } client.recordLog(message: message, severity: severity, attributes: attributes) } - public func startSpan(name: String, attributes: [String : AttributeValue]) -> any Span { + public func startSpan(name: String, attributes: [String : AttributeValue] = [:]) -> any Span { lock.lock() defer { lock.unlock() } return client.startSpan(name: name, attributes: attributes) diff --git a/Sources/Observability/InstrumentationManager.swift b/Sources/Observability/InstrumentationManager.swift index 2f1ee46f..2f21c37d 100644 --- a/Sources/Observability/InstrumentationManager.swift +++ b/Sources/Observability/InstrumentationManager.swift @@ -7,6 +7,8 @@ import URLSessionInstrumentation import Common import API +import Sampling +import SamplingLive private let tracesPath = "/v1/traces" private let logsPath = "/v1/logs" @@ -26,18 +28,30 @@ final class InstrumentationManager { private var cachedLongCounters = AtomicDictionary() private var cachedHistograms = AtomicDictionary() private var cachedUpDownCounters = AtomicDictionary() + private let sampler: ExportSampler public init(sdkKey: String, options: Options, sessionManager: SessionManager) { self.sdkKey = sdkKey self.options = options self.sessionManager = sessionManager + + let sampler = ExportSampler.customSampler() + /// Here is how we will inject the sampling config coming from backend, if the next lines + /// are uncommented, them will inject a example local config for testing purposes only. + ///let samplingConfig = loadSampleConfig() + ///sampler.setConfig(samplingConfig) + + let processorAndProvider = URL(string: options.otlpEndpoint) .flatMap { $0.appending(path: logsPath) } .map { url in - OtlpHttpLogExporter( - endpoint: url, - envVarHeaders: options.customHeaders + SamplingLogExporterDecorator( + exporter: OtlpHttpLogExporter( + endpoint: url, + envVarHeaders: options.customHeaders + ), + sampler: sampler ) } .map { exporter in @@ -71,9 +85,12 @@ final class InstrumentationManager { URL(string: options.otlpEndpoint) .flatMap { $0.appending(path: tracesPath) } .map { url in - OtlpHttpTraceExporter( - endpoint: url, - envVarHeaders: options.customHeaders + SamplingTraceExporterDecorator( + exporter: OtlpHttpTraceExporter( + endpoint: url, + envVarHeaders: options.customHeaders + ), + sampler: sampler ) } .map { exporter in @@ -153,6 +170,8 @@ final class InstrumentationManager { tracer: self.otelTracer ) ) + + self.sampler = sampler } func recordMetric(metric: Metric) { @@ -206,7 +225,7 @@ final class InstrumentationManager { var attributes = attributes let sessionId = sessionManager.sessionInfo.id if !sessionId.isEmpty { - attributes[SemanticConvention.highlightSessionId.rawValue] = .string(sessionId) + attributes[SemanticConvention.highlightSessionId] = .string(sessionId) } otelLogger?.logRecordBuilder() .setBody(.string(message)) @@ -230,8 +249,8 @@ final class InstrumentationManager { let sessionId = sessionManager.sessionInfo.id if !sessionId.isEmpty { - builder?.setAttribute(key: SemanticConvention.highlightSessionId.rawValue, value: sessionId) - attributes[SemanticConvention.highlightSessionId.rawValue] = .string(sessionId) + builder?.setAttribute(key: SemanticConvention.highlightSessionId, value: sessionId) + attributes[SemanticConvention.highlightSessionId] = .string(sessionId) } @@ -265,3 +284,17 @@ final class InstrumentationManager { return builder.startSpan() } } + +func loadSampleConfig() -> SamplingConfig? { + guard let url = Bundle.module.url(forResource: "Config", withExtension: "json") else { + return nil + } + do { + let data = try Data(contentsOf: url) + let decoder = JSONDecoder() + let root = try decoder.decode(Root.self, from: data) + return root.data.sampling + } catch { + return nil + } +} diff --git a/Sources/Observability/Resources/Config.json b/Sources/Observability/Resources/Config.json new file mode 100644 index 00000000..4bc24304 --- /dev/null +++ b/Sources/Observability/Resources/Config.json @@ -0,0 +1,155 @@ +{ + "data": { + "sampling": { + "logs": [ + { + "severityText": { + "matchValue": "ERROR" + }, + "samplingRatio": 2 + }, + { + "severityText": { + "regexValue": "ERROR: .*" + }, + "samplingRatio": 2 + }, + { + "message": { + "matchValue": "Connection failed" + }, + "samplingRatio": 2 + }, + { + "message": { + "regexValue": "Error: .*" + }, + "samplingRatio": 2 + }, + { + "message": { + "regexValue": "Database connection .*" + }, + "severityText": { + "matchValue": "ERROR" + }, + "attributes": [ + { + "key": { + "regexValue": "service.*" + }, + "attribute": { + "regexValue": "db-.*" + } + }, + { + "key": { + "matchValue": "retry.enabled" + }, + "attribute": { + "matchValue": true + } + }, + { + "key": { + "matchValue": "retry.count" + }, + "attribute": { + "matchValue": 3 + } + }, + { + "key": { + "matchValue": "retry.timeout" + }, + "attribute": { + "matchValue": 15.5 + } + } + ], + "samplingRatio": 2 + } + ], + "spans": [ + { + "name": { + "matchValue": "test-span" + }, + "samplingRatio": 2 + }, + { + "name": { + "regexValue": "test-span-\\d+" + }, + "samplingRatio": 2 + }, + { + "events": [ + { + "name": { + "matchValue": "test-event" + } + } + ], + "samplingRatio": 2 + }, + { + "events": [ + { + "name": { + "regexValue": "test-event-\\d+" + } + } + ], + "samplingRatio": 2 + }, + { + "attributes": [ + { + "key": { + "matchValue": "http.method" + }, + "attribute": { + "matchValue": "POST" + } + } + ], + "samplingRatio": 2 + }, + { + "events": [ + { + "attributes": [ + { + "key": { + "matchValue": "error.type" + }, + "attribute": { + "matchValue": "network" + } + }, + { + "key": { + "matchValue": "db.error" + }, + "attribute": { + "regexValue": "Database connection .*" + } + }, + { + "key": { + "matchValue": "error.code" + }, + "attribute": { + "matchValue": 503 + } + } + ] + } + ], + "samplingRatio": 2 + } + ] + } + } +} diff --git a/Sources/Sampling/ExportSampler.swift b/Sources/Sampling/ExportSampler.swift new file mode 100644 index 00000000..6235e94e --- /dev/null +++ b/Sources/Sampling/ExportSampler.swift @@ -0,0 +1,21 @@ +import OpenTelemetryApi +import OpenTelemetrySdk + +public struct ExportSampler { + public var sampleSpan: (SpanData) -> SamplingResult + public var sampleLog: (ReadableLogRecord) -> SamplingResult + public var isSamplingEnabled: () -> Bool + public var setConfig: (_ config: SamplingConfig?) -> Void + + public init( + sampleSpan: @escaping (SpanData) -> SamplingResult, + sampleLog: @escaping (ReadableLogRecord) -> SamplingResult, + isSamplingEnabled: @escaping () -> Bool, + setConfig: @escaping (_: SamplingConfig?) -> Void + ) { + self.sampleSpan = sampleSpan + self.sampleLog = sampleLog + self.isSamplingEnabled = isSamplingEnabled + self.setConfig = setConfig + } +} diff --git a/Sources/Sampling/SamplingConfig.swift b/Sources/Sampling/SamplingConfig.swift new file mode 100644 index 00000000..d2cd8feb --- /dev/null +++ b/Sources/Sampling/SamplingConfig.swift @@ -0,0 +1,122 @@ +import Foundation +import OpenTelemetryApi + +public enum MatchConfig: Codable { + case basic(value: AttributeValue) + case regex(expression: String) + + enum CodingKeys: String, CodingKey { + case basic = "matchValue" + case regex = "regexValue" + } + + public init(from decoder: any Decoder) throws { + let container = try decoder.container(keyedBy: CodingKeys.self) + if let regex = try? container.decode(String.self, forKey: .regex) { + self = .regex(expression: regex) + } else if let value = try? container.decode(String.self, forKey: .basic) { + self = .basic(value: .string(value)) + } else if let value = try? container.decode(Int.self, forKey: .basic) { + self = .basic(value: .int(value)) + } else if let value = try? container.decode(Bool.self, forKey: .basic) { + self = .basic(value: .bool(value)) + } else if let value = try? container.decode(Double.self, forKey: .basic) { + self = .basic(value: .double(value)) + } else if let value = try? container.decode(AttributeArray.self, forKey: .basic) { + self = .basic(value: .array(value)) + } else if let value = try? container.decode(AttributeSet.self, forKey: .basic) { + self = .basic(value: .set(value)) + } + else { + throw DecodingError.typeMismatch( + Any.self, + .init(codingPath: decoder.codingPath, debugDescription: "Unsupported type for MatchConfig") + ) + } + } + + public func encode(to encoder: any Encoder) throws { + var container = encoder.container(keyedBy: CodingKeys.self) + switch self { + case .basic(let value): + try container.encode(value, forKey: .basic) + case .regex(let expression): + try container.encode(expression, forKey: .regex) + } + } +} + +public struct AttributeMatchConfig: Codable { + public let key: MatchConfig + public let attribute: MatchConfig + + public init(key: MatchConfig, attribute: MatchConfig) { + self.key = key + self.attribute = attribute + } +} + +public struct SpanEventMatchConfig: Codable { + public let name: MatchConfig? + public let attributes: [AttributeMatchConfig]? + + public init(name: MatchConfig? = nil, attributes: [AttributeMatchConfig] = []) { + self.name = name + self.attributes = attributes + } +} + +public struct SpanSamplingConfig: Codable { + public let name: MatchConfig? + public let attributes: [AttributeMatchConfig]? + public let events: [SpanEventMatchConfig]? + public let samplingRatio: Int + + public init( + name: MatchConfig? = nil, + attributes: [AttributeMatchConfig] = [], + events: [SpanEventMatchConfig] = [], + samplingRatio: Int + ) { + self.name = name + self.attributes = attributes + self.events = events + self.samplingRatio = samplingRatio + } +} + +public struct LogSamplingConfig: Codable { + public let message: MatchConfig? + public let severityText: MatchConfig? + public let attributes: [AttributeMatchConfig]? + public let samplingRatio: Int + + public init( + message: MatchConfig? = nil, + severityText: MatchConfig? = nil, + attributes: [AttributeMatchConfig] = [], + samplingRatio: Int + ) { + self.message = message + self.severityText = severityText + self.attributes = attributes + self.samplingRatio = samplingRatio + } +} + +public struct SamplingConfig: Codable { + public let spans: [SpanSamplingConfig]? + public let logs: [LogSamplingConfig]? + + public init(spans: [SpanSamplingConfig] = [], logs: [LogSamplingConfig] = []) { + self.spans = spans + self.logs = logs + } +} + +public struct Root: Codable { + public struct SamplingData: Codable { + public let sampling: SamplingConfig + } + public let data: SamplingData +} diff --git a/Sources/Sampling/SamplingResult.swift b/Sources/Sampling/SamplingResult.swift new file mode 100644 index 00000000..75e50c18 --- /dev/null +++ b/Sources/Sampling/SamplingResult.swift @@ -0,0 +1,12 @@ +import OpenTelemetryApi + + +public struct SamplingResult { + public let sample: Bool + public let attributes: [String: AttributeValue]? + + public init(sample: Bool, attributes: [String : AttributeValue]? = nil) { + self.sample = sample + self.attributes = attributes + } +} diff --git a/Sources/SamplingLive/CustomSampler.swift b/Sources/SamplingLive/CustomSampler.swift new file mode 100644 index 00000000..5c7bf842 --- /dev/null +++ b/Sources/SamplingLive/CustomSampler.swift @@ -0,0 +1,217 @@ +import Foundation + +import OpenTelemetryApi +import OpenTelemetrySdk + +import Common +import Sampling + +extension ExportSampler { + public static func customSampler( + sampler: ((Int) -> Bool)? = nil + ) -> Self { + final class CustomSampler { + private let sampler: (Int) -> Bool + private var regexCache = AtomicDictionary>() + private let queue = DispatchQueue(label: "com.launchdarkly.sampler.custom", attributes: .concurrent) + private var _config: SamplingConfig? + private var config: SamplingConfig? { + get { + queue.sync { [weak self] in self?._config } + } + set { + queue.async(flags: .barrier) { [weak self] in self?._config = newValue } + } + } + + init(config: SamplingConfig? = nil, sampler: ((Int) -> Bool)? = nil) { + self._config = config + self.sampler = sampler ?? ThreadSafeSampler.shared.sample(_:) + } + + func sampleSpan(_ spanData: SpanData) -> SamplingResult { + guard let config else { + return .init(sample: true) + } + + for spanConfig in config.spans ?? [] { + if matchesSpanConfig(config: spanConfig, span: spanData) { + return .init( + sample: sampler(spanConfig.samplingRatio), + attributes: [LDSemanticAttribute.ATTR_SAMPLING_RATIO: .int(spanConfig.samplingRatio)] + ) + } + } + + /// Didn't match any sampling config, or there were no configs, so we sample it. + return .init(sample: true) + } + + func sampleLog(_ logData: ReadableLogRecord) -> SamplingResult { + guard let config, !(config.logs?.isEmpty ?? true) else { + return .init(sample: true) + } + + for logConfig in config.logs ?? [] { + if matchesLogConfig(config: logConfig, record: logData) { + return .init( + sample: sampler(logConfig.samplingRatio), + attributes: [LDSemanticAttribute.ATTR_SAMPLING_RATIO: .int(logConfig.samplingRatio)] + ) + } + } + + /// Didn't match any sampling config, or there were no configs, so we sample it. + return .init(sample: true) + } + + func setConfig(_ config: SamplingConfig?) { + self.config = config + } + + /// Check if sampling is enabled. + /// Sampling is enabled if there is at least one configuration in either the log or span sampling. + /// - Returns: true if sampling is enabled + func isSamplingEnabled() -> Bool { + guard let config else { + return false + } + + return config.spans?.isEmpty == false || config.logs?.isEmpty == false + } + + private func matchesValue( + matchConfig: MatchConfig?, + value: AttributeValue + ) -> Bool { + guard let matchConfig else { return false } + switch matchConfig { + case .basic(let configValue): + return configValue == value + case .regex(let pattern): + guard case .string(let valueString) = value else { return false } + var regex = regexCache[pattern] + if regex == nil { + regex = try? Regex(pattern) + regexCache[pattern] = regex + } + return (try? regex?.firstMatch(in: valueString) != nil) ?? false + } + } + + private func matchesAttributes( + attributeConfigs: [AttributeMatchConfig]?, + attributes: [String: AttributeValue]? + ) -> Bool { + guard let attributeConfigs else { + return true + } + guard !attributeConfigs.isEmpty else { + return true + } + + // No attributes, so they cannot match. + guard let attributes else { + return false + } + + return attributeConfigs.allSatisfy { config in + let result = attributes.contains { key, value in + let match = matchesValue(matchConfig: config.key, value: .string(key)) && matchesValue(matchConfig: config.attribute, value: value) + return match + } + return result + } + } + + private func matchEvent( + eventConfig: SpanEventMatchConfig, + event: SpanData.Event + ) -> Bool { + if let eventConfigName = eventConfig.name { + // Match by Event name + if !matchesValue(matchConfig: eventConfigName, value: .string(event.name)) { + return false + } + } + + // Match by event attributes if specified + if !matchesAttributes(attributeConfigs: eventConfig.attributes, attributes: event.attributes) { + return false + } + + return true + } + + private func matchesEvents( + eventConfigs: [SpanEventMatchConfig]?, + events: [SpanData.Event] + ) -> Bool { + guard let eventConfigs else { + return true + } + + guard !eventConfigs.isEmpty else { + return true + } + + guard !events.isEmpty else { + return false + } + + + return eventConfigs.allSatisfy { eventConfig in + events.contains { event in + matchEvent(eventConfig: eventConfig, event: event) + } + } + } + + private func matchesSpanConfig( + config: SpanSamplingConfig, + span: SpanData + ) -> Bool { + // Check span name if it's defined in the config + if let configName = config.name { + if !matchesValue(matchConfig: configName, value: .string(span.name)) { + return false + } + } + + if !matchesAttributes(attributeConfigs: config.attributes, attributes: span.attributes) { + return false + } + + return matchesEvents(eventConfigs: config.events, events: span.events) + } + + private func matchesLogConfig( + config: LogSamplingConfig, + record: ReadableLogRecord + ) -> Bool { + if let severityText = config.severityText, let severity = record.severity?.description { + if !matchesValue(matchConfig: severityText, value: .string(severity)) { + return false + } + } + if let configName = config.message, let body = record.body { + if !matchesValue(matchConfig: configName, value: body) { + return false + } + } + + return matchesAttributes(attributeConfigs: config.attributes, attributes: record.attributes) + } + + } + + let sampler = CustomSampler(sampler: sampler) + + return .init( + sampleSpan: { sampler.sampleSpan($0) }, + sampleLog: { sampler.sampleLog($0) }, + isSamplingEnabled: { sampler.isSamplingEnabled() }, + setConfig: { sampler.setConfig($0) } + ) + } +} diff --git a/Sources/SamplingLive/ExportSampler+SampleLogs.swift b/Sources/SamplingLive/ExportSampler+SampleLogs.swift new file mode 100644 index 00000000..a61d11c6 --- /dev/null +++ b/Sources/SamplingLive/ExportSampler+SampleLogs.swift @@ -0,0 +1,30 @@ +import OpenTelemetrySdk + +import Sampling + +extension ExportSampler { + func sampleLogs( + items: [ReadableLogRecord] + ) -> [ReadableLogRecord] { + guard isSamplingEnabled() else { + return items + } + + return items.compactMap { item in + let sampleResult = sampleLog(item) + guard sampleResult.sample else { + return nil + } + return ReadableLogRecord( + resource: item.resource, + instrumentationScopeInfo: item.instrumentationScopeInfo, + timestamp: item.timestamp, + observedTimestamp: item.observedTimestamp, + spanContext: item.spanContext, + severity: item.severity, + body: item.body, + attributes: item.attributes.merging(sampleResult.attributes ?? [:], uniquingKeysWith: { current, new in current }) // Merge, prioritizing values from logRecord for duplicate keys + ) + } + } +} diff --git a/Sources/SamplingLive/ExportSampler+SampleSpans.swift b/Sources/SamplingLive/ExportSampler+SampleSpans.swift new file mode 100644 index 00000000..86789774 --- /dev/null +++ b/Sources/SamplingLive/ExportSampler+SampleSpans.swift @@ -0,0 +1,51 @@ +import OpenTelemetrySdk +import OpenTelemetryApi + +import Sampling + +extension ExportSampler { + func sampleSpans( + items: [SpanData] + ) -> [SpanData] { + if !isSamplingEnabled() { + return items + } + + var omittedSpansIds = Set() + var spanById = [SpanId: SpanData]() + var childrenByParentId = [SpanId: [SpanId]]() + + /// The first pass we sample items which are directly impacted by a sampling decision. + /// We also build a map of children spans by parent span id, which allows us to quickly traverse the span tree. + for item in items { + if let parentSpanId = item.parentSpanId { + childrenByParentId[parentSpanId, default: []].append(item.spanId) + } + + let sampleResult = sampleSpan(item) + if sampleResult.sample { + var mutableSpanData = item + mutableSpanData.settingAttributes( + item.attributes.merging(sampleResult.attributes ?? [:], uniquingKeysWith: { current, new in current }) + ) // Merge, prioritizing values from spanData for duplicate keys + spanById[item.spanId] = mutableSpanData + } else { + omittedSpansIds.insert(item.spanId) + } + } + + /// Find all children of spans that have been sampled out and remove them. + /// Repeat until there are no more children to remove. + while let omittedSpanId = omittedSpansIds.popFirst() { + + guard let affectedSpans = childrenByParentId[omittedSpanId] else { continue } + + for spanIdToRemove in affectedSpans { + spanById.removeValue(forKey: spanIdToRemove) + omittedSpansIds.insert(spanIdToRemove) + } + } + + return items.compactMap { spanById[$0.spanId] } + } +} diff --git a/Sources/SamplingLive/SamplingLogExporter.swift b/Sources/SamplingLive/SamplingLogExporter.swift new file mode 100644 index 00000000..03de8dd5 --- /dev/null +++ b/Sources/SamplingLive/SamplingLogExporter.swift @@ -0,0 +1,44 @@ +import Foundation + +import OpenTelemetrySdk +import OpenTelemetryApi + +import Sampling + +public final class SamplingLogExporterDecorator: LogRecordExporter { + private let exporter: LogRecordExporter + private let sampler: ExportSampler + + public init(exporter: LogRecordExporter, sampler: ExportSampler) { + self.exporter = exporter + self.sampler = sampler + } + + public func forceFlush( + explicitTimeout: TimeInterval? + ) -> ExportResult { + exporter.forceFlush(explicitTimeout: explicitTimeout) + } + + public func shutdown( + explicitTimeout: TimeInterval? + ) { + exporter.shutdown(explicitTimeout: explicitTimeout) + } + + public func export( + logRecords: [ReadableLogRecord], + explicitTimeout: TimeInterval? + ) -> ExportResult { + let sampledItems = sampler.sampleLogs( + items: logRecords + ) + guard !sampledItems.isEmpty else { + return .success + } + + return exporter.export(logRecords: sampledItems, explicitTimeout: explicitTimeout) + } + + +} diff --git a/Sources/SamplingLive/SamplingTraceExporter.swift b/Sources/SamplingLive/SamplingTraceExporter.swift new file mode 100644 index 00000000..7d1142cc --- /dev/null +++ b/Sources/SamplingLive/SamplingTraceExporter.swift @@ -0,0 +1,38 @@ +import Foundation + +import OpenTelemetrySdk +import OpenTelemetryApi + +import Sampling + +public final class SamplingTraceExporterDecorator: SpanExporter { + private let exporter: SpanExporter + private let sampler: ExportSampler + + public init(exporter: SpanExporter, sampler: ExportSampler) { + self.exporter = exporter + self.sampler = sampler + } + + public func shutdown(explicitTimeout: TimeInterval?) { + exporter.shutdown(explicitTimeout: explicitTimeout) + } + + public func flush( + explicitTimeout: TimeInterval? + ) -> SpanExporterResultCode { + exporter.flush(explicitTimeout: explicitTimeout) + } + + public func export( + spans: [SpanData], + explicitTimeout: TimeInterval? + ) -> SpanExporterResultCode { + let sampledItems = sampler.sampleSpans(items: spans) + guard !sampledItems.isEmpty else { + return .success + } + + return exporter.export(spans: sampledItems, explicitTimeout: explicitTimeout) + } +} diff --git a/Sources/SamplingLive/ThreadSafeSampler.swift b/Sources/SamplingLive/ThreadSafeSampler.swift new file mode 100644 index 00000000..fd9ac510 --- /dev/null +++ b/Sources/SamplingLive/ThreadSafeSampler.swift @@ -0,0 +1,22 @@ +import Foundation + +final class ThreadSafeSampler { + static let shared = ThreadSafeSampler() + private var generator = SystemRandomNumberGenerator() + private let queue = DispatchQueue(label: "com.launchdarkly.sampler") + + private init() {} + + private func nextInt(in range: Range) -> Int { + queue.sync { + .random(in: range, using: &generator) + } + } + + func sample(_ ratio: Int) -> Bool { + if ratio <= 0 { return false } + if ratio == 1 { return true } + + return nextInt(in: 0.. SamplingResult = { _ in .init(sample: true) }, + sampleLog: @escaping (ReadableLogRecord) -> SamplingResult = { _ in .init(sample: true) }, + isSamplingEnabled: Bool = true + ) -> ExportSampler { + final class FakeExportSampler { + var config: SamplingConfig? + + func setConfig(_ config: SamplingConfig?) { + self.config = config + } + } + let fakeExportSampler = FakeExportSampler() + return .init( + sampleSpan: sampleSpan, + sampleLog: sampleLog, + isSamplingEnabled: { isSamplingEnabled }, + setConfig: { fakeExportSampler.setConfig($0) } + ) + } +} diff --git a/Tests/SamplingLiveTests/Mocks.swift b/Tests/SamplingLiveTests/Mocks.swift new file mode 100644 index 00000000..c8550e7a --- /dev/null +++ b/Tests/SamplingLiveTests/Mocks.swift @@ -0,0 +1,51 @@ +import Foundation + +@testable import OpenTelemetrySdk +import OpenTelemetryApi + +func makeMockReadableLogRecord( + body: AttributeValue? = nil, + severity: Severity? = nil, + attributes: [String: AttributeValue] = .init() +) -> ReadableLogRecord { + .init( + resource: .empty, + instrumentationScopeInfo: .init(), + timestamp: .now, + severity: severity, + body: body, + attributes: attributes + ) +} + +func makeMockSpanData( + name: String, + spanId: SpanId = .random(), + parentSpanId: SpanId? = nil, + events: [SpanData.Event] = [], + attributes: [String: AttributeValue] = [:] +) -> SpanData { + SpanData( + traceId: .random(), + spanId: spanId, + parentSpanId: parentSpanId, + name: name, + kind: .client, + startTime: .now, + attributes: attributes, + events: events, + endTime: .now.addingTimeInterval(60 * 2) + ) +} + +func makeMockSpanEvent( + name: String, + timestamp: Date = .now, + attributes: [String: AttributeValue]? = nil +) -> SpanData.Event { + SpanData.Event( + name: name, + timestamp: timestamp, + attributes: attributes + ) +} diff --git a/Tests/SamplingLiveTests/Resources/Stubs/Config.json b/Tests/SamplingLiveTests/Resources/Stubs/Config.json new file mode 100644 index 00000000..826c2b90 --- /dev/null +++ b/Tests/SamplingLiveTests/Resources/Stubs/Config.json @@ -0,0 +1,155 @@ +{ + "data": { + "sampling": { + "logs": [ + { + "severityText": { + "matchValue": "ERROR" + }, + "samplingRatio": 0 + }, + { + "severityText": { + "regexValue": "ERROR: .*" + }, + "samplingRatio": 0 + }, + { + "message": { + "matchValue": "Connection failed" + }, + "samplingRatio": 0 + }, + { + "message": { + "regexValue": "Error: .*" + }, + "samplingRatio": 0 + }, + { + "message": { + "regexValue": "Database connection .*" + }, + "severityText": { + "matchValue": "ERROR" + }, + "attributes": [ + { + "key": { + "regexValue": "service.*" + }, + "attribute": { + "regexValue": "db-.*" + } + }, + { + "key": { + "matchValue": "retry.enabled" + }, + "attribute": { + "matchValue": true + } + }, + { + "key": { + "matchValue": "retry.count" + }, + "attribute": { + "matchValue": 3 + } + }, + { + "key": { + "matchValue": "retry.timeout" + }, + "attribute": { + "matchValue": 15.5 + } + } + ], + "samplingRatio": 0 + } + ], + "spans": [ + { + "name": { + "matchValue": "test-span" + }, + "samplingRatio": 0 + }, + { + "name": { + "regexValue": "test-span-\\d+" + }, + "samplingRatio": 0 + }, + { + "events": [ + { + "name": { + "matchValue": "test-event" + } + } + ], + "samplingRatio": 0 + }, + { + "events": [ + { + "name": { + "regexValue": "test-event-\\d+" + } + } + ], + "samplingRatio": 0 + }, + { + "attributes": [ + { + "key": { + "matchValue": "http.method" + }, + "attribute": { + "matchValue": "POST" + } + } + ], + "samplingRatio": 0 + }, + { + "events": [ + { + "attributes": [ + { + "key": { + "matchValue": "error.type" + }, + "attribute": { + "matchValue": "network" + } + }, + { + "key": { + "matchValue": "db.error" + }, + "attribute": { + "regexValue": "Database connection .*" + } + }, + { + "key": { + "matchValue": "error.code" + }, + "attribute": { + "matchValue": 503 + } + } + ] + } + ], + "samplingRatio": 0 + } + ] + } + } +} diff --git a/Tests/SamplingLiveTests/Resources/Stubs/MinConfig.json b/Tests/SamplingLiveTests/Resources/Stubs/MinConfig.json new file mode 100644 index 00000000..54a856e9 --- /dev/null +++ b/Tests/SamplingLiveTests/Resources/Stubs/MinConfig.json @@ -0,0 +1,14 @@ +{ + "data": { + "sampling": { + "spans": [ + { + "name": { + "matchValue": "test-span" + }, + "samplingRatio": 0 + } + ] + } + } +} diff --git a/Tests/SamplingLiveTests/SampleSpansParentRelationship.swift b/Tests/SamplingLiveTests/SampleSpansParentRelationship.swift new file mode 100644 index 00000000..103411f9 --- /dev/null +++ b/Tests/SamplingLiveTests/SampleSpansParentRelationship.swift @@ -0,0 +1,123 @@ +import Testing + +@testable import OpenTelemetrySdk +import OpenTelemetryApi + +import Common +import Sampling +@testable import SamplingLive + +struct SampleSpansParentRelationship { + @Test("should remove child and grandchild spans when parent is not sampled") + func shouldRemoveChildAndGrandchildSpansWhenParentIsNotSampled() { + let sampler = ExportSampler.fake( + sampleSpan: { span in + if span.name == "parent" { + return .init(sample: false) + } else { + return .init(sample: true) + } + }, + isSamplingEnabled: true + ) + + let parentSpan = makeMockSpanData(name: "parent") + let childSpan = makeMockSpanData(name: "child", parentSpanId: parentSpan.spanId) + let grandchildSpan = makeMockSpanData(name: "grandchild", parentSpanId: childSpan.spanId) + let unrelatedSpan = makeMockSpanData(name: "unrelated") + + let spans = [parentSpan, childSpan, grandchildSpan, unrelatedSpan] + + let result = sampler.sampleSpans(items: spans) + + #expect(result.count == 1) + #expect(result[0] == unrelatedSpan) + } + + @Test("should keep child spans when parent is sampled") + func shouldKeepChildSpansWhenParentIsSampled() { + let sampler = ExportSampler.fake( + sampleSpan: { span in + return .init(sample: true) + }, + isSamplingEnabled: true + ) + + let parentSpan = makeMockSpanData(name: "parent") + let childSpan1 = makeMockSpanData(name: "child1", parentSpanId: parentSpan.spanId) + let childSpan2 = makeMockSpanData(name: "child2", parentSpanId: parentSpan.spanId) + + let spans = [parentSpan, childSpan1, childSpan2] + + let result = sampler.sampleSpans(items: spans) + + #expect(result.count == 3) + #expect(result.allSatisfy { spans.contains($0) }) + } + + @Test("should remove child spans even if parent is sampled but child is not") + func shouldRemoveChildSpansEvenIfParentIsSampledButChildIsNot() { + let sampler = ExportSampler.fake( + sampleSpan: { span in + if span.name == "child" { + return .init(sample: false) + } else { + return .init(sample: true) + } + }, + isSamplingEnabled: true + ) + + let parentSpan = makeMockSpanData(name: "parent") + let childSpan = makeMockSpanData(name: "child", parentSpanId: parentSpan.spanId) + let grandchildSpan = makeMockSpanData(name: "grandchild", parentSpanId: childSpan.spanId) + + let spans = [parentSpan, childSpan, grandchildSpan] + + let result = sampler.sampleSpans(items: spans) + + #expect(result.count == 1) + #expect(result[0] == parentSpan) + } + + @Test("should handle complex span hierarchy with mixed sampling") + func shouldHandleComplexSpanHierarchyWithMixedSampling() { + /* + Create a complex hierarchy: + parent1 (sampled) -> child1 (not sampled) -> grandchild1 (sampled) + parent2 (not sampled) -> child2 (sampled) -> grandchild2 (sampled) + unrelated (sampled) + */ + + let sampler = ExportSampler.fake( + sampleSpan: { span in + if ["parent1", "grandchild1", "child2", "grandchild2", "unrelated"].contains(span.name) { + return .init(sample: true) + } else if ["child1", "parent2"].contains(span.name) { + return .init(sample: false) + } else { + return .init(sample: false) + } + }, + isSamplingEnabled: true + ) + + let parent1 = makeMockSpanData(name: "parent1") + let child1 = makeMockSpanData(name: "child1", parentSpanId: parent1.spanId) + let grandchild1 = makeMockSpanData(name: "grandchild1", parentSpanId: child1.spanId) + + let parent2 = makeMockSpanData(name: "parent2") + let child2 = makeMockSpanData(name: "child2", parentSpanId: parent2.spanId) + let grandchild2 = makeMockSpanData(name: "grandchild2", parentSpanId: child2.spanId) + + let unrelated = makeMockSpanData(name: "unrelated") + + let spans = [parent1, child1, grandchild1, parent2, child2, grandchild2, unrelated] + + let result = sampler.sampleSpans(items: spans) + + #expect(result.count == 2) + #expect(result.contains(parent1)) + #expect(result.contains(unrelated)) + } +} diff --git a/Tests/SamplingLiveTests/SamplerTests.swift b/Tests/SamplingLiveTests/SamplerTests.swift new file mode 100644 index 00000000..989aa71e --- /dev/null +++ b/Tests/SamplingLiveTests/SamplerTests.swift @@ -0,0 +1,17 @@ +import Testing + +import Sampling +@testable import SamplingLive + +struct SamplerTests { + + @Test("defaultSampler should always return true for ratio 1") + func defaultSamplerTrueForRatioOne() { + #expect(ThreadSafeSampler.shared.sample(1)) + } + + @Test("defaultSampler should always return false for ratio 0") + func defaultSamplerFalseForRatioZero() { + #expect(!ThreadSafeSampler.shared.sample(0)) + } +} diff --git a/Tests/SamplingLiveTests/SamplingLogsTests.swift b/Tests/SamplingLiveTests/SamplingLogsTests.swift new file mode 100644 index 00000000..5304176e --- /dev/null +++ b/Tests/SamplingLiveTests/SamplingLogsTests.swift @@ -0,0 +1,171 @@ +import Testing + +@testable import OpenTelemetrySdk +import OpenTelemetryApi + +import Common +import Sampling +@testable import SamplingLive + +struct SamplingLogsTests { + @Test func threadSafeSampler() { + let sampler = ThreadSafeSampler.shared + + #expect(sampler.sample(1) == true) + #expect(sampler.sample(0) == false) + } + + @Test("Given a list of logs, when sampling is disabled, then exporter should return all logs") + func disablingSampling() { + let sampler = ExportSampler.fake(isSamplingEnabled: false) + let logs = (1...3).map { makeMockReadableLogRecord(body: .string("log\($0)")) } + + let result = sampler.sampleLogs(items: logs) + + #expect(logs.count == result.count) + + for (r, l) in zip(result, logs) { + #expect(r.body == l.body) + } + } + + @Test("Given a empty list of logs, when sampling is enabled, then exporter should return empty result") + func enablingSampling() { + let sampler = ExportSampler.fake(isSamplingEnabled: true) + + let result = sampler.sampleLogs(items: []) + + #expect(result.isEmpty) + } + + @Test("Given a list of logs, when sampling is enabled and no logs are sampled, then should return empty") + func enablingNoLogsSampled() { + let sampler = ExportSampler.fake( + sampleLog: { _ in .init(sample: false) }, + isSamplingEnabled: true + ) + let logs = (1...3).map { makeMockReadableLogRecord(body: .string("log\($0)")) } + + let result = sampler.sampleLogs(items: logs) + + #expect(result.isEmpty) + } + + @Test( + """ + Given a list of logs without additional attributes, + when sampling is enabled and all logs are sampled, then should return all + """ + ) + func enablingAllLogsSampledWithoutAdditionalAttributes() { + let sampler = ExportSampler.fake( + sampleLog: { _ in .init(sample: true, attributes: nil) }, + isSamplingEnabled: true + ) + let logs = (1...3).map { makeMockReadableLogRecord(body: .string("log\($0)")) } + + let result = sampler.sampleLogs(items: logs) + + #expect(logs.count == result.count) + + for (r, l) in zip(result, logs) { + #expect(r.body == l.body) + } + } + + @Test("Given a list of logs, when sampling is enabled, then should return only sampled ones") + func enablingSamplingSomeAreSampled() { + let lowerBound = 1 + let upperBound = 4 + let logs = (lowerBound...upperBound).map { makeMockReadableLogRecord(body: .string("log\($0)")) } + + let sampler = ExportSampler.fake( + sampleLog: { logData in + if logData.body == logs[2].body { + return .init(sample: false) + } else { + return .init(sample: true) + } + }, + isSamplingEnabled: true + ) + + + let result = sampler.sampleLogs(items: logs) + + + let logAtIndex2ThatIsNotSampledIsInTheArray = { (log: ReadableLogRecord) -> Bool in + log.body == logs[2].body + } + #expect(result.count == logs.count - 1) + + #expect(result.contains(where: logAtIndex2ThatIsNotSampledIsInTheArray) == false) + } + + @Test("Given a log with attributes, when sampling is enabled and sampling attributes is provided, then should add sampling attributes and preserve the original ones") + func enablingSamplingProvidingSamplingAttributes() { + let originalAttributes = [ + "service.name": AttributeValue.string("api-service"), + "environment": AttributeValue.string("production") + ] + let samplingAttributes = [ + LDSemanticAttribute.ATTR_SAMPLING_RATIO: AttributeValue.int(42) + ] + let mockLog = makeMockReadableLogRecord(body: .string("test-log"), attributes: originalAttributes) + let sampler = ExportSampler.fake( + sampleLog: { logData in + if let logBody = logData.body, case AttributeValue.string(let body) = logBody, body == "test-log" { + return .init(sample: true, attributes: samplingAttributes) + } else { + return .init(sample: false) + } + }, + isSamplingEnabled: true + ) + + // When + let items = [mockLog] + let result = sampler.sampleLogs(items: items) + #expect(result.count == items.count) + + #expect(result[0].attributes["service.name"] == originalAttributes["service.name"]) + #expect(result[0].attributes["environment"] == originalAttributes["environment"]) + #expect(result[0].attributes[LDSemanticAttribute.ATTR_SAMPLING_RATIO] == samplingAttributes[LDSemanticAttribute.ATTR_SAMPLING_RATIO]) + } + + @Test("Given a set of logs, when having a mixed sampling results with and without attributes, then should handle them correctly") + func mixedLogsSamplingResults() { + let lowerBound = 1 + let upperBound = 4 + let logs = (lowerBound...upperBound).map { makeMockReadableLogRecord(body: .string("log\($0)")) } + + let samplingAttributes = [ + LDSemanticAttribute.ATTR_SAMPLING_RATIO: AttributeValue.int(50) + ] + + let sampler = ExportSampler.fake( + sampleLog: { logData in + if logData.body == logs[0].body { + return .init(sample: true, attributes: nil) + } else if logData.body == logs[2].body { + return .init(sample: true, attributes: samplingAttributes) + } else { + return .init(sample: false) + } + }, + isSamplingEnabled: true + ) + + let result = sampler.sampleLogs(items: logs) + + let sampled = [logs[0], logs[2]] + #expect(result.count == sampled.count) + + #expect(result[0].body == logs[0].body) + #expect(result[0].attributes == logs[0].attributes) + + #expect(result[1].body == logs[2].body) + #expect(result[1].attributes != logs[2].attributes) // result[1] has modified attributes with sampling attr. added + #expect(result[1].attributes[LDSemanticAttribute.ATTR_SAMPLING_RATIO] == AttributeValue.int(50)) + } +} diff --git a/Tests/SamplingLiveTests/SamplingSpansTests.swift b/Tests/SamplingLiveTests/SamplingSpansTests.swift new file mode 100644 index 00000000..2cc9bc0f --- /dev/null +++ b/Tests/SamplingLiveTests/SamplingSpansTests.swift @@ -0,0 +1,131 @@ +import Testing + +@testable import OpenTelemetrySdk +import OpenTelemetryApi + +import Common +import Sampling +@testable import SamplingLive + +struct SamplingSpansTests { + + @Test("Given a set of spans and a sampler, when sampling is disabled, then should return all spans") + func samplingDisabled() { + let spans = (1...3).map { makeMockSpanData(name: "span\($0)") } + let sampler = ExportSampler.fake(isSamplingEnabled: false) + + let result = sampler.sampleSpans(items: spans) + + #expect(result.count == spans.count) + for (r, sd) in zip(result, spans) { + #expect(r.spanId == sd.spanId) + #expect(r.name == sd.name) + } + } + + @Test("Given an empty array of spans and a sampler, when sampling is enabled, then should return handle empty array") + func samplingEnabledEmptySpansArray() { + let spans = [SpanData]() + let sampler = ExportSampler.fake(isSamplingEnabled: true) + + let result = sampler.sampleSpans(items: spans) + + #expect(result.isEmpty) + } + + @Test("Given an array of spans and a sampler, when sampling is enabled and no span is sampled, then should return empty result") + func samplingEnabledNoSamplingThenEmptyResult() { + let spans = (1...3).map { makeMockSpanData(name: "span\($0)") } + let sampler = ExportSampler.fake( + sampleSpan: { _ in .init(sample: false) }, + isSamplingEnabled: true + ) + + let result = sampler.sampleSpans(items: spans) + + #expect(result.isEmpty) + } + + @Test("Given an array of spans and a sample, when sampling is enabled and all are sampled without additional attributes, then should return all") + func samplingEnabledAllSampledWithoutAdditionalAttributes() { + let sampler = ExportSampler.fake( + sampleSpan: { _ in .init(sample: true, attributes: nil) }, + isSamplingEnabled: true + ) + let spans = (1...3).map { makeMockSpanData(name: "span\($0)") } + + let result = sampler.sampleSpans(items: spans) + + #expect(result.count == spans.count) + #expect(result.contains(spans)) + } + + @Test("Given an array of spans, when some spans are sampled, then should return a subset of the spans") + func samplingEnabledSomeSampled() { + let sampler = ExportSampler.fake( + sampleSpan: { span in + if span.name == "span2" { + return .init(sample: false) + } else { + return .init(sample: true) + } + }, + isSamplingEnabled: true + ) + let span1 = makeMockSpanData(name: "span1") + let span2 = makeMockSpanData(name: "span2") + let span3 = makeMockSpanData(name: "span1") + let spans = [span1, span2, span3] + + let result = sampler.sampleSpans(items: spans) + + + #expect(result.count == 2) + #expect(span1 == result[0]) + #expect(span3 == result[1]) + } + + @Test( + """ + Given an array of spans, + when some spans are sampled and sampling attributes are provided, + then should add the sampling attributes and preserve the original ones + """ + ) + func samplingEnabledWithSamplingAttributes() { + let originalAttributes = [ + "service.name": AttributeValue.string("api-service"), + "environment": AttributeValue.string("production") + ] + + let originalSpan = makeMockSpanData(name: "test-span", parentSpanId: .random(), attributes: originalAttributes) + + let spans = [originalSpan] + + let samplingAttributes = [ + LDSemanticAttribute.ATTR_SAMPLING_RATIO: AttributeValue.int(42) + ] + + let sampler = ExportSampler.fake( + sampleSpan: { span in + if span.name == "test-span" { + return .init(sample: true, attributes: samplingAttributes) + } else { + return .init(sample: true) + } + }, + isSamplingEnabled: true + ) + + let result = sampler.sampleSpans(items: spans) + + #expect(result.count == 1) + #expect(originalSpan.attributes != result[0].attributes) + #expect(result[0].spanId == originalSpan.spanId) + #expect(result[0].traceId == originalSpan.traceId) + #expect(result[0].parentSpanId == originalSpan.parentSpanId) + #expect(result[0].instrumentationScope == originalSpan.instrumentationScope) + } +} + +