Skip to content

Commit df1ae9f

Browse files
matsumo-andmatsumo-and
committed
Improve standalone sampler test coverage
- Replace individual SamplerTest cases with @unroll matrix that covers all product-flag combinations and asserts activeProducts contents directly; add package-private getActiveProducts() getter to StandaloneSampler to enable this - Add StandaloneSamplerTest case for spans with both LLMOBS and ASM bits set simultaneously, verifying LLMOBS wins via list ordering - Add TraceCollectorTest exercising the full span-finish → CoreTracer write path to verify that setSamplingPriorityIfNecessary skips the sampler when APM is disabled, a standalone product flag is set, and priority is already non-UNSET Signed-off-by: matsumo-and <yh134.toisanda@gmail.com>
1 parent 9645126 commit df1ae9f

5 files changed

Lines changed: 124 additions & 90 deletions

File tree

‎dd-trace-core/src/main/java/datadog/trace/common/sampling/StandaloneSampler.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,10 @@ public StandaloneSampler(final List<StandaloneProduct> activeProducts, final Clo
5353
this.lastSampleTime = new AtomicLong(clock.millis() - RATE_IN_MILLISECONDS);
5454
}
5555

56+
List<StandaloneProduct> getActiveProducts() {
57+
return activeProducts;
58+
}
59+
5660
@Override
5761
public <T extends CoreSpan<T>> boolean sample(final T span) {
5862
// Priority sampling sends all traces to the core agent, including traces marked dropped,

‎dd-trace-core/src/test/groovy/datadog/trace/common/sampling/SamplerTest.groovy‎

Lines changed: 21 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -2,95 +2,34 @@ package datadog.trace.common.sampling
22

33
import datadog.trace.api.Config
44
import datadog.trace.test.util.DDSpecification
5+
import spock.lang.Unroll
56

67
class SamplerTest extends DDSpecification {
78

8-
void "test that StandaloneSampler is selected when apm tracing disabled and appsec enabled"() {
9+
@Unroll
10+
def 'sampler selection: apmEnabled=#apmEnabled llmobs=#llmobs appsec=#appsec iast=#iast sca=#sca → #expectedType.simpleName with activeProducts=#expectedProducts'() {
911
setup:
10-
System.setProperty("dd.apm.tracing.enabled", "false")
11-
System.setProperty("dd.appsec.enabled", "true")
12-
Config config = new Config()
12+
if (!apmEnabled) injectSysConfig("dd.apm.tracing.enabled", "false")
13+
if (llmobs) injectSysConfig("dd.llmobs.enabled", "true")
14+
if (appsec) injectSysConfig("dd.appsec.enabled", "true")
15+
if (iast) injectSysConfig("dd.iast.enabled", "true")
16+
if (sca) injectSysConfig("dd.appsec.sca.enabled", "true")
1317

1418
when:
15-
Sampler sampler = Sampler.Builder.forConfig(config, null)
19+
Sampler sampler = Sampler.Builder.forConfig(Config.get(), null)
1620

1721
then:
18-
sampler instanceof StandaloneSampler
19-
}
20-
21-
void "test that StandaloneSampler is selected when apm tracing disabled and iast enabled"() {
22-
setup:
23-
System.setProperty("dd.apm.tracing.enabled", "false")
24-
System.setProperty("dd.iast.enabled", "true")
25-
Config config = new Config()
26-
27-
when:
28-
Sampler sampler = Sampler.Builder.forConfig(config, null)
29-
30-
then:
31-
sampler instanceof StandaloneSampler
32-
}
33-
34-
void "test that StandaloneSampler is selected when apm tracing disabled and sca enabled"() {
35-
setup:
36-
System.setProperty("dd.apm.tracing.enabled", "false")
37-
System.setProperty("dd.appsec.sca.enabled", "true")
38-
Config config = new Config()
39-
40-
when:
41-
Sampler sampler = Sampler.Builder.forConfig(config, null)
42-
43-
then:
44-
sampler instanceof StandaloneSampler
45-
}
46-
47-
void "test that StandaloneSampler is selected when apm tracing disabled and llmobs enabled"() {
48-
setup:
49-
System.setProperty("dd.apm.tracing.enabled", "false")
50-
System.setProperty("dd.llmobs.enabled", "true")
51-
Config config = new Config()
52-
53-
when:
54-
Sampler sampler = Sampler.Builder.forConfig(config, null)
55-
56-
then:
57-
sampler instanceof StandaloneSampler
58-
}
59-
60-
void "test that StandaloneSampler is selected when apm tracing disabled and both llmobs and asm enabled"() {
61-
setup:
62-
System.setProperty("dd.apm.tracing.enabled", "false")
63-
System.setProperty("dd.llmobs.enabled", "true")
64-
System.setProperty("dd.appsec.enabled", "true")
65-
Config config = new Config()
66-
67-
when:
68-
Sampler sampler = Sampler.Builder.forConfig(config, null)
69-
70-
then:
71-
sampler instanceof StandaloneSampler
72-
}
73-
74-
void "test that ForcePrioritySampler with SAMPLER_DROP is selected when apm tracing disabled and no other products enabled"() {
75-
setup:
76-
System.setProperty("dd.apm.tracing.enabled", "false")
77-
Config config = new Config()
78-
79-
when:
80-
Sampler sampler = Sampler.Builder.forConfig(config, null)
81-
82-
then:
83-
sampler instanceof ForcePrioritySampler
84-
}
85-
86-
void "test that StandaloneSampler is not selected when apm tracing enabled"() {
87-
setup:
88-
Config config = new Config()
89-
90-
when:
91-
Sampler sampler = Sampler.Builder.forConfig(config, null)
92-
93-
then:
94-
!(sampler instanceof StandaloneSampler)
22+
expectedType.isInstance(sampler)
23+
expectedProducts == null || (sampler as StandaloneSampler).getActiveProducts() == expectedProducts
24+
25+
where:
26+
apmEnabled | llmobs | appsec | iast | sca || expectedType | expectedProducts
27+
true | false | false | false | false || RateByServiceTraceSampler | null
28+
false | true | false | false | false || StandaloneSampler | [StandaloneProduct.LLMOBS]
29+
false | false | true | false | false || StandaloneSampler | [StandaloneProduct.ASM]
30+
false | false | false | true | false || StandaloneSampler | [StandaloneProduct.ASM]
31+
false | false | false | false | true || StandaloneSampler | [StandaloneProduct.ASM]
32+
false | true | true | false | false || StandaloneSampler | [StandaloneProduct.LLMOBS, StandaloneProduct.ASM]
33+
false | false | false | false | false || ForcePrioritySampler | null
9534
}
9635
}

‎dd-trace-core/src/test/groovy/datadog/trace/common/sampling/StandaloneSamplerTest.groovy‎

Lines changed: 44 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,9 @@ class StandaloneSamplerTest extends DDCoreSpecification {
7474
setup:
7575
def current = new AtomicLong(System.currentTimeMillis())
7676
final Clock clock = Mock(Clock) {
77-
millis() >> { current.get() }
77+
millis() >> {
78+
current.get()
79+
}
7880
}
7981
def sampler = new StandaloneSampler([StandaloneProduct.ASM], clock)
8082
def tracer = tracerBuilder().writer(writer).sampler(sampler).build()
@@ -84,7 +86,9 @@ class StandaloneSamplerTest extends DDCoreSpecification {
8486
sampler.setSamplingPriority(span1)
8587

8688
then:
87-
1 * clock.millis() >> { current.updateAndGet(v -> v + 1000) }
89+
1 * clock.millis() >> {
90+
current.updateAndGet(v -> v + 1000)
91+
}
8892
span1.getSamplingPriority() == PrioritySampling.SAMPLER_KEEP
8993
span1.context().getPropagationTags().createTagMap().get("_dd.p.dm") == "-5"
9094

@@ -93,7 +97,9 @@ class StandaloneSamplerTest extends DDCoreSpecification {
9397
sampler.setSamplingPriority(span2)
9498

9599
then:
96-
1 * clock.millis() >> { current.updateAndGet(v -> v + 1000) }
100+
1 * clock.millis() >> {
101+
current.updateAndGet(v -> v + 1000)
102+
}
97103
span2.getSamplingPriority() == PrioritySampling.SAMPLER_DROP
98104
!span2.context().getPropagationTags().createTagMap().containsKey("_dd.p.dm")
99105

@@ -102,7 +108,9 @@ class StandaloneSamplerTest extends DDCoreSpecification {
102108
sampler.setSamplingPriority(span3)
103109

104110
then:
105-
1 * clock.millis() >> { current.updateAndGet(v -> v + 60000) }
111+
1 * clock.millis() >> {
112+
current.updateAndGet(v -> v + 60000)
113+
}
106114
span3.getSamplingPriority() == PrioritySampling.SAMPLER_KEEP
107115
span3.context().getPropagationTags().createTagMap().get("_dd.p.dm") == "-5"
108116

@@ -150,11 +158,33 @@ class StandaloneSamplerTest extends DDCoreSpecification {
150158
tracer.close()
151159
}
152160

161+
void "LLMOBS+ASM: span with both LLMOBS and ASM bits set is kept with DEFAULT mechanism (LLMOBS wins)"() {
162+
setup:
163+
def sampler = new StandaloneSampler([StandaloneProduct.LLMOBS, StandaloneProduct.ASM], Clock.systemUTC())
164+
def tracer = tracerBuilder().writer(writer).sampler(sampler).build()
165+
166+
when:
167+
def span = tracer.buildSpan("testInstrumentation", "waf-llm-request").start()
168+
def scope = tracer.activateSpan(span)
169+
tracer.getTraceSegment().setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.LLMOBS | ProductTraceSource.ASM)
170+
sampler.setSamplingPriority(span)
171+
scope.close()
172+
173+
then:
174+
span.getSamplingPriority() == PrioritySampling.SAMPLER_KEEP
175+
span.context().getPropagationTags().createTagMap().get("_dd.p.dm") == "-0"
176+
177+
cleanup:
178+
tracer.close()
179+
}
180+
153181
void "LLMOBS+ASM: APM-only spans are rate-limited with APPSEC mechanism"() {
154182
setup:
155183
def current = new AtomicLong(System.currentTimeMillis())
156184
final Clock clock = Mock(Clock) {
157-
millis() >> { current.get() }
185+
millis() >> {
186+
current.get()
187+
}
158188
}
159189
def sampler = new StandaloneSampler([StandaloneProduct.LLMOBS, StandaloneProduct.ASM], clock)
160190
def tracer = tracerBuilder().writer(writer).sampler(sampler).build()
@@ -164,23 +194,29 @@ class StandaloneSamplerTest extends DDCoreSpecification {
164194
sampler.setSamplingPriority(span1)
165195
166196
then:
167-
1 * clock.millis() >> { current.updateAndGet(v -> v + 1000) }
197+
1 * clock.millis() >> {
198+
current.updateAndGet(v -> v + 1000)
199+
}
168200
span1.getSamplingPriority() == PrioritySampling.SAMPLER_KEEP
169201
170202
when: "second APM span within the same minute"
171203
def span2 = tracer.buildSpan("testInstrumentation", "apm-request2").start()
172204
sampler.setSamplingPriority(span2)
173205
174206
then:
175-
1 * clock.millis() >> { current.updateAndGet(v -> v + 1000) }
207+
1 * clock.millis() >> {
208+
current.updateAndGet(v -> v + 1000)
209+
}
176210
span2.getSamplingPriority() == PrioritySampling.SAMPLER_DROP
177211
178212
when: "third APM span after 1 minute"
179213
def span3 = tracer.buildSpan("testInstrumentation", "apm-request3").start()
180214
sampler.setSamplingPriority(span3)
181215
182216
then:
183-
1 * clock.millis() >> { current.updateAndGet(v -> v + 60000) }
217+
1 * clock.millis() >> {
218+
current.updateAndGet(v -> v + 60000)
219+
}
184220
span3.getSamplingPriority() == PrioritySampling.SAMPLER_KEEP
185221
186222
cleanup:

‎dd-trace-core/src/test/groovy/datadog/trace/core/TraceCollectorTest.groovy‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
package datadog.trace.core
2+
3+
import datadog.trace.api.ProductTraceSource
4+
import datadog.trace.api.sampling.PrioritySampling
5+
import datadog.trace.api.sampling.SamplingMechanism
6+
import datadog.trace.bootstrap.instrumentation.api.Tags
7+
import datadog.trace.common.sampling.StandaloneProduct
8+
import datadog.trace.common.sampling.StandaloneSampler
9+
import datadog.trace.common.writer.ListWriter
10+
import datadog.trace.core.test.DDCoreSpecification
11+
12+
import java.time.Clock
13+
14+
class TraceCollectorTest extends DDCoreSpecification {
15+
16+
def writer = new ListWriter()
17+
18+
void "setSamplingPriorityIfNecessary: sampler is skipped when APM disabled, standalone product flag set, and priority already non-UNSET"() {
19+
setup:
20+
injectSysConfig("dd.apm.tracing.enabled", "false")
21+
def sampler = Spy(StandaloneSampler, constructorArgs: [[StandaloneProduct.LLMOBS], Clock.systemUTC()])
22+
def tracer = tracerBuilder().writer(writer).sampler(sampler).build()
23+
24+
when:
25+
def span = tracer.buildSpan("testInstrumentation", "llm-request").start()
26+
def scope = tracer.activateSpan(span)
27+
tracer.getTraceSegment().setTagTop(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.LLMOBS)
28+
span.setSamplingPriority(PrioritySampling.USER_KEEP, SamplingMechanism.MANUAL)
29+
scope.close()
30+
span.finish()
31+
writer.waitForTraces(1)
32+
33+
then:
34+
0 * sampler.setSamplingPriority(_)
35+
span.getSamplingPriority() == PrioritySampling.USER_KEEP
36+
37+
cleanup:
38+
tracer.close()
39+
}
40+
}

‎internal-api/src/test/groovy/datadog/trace/api/ProductTraceSourceTest.groovy‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,21 @@ class ProductTraceSourceTest extends DDSpecification {
3030
ProductTraceSource.DSM | ProductTraceSource.ASM | false
3131
}
3232

33+
void 'test isAnyStandaloneProductMarked'(){
34+
when:
35+
final result = ProductTraceSource.isAnyStandaloneProductMarked(value)
36+
37+
then:
38+
result == expected
39+
40+
where:
41+
value | expected
42+
ProductTraceSource.UNSET | false
43+
ProductTraceSource.APM | false
44+
ProductTraceSource.ASM | true
45+
ProductTraceSource.LLMOBS | true
46+
}
47+
3348
void 'test getBitfieldHex'(){
3449
when:
3550
final result = ProductTraceSource.getBitfieldHex(value)

0 commit comments

Comments
 (0)