diff --git a/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/build.gradle b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/build.gradle index 83caaeb1324..ba64ade17f7 100644 --- a/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/build.gradle +++ b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/build.gradle @@ -13,6 +13,7 @@ tasks.named("compileJava") { addTestSuiteForDir('latestDepTest', 'test') dependencies { + testImplementation project(':dd-java-agent:agent-llmobs') testImplementation project(':dd-java-agent:instrumentation:datadog:tracing:trace-annotation') testImplementation libs.guava testImplementation group: 'io.netty', name: 'netty-all', version: '4.1.9.Final' diff --git a/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/test/java/datadog/trace/instrumentation/java/concurrent/LLMObsCompletableFuturePropagationTest.java b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/test/java/datadog/trace/instrumentation/java/concurrent/LLMObsCompletableFuturePropagationTest.java new file mode 100644 index 00000000000..c8cecdcf104 --- /dev/null +++ b/dd-java-agent/instrumentation/java/java-concurrent/java-concurrent-1.8/src/test/java/datadog/trace/instrumentation/java/concurrent/LLMObsCompletableFuturePropagationTest.java @@ -0,0 +1,182 @@ +package datadog.trace.instrumentation.java.concurrent; + +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; + +import datadog.trace.agent.test.AbstractInstrumentationTest; +import datadog.trace.api.DDSpanTypes; +import datadog.trace.api.WellKnownTags; +import datadog.trace.api.llmobs.LLMObs; +import datadog.trace.api.llmobs.LLMObsSpan; +import datadog.trace.api.llmobs.LLMObsTags; +import datadog.trace.api.llmobs.noop.NoOpLLMObsEvalProcessor; +import datadog.trace.api.llmobs.noop.NoOpLLMObsSpanFactory; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.Tags; +import datadog.trace.core.DDSpan; +import datadog.trace.junit.utils.config.WithConfig; +import datadog.trace.llmobs.domain.DDLLMObsSpan; +import datadog.trace.llmobs.domain.LLMObsInternal; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; + +@WithConfig(key = "llmobs.enabled", value = "true") +class LLMObsCompletableFuturePropagationTest extends AbstractInstrumentationTest { + private static final String LLMOBS_TAG_PREFIX = "_ml_obs_tag."; + private static final String SESSION_ID = "session-from-workflow"; + + @BeforeAll + static void setupLlmObsApi() { + LLMObsInternal.setLLMObsSpanFactory(new TestLLMObsSpanFactory()); + LLMObsInternal.setLLMObsEvalProcessor(NoOpLLMObsEvalProcessor.INSTANCE); + } + + @AfterAll + static void resetLlmObsApi() { + LLMObsInternal.setLLMObsSpanFactory(NoOpLLMObsSpanFactory.INSTANCE); + LLMObsInternal.setLLMObsEvalProcessor(NoOpLLMObsEvalProcessor.INSTANCE); + } + + @Test + void activeLlmObsContextPropagatesThroughCompletableFutureExecutor() throws Exception { + ExecutorService executor = Executors.newSingleThreadExecutor(); + AgentSpan request = startSpan("test", "request"); + LLMObsSpan workflow = null; + TaskSpanResult task = null; + + try (AgentScope scope = activateSpan(request)) { + workflow = LLMObs.startWorkflowSpan("workflow", null, SESSION_ID); + String workflowTraceId = workflow.getTraceId().toString(); + + CompletableFuture future = + CompletableFuture.supplyAsync( + () -> { + LLMObsSpan child = LLMObs.startTaskSpan("task", null, null); + try { + assertEquals(workflowTraceId, child.getTraceId().toString()); + return new TaskSpanResult(child.getTraceId().toString(), child.getSpanId()); + } finally { + child.finish(); + } + }, + executor); + + task = future.get(5, TimeUnit.SECONDS); + } finally { + if (workflow != null) { + workflow.finish(); + } + request.finish(); + executor.shutdownNow(); + } + + writer.waitForTraces(1); + List trace = writer.get(0); + + DDSpan requestSpan = findSpan(trace, "request"); + DDSpan workflowSpan = findSpan(trace, "workflow"); + DDSpan taskSpan = findSpan(trace, "task"); + + assertEquals(0L, requestSpan.getParentId()); + assertEquals(task.traceId, taskSpan.getTraceId().toString()); + assertEquals(task.spanId, taskSpan.getSpanId()); + assertEquals(requestSpan.getTraceId(), workflowSpan.getTraceId()); + assertEquals(requestSpan.getTraceId(), taskSpan.getTraceId()); + assertNotEquals(0L, workflowSpan.getSpanId()); + assertEquals(DDSpanTypes.LLMOBS, workflowSpan.getType()); + assertEquals(DDSpanTypes.LLMOBS, taskSpan.getType()); + assertEquals(String.valueOf(workflowSpan.getSpanId()), taskSpan.getTag(parentIdTag())); + assertEquals(SESSION_ID, workflowSpan.getTag(sessionIdTag())); + assertEquals(SESSION_ID, taskSpan.getTag(sessionIdTag())); + } + + private static DDSpan findSpan(List trace, String operationName) { + return trace.stream() + .filter(span -> operationName.contentEquals(span.getOperationName())) + .findFirst() + .orElseThrow(() -> new AssertionError("Missing span: " + operationName)); + } + + private static String parentIdTag() { + return LLMOBS_TAG_PREFIX + "parent_id"; + } + + private static String sessionIdTag() { + return LLMOBS_TAG_PREFIX + LLMObsTags.SESSION_ID; + } + + private static final class TaskSpanResult { + private final String traceId; + private final long spanId; + + private TaskSpanResult(String traceId, long spanId) { + this.traceId = traceId; + this.spanId = spanId; + } + } + + private static final class TestLLMObsSpanFactory implements LLMObs.LLMObsSpanFactory { + private static final String ML_APP = "test-ml-app"; + private static final String SERVICE_NAME = "test-service"; + private static final WellKnownTags WELL_KNOWN_TAGS = + new WellKnownTags("runtime-id", "host", "env", SERVICE_NAME, "version", "java"); + + @Override + public LLMObsSpan startLLMSpan( + String spanName, String modelName, String modelProvider, String mlApp, String sessionId) { + DDLLMObsSpan span = start(Tags.LLMOBS_LLM_SPAN_KIND, spanName, mlApp, sessionId); + span.setTag(LLMObsTags.MODEL_NAME, modelName); + span.setTag(LLMObsTags.MODEL_PROVIDER, modelProvider); + return span; + } + + @Override + public LLMObsSpan startAgentSpan(String spanName, String mlApp, String sessionId) { + return start(Tags.LLMOBS_AGENT_SPAN_KIND, spanName, mlApp, sessionId); + } + + @Override + public LLMObsSpan startToolSpan(String spanName, String mlApp, String sessionId) { + return start(Tags.LLMOBS_TOOL_SPAN_KIND, spanName, mlApp, sessionId); + } + + @Override + public LLMObsSpan startTaskSpan(String spanName, String mlApp, String sessionId) { + return start(Tags.LLMOBS_TASK_SPAN_KIND, spanName, mlApp, sessionId); + } + + @Override + public LLMObsSpan startWorkflowSpan(String spanName, String mlApp, String sessionId) { + return start(Tags.LLMOBS_WORKFLOW_SPAN_KIND, spanName, mlApp, sessionId); + } + + @Override + public LLMObsSpan startEmbeddingSpan( + String spanName, String mlApp, String modelProvider, String modelName, String sessionId) { + DDLLMObsSpan span = start(Tags.LLMOBS_EMBEDDING_SPAN_KIND, spanName, mlApp, sessionId); + span.setTag(LLMObsTags.MODEL_NAME, modelName); + span.setTag(LLMObsTags.MODEL_PROVIDER, modelProvider); + return span; + } + + @Override + public LLMObsSpan startRetrievalSpan(String spanName, String mlApp, String sessionId) { + return start(Tags.LLMOBS_RETRIEVAL_SPAN_KIND, spanName, mlApp, sessionId); + } + + private static DDLLMObsSpan start( + String kind, String spanName, String mlApp, String sessionId) { + return new DDLLMObsSpan( + kind, spanName, mlApp == null ? ML_APP : mlApp, sessionId, SERVICE_NAME, WELL_KNOWN_TAGS); + } + } +}