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 + "]";
+ }
+ };
+ }
+}