@@ -30,32 +30,48 @@ public class AiSampler implements Sampler {
3030 private static final double SAMPLE_RATE_TO_DISABLE_INGESTION_SAMPLING = 99.99 ;
3131
3232 private final boolean ingestionSamplingEnabled ;
33- private final boolean localParentBased ;
33+ private final boolean sampleWhenLocalParentSampled ;
34+ private final boolean dropWhenLocalParentDropped ;
3435 private final SamplingPercentage requestSamplingPercentage ;
3536 // when localParentBased=false, then this applies to all dependencies, not only parentless
3637 private final SamplingPercentage parentlessDependencySamplingPercentage ;
3738 private final Cache <Double , SamplingResult > recordAndSampleWithSampleRateMap = Cache .bounded (100 );
3839
39- public AiSampler (
40+ public static AiSampler create (
4041 SamplingPercentage requestSamplingPercentage ,
4142 SamplingPercentage parentlessDependencySamplingPercentage ,
4243 boolean ingestionSamplingEnabled ) {
43- this (
44+ return new AiSampler (
4445 requestSamplingPercentage ,
4546 parentlessDependencySamplingPercentage ,
4647 ingestionSamplingEnabled ,
48+ true ,
4749 true );
4850 }
4951
50- public AiSampler (
52+ public static AiSampler createSamplingOverride (
53+ SamplingPercentage samplingPercentage ,
54+ boolean sampleWhenLocalParentSampled ,
55+ boolean dropWhenLocalParentDropped ) {
56+ return new AiSampler (
57+ samplingPercentage ,
58+ samplingPercentage ,
59+ false ,
60+ sampleWhenLocalParentSampled ,
61+ dropWhenLocalParentDropped );
62+ }
63+
64+ private AiSampler (
5165 SamplingPercentage requestSamplingPercentage ,
5266 SamplingPercentage parentlessDependencySamplingPercentage ,
5367 boolean ingestionSamplingEnabled ,
54- boolean localParentBased ) {
68+ boolean sampleWhenLocalParentSampled ,
69+ boolean dropWhenLocalParentDropped ) {
5570 this .requestSamplingPercentage = requestSamplingPercentage ;
5671 this .parentlessDependencySamplingPercentage = parentlessDependencySamplingPercentage ;
5772 this .ingestionSamplingEnabled = ingestionSamplingEnabled ;
58- this .localParentBased = localParentBased ;
73+ this .sampleWhenLocalParentSampled = sampleWhenLocalParentSampled ;
74+ this .dropWhenLocalParentDropped = dropWhenLocalParentDropped ;
5975 }
6076
6177 @ Override
@@ -66,8 +82,42 @@ public SamplingResult shouldSample(
6682 SpanKind spanKind ,
6783 Attributes attributes ,
6884 List <LinkData > parentLinks ) {
69- if (localParentBased ) {
70- SamplingResult samplingResult = useLocalParentDecisionIfPossible (parentContext );
85+
86+ Span parentSpan = Span .fromContext (parentContext );
87+ SpanContext parentSpanContext = parentSpan .getSpanContext ();
88+ Double parentSpanSampleRate = null ;
89+ if (parentSpan instanceof ReadableSpan ) {
90+ parentSpanSampleRate =
91+ ((ReadableSpan ) parentSpan ).getAttribute (AiSemanticAttributes .SAMPLE_RATE );
92+ }
93+
94+ return internalShouldSample (
95+ parentSpanContext , parentSpanSampleRate , traceId , spanKind , attributes );
96+ }
97+
98+ public SamplingResult shouldSampleLog (SpanContext spanContext , @ Nullable Double spanSampleRate ) {
99+ return internalShouldSample (
100+ spanContext ,
101+ spanSampleRate ,
102+ spanContext .getTraceId (),
103+ SpanKind .INTERNAL , // unused
104+ Attributes .empty ());
105+ }
106+
107+ public SamplingPercentage getParentlessDependencySamplingPercentage () {
108+ return parentlessDependencySamplingPercentage ;
109+ }
110+
111+ private SamplingResult internalShouldSample (
112+ SpanContext parentSpanContext ,
113+ @ Nullable Double parentSpanSampleRate ,
114+ String traceId ,
115+ SpanKind spanKind ,
116+ Attributes attributes ) {
117+
118+ if (sampleWhenLocalParentSampled || dropWhenLocalParentDropped ) {
119+ SamplingResult samplingResult =
120+ useLocalParentDecisionIfPossible (parentSpanContext , parentSpanSampleRate );
71121 if (samplingResult != null ) {
72122 return samplingResult ;
73123 }
@@ -78,7 +128,6 @@ public SamplingResult shouldSample(
78128 // optimization for fixed-rate sampling
79129 sp = requestSamplingPercentage .get ();
80130 } else {
81- SpanContext parentSpanContext = Span .fromContext (parentContext ).getSpanContext ();
82131 boolean isRequest = RequestChecker .isRequest (spanKind , parentSpanContext , attributes ::get );
83132 sp =
84133 isRequest
@@ -113,26 +162,26 @@ public SamplingResult shouldSample(
113162 return samplingResult ;
114163 }
115164
165+ @ SuppressWarnings ("SystemOut" )
116166 @ Nullable
117- private static SamplingResult useLocalParentDecisionIfPossible (Context parentContext ) {
167+ private SamplingResult useLocalParentDecisionIfPossible (
168+ SpanContext parentSpanContext , @ Nullable Double parentSpanSampleRate ) {
118169 // remote parent-based sampling messes up item counts since item count is not propagated in
119170 // tracestate (yet), but local parent-based sampling doesn't have this issue since we are
120171 // propagating item count locally
121- Span parentSpan = Span .fromContext (parentContext );
122- SpanContext parentSpanContext = parentSpan .getSpanContext ();
123172 if (!parentSpanContext .isValid () || parentSpanContext .isRemote ()) {
173+ System .out .println ("useLocalParentDecisionIfPossible1" );
124174 return null ;
125175 }
126176 if (!parentSpanContext .isSampled ()) {
127- return SamplingResult .drop ();
177+ System .out .println ("useLocalParentDecisionIfPossible2: " + dropWhenLocalParentDropped );
178+ return dropWhenLocalParentDropped ? SamplingResult .drop () : null ;
128179 }
129- if (parentSpan instanceof ReadableSpan ) {
130- Double parentSampleRate =
131- ((ReadableSpan ) parentSpan ).getAttribute (AiSemanticAttributes .SAMPLE_RATE );
132- if (parentSampleRate != null ) {
133- return new RecordAndSampleWithItemCount (parentSampleRate );
134- }
180+ if (sampleWhenLocalParentSampled && parentSpanSampleRate != null ) {
181+ System .out .println ("useLocalParentDecisionIfPossible3" );
182+ return new RecordAndSampleWithItemCount (parentSpanSampleRate );
135183 }
184+ System .out .println ("useLocalParentDecisionIfPossible5" );
136185 return null ;
137186 }
138187
0 commit comments