Skip to content

Commit ec93bc5

Browse files
nikita-tkachenko-datadognikita-tkachenko-datadog
committed
Flush on tracer close if CI Visibility enabled
1 parent 2500e38 commit ec93bc5

2 files changed

Lines changed: 61 additions & 4 deletions

File tree

‎dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -189,7 +189,8 @@ public static CoreTracerBuilder builder() {
189189
private final boolean disableSamplingMechanismValidation;
190190
private final TimeSource timeSource;
191191
private final ProfilingContextIntegration profilingContextIntegration;
192-
private boolean injectBaggageAsTags;
192+
private final boolean injectBaggageAsTags;
193+
private final boolean flushOnClose;
193194

194195
/**
195196
* JVM shutdown callback, keeping a reference to it to remove this if DDTracer gets destroyed
@@ -280,6 +281,7 @@ public static class CoreTracerBuilder {
280281
private boolean pollForTracerFlareRequests;
281282
private boolean pollForTracingConfiguration;
282283
private boolean injectBaggageAsTags;
284+
private boolean flushOnClose;
283285

284286
public CoreTracerBuilder serviceName(String serviceName) {
285287
this.serviceName = serviceName;
@@ -413,6 +415,11 @@ public CoreTracerBuilder injectBaggageAsTags(boolean injectBaggageAsTags) {
413415
return this;
414416
}
415417

418+
public CoreTracerBuilder flushOnClose(boolean flushOnClose) {
419+
this.flushOnClose = flushOnClose;
420+
return this;
421+
}
422+
416423
public CoreTracerBuilder() {
417424
// Apply the default values from config.
418425
config(Config.get());
@@ -443,6 +450,7 @@ public CoreTracerBuilder config(final Config config) {
443450
partialFlushMinSpans(config.getPartialFlushMinSpans());
444451
strictTraceWrites(config.isTraceStrictWritesEnabled());
445452
injectBaggageAsTags(config.isInjectBaggageAsTagsEnabled());
453+
flushOnClose(config.isCiVisibilityEnabled());
446454
return this;
447455
}
448456

@@ -473,7 +481,8 @@ public CoreTracer build() {
473481
profilingContextIntegration,
474482
pollForTracerFlareRequests,
475483
pollForTracingConfiguration,
476-
injectBaggageAsTags);
484+
injectBaggageAsTags,
485+
flushOnClose);
477486
}
478487
}
479488

@@ -504,7 +513,8 @@ private CoreTracer(
504513
final ProfilingContextIntegration profilingContextIntegration,
505514
final boolean pollForTracerFlareRequests,
506515
final boolean pollForTracingConfiguration,
507-
final boolean injectBaggageAsTags) {
516+
final boolean injectBaggageAsTags,
517+
final boolean flushOnClose) {
508518

509519
assert localRootSpanTags != null;
510520
assert defaultSpanTags != null;
@@ -712,6 +722,7 @@ private CoreTracer(
712722
propagationTagsFactory = PropagationTags.factory(config);
713723
this.profilingContextIntegration = profilingContextIntegration;
714724
this.injectBaggageAsTags = injectBaggageAsTags;
725+
this.flushOnClose = flushOnClose;
715726
this.allowInferredServices = SpanNaming.instance().namingSchema().allowInferredServices();
716727
if (profilingContextIntegration != ProfilingContextIntegration.NoOp.INSTANCE) {
717728
Map<String, Object> tmp = new HashMap<>(localRootSpanTags);
@@ -1087,6 +1098,9 @@ public CallbackProvider getUniversalCallbackProvider() {
10871098

10881099
@Override
10891100
public void close() {
1101+
if (flushOnClose) {
1102+
flush();
1103+
}
10901104
tracingConfigPoller.stop();
10911105
pendingTraceBuffer.close();
10921106
writer.close();

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

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,11 @@ import datadog.trace.api.Config
1111
import datadog.trace.api.StatsDClient
1212
import datadog.trace.api.remoteconfig.ServiceNameCollector
1313
import datadog.trace.api.sampling.PrioritySampling
14+
import datadog.trace.api.sampling.SamplingMechanism
1415
import datadog.trace.common.sampling.AllSampler
1516
import datadog.trace.common.sampling.PrioritySampler
1617
import datadog.trace.common.sampling.RateByServiceTraceSampler
1718
import datadog.trace.common.sampling.Sampler
18-
import datadog.trace.api.sampling.SamplingMechanism
1919
import datadog.trace.common.writer.DDAgentWriter
2020
import datadog.trace.common.writer.ListWriter
2121
import datadog.trace.common.writer.LoggingWriter
@@ -27,6 +27,7 @@ import okhttp3.OkHttpClient
2727
import spock.lang.Timeout
2828

2929
import java.nio.charset.StandardCharsets
30+
import java.util.concurrent.CopyOnWriteArrayList
3031

3132
import static datadog.trace.api.config.GeneralConfig.ENV
3233
import static datadog.trace.api.config.GeneralConfig.HEALTH_METRICS_ENABLED
@@ -670,6 +671,48 @@ class CoreTracerTest extends DDCoreSpecification {
670671
"service" | "env" | "SERVICE" | "ENV"
671672
"SERVICE" | "ENV" | "service" | "env"
672673
}
674+
675+
def "flushes on tracer close if configured to do so"() {
676+
given:
677+
def writer = new WriterWithExplicitFlush()
678+
def tracer = tracerBuilder().writer(writer).flushOnClose(true).build()
679+
680+
when:
681+
tracer.buildSpan('my_span').start().finish()
682+
tracer.close()
683+
684+
then:
685+
!writer.flushedTraces.empty
686+
}
687+
}
688+
689+
class WriterWithExplicitFlush implements datadog.trace.common.writer.Writer {
690+
final List<List<DDSpan>> writtenTraces = new CopyOnWriteArrayList<>()
691+
final List<List<DDSpan>> flushedTraces = new CopyOnWriteArrayList<>()
692+
693+
@Override
694+
void write(List<DDSpan> trace) {
695+
writtenTraces.add(trace)
696+
}
697+
698+
@Override
699+
void start() {
700+
}
701+
702+
@Override
703+
boolean flush() {
704+
flushedTraces.addAll(writtenTraces)
705+
writtenTraces.clear()
706+
return true
707+
}
708+
709+
@Override
710+
void close() {
711+
}
712+
713+
@Override
714+
void incrementDropCounts(int spanCount) {
715+
}
673716
}
674717

675718
class ControllableSampler implements Sampler, PrioritySampler {

0 commit comments

Comments
 (0)