Instrument spring-batch: chunk
MODIFY
instrumentation/spring/spring-batch-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/batch/SpringBatchInstrumentationConfig.java
instrumentation/spring/spring-batch-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/batch/SpringBatchInstrumentationConfig.java
@@ -24,5 +24,10 @@ public final class SpringBatchInstrumentationConfig {
.isInstrumentationPropertyEnabled(instrumentationNames(), type + ".enabled", true);
}
+ public static boolean shouldCreateRootSpanForChunk() {
+ return Config.get()
+ .isInstrumentationPropertyEnabled(instrumentationNames(), "chunk.root-span", false);
+ }
+
private SpringBatchInstrumentationConfig() {}
}

MODIFY
instrumentation/spring/spring-batch-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/batch/SpringBatchInstrumentationModule.java
instrumentation/spring/spring-batch-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/batch/SpringBatchInstrumentationModule.java
@@ -10,6 +10,7 @@ import static io.opentelemetry.javaagent.instrumentation.spring.batch.SpringBatc
import static io.opentelemetry.javaagent.tooling.ClassLoaderMatcher.hasClassesNamed;
import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.instrumentation.spring.batch.chunk.StepBuilderInstrumentation;
import io.opentelemetry.javaagent.instrumentation.spring.batch.job.JobBuilderHelperInstrumentation;
import io.opentelemetry.javaagent.instrumentation.spring.batch.job.JobFactoryBeanInstrumentation;
import io.opentelemetry.javaagent.instrumentation.spring.batch.job.JobParserJobFactoryBeanInstrumentation;
@@ -40,6 +41,7 @@ public class SpringBatchInstrumentationModule extends InstrumentationModule {
"io.opentelemetry.javaagent.instrumentation.spring.batch.ContextAndScope";
context.put("org.springframework.batch.core.JobExecution", contextAndScope);
context.put("org.springframework.batch.core.StepExecution", contextAndScope);
+ context.put("org.springframework.batch.core.scope.context.ChunkContext", contextAndScope);
return context;
}
@@ -54,6 +56,9 @@ public class SpringBatchInstrumentationModule extends InstrumentationModule {
if (isTracingEnabled("step")) {
instrumentations.add(new StepBuilderHelperInstrumentation());
}
+ if (isTracingEnabled("chunk")) {
+ instrumentations.add(new StepBuilderInstrumentation());
+ }
return instrumentations;
}

ADD
instrumentation/spring/spring-batch-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/batch/chunk/ChunkExecutionTracer.java
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.spring.batch.chunk;
+
+import static io.opentelemetry.api.trace.Span.Kind.INTERNAL;
+import static io.opentelemetry.javaagent.instrumentation.spring.batch.SpringBatchInstrumentationConfig.shouldCreateRootSpanForChunk;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanBuilder;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.instrumentation.api.tracer.BaseTracer;
+import org.springframework.batch.core.scope.context.ChunkContext;
+
+public class ChunkExecutionTracer extends BaseTracer {
+ private static final ChunkExecutionTracer TRACER = new ChunkExecutionTracer();
+
+ public static ChunkExecutionTracer tracer() {
+ return TRACER;
+ }
+
+ public Context startSpan(ChunkContext chunkContext) {
+ String jobName = chunkContext.getStepContext().getJobName();
+ String stepName = chunkContext.getStepContext().getStepName();
+ SpanBuilder spanBuilder =
+ tracer.spanBuilder("BatchJob " + jobName + "." + stepName + ".Chunk").setSpanKind(INTERNAL);
+ if (shouldCreateRootSpanForChunk()) {
+ linkParentSpan(spanBuilder);
+ }
+ Span span = spanBuilder.startSpan();
+ return Context.current().with(span);
+ }
+
+ private void linkParentSpan(SpanBuilder spanBuilder) {
+ spanBuilder.setNoParent();
+
+ SpanContext parentSpanContext = Span.current().getSpanContext();
+ if (parentSpanContext.isValid()) {
+ spanBuilder.addLink(parentSpanContext);
+ }
+ }
+
+ public void end(Context context) {
+ end(Span.fromContext(context));
+ }
+
+ public void endExceptionally(Context context, Throwable throwable) {
+ endExceptionally(Span.fromContext(context), throwable);
+ }
+
+ @Override
+ protected String getInstrumentationName() {
+ return "io.opentelemetry.javaagent.spring-batch";
+ }
+}

