Spring Batch에서 Step간 데이터 공유하는 방법

Gongmeda·2022년 12월 25일

Spring Batch Context에는 두 가지 Context가 있습니다.

Job ExecutionContext
- Job 전체에 걸쳐 유지
- 각 Step이 종료되는 시점에 업데이트
Step ExecutionContext
- 각 Step에서만 유지
- 각 chunk가 commit되는 시점에 업데이트

위와 같은 특성 때문에 Step간 데이터를 공유하기 위해서는 Job ExecutionContext에 데이터를 저장해 공유해야 합니다.

Job ExecutionContext에 데이터 저장

Tasklet

RepeatStatus execute(StepContribution contribution, ChunkContext chunkContext) {
	// contribution 사용
    contribution
    	.getStepExecution()
        .getJobExecution()
        .getExecutionContext()
        .put("key", value);
    
    // chunkContext 사용
    chunkContext
    	.getStepContext()
        .getStepExecution()
        .getJobExecution()
        .getExecutionContext()
        .put("key", value);
}

Tasklet을 구현할 때 사용할 수 있는 두 파라미터 모두 StepExecution 객체를 통해 Job ExecutionContext에 접근할 수 있습니다.

이때, 두 StepExecution 객체가 동일한지는 아래의 코드로 확인해볼 수 있습니다.

contribution.getStepExecution() == chunkContext.getStepContext().getStepExecution();
// True

Reader/Processor/Writer

public class SavingItemWriter implements ItemWriter<Object> {
    private StepExecution stepExecution;

    public void write(Chunk<? extends Object> items) throws Exception {
        // ...

        ExecutionContext stepContext = this.stepExecution.getExecutionContext();
        stepContext.put("key", value);
    }

    @BeforeStep
    public void saveStepExecution(StepExecution stepExecution) {
        this.stepExecution = stepExecution;
    }
}

ItemReader, ItemProcessor 또는 ItemWriter 를 구현해 사용할 때는 StepExecutionListener 인터페이스를 구현하거나 @BeforeStep 어노테이션을 활용해서 StepExecution 객체를 주입받아 사용할 수 있습니다.

만약 JpaPagingItemReader 와 같이 제공되는 구현체에서 ExecutionContext에 접근해야 한다면 아래와 같이 구현체 클래스를 extend해서 동일하게 StepExecution 객체를 주입받아 사용할 수 있습니다.

public class CustomJpaPagingItemReader<T> extends JpaPagingItemReader<T> {

    private StepExecution stepExecution;

    @Override
    protected void doReadPage() {
        super.doReadPage();
        if (!results.isEmpty()) {
            // stepExecution으로 ExecutionContext 접근
        }
    }

    @BeforeStep
    public void saveStepExecution(StepExecution stepExecution) {
        this.stepExecution = stepExecution;
    }
}

지금까지 코드들에서는 Job ExecutionContext에 직접 접근해서 데이터를 저장했습니다.

하지만 Job ExecutionContext에 직접 저장하는 것은 의도치 않게 데이터 유실을 겪을 수 있습니다. 이는 Step이 실행중 Job ExecutionContext에 데이터를 저장한다 하더라도 해당 데이터는 Step 실행중에는 저장(영속화)되지 않으며, Step이 실패하는 경우에는 데이터가 유실될 수 있기 때문입니다.

또한, 다른 Job에서 Step 구현체를 재사용하는 경우에는 해당 데이터가 필요 없는 경우가 생길 수도 있는데, 구현체에서 Job ExecutionContext에 직접 접근하여 데이터를 저장하는 것은 Job과 Step간의 강한 결합을 형성합니다. 즉, 새로운 Job에서 해당 Step 구현체를 Job ExecutionContext의 변경 없이 사용하기 위해서는 구현체 코드를 수정해야 하는 것 입니다.

이러한 문제들로 인해 Spring 공식 문서에서는 다음과 같은 방식을 권장합니다.

Step ExecutionContext에 데이터 저장 + ExecutionContextPromotionListener

"Listener는 배치 흐름 중에 Job, Step, Chunk 실행 전후에 어떤 일을 하도록 하는 Interceptor 개념의 클래스다."

Spring Batch에는 Listener라는 개념이 존재합니다.

이를 활용해서 Step 전후에 특정 동작을 하도록 설정할 수 있는데 위에서 언급되었던 StepExecutionListener 인터페이스 또한 Spring Batch가 제공하는 Listener 구현체 중 하나 입니다.

Spring Batch가 제공하는 ExecutionContextPromotionListener 는 설정된 키에 대해 Step ExecutionContext의 데이터들을 Step이 완료되는 시점에 자동으로 Job ExecutionContext로 승격(promote) 시켜주는 Listener입니다.

해당 Listener를 사용하여 Step ExecutionContext에 데이터를 저장하고 자동 승격 설정을 통해 Step간 데이터를 공유하면 직접 Job ExecutionContext에 접근하여 생기는 문제들을 해결할 수 있습니다.

@Bean
public Step step() {
	return new stepBuilderFactory.get("step")
    	.<String, String>chunk(10)
        .reader(reader())
        .writer(writer())
        .listener(promotionListener())
        .build();
}

@Bean
public ExecutionContextPromotionListener promotionListener() {
	ExecutionContextPromotionListener listener = new ExecutionContextPromotionListener();
    listener.setKeys(new String[]{"key"}); // 자동으로 승격시킬 키 목록
    return listener;
}

@JobScope를 사용한 Data Holder Bean

@JobScope
@Component
public class JobDataHolder {
	// 공유할 내용
}

@RequiredArgsConstructor
@StepScope
@Component
public MyTasklet implements Tasklet {
	
    private final JobDataHolder jobDataHolder;
    
    // ...
}