@@ -25,24 +25,20 @@ import datadog.trace.agent.tooling.TracerInstaller
2525import datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers
2626import datadog.trace.agent.tooling.bytebuddy.matcher.GlobalIgnores
2727import datadog.trace.api.Config
28- import datadog.trace.api.DDSpanId
2928import datadog.trace.api.IdGenerationStrategy
3029import datadog.trace.api.ProcessTags
3130import datadog.trace.api.StatsDClient
3231import datadog.trace.api.TraceConfig
3332import datadog.trace.api.config.GeneralConfig
3433import datadog.trace.api.config.TracerConfig
35- import datadog.trace.api.gateway.RequestContext
36- import datadog.trace.api.internal.TraceSegment
34+ import datadog.trace.api.datastreams.AgentDataStreamsMonitoring
3735import datadog.trace.api.sampling.SamplingRule
3836import datadog.trace.api.time.SystemTimeSource
39- import datadog.trace.api.datastreams.AgentDataStreamsMonitoring
4037import datadog.trace.bootstrap.ActiveSubsystems
41- import datadog.trace.bootstrap.CallDepthThreadLocalMap
4238import datadog.trace.bootstrap.InstrumentationErrors
4339import datadog.trace.bootstrap.debugger.DebuggerContext
4440import datadog.trace.bootstrap.instrumentation.api.AgentSpan
45- import datadog.trace.bootstrap.instrumentation.api.AgentTracer.TracerAPI
41+ import datadog.trace.bootstrap.instrumentation.api.AgentTracer
4642import datadog.trace.common.metrics.EventListener
4743import datadog.trace.common.metrics.Sink
4844import datadog.trace.common.writer.DDAgentWriter
@@ -70,7 +66,6 @@ import org.junit.jupiter.api.extension.ExtendWith
7066import org.slf4j.Logger
7167import org.slf4j.LoggerFactory
7268import org.spockframework.mock.MockUtil
73- import org.spockframework.mock.runtime.MockInvocation
7469import spock.lang.Shared
7570
7671import java.lang.instrument.ClassFileTransformer
@@ -152,7 +147,7 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
152147
153148 @SuppressWarnings (' PropertyName' )
154149 @Shared
155- TracerAPI TEST_TRACER
150+ AgentTracer. TracerAPI TEST_TRACER
156151
157152 @SuppressWarnings (' PropertyName' )
158153 @Shared
@@ -175,10 +170,6 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
175170 @Shared
176171 TestProfilingContextIntegration TEST_PROFILING_CONTEXT_INTEGRATION = new TestProfilingContextIntegration ()
177172
178- @SuppressWarnings (' PropertyName' )
179- @Shared
180- Set<DDSpanId > TEST_SPANS = Sets . newHashSet()
181-
182173 @SuppressWarnings (' PropertyName' )
183174 @Shared
184175 RecordingDatastreamsPayloadWriter TEST_DATA_STREAMS_WRITER
@@ -264,7 +255,7 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
264255 @Shared
265256 ConcurrentHashMap<DDSpan , List<Exception > > spanFinishLocations = new ConcurrentHashMap<> ()
266257 @Shared
267- ConcurrentHashMap<DDSpan , DDSpan > originalToSpySpan = new ConcurrentHashMap<> ()
258+ ConcurrentHashMap<DDSpan , DDSpan > originalToTrackingSpan = new ConcurrentHashMap<> ()
268259
269260 protected boolean enabledFinishTimingChecks () {
270261 false
@@ -399,67 +390,13 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
399390 boolean enabledFinishTimingChecks = this . enabledFinishTimingChecks()
400391 TEST_TRACER . startSpan(* _) >> {
401392 AgentSpan agentSpan = callRealMethod()
402- TEST_SPANS . add(agentSpan. spanId)
403393 if (! enabledFinishTimingChecks) {
404394 return agentSpan
405395 }
406396
407- // rest of closure if for checking duplicate finishes and tags set after finish
408- AgentSpan spiedAgentSpan = Spy (agentSpan)
409- originalToSpySpan[agentSpan] = spiedAgentSpan
410- def handleFinish = { MockInvocation mi ->
411- def depth = CallDepthThreadLocalMap . incrementCallDepth(DDSpan )
412- try {
413- if (depth > 0 ) {
414- return
415- }
416- List<Exception > locations
417- List<Exception > newLocations
418- do {
419- locations = spanFinishLocations. get(agentSpan)
420- newLocations = (locations ?: []) + new Exception ()
421- } while (! (locations == null ?
422- spanFinishLocations. putIfAbsent(agentSpan, newLocations) == null :
423- spanFinishLocations. replace(agentSpan, locations, newLocations)))
424- mi. callRealMethod()
425- } finally {
426- CallDepthThreadLocalMap . decrementCallDepth(DDSpan )
427- }
428- }
429- spiedAgentSpan. finish() >> {
430- handleFinish(delegate)
431- }
432- spiedAgentSpan. finish(_ as long ) >> {
433- handleFinish(delegate)
434- }
435- spiedAgentSpan. finishWithDuration() >> {
436- handleFinish(delegate)
437- }
438- spiedAgentSpan. finishWithEndToEnd() >> {
439- handleFinish(delegate)
440- }
441- spiedAgentSpan. getLocalRootSpan() >> {
442- DDSpan unwrappedSpan = callRealMethod()
443- originalToSpySpan. getOrDefault(unwrappedSpan, unwrappedSpan)
444- }
445- RequestContext requestContext = agentSpan. getRequestContext()
446- TraceSegment segment = requestContext. getTraceSegment()
447- RequestContext spiedReqCtx = Spy (requestContext)
448- TraceSegment checkedSegment = new PreconditionCheckTraceSegment (
449- check : {
450- -> if (useStrictTraceWrites() && spiedAgentSpan. localRootSpan. isFinished()) {
451- throw new AssertionError (" Interaction with TraceSegment after root span has already finished: $spiedAgentSpan " )
452- }},
453- delegate : segment
454- )
455- spiedAgentSpan. getRequestContext() >> {
456- spiedReqCtx
457- }
458- spiedReqCtx. getTraceSegment() >> {
459- checkedSegment
460- }
461-
462- spiedAgentSpan
397+ def trackingSpan = new TrackingSpanDecorator (agentSpan, spanFinishLocations, originalToTrackingSpan, useStrictTraceWrites())
398+ originalToTrackingSpan[agentSpan] = trackingSpan
399+ return trackingSpan
463400 }
464401
465402 // if a test enables the instrumentation it verifies,
@@ -498,7 +435,6 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
498435
499436 println " Starting test: ${ getSpecificationContext().getCurrentIteration().getName()} from ${ specificationContext.currentSpec.name} "
500437 TEST_TRACER . flush()
501- TEST_SPANS . clear()
502438
503439 if (isTestAgentEnabled()) {
504440 TEST_AGENT_WRITER . flush()
@@ -561,7 +497,7 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
561497 }
562498 } finally {
563499 spanFinishLocations. clear()
564- originalToSpySpan . clear()
500+ originalToTrackingSpan . clear()
565501 }
566502 assert InstrumentationErrors . errorCount == 0
567503 }
@@ -578,7 +514,7 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
578514 pw. write " Location $i :\n "
579515 def st = e. stackTrace
580516 int loc = st. findIndexOf {
581- it. className. startsWith(' datadog.trace.core.DDSpan$SpockMock$ ' ) &&
517+ it. className. startsWith(TrackingSpanDecorator . class . name ) &&
582518 it. methodName. startsWith(' finish' )
583519 }
584520 for (int j = loc == -1 ? 0 : loc; j < st. length; j++ ) {
@@ -590,77 +526,6 @@ abstract class InstrumentationSpecification extends DDSpecification implements A
590526 }
591527 }
592528
593- class PreconditionCheckTraceSegment implements TraceSegment {
594- Closure check
595- TraceSegment delegate
596-
597- @Override
598- void setTagTop (String key , Object value , boolean sanitize ) {
599- check()
600- delegate. setTagTop(key, value, sanitize)
601- }
602-
603- @Override
604- void setTagCurrent (String key , Object value , boolean sanitize ) {
605- check()
606- delegate. setTagCurrent(key, value, sanitize)
607- }
608-
609- @Override
610- Object getTagTop (String key , boolean sanitize ) {
611- check()
612- return delegate. getTagTop(key, sanitize)
613- }
614-
615- @Override
616- Object getTagCurrent (String key , boolean sanitize ) {
617- check()
618- return delegate. getTagCurrent(key, sanitize)
619- }
620-
621- @Override
622- void setDataTop (String key , Object value ) {
623- check()
624- delegate. setDataTop(key, value)
625- }
626-
627- @Override
628- Object getDataTop (String key ) {
629- check()
630- return delegate. getDataTop(key)
631- }
632-
633- @Override
634- void setMetaStructTop (String key , Object value ) {
635- check()
636- delegate. setMetaStructTop(key, value)
637- }
638-
639- @Override
640- void setMetaStructCurrent (String key , Object value ) {
641- check()
642- delegate. setMetaStructCurrent(key, value)
643- }
644-
645- @Override
646- void effectivelyBlocked () {
647- check()
648- delegate. effectivelyBlocked()
649- }
650-
651- @Override
652- void setDataCurrent (String key , Object value ) {
653- check()
654- delegate. setDataCurrent(key, value)
655- }
656-
657- @Override
658- Object getDataCurrent (String key ) {
659- check()
660- return delegate. getDataCurrent(key)
661- }
662- }
663-
664529 /* * Override to clean up things after the agent is removed */
665530 protected void cleanupAfterAgent () {}
666531
@@ -793,5 +658,3 @@ class AbortTransformationException extends RuntimeException {
793658 super (message)
794659 }
795660}
796-
797-
0 commit comments