ADD
instrumentation/spring/spring-batch-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/batch/chunk/StepBuilderInstrumentation.java
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.spring.batch.chunk;
+
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
+import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
+import io.opentelemetry.javaagent.instrumentation.spring.batch.ContextAndScope;
+import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import org.springframework.batch.core.scope.context.ChunkContext;
+import org.springframework.batch.core.step.builder.AbstractTaskletStepBuilder;
+
+public class StepBuilderInstrumentation implements TypeInstrumentation {
+ @Override
+ public ElementMatcher<? super TypeDescription> typeMatcher() {
+ // Spring Batch Java DSL and XML config
+ return named("org.springframework.batch.core.step.builder.AbstractTaskletStepBuilder")
+ // JSR-352 XML config
+ .or(named("org.springframework.batch.core.jsr.step.builder.JsrSimpleStepBuilder"))
+ .or(named("org.springframework.batch.core.jsr.step.builder.JsrBatchletStepBuilder"));
+ }
+
+ @Override
+ public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+ return singletonMap(
+ named("build").and(isPublic()).and(takesArguments(0)),
+ this.getClass().getName() + "$BuildAdvice");
+ }
+
+ public static class BuildAdvice {
+ @Advice.OnMethodEnter(suppress = Throwable.class)
+ public static void onEnter(@Advice.This AbstractTaskletStepBuilder<?> stepBuilder) {
+ ContextStore<ChunkContext, ContextAndScope> chunkExecutionContextStore =
+ InstrumentationContext.get(ChunkContext.class, ContextAndScope.class);
+ stepBuilder.listener(new TracingChunkExecutionListener(chunkExecutionContextStore));
+ }
+ }
+}

ADD
instrumentation/spring/spring-batch-3.0/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/spring/batch/chunk/TracingChunkExecutionListener.java
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.spring.batch.chunk;
+
+import static io.opentelemetry.javaagent.instrumentation.spring.batch.chunk.ChunkExecutionTracer.tracer;
+
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
+import io.opentelemetry.javaagent.instrumentation.spring.batch.ContextAndScope;
+import org.springframework.batch.core.ChunkListener;
+import org.springframework.batch.core.scope.context.ChunkContext;
+import org.springframework.core.Ordered;
+
+public class TracingChunkExecutionListener implements ChunkListener, Ordered {
+ private final ContextStore<ChunkContext, ContextAndScope> executionContextStore;
+
+ public TracingChunkExecutionListener(
+ ContextStore<ChunkContext, ContextAndScope> executionContextStore) {
+ this.executionContextStore = executionContextStore;
+ }
+
+ @Override
+ public void beforeChunk(ChunkContext chunkContext) {
+ Context context = tracer().startSpan(chunkContext);
+ // beforeJob & afterJob always execute on the same thread
+ Scope scope = context.makeCurrent();
+ executionContextStore.put(chunkContext, new ContextAndScope(context, scope));
+ }
+
+ @Override
+ public void afterChunk(ChunkContext chunkContext) {
+ ContextAndScope contextAndScope = executionContextStore.get(chunkContext);
+ if (contextAndScope != null) {
+ executionContextStore.put(chunkContext, null);
+ contextAndScope.closeScope();
+ tracer().end(contextAndScope.getContext());
+ }
+ }
+
+ @Override
+ public void afterChunkError(ChunkContext chunkContext) {
+ ContextAndScope contextAndScope = executionContextStore.get(chunkContext);
+ if (contextAndScope != null) {
+ executionContextStore.put(chunkContext, null);
+ contextAndScope.closeScope();
+ Throwable throwable =
+ (Throwable) chunkContext.getAttribute(ChunkListener.ROLLBACK_EXCEPTION_KEY);
+ tracer().endExceptionally(contextAndScope.getContext(), throwable);
+ }
+ }
+
+ @Override
+ public int getOrder() {
+ return HIGHEST_PRECEDENCE;
+ }
+
+ // equals() and hashCode() methods guarantee that only one instance of
+ // TracingJobExecutionListener will be present in an ordered set of listeners
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ return o != null && getClass() == o.getClass();
+ }
+
+ @Override
+ public int hashCode() {
+ return getClass().hashCode();
+ }
+}

