Skip to content

Instantly share code, notes, and snippets.

@monosoul
Last active February 8, 2023 10:46
Show Gist options
  • Save monosoul/048e0456851fdc4bafe708d1807c6a87 to your computer and use it in GitHub Desktop.
Save monosoul/048e0456851fdc4bafe708d1807c6a87 to your computer and use it in GitHub Desktop.
Use DD agent 1.7.0
package datadog.trace.core.scopemanager
import datadog.trace.bootstrap.instrumentation.api.ScopeSource.MANUAL
import datadog.trace.core.scopeManager
import datadog.trace.core.scopemanager.ContinuableScopeManager.ScopeStack
import kotlinx.coroutines.ThreadContextElement
import org.slf4j.LoggerFactory
import kotlin.coroutines.CoroutineContext
import datadog.trace.api.GlobalTracer as DDGlobalTracer
private class ScopeStackCoroutineContext : ThreadContextElement<ScopeStack?> {
private val scopeManager = DDGlobalTracer.get().scopeManager
private val span = scopeManager?.activeSpan()
private val scopeStack = ScopeStack()
private val logger = LoggerFactory.getLogger(ScopeStackCoroutineContext::class.java)
init {
if (scopeManager == null) {
logger.warn("Unexpected scope manager or tracer implementation, unable to propagate tracing scope")
}
}
override val key: CoroutineContext.Key<*>
get() = Key
override fun updateThreadContext(context: CoroutineContext): ScopeStack? = scopeManager?.let {
val oldScopeStack = it.tlsScopeStack.get()
it.tlsScopeStack.set(scopeStack)
if (scopeStack.depth() == 0 && span != null) {
scopeManager.activate(span, MANUAL)
}
oldScopeStack
}
override fun restoreThreadContext(context: CoroutineContext, oldState: ScopeStack?) {
scopeManager?.tlsScopeStack?.set(oldState)
}
companion object Key : CoroutineContext.Key<ScopeStackCoroutineContext>
}
fun scopeStackCoroutineContext(): ThreadContextElement<*> = ScopeStackCoroutineContext()
package datadog.trace.core
import datadog.trace.api.Tracer
import datadog.trace.core.scopemanager.ContinuableScopeManager
val Tracer.scopeManager: ContinuableScopeManager?
get() = (this as? CoreTracer)?.scopeManager as? ContinuableScopeManager
import datadog.opentracing.DDTracer
import datadog.trace.core.CoreTracer.CoreTracerBuilder
import datadog.trace.core.scopemanager.scopeStackCoroutineContext
import io.opentracing.util.GlobalTracer
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.time.delay
import org.junit.jupiter.api.BeforeEach
import org.junit.jupiter.api.Test
import org.slf4j.LoggerFactory
import java.time.Duration
internal class TracingExtensionsKtTest {
private val logger = LoggerFactory.getLogger("testLogger")
@BeforeEach
internal fun setUp() {
CoreTracerBuilder()
.build()
.also(datadog.trace.api.GlobalTracer::registerIfAbsent)
.let { DDTracer(it) }
.also(GlobalTracer::registerIfAbsent)
}
@Test
internal fun test() {
val topLevelSpan = GlobalTracer.get().buildSpan("top-level-span").start()
GlobalTracer.get().activateSpan(topLevelSpan).use {
runBlocking(scopeStackCoroutineContext()) {
val uninterruptedSpan = GlobalTracer.get().buildSpan("uninterrupted-span").start()
GlobalTracer.get().activateSpan(uninterruptedSpan).use {
logger.info("Hello from uninterrupted span!")
uninterruptedSpan.finish()
}
val interruptedSpan = GlobalTracer.get().buildSpan("interrupted-span").start()
GlobalTracer.get().activateSpan(interruptedSpan).use {
logger.info("Hello from interrupted span!")
delay(Duration.ofMillis(100))
interruptedSpan.finish()
}
}
topLevelSpan.finish()
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment