From e78c76687e50b624a500a0fb89d3ec0f891acef7 Mon Sep 17 00:00:00 2001 From: Daijiro Wachi Date: Thu, 15 Oct 2026 22:56:48 +0900 Subject: [PATCH] Run each batch job in its own OpenTelemetry trace. Jobs run on batchTaskExecutor, and an attached OpenTelemetry Java agent carries the submitting request's context into that pool. Cohort generation SQL therefore joined the HTTP request's trace, which could then last hours. The agent's Spring Batch instrumentation does not support Spring Batch 5 (opentelemetry-java-instrumentation#7306). TracingJobLauncher wraps each launched job so that execute() starts a new root span "BatchJob ", linked to the launching span, with the job's own work as children. The launching span records the job execution id. Unsuccessful jobs are marked as errors without the exit description, which can carry values from failing SQL; STOPPED (cancel) is not an error. Only opentelemetry-api is added (version from the Spring Boot BOM). It is a no-op unless an agent or SDK is present. --- pom.xml | 12 + src/main/java/org/ohdsi/webapi/JobConfig.java | 6 +- .../ohdsi/webapi/job/TracingJobLauncher.java | 144 ++++++++++++ .../webapi/job/TracingJobLauncherTest.java | 207 ++++++++++++++++++ 4 files changed, 367 insertions(+), 2 deletions(-) create mode 100644 src/main/java/org/ohdsi/webapi/job/TracingJobLauncher.java create mode 100644 src/test/java/org/ohdsi/webapi/job/TracingJobLauncherTest.java diff --git a/pom.xml b/pom.xml index e10736be82..f44a6407e0 100644 --- a/pom.xml +++ b/pom.xml @@ -397,6 +397,13 @@ spring-boot-starter-batch + + + io.opentelemetry + opentelemetry-api + + org.springframework.boot spring-boot-starter-jdbc @@ -757,6 +764,11 @@ + + io.opentelemetry + opentelemetry-sdk-testing + test + org.junit.vintage junit-vintage-engine diff --git a/src/main/java/org/ohdsi/webapi/JobConfig.java b/src/main/java/org/ohdsi/webapi/JobConfig.java index 186368dce3..12fca7eae8 100644 --- a/src/main/java/org/ohdsi/webapi/JobConfig.java +++ b/src/main/java/org/ohdsi/webapi/JobConfig.java @@ -8,6 +8,7 @@ import org.ohdsi.webapi.job.JdbcSearchableJobExecutionDao; import org.ohdsi.webapi.job.JobTemplate; import org.ohdsi.webapi.job.SearchableJobExecutionDao; +import org.ohdsi.webapi.job.TracingJobLauncher; import org.ohdsi.webapi.security.authz.AuthorizationService; import org.ohdsi.webapi.util.ManagedThreadPoolTaskExecutor; import org.slf4j.Logger; @@ -124,11 +125,12 @@ public TaskExecutor taskExecutor() { return taskExecutor; } - // Spring Batch 5: Custom JobLauncher for async execution + // Spring Batch 5: Custom JobLauncher for async execution. Each job runs in its own + // OpenTelemetry trace, linked to the request that launched it (no-op without an agent). @Bean @Primary public JobLauncher asyncJobLauncher(JobRepository jobRepository, @Qualifier("batchTaskExecutor") TaskExecutor batchTaskExecutor) throws Exception { - TaskExecutorJobLauncher jobLauncher = new TaskExecutorJobLauncher(); + TaskExecutorJobLauncher jobLauncher = new TracingJobLauncher(); jobLauncher.setJobRepository(jobRepository); jobLauncher.setTaskExecutor(batchTaskExecutor); jobLauncher.afterPropertiesSet(); diff --git a/src/main/java/org/ohdsi/webapi/job/TracingJobLauncher.java b/src/main/java/org/ohdsi/webapi/job/TracingJobLauncher.java new file mode 100644 index 0000000000..1a028cc087 --- /dev/null +++ b/src/main/java/org/ohdsi/webapi/job/TracingJobLauncher.java @@ -0,0 +1,144 @@ +package org.ohdsi.webapi.job; + +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanBuilder; +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.api.trace.StatusCode; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Scope; +import org.springframework.batch.core.BatchStatus; +import org.springframework.batch.core.Job; +import org.springframework.batch.core.JobExecution; +import org.springframework.batch.core.JobParameters; +import org.springframework.batch.core.JobParametersIncrementer; +import org.springframework.batch.core.JobParametersInvalidException; +import org.springframework.batch.core.JobParametersValidator; +import org.springframework.batch.core.launch.support.TaskExecutorJobLauncher; +import org.springframework.batch.core.repository.JobExecutionAlreadyRunningException; +import org.springframework.batch.core.repository.JobInstanceAlreadyCompleteException; +import org.springframework.batch.core.repository.JobRestartException; + +/** + * Runs every batch job in its own trace. + * + *