ADD
instrumentation/spring/spring-batch-3.0/javaagent/src/test/groovy/ChunkRootSpanTest.groovy
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import static io.opentelemetry.api.trace.Span.Kind.INTERNAL
+import static io.opentelemetry.instrumentation.test.utils.ConfigUtils.setConfig
+import static io.opentelemetry.instrumentation.test.utils.ConfigUtils.updateConfigAndResetInstrumentation
+import static java.util.Collections.emptyMap
+
+import io.opentelemetry.instrumentation.test.AgentTestRunner
+import org.springframework.batch.core.JobParameter
+import org.springframework.context.ConfigurableApplicationContext
+import org.springframework.context.annotation.AnnotationConfigApplicationContext
+import org.springframework.context.support.ClassPathXmlApplicationContext
+
+abstract class ChunkRootSpanTest extends AgentTestRunner {
+
+ abstract runJob(String jobName, Map<String, JobParameter> params = emptyMap())
+
+ def "should create separate traces for each chunk"() {
+ when:
+ runJob("itemsAndTaskletJob")
+
+ then:
+ assertTraces(5) {
+ def itemStepSpan = null
+ def taskletStepSpan = null
+
+ trace(0, 3) {
+ itemStepSpan = span(1)
+ taskletStepSpan = span(2)
+
+ span(0) {
+ name "BatchJob itemsAndTaskletJob"
+ kind INTERNAL
+ }
+ span(1) {
+ name "BatchJob itemsAndTaskletJob.itemStep"
+ kind INTERNAL
+ childOf span(0)
+ }
+ span(2) {
+ name "BatchJob itemsAndTaskletJob.taskletStep"
+ kind INTERNAL
+ childOf span(0)
+ }
+ }
+ trace(1, 1) {
+ span(0) {
+ name "BatchJob itemsAndTaskletJob.itemStep.Chunk"
+ kind INTERNAL
+ hasLink itemStepSpan
+ }
+ }
+ trace(2, 1) {
+ span(0) {
+ name "BatchJob itemsAndTaskletJob.itemStep.Chunk"
+ kind INTERNAL
+ hasLink itemStepSpan
+ }
+ }
+ trace(3, 1) {
+ span(0) {
+ name "BatchJob itemsAndTaskletJob.itemStep.Chunk"
+ kind INTERNAL
+ hasLink itemStepSpan
+ }
+ }
+ trace(4, 1) {
+ span(0) {
+ name "BatchJob itemsAndTaskletJob.taskletStep.Chunk"
+ kind INTERNAL
+ hasLink taskletStepSpan
+ }
+ }
+ }
+ }
+}
+
+class JavaConfigChunkRootSpanTest extends ChunkRootSpanTest implements ApplicationConfigTrait {
+ static final PREVIOUS_CONFIG = updateConfigAndResetInstrumentation({
+ it.setProperty("otel.instrumentation.spring-batch.enabled", "true")
+ it.setProperty("otel.instrumentation.spring-batch.chunk.root-span", "true")
+ })
+
+ @Override
+ ConfigurableApplicationContext createApplicationContext() {
+ new AnnotationConfigApplicationContext(SpringBatchApplication)
+ }
+
+ @Override
+ def additionalCleanup() {
+ setConfig(PREVIOUS_CONFIG)
+ }
+}
+
+class XmlConfigChunkRootSpanTest extends ChunkRootSpanTest implements ApplicationConfigTrait {
+ static final PREVIOUS_CONFIG = updateConfigAndResetInstrumentation({
+ it.setProperty("otel.instrumentation.spring-batch.enabled", "true")
+ it.setProperty("otel.instrumentation.spring-batch.chunk.root-span", "true")
+ })
+
+ @Override
+ ConfigurableApplicationContext createApplicationContext() {
+ new ClassPathXmlApplicationContext("spring-batch.xml")
+ }
+
+ @Override
+ def additionalCleanup() {
+ setConfig(PREVIOUS_CONFIG)
+ }
+}
+
+class JsrConfigChunkRootSpanTest extends ChunkRootSpanTest implements JavaxBatchConfigTrait {
+ static final PREVIOUS_CONFIG = updateConfigAndResetInstrumentation({
+ it.setProperty("otel.instrumentation.spring-batch.enabled", "true")
+ it.setProperty("otel.instrumentation.spring-batch.chunk.root-span", "true")
+ })
+
+ @Override
+ def additionalCleanup() {
+ setConfig(PREVIOUS_CONFIG)
+ }
+}

