diff --git a/spring-batch-integration/src/main/java/org/springframework/batch/integration/chunk/ChunkTaskExecutorItemWriter.java b/spring-batch-integration/src/main/java/org/springframework/batch/integration/chunk/ChunkTaskExecutorItemWriter.java index c70ef8ed85..d4cd648878 100644 --- a/spring-batch-integration/src/main/java/org/springframework/batch/integration/chunk/ChunkTaskExecutorItemWriter.java +++ b/spring-batch-integration/src/main/java/org/springframework/batch/integration/chunk/ChunkTaskExecutorItemWriter.java @@ -17,6 +17,7 @@ import org.springframework.batch.core.BatchStatus; import org.springframework.batch.core.ExitStatus; +import org.springframework.batch.core.scope.context.StepSynchronizationManager; import org.springframework.batch.core.step.StepContribution; import org.springframework.batch.core.step.StepExecution; import org.springframework.batch.core.listener.StepExecutionListener; @@ -94,8 +95,15 @@ public ChunkTaskExecutorItemWriter(ChunkProcessor chunkRequestProcessor, Task public void write(Chunk chunk) { ChunkRequest request = new ChunkRequest<>(++sequence, chunk, this.stepExecution.getJobExecutionId(), this.stepExecution.createStepContribution()); - FutureTask chunkResponseFutureTask = new FutureTask<>( - () -> this.chunkProcessorChunkHandler.handle(request)); + FutureTask chunkResponseFutureTask = new FutureTask<>(() -> { + try { + StepSynchronizationManager.register(this.stepExecution); + return this.chunkProcessorChunkHandler.handle(request); + } + finally { + StepSynchronizationManager.close(); + } + }); this.responses.add(chunkResponseFutureTask); this.taskExecutor.execute(chunkResponseFutureTask); } diff --git a/spring-batch-integration/src/test/java/org/springframework/batch/integration/chunk/ChunkTaskExecutorItemWriterTests.java b/spring-batch-integration/src/test/java/org/springframework/batch/integration/chunk/ChunkTaskExecutorItemWriterTests.java new file mode 100644 index 0000000000..a7ead08a9b --- /dev/null +++ b/spring-batch-integration/src/test/java/org/springframework/batch/integration/chunk/ChunkTaskExecutorItemWriterTests.java @@ -0,0 +1,70 @@ +/* + * Copyright 2025-present the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.batch.integration.chunk; + +import java.util.concurrent.atomic.AtomicReference; + +import org.junit.jupiter.api.Test; + +import org.springframework.batch.core.ExitStatus; +import org.springframework.batch.core.scope.context.StepContext; +import org.springframework.batch.core.scope.context.StepSynchronizationManager; +import org.springframework.batch.core.step.StepExecution; +import org.springframework.batch.core.step.item.ChunkProcessor; +import org.springframework.batch.infrastructure.item.Chunk; +import org.springframework.batch.test.MetaDataInstanceFactory; +import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; + +class ChunkTaskExecutorItemWriterTests { + + @Test + void stepContextIsPropagatedToWorkerThread() throws Exception { + // given + StepExecution stepExecution = MetaDataInstanceFactory.createStepExecution(); + AtomicReference capturedContext = new AtomicReference<>(); + ChunkProcessor chunkProcessor = (chunk, contribution) -> { + capturedContext.set(StepSynchronizationManager.getContext()); + contribution.setExitStatus(ExitStatus.COMPLETED); + }; + ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor(); + taskExecutor.setCorePoolSize(1); + taskExecutor.setThreadNamePrefix("worker-thread-"); + taskExecutor.setWaitForTasksToCompleteOnShutdown(true); + taskExecutor.afterPropertiesSet(); + try { + ChunkTaskExecutorItemWriter itemWriter = new ChunkTaskExecutorItemWriter<>(chunkProcessor, + taskExecutor); + itemWriter.beforeStep(stepExecution); + + // when + itemWriter.write(Chunk.of("foo", "bar")); + ExitStatus exitStatus = itemWriter.afterStep(stepExecution); + + // then + assertEquals(ExitStatus.COMPLETED.getExitCode(), exitStatus.getExitCode()); + StepContext context = capturedContext.get(); + assertNotNull(context, "StepContext should be available on the worker thread"); + assertEquals(stepExecution, context.getStepExecution()); + } + finally { + taskExecutor.shutdown(); + } + } + +}