Skip to content

Commit d4ec25a

Browse files
mariedmmariedm
committed
RUM-12642 Fix distributed trace sampling to be session consistent
1 parent da479bf commit d4ec25a

File tree

4 files changed

+98
-29
lines changed

4 files changed

+98
-29
lines changed

‎DatadogRUM/Sources/Instrumentation/Resources/URLSessionRUMResourcesHandler.swift‎

Lines changed: 39 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,8 @@ import Foundation
88
import DatadogInternal
99

1010
internal struct DistributedTracing {
11-
/// Tracing sampler used to sample traces generated by the SDK.
12-
let sampler: Sampler
11+
/// The sampling rate for tracing. Value between `0.0` and `100.0`, where `0.0` means NO trace will be sent and `100.0` means ALL traces will be sent.
12+
let samplingRate: SampleRate
1313
/// The distributed tracing ID generator.
1414
let traceIDGenerator: TraceIDGenerator
1515
let spanIDGenerator: SpanIDGenerator
@@ -19,13 +19,13 @@ internal struct DistributedTracing {
1919
let traceContextInjection: TraceContextInjection
2020

2121
init(
22-
sampler: Sampler,
22+
samplingRate: SampleRate,
2323
firstPartyHosts: FirstPartyHosts,
2424
traceIDGenerator: TraceIDGenerator,
2525
spanIDGenerator: SpanIDGenerator,
2626
traceContextInjection: TraceContextInjection
2727
) {
28-
self.sampler = sampler
28+
self.samplingRate = samplingRate
2929
self.traceIDGenerator = traceIDGenerator
3030
self.spanIDGenerator = spanIDGenerator
3131
self.firstPartyHosts = firstPartyHosts
@@ -184,11 +184,12 @@ extension DistributedTracing {
184184
payload: request.value(forHTTPHeaderField: GraphQLHeaders.payload)
185185
)
186186

187+
let sampler = sampler(sessionID: rumSessionId, traceID: traceID.idLo)
187188
let injectedSpanContext = TraceContext(
188189
traceID: traceID,
189190
spanID: spanID,
190191
parentSpanID: nil,
191-
sampleRate: sampler.samplingRate,
192+
sampleRate: samplingRate,
192193
isKept: sampler.sample(),
193194
rumSessionId: rumSessionId,
194195
userId: userId,
@@ -254,10 +255,42 @@ extension DistributedTracing {
254255
.init(
255256
traceID: $0.traceID,
256257
spanID: $0.spanID,
257-
samplingRate: Double(sampler.samplingRate.percentageProportion)
258+
samplingRate: Double(samplingRate.percentageProportion)
258259
)
259260
}
260261
}
262+
263+
/// Creates a sampler that makes consistent sampling decisions per session.
264+
///
265+
/// This method implements deterministic sampling based on the RUM session ID.
266+
/// When a session ID is available, it uses the last 48 bits of the session UUID as a `seed`
267+
/// to create a `DeterministicSampler`, ensuring all resources within the same session
268+
/// have the same sampling decision.
269+
///
270+
/// Fallback chain:
271+
/// 1. Session ID (preferred) → session-consistent sampling
272+
/// 2. Trace ID → trace-consistent sampling
273+
/// 3. Random sampler → fallback if neither is available
274+
///
275+
/// - Parameters:
276+
/// - sessionID: The RUM session ID
277+
/// - traceID: The trace ID as a fallback
278+
/// - Returns: A `Sampling` instance that will consistently sample based on the provided seed
279+
private func sampler(sessionID: String?, traceID: UInt64?) -> Sampling {
280+
if let sessionID,
281+
// For a UUID with value aaaaaaaa-bbbb-Mccc-Nddd-1234567890ab
282+
// we use as the base id the last part: 0x1234567890ab
283+
let seed = sessionID
284+
.split(separator: "-")
285+
.last
286+
.flatMap({ UInt64($0, radix: 16) }) {
287+
return DeterministicSampler(seed: seed, samplingRate: samplingRate)
288+
} else if let traceID {
289+
return DeterministicSampler(seed: traceID, samplingRate: samplingRate)
290+
}
291+
292+
return Sampler(samplingRate: samplingRate)
293+
}
261294
}
262295

263296
private extension HTTPURLResponse {

‎DatadogRUM/Sources/RUM+Internal.swift‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,15 +51,15 @@ extension InternalExtension where ExtendedType == RUM {
5151
switch configuration.firstPartyHostsTracing {
5252
case let .trace(hosts, sampleRate, traceContextInjection):
5353
distributedTracing = DistributedTracing(
54-
sampler: Sampler(samplingRate: rumConfiguration.debugSDK ? 100 : sampleRate),
54+
samplingRate: rumConfiguration.debugSDK ? 100 : sampleRate,
5555
firstPartyHosts: FirstPartyHosts(hosts),
5656
traceIDGenerator: rumConfiguration.traceIDGenerator,
5757
spanIDGenerator: rumConfiguration.spanIDGenerator,
5858
traceContextInjection: traceContextInjection
5959
)
6060
case let .traceWithHeaders(hostsWithHeaders, sampleRate, traceContextInjection):
6161
distributedTracing = DistributedTracing(
62-
sampler: Sampler(samplingRate: rumConfiguration.debugSDK ? 100 : sampleRate),
62+
samplingRate: rumConfiguration.debugSDK ? 100 : sampleRate,
6363
firstPartyHosts: FirstPartyHosts(hostsWithHeaders),
6464
traceIDGenerator: rumConfiguration.traceIDGenerator,
6565
spanIDGenerator: rumConfiguration.spanIDGenerator,

‎DatadogRUM/Tests/Instrumentation/Resources/URLSessionRUMResourcesHandlerTests.swift‎

Lines changed: 56 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
3232
// Given
3333
let handler = createHandler(
3434
distributedTracing: .init(
35-
sampler: .mockKeepAll(),
35+
samplingRate: .maxSampleRate,
3636
firstPartyHosts: .init(),
3737
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
3838
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 0),
@@ -72,7 +72,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
7272
// Given
7373
let handler = createHandler(
7474
distributedTracing: .init(
75-
sampler: .mockKeepAll(),
75+
samplingRate: .maxSampleRate,
7676
firstPartyHosts: .init(),
7777
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
7878
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 0),
@@ -108,7 +108,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
108108
// Given
109109
let handler = createHandler(
110110
distributedTracing: .init(
111-
sampler: .mockKeepAll(),
111+
samplingRate: .maxSampleRate,
112112
firstPartyHosts: .init(),
113113
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
114114
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 0),
@@ -147,7 +147,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
147147
// Given
148148
let handler = createHandler(
149149
distributedTracing: .init(
150-
sampler: .mockKeepAll(),
150+
samplingRate: .maxSampleRate,
151151
firstPartyHosts: .init(),
152152
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
153153
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 0),
@@ -183,7 +183,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
183183
/// Given
184184
let handler = createHandler(
185185
distributedTracing: .init(
186-
sampler: .mockRejectAll(),
186+
samplingRate: 0,
187187
firstPartyHosts: .init(),
188188
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
189189
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 0),
@@ -215,7 +215,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
215215
/// Given
216216
let handler = createHandler(
217217
distributedTracing: .init(
218-
sampler: .mockRejectAll(),
218+
samplingRate: 0,
219219
firstPartyHosts: .init(),
220220
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
221221
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 0),
@@ -245,7 +245,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
245245
/// Given
246246
let handler = createHandler(
247247
distributedTracing: .init(
248-
sampler: .mockRejectAll(),
248+
samplingRate: 0,
249249
firstPartyHosts: .init(),
250250
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
251251
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 0),
@@ -278,7 +278,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
278278
/// Given
279279
let handler = createHandler(
280280
distributedTracing: .init(
281-
sampler: .mockRejectAll(),
281+
samplingRate: 0,
282282
firstPartyHosts: .init(),
283283
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
284284
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 0),
@@ -308,7 +308,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
308308
// Given
309309
let handler = createHandler(
310310
distributedTracing: .init(
311-
sampler: .mockKeepAll(),
311+
samplingRate: .maxSampleRate,
312312
firstPartyHosts: .init(),
313313
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
314314
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 0),
@@ -368,7 +368,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
368368
// Given
369369
let handler = createHandler(
370370
distributedTracing: .init(
371-
sampler: .mockKeepAll(),
371+
samplingRate: .maxSampleRate,
372372
firstPartyHosts: .init(),
373373
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
374374
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 0),
@@ -427,7 +427,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
427427
// Given
428428
let handler = createHandler(
429429
distributedTracing: .init(
430-
sampler: .mockKeepAll(),
430+
samplingRate: .maxSampleRate,
431431
firstPartyHosts: .init(),
432432
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
433433
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 0),
@@ -511,11 +511,11 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
511511
commandSubscriber.onCommandReceived = { _ in receiveCommand.fulfill() }
512512

513513
// Given
514-
let traceSamplingRate: Double = .mockRandom(min: 0, max: 100)
514+
let traceSamplingRate: Float = .mockRandom(min: 0, max: 100)
515515

516516
let handler = createHandler(
517517
distributedTracing: .init(
518-
sampler: Sampler(samplingRate: Float(traceSamplingRate)),
518+
samplingRate: traceSamplingRate,
519519
firstPartyHosts: .init(),
520520
traceIDGenerator: DefaultTraceIDGenerator(),
521521
spanIDGenerator: DefaultSpanIDGenerator(),
@@ -544,7 +544,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
544544
let spanContext = try XCTUnwrap(resourceStartCommand.spanContext)
545545
XCTAssertEqual(spanContext.traceID, .init(idLo: 100))
546546
XCTAssertEqual(spanContext.spanID, .init(rawValue: 200))
547-
XCTAssertEqual(spanContext.samplingRate, traceSamplingRate / 100, accuracy: 0.01)
547+
XCTAssertEqual(spanContext.samplingRate, Double(traceSamplingRate / 100), accuracy: 0.01)
548548
}
549549

550550
func testGivenTaskInterceptionWithMetricsAndResponse_whenInterceptionCompletes_itStopsRUMResourceWithMetrics() throws {
@@ -700,7 +700,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
700700
func testGivenAllTracingHeaderTypes_itUsesTheSameIds() throws {
701701
let handler = createHandler(
702702
distributedTracing: .init(
703-
sampler: .mockKeepAll(),
703+
samplingRate: .maxSampleRate,
704704
firstPartyHosts: .init(),
705705
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
706706
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 0),
@@ -744,7 +744,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
744744
// Given
745745
let handler = createHandler(
746746
distributedTracing: .init(
747-
sampler: .mockKeepAll(),
747+
samplingRate: .maxSampleRate,
748748
firstPartyHosts: .init(),
749749
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
750750
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 0),
@@ -781,7 +781,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
781781
// Given
782782
let handler = createHandler(
783783
distributedTracing: .init(
784-
sampler: .mockKeepAll(),
784+
samplingRate: .maxSampleRate,
785785
firstPartyHosts: .init(),
786786
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
787787
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 0),
@@ -819,7 +819,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
819819
// Given
820820
let handler = createHandler(
821821
distributedTracing: .init(
822-
sampler: .mockKeepAll(),
822+
samplingRate: .maxSampleRate,
823823
firstPartyHosts: .init(),
824824
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
825825
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 0),
@@ -872,7 +872,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
872872
// Given
873873
let handler = createHandler(
874874
distributedTracing: .init(
875-
sampler: .mockKeepAll(),
875+
samplingRate: .maxSampleRate,
876876
firstPartyHosts: .init(),
877877
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
878878
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 0),
@@ -904,7 +904,7 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
904904
// Given
905905
let handler = createHandler(
906906
distributedTracing: .init(
907-
sampler: .mockKeepAll(),
907+
samplingRate: .maxSampleRate,
908908
firstPartyHosts: .init(),
909909
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
910910
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 0),
@@ -1072,6 +1072,42 @@ class URLSessionRUMResourcesHandlerTests: XCTestCase {
10721072
XCTAssertEqual(attributes[CrossPlatformAttributes.graphqlOperationType] as? String, "mutation")
10731073
}
10741074

1075+
// MARK: - Deterministic Sampling Tests
1076+
1077+
func testGivenSameSessionID_withDeterministicSampling_itProducesConsistentSamplingDecision() throws {
1078+
// Given
1079+
let sessionID = "12345678-1234-4abc-9def-123456789abc"
1080+
let handler = createHandler(
1081+
distributedTracing: .init(
1082+
samplingRate: 50,
1083+
firstPartyHosts: .init(),
1084+
traceIDGenerator: RelativeTracingUUIDGenerator(startingFrom: .init(idHi: 10, idLo: 100)),
1085+
spanIDGenerator: RelativeSpanIDGenerator(startingFrom: 100, advancingByCount: 1),
1086+
traceContextInjection: .all
1087+
)
1088+
)
1089+
1090+
// When - Make multiple requests with the same session ID
1091+
var samplingDecisions: [Bool] = []
1092+
for _ in 1...10 {
1093+
let (_, traceContext) = handler.modify(
1094+
request: .mockWith(url: "https://www.example.com"),
1095+
headerTypes: [.datadog],
1096+
networkContext: NetworkContext(
1097+
rumContext: .init(
1098+
applicationID: .mockRandom(),
1099+
sessionID: sessionID
1100+
)
1101+
)
1102+
)
1103+
samplingDecisions.append(traceContext?.isKept ?? false)
1104+
}
1105+
1106+
// Then - All sampling decisions should be the same
1107+
let firstDecision = samplingDecisions.first!
1108+
XCTAssertTrue(samplingDecisions.allSatisfy { $0 == firstDecision }, "All sampling decisions for the same session should be identical")
1109+
}
1110+
10751111
// MARK: - Helper Methods
10761112

10771113
private func extractBaggageKeyValuePairs(from header: String) -> [String: String] {

‎DatadogRUM/Tests/RUMTests.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -492,6 +492,6 @@ class RUMTests: XCTestCase {
492492
let feature = try XCTUnwrap(core.get(feature: NetworkInstrumentationFeature.self))
493493
let urlSessionHandler = try XCTUnwrap(feature.handlers.first as? URLSessionRUMResourcesHandler)
494494
XCTAssertEqual(urlSessionHandler.distributedTracing?.firstPartyHosts.hosts, hosts)
495-
XCTAssertEqual(urlSessionHandler.distributedTracing?.sampler.samplingRate, debugSDK ? 100.0 : sampleRate)
495+
XCTAssertEqual(urlSessionHandler.distributedTracing?.samplingRate, debugSDK ? 100.0 : sampleRate)
496496
}
497497
}

0 commit comments

Comments
 (0)