MODIFY
instrumentation/spring/spring-batch-3.0/javaagent/src/test/groovy/SpringBatchTest.groovy
instrumentation/spring/spring-batch-3.0/javaagent/src/test/groovy/SpringBatchTest.groovy
import static io.opentelemetry.api.trace.Span.Kind.INTERNAL
import static io.opentelemetry.instrumentation.test.utils.ConfigUtils.setConfig
-import static io.opentelemetry.instrumentation.test.utils.ConfigUtils.updateConfig
+import static io.opentelemetry.instrumentation.test.utils.ConfigUtils.updateConfigAndResetInstrumentation
import static java.util.Collections.emptyMap
import io.opentelemetry.instrumentation.api.config.Config
@@ -25,7 +25,7 @@ abstract class SpringBatchTest extends AgentTestRunner {
then:
assertTraces(1) {
- trace(0, 2) {
+ trace(0, 3) {
span(0) {
name "BatchJob taskletJob"
kind INTERNAL
@@ -35,17 +35,49 @@ abstract class SpringBatchTest extends AgentTestRunner {
kind INTERNAL
childOf span(0)
}
+ span(2) {
+ name "BatchJob taskletJob.step.Chunk"
+ kind INTERNAL
+ childOf span(1)
+ }
}
}
}
- def "should trace job with multiple steps"() {
+ def "should handle exception in tasklet job+step"() {
when:
- runJob("itemsAndTaskletJob")
+ runJob("taskletJob", ["fail": new JobParameter(1)])
then:
assertTraces(1) {
trace(0, 3) {
+ span(0) {
+ name "BatchJob taskletJob"
+ kind INTERNAL
+ }
+ span(1) {
+ name "BatchJob taskletJob.step"
+ kind INTERNAL
+ childOf span(0)
+ }
+ span(2) {
+ name "BatchJob taskletJob.step.Chunk"
+ kind INTERNAL
+ childOf span(1)
+ errored true
+ errorEvent RuntimeException, "fail"
+ }
+ }
+ }
+ }
+
+ def "should trace chunked items job"() {
+ when:
+ runJob("itemsAndTaskletJob")
+
+ then:
+ assertTraces(1) {
+ trace(0, 7) {
span(0) {
name "BatchJob itemsAndTaskletJob"
kind INTERNAL
@@ -56,17 +88,37 @@ abstract class SpringBatchTest extends AgentTestRunner {
childOf span(0)
}
span(2) {
+ name "BatchJob itemsAndTaskletJob.itemStep.Chunk"
+ kind INTERNAL
+ childOf span(1)
+ }
+ span(3) {
+ name "BatchJob itemsAndTaskletJob.itemStep.Chunk"
+ kind INTERNAL
+ childOf span(1)
+ }
+ span(4) {
+ name "BatchJob itemsAndTaskletJob.itemStep.Chunk"
+ kind INTERNAL
+ childOf span(1)
+ }
+ span(5) {
name "BatchJob itemsAndTaskletJob.taskletStep"
kind INTERNAL
childOf span(0)
}
+ span(6) {
+ name "BatchJob itemsAndTaskletJob.taskletStep.Chunk"
+ kind INTERNAL
+ childOf span(5)
+ }
}
}
}
}
class JavaConfigBatchJobTest extends SpringBatchTest implements ApplicationConfigTrait {
- static final Config PREVIOUS_CONFIG = updateConfig {
+ static final Config PREVIOUS_CONFIG = updateConfigAndResetInstrumentation {
it.setProperty("otel.instrumentation.spring-batch.enabled", "true")
}
@@ -81,7 +133,7 @@ class JavaConfigBatchJobTest extends SpringBatchTest implements ApplicationConfi
}
class XmlConfigBatchJobTest extends SpringBatchTest implements ApplicationConfigTrait {
- static final Config PREVIOUS_CONFIG = updateConfig {
+ static final Config PREVIOUS_CONFIG = updateConfigAndResetInstrumentation {
it.setProperty("otel.instrumentation.spring-batch.enabled", "true")
}
@@ -96,7 +148,7 @@ class XmlConfigBatchJobTest extends SpringBatchTest implements ApplicationConfig
}
class JsrConfigBatchJobTest extends SpringBatchTest implements JavaxBatchConfigTrait {
- static final Config PREVIOUS_CONFIG = updateConfig {
+ static final Config PREVIOUS_CONFIG = updateConfigAndResetInstrumentation {
it.setProperty("otel.instrumentation.spring-batch.enabled", "true")
}

MODIFY
instrumentation/spring/spring-batch-3.0/javaagent/src/test/groovy/springbatch/TestTasklet.groovy
instrumentation/spring/spring-batch-3.0/javaagent/src/test/groovy/springbatch/TestTasklet.groovy
@@ -13,6 +13,9 @@ import org.springframework.batch.repeat.RepeatStatus
class TestTasklet implements Tasklet {
@Override
RepeatStatus execute(StepContribution contribution, ChunkContext chunkContext) throws Exception {
+ if (chunkContext.stepContext.stepExecution.jobParameters.getLong("fail") == 1) {
+ throw new RuntimeException("fail")
+ }
RepeatStatus.FINISHED
}
}
\ No newline at end of file

Instrument spring-batch: chunk (#1850)