Jobs execute on {@code batchTaskExecutor}. When an OpenTelemetry Java agent is attached, it + * carries the submitting HTTP request's context into that pool, so the SQL of a cohort generation + * would otherwise join the request's trace and stretch it over minutes or hours. Instead, each job + * starts a new root span, linked to the span that launched it, and the launching span records the + * job execution id. Short request traces and per-job traces are also what make tail sampling + * (keep errors and slow traces) practical in a collector. + * + *

Only the OpenTelemetry API is used, which is a no-op unless an agent or SDK is present, so + * deployments without telemetry are unaffected. The agent's own Spring Batch instrumentation does + * not support Spring Batch 5 (opentelemetry-java-instrumentation#7306). + */ +public class TracingJobLauncher extends TaskExecutorJobLauncher { + + static final String INSTRUMENTATION_SCOPE = "org.ohdsi.webapi.job"; + + static final AttributeKey JOB_SYSTEM = AttributeKey.stringKey("job.system"); + static final AttributeKey JOB_NAME = AttributeKey.stringKey("spring_batch.job.name"); + static final AttributeKey JOB_EXECUTION_ID = AttributeKey.longKey("spring_batch.job.execution.id"); + static final AttributeKey JOB_STATUS = AttributeKey.stringKey("spring_batch.job.status"); + + private final Tracer tracer; + + public TracingJobLauncher() { + this(GlobalOpenTelemetry.getTracer(INSTRUMENTATION_SCOPE)); + } + + TracingJobLauncher(Tracer tracer) { + this.tracer = tracer; + } + + @Override + public JobExecution run(Job job, JobParameters jobParameters) + throws JobExecutionAlreadyRunningException, JobRestartException, + JobInstanceAlreadyCompleteException, JobParametersInvalidException { + Span launching = Span.current(); + JobExecution execution = super.run(new TracedJob(job, launching.getSpanContext(), tracer), jobParameters); + if (execution.getId() != null) { + launching.setAttribute(JOB_EXECUTION_ID, execution.getId()); + } + return execution; + } + + /** + * Delegates everything to the real job, and wraps {@link #execute} in the job's root span. It is + * created on the launching thread, which is where the launching span context is captured; the + * span itself starts on the executor thread. + */ + static final class TracedJob implements Job { + + private final Job delegate; + private final SpanContext launchedBy; + private final Tracer tracer; + + TracedJob(Job delegate, SpanContext launchedBy, Tracer tracer) { + this.delegate = delegate; + this.launchedBy = launchedBy; + this.tracer = tracer; + } + + @Override + public String getName() { + return delegate.getName(); + } + + @Override + public boolean isRestartable() { + return delegate.isRestartable(); + } + + @Override + public JobParametersIncrementer getJobParametersIncrementer() { + return delegate.getJobParametersIncrementer(); + } + + @Override + public JobParametersValidator getJobParametersValidator() { + return delegate.getJobParametersValidator(); + } + + @Override + public void execute(JobExecution execution) { + SpanBuilder builder = tracer.spanBuilder("BatchJob " + delegate.getName()) + .setNoParent() + .setSpanKind(SpanKind.INTERNAL) + .setAttribute(JOB_SYSTEM, "spring_batch") + .setAttribute(JOB_NAME, delegate.getName()); + if (execution.getId() != null) { + builder.setAttribute(JOB_EXECUTION_ID, execution.getId()); + } + if (launchedBy.isValid()) { + builder.addLink(launchedBy); + } + Span span = builder.startSpan(); + try (Scope ignored = span.makeCurrent()) { + delegate.execute(execution); + } catch (RuntimeException | Error e) { + span.recordException(e); + span.setStatus(StatusCode.ERROR); + throw e; + } finally { + BatchStatus status = execution.getStatus(); + if (status != null) { + span.setAttribute(JOB_STATUS, status.name()); + // No description: exit messages can carry values from the failing SQL. + if (status.isUnsuccessful()) { + span.setStatus(StatusCode.ERROR); + } + } + span.end(); + } + } + + @Override + public String toString() { + return delegate.toString(); + } + } +} diff --git a/src/test/java/org/ohdsi/webapi/job/TracingJobLauncherTest.java b/src/test/java/org/ohdsi/webapi/job/TracingJobLauncherTest.java new file mode 100644 index 0000000000..e7b1874c0e --- /dev/null +++ b/src/test/java/org/ohdsi/webapi/job/TracingJobLauncherTest.java @@ -0,0 +1,207 @@ +package org.ohdsi.webapi.job; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.api.trace.StatusCode; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Scope; +import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule; +import io.opentelemetry.sdk.trace.data.SpanData; +import java.util.List; +import java.util.function.Consumer; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.springframework.batch.core.BatchStatus; +import org.springframework.batch.core.Job; +import org.springframework.batch.core.JobExecution; +import org.springframework.batch.core.JobInstance; +import org.springframework.batch.core.JobParameters; +import org.springframework.batch.core.JobParametersBuilder; +import org.springframework.batch.core.JobParametersIncrementer; +import org.springframework.batch.core.JobParametersValidator; +import org.springframework.batch.core.repository.JobRepository; +import org.springframework.core.task.SyncTaskExecutor; + +/** + * A batch job must not join the trace of the request that launched it: it runs in its own root + * span, linked to the launching span, with SQL and everything else inside the job as children. + */ +public class TracingJobLauncherTest { + + private static final long EXECUTION_ID = 42L; + + @Rule + public OpenTelemetryRule otel = OpenTelemetryRule.create(); + + private Tracer tracer; + private TracingJobLauncher launcher; + private final JobParameters parameters = new JobParametersBuilder().addLong("cohort_definition_id", 7L).toJobParameters(); + + @Before + public void setUp() throws Exception { + tracer = otel.getOpenTelemetry().getTracer("test"); + + JobRepository repository = mock(JobRepository.class); + when(repository.getLastJobExecution(anyString(), any())).thenReturn(null); + when(repository.createJobExecution(anyString(), any())) + .thenAnswer(inv -> new JobExecution(new JobInstance(1L, inv.getArgument(0)), EXECUTION_ID, inv.getArgument(1))); + + launcher = new TracingJobLauncher(otel.getOpenTelemetry().getTracer(TracingJobLauncher.INSTRUMENTATION_SCOPE)); + launcher.setJobRepository(repository); + // Synchronous on purpose: the job then runs with the request's context still current, + // which is exactly what the agent's executor instrumentation reproduces on a pool thread. + launcher.setTaskExecutor(new SyncTaskExecutor()); + launcher.afterPropertiesSet(); + } + + @Test + public void jobRunsInItsOwnTraceLinkedToTheLaunchingRequest() throws Exception { + SpanContext[] insideJob = new SpanContext[1]; + SpanData request = launchFromRequest(job("generateCohort", execution -> { + insideJob[0] = Span.current().getSpanContext(); + tracer.spanBuilder("SELECT cohort").startSpan().end(); + execution.setStatus(BatchStatus.COMPLETED); + })); + + SpanData jobSpan = span("BatchJob generateCohort"); + SpanData sql = span("SELECT cohort"); + + assertFalse("job span must be a root", jobSpan.getParentSpanContext().isValid()); + assertNotEquals(request.getTraceId(), jobSpan.getTraceId()); + assertEquals(1, jobSpan.getLinks().size()); + assertEquals(request.getSpanContext(), jobSpan.getLinks().get(0).getSpanContext()); + + assertEquals("work inside the job is current under the job span", jobSpan.getSpanContext(), insideJob[0]); + assertEquals(jobSpan.getSpanId(), sql.getParentSpanId()); + assertEquals(jobSpan.getTraceId(), sql.getTraceId()); + } + + @Test + public void recordsJobIdentityAndOutcome() throws Exception { + launchFromRequest(job("generateCohort", execution -> execution.setStatus(BatchStatus.COMPLETED))); + + SpanData jobSpan = span("BatchJob generateCohort"); + assertEquals("spring_batch", jobSpan.getAttributes().get(TracingJobLauncher.JOB_SYSTEM)); + assertEquals("generateCohort", jobSpan.getAttributes().get(TracingJobLauncher.JOB_NAME)); + assertEquals(Long.valueOf(EXECUTION_ID), jobSpan.getAttributes().get(TracingJobLauncher.JOB_EXECUTION_ID)); + assertEquals("COMPLETED", jobSpan.getAttributes().get(TracingJobLauncher.JOB_STATUS)); + assertEquals(StatusCode.UNSET, jobSpan.getStatus().getStatusCode()); + } + + @Test + public void launchingSpanRecordsTheExecutionIdSoTheJobTraceCanBeFound() throws Exception { + SpanData request = launchFromRequest(job("generateCohort", execution -> execution.setStatus(BatchStatus.COMPLETED))); + assertEquals(Long.valueOf(EXECUTION_ID), request.getAttributes().get(TracingJobLauncher.JOB_EXECUTION_ID)); + } + + @Test + public void failedJobIsAnErrorWithoutForwardingTheExitMessage() throws Exception { + launchFromRequest(job("generateCohort", execution -> { + execution.setStatus(BatchStatus.FAILED); + execution.setExitStatus(execution.getExitStatus().addExitDescription("Key (person_id)=(123) already exists")); + })); + + SpanData jobSpan = span("BatchJob generateCohort"); + assertEquals(StatusCode.ERROR, jobSpan.getStatus().getStatusCode()); + assertEquals("", jobSpan.getStatus().getDescription()); + assertEquals("FAILED", jobSpan.getAttributes().get(TracingJobLauncher.JOB_STATUS)); + } + + @Test + public void stoppedJobIsNotAnError() throws Exception { + launchFromRequest(job("generateCohort", execution -> execution.setStatus(BatchStatus.STOPPED))); + assertEquals(StatusCode.UNSET, span("BatchJob generateCohort").getStatus().getStatusCode()); + } + + @Test + public void exceptionFromTheJobIsRecordedAndRethrown() throws Exception { + // TaskExecutorJobLauncher rethrows what the job throws, after marking the execution failed. + try { + launchFromRequest(job("generateCohort", execution -> { + throw new IllegalStateException("boom"); + })); + } catch (IllegalStateException expected) { + // fall through to the assertions + } + + SpanData jobSpan = span("BatchJob generateCohort"); + assertEquals(StatusCode.ERROR, jobSpan.getStatus().getStatusCode()); + assertTrue(jobSpan.getEvents().stream().anyMatch(e -> e.getName().equals("exception"))); + } + + @Test + public void jobLaunchedOutsideARequestIsARootWithoutLinks() throws Exception { + launcher.run(job("warmCache", execution -> execution.setStatus(BatchStatus.COMPLETED)), parameters); + + SpanData jobSpan = span("BatchJob warmCache"); + assertFalse(jobSpan.getParentSpanContext().isValid()); + assertTrue(jobSpan.getLinks().isEmpty()); + } + + @Test + public void wrapperDelegatesTheJobContract() { + // Job's default methods return a new validator per call, so pin real instances. + JobParametersValidator validator = mock(JobParametersValidator.class); + JobParametersIncrementer incrementer = mock(JobParametersIncrementer.class); + Job delegate = mock(Job.class); + when(delegate.getName()).thenReturn("generateCohort"); + when(delegate.isRestartable()).thenReturn(false); + when(delegate.getJobParametersValidator()).thenReturn(validator); + when(delegate.getJobParametersIncrementer()).thenReturn(incrementer); + Job traced = new TracingJobLauncher.TracedJob(delegate, SpanContext.getInvalid(), tracer); + + assertEquals("generateCohort", traced.getName()); + assertFalse(traced.isRestartable()); + assertSame(validator, traced.getJobParametersValidator()); + assertSame(incrementer, traced.getJobParametersIncrementer()); + assertEquals(delegate.toString(), traced.toString()); + } + + private SpanData launchFromRequest(Job job) throws Exception { + Span request = tracer.spanBuilder("POST /WebAPI/cohortdefinition/{id}/generate/{sourceKey}").startSpan(); + try (Scope ignored = request.makeCurrent()) { + launcher.run(job, parameters); + } finally { + request.end(); + } + return span("POST /WebAPI/cohortdefinition/{id}/generate/{sourceKey}"); + } + + private SpanData span(String name) { + List spans = otel.getSpans(); + return spans.stream() + .filter(s -> s.getName().equals(name)) + .findFirst() + .orElseThrow(() -> new AssertionError("no span named " + name + " in " + spans)); + } + + private static Job job(String name, Consumer body) { + return new Job() { + @Override + public String getName() { + return name; + } + + @Override + public void execute(JobExecution execution) { + body.accept(execution); + } + + @Override + public String toString() { + return "job[" + name + "]"; + } + }; + } +}