Is it possible to combine partition and parallel steps in spring batch?

Question

I am just wondering is it doable in Spring Batch?

Step1
Step2 (flow) -> flow1, flow2, flow3
Step3

Where each
flow1 -> partition into 5 GridSize
flow2 -> partition into 5 GridSize
flow3 -> partition into 5 GridSize

return jobBuilderFactory.get("dataLoad")
            .incrementer(new RunIdIncrementer())
            .listener(listener)
            .start(step1())
            .next(step2())
            .next(step3())
            .build()
            .build();
@Bean
public Flow step2() {
    Flow subflow1 = new FlowBuilder<Flow>("readTable1Flow").from(readTable1()).end();
    Flow subflow2 = new FlowBuilder<Flow>("readTable2Flow").from(readTable2()).end();
    Flow subflow3 = new FlowBuilder<Flow>("readTable3Flow").from(readTable3()).end();

    return new FlowBuilder<Flow>("splitflow").split(new SimpleAsyncTaskExecutor())
            .add(subflow1, subflow2, subflow3).build();
}
@Bean
public Step readTable1() {
    return stepBuilderFactory.get("readTable1Step")
            .partitioner(slaveStep1().getName(), partitioner1())
            .partitionHandler(slaveStep1Handler())
            .build();
}

@Bean
public Step readTable2() {
    return stepBuilderFactory.get("readTable2Step")
            .partitioner(slaveStep2().getName(), partitioner2())
            .partitionHandler(slaveStep2Handler())
            .build();
}
@Bean
public Step readTable3() {
    return stepBuilderFactory.get("readTable3Step")
            .partitioner(slaveStep3().getName(), partitioner2())
            .partitionHandler(slaveStep3Handler())
            .build();
}

Answer 1

This is possible. You can have a split flow where each step running in parallel is a partitioned step. Here is an example:

import java.util.HashMap;
import java.util.Map;

import org.springframework.batch.core.Job;
import org.springframework.batch.core.JobParameters;
import org.springframework.batch.core.Step;
import org.springframework.batch.core.configuration.annotation.EnableBatchProcessing;
import org.springframework.batch.core.configuration.annotation.JobBuilderFactory;
import org.springframework.batch.core.configuration.annotation.StepBuilderFactory;
import org.springframework.batch.core.configuration.annotation.StepScope;
import org.springframework.batch.core.job.builder.FlowBuilder;
import org.springframework.batch.core.job.flow.Flow;
import org.springframework.batch.core.launch.JobLauncher;
import org.springframework.batch.core.partition.support.Partitioner;
import org.springframework.batch.core.step.tasklet.Tasklet;
import org.springframework.batch.item.ExecutionContext;
import org.springframework.batch.repeat.RepeatStatus;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.ApplicationContext;
import org.springframework.context.annotation.AnnotationConfigApplicationContext;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.SimpleAsyncTaskExecutor;

@Configuration
@EnableBatchProcessing
public class MyJob {

    @Autowired
    private JobBuilderFactory jobs;

    @Autowired
    private StepBuilderFactory steps;

    @Bean
    public Step step1() {
        return steps.get("step1")
                .tasklet((contribution, chunkContext) -> {
                    System.out.println(Thread.currentThread().getName() + ": step1");
                    return RepeatStatus.FINISHED;
                })
                .build();
    }

    @Bean
    public Flow step2() {
        Flow subflow1 = new FlowBuilder<Flow>("step21_master").from(step21_master()).end();
        Flow subflow2 = new FlowBuilder<Flow>("step22_master").from(step22_master()).end();
        Flow subflow3 = new FlowBuilder<Flow>("step23_master").from(step23_master()).end();

        return new FlowBuilder<Flow>("splitflow").split(taskExecutor())
                .add(subflow1, subflow2, subflow3).build();
    }

    @Bean
    public Step step21_master() {
        return steps.get("step21_master")
                .partitioner("workerStep", partitioner("step21_master"))
                .step(workerStep())
                .gridSize(3)
                .taskExecutor(taskExecutor())
                .build();
    }

    @Bean
    public Step step22_master() {
        return steps.get("step22_master")
                .partitioner("workerStep", partitioner("step22_master"))
                .step(workerStep())
                .gridSize(3)
                .taskExecutor(taskExecutor())
                .build();
    }

    @Bean
    public Step step23_master() {
        return steps.get("step23_master")
                .partitioner("workerStep", partitioner("step23_master"))
                .step(workerStep())
                .gridSize(3)
                .taskExecutor(taskExecutor())
                .build();
    }

    @Bean
    public Step step3() {
        return steps.get("step3")
                .tasklet((contribution, chunkContext) -> {
                    System.out.println(Thread.currentThread().getName() + ": step3");
                    return RepeatStatus.FINISHED;
                })
                .build();
    }

    @Bean
    public Step workerStep() {
        return steps.get("workerStep")
                .tasklet(getTasklet(null))
                .build();
    }

    @Bean
    @StepScope
    public Tasklet getTasklet(@Value("#{stepExecutionContext['data']}") String partitionData) {
        return (contribution, chunkContext) -> {
            System.out.println(Thread.currentThread().getName() + " processing partitionData = " + partitionData);
            return RepeatStatus.FINISHED;
        };
    }

    @Bean
    public Job job() {
        return jobs.get("job")
                .flow(step1()).on("*").to(step2())
                .next(step3())
                .build()
                .build();
    }

    @Bean
    public SimpleAsyncTaskExecutor taskExecutor() {
        return new SimpleAsyncTaskExecutor();
    }

    public Partitioner partitioner(String stepName) {
        return gridSize -> {
            Map<String, ExecutionContext> map = new HashMap<>(gridSize);
            for (int i = 0; i < gridSize; i++) {
                ExecutionContext executionContext = new ExecutionContext();
                executionContext.put("data", stepName + ":data" + i);
                String key = stepName + ":partition" + i;
                map.put(key, executionContext);
            }
            return map;
        };
    }

    public static void main(String[] args) throws Exception {
        ApplicationContext context = new AnnotationConfigApplicationContext(MyJob.class);
        JobLauncher jobLauncher = context.getBean(JobLauncher.class);
        Job job = context.getBean(Job.class);
        jobLauncher.run(job, new JobParameters());
    }

}

In this example, step2 is a split flow (as in your example) where each sub flow is a partitioned step (a master step) with 3 worker steps. If you run this sample, you should see something like:

[main] INFO org.springframework.batch.core.launch.support.SimpleJobLauncher - Job: [FlowJob: [name=job]] launched with the following parameters: [{}]
[main] INFO org.springframework.batch.core.job.SimpleStepHandler - Executing step: [step1]
main: step1
[SimpleAsyncTaskExecutor-1] INFO org.springframework.batch.core.job.SimpleStepHandler - Executing step: [step21_master]
[SimpleAsyncTaskExecutor-3] INFO org.springframework.batch.core.job.SimpleStepHandler - Executing step: [step23_master]
[SimpleAsyncTaskExecutor-2] INFO org.springframework.batch.core.job.SimpleStepHandler - Executing step: [step22_master]
SimpleAsyncTaskExecutor-4 processing partitionData = step21_master: data2
SimpleAsyncTaskExecutor-12 processing partitionData = step22_master: data2
SimpleAsyncTaskExecutor-11 processing partitionData = step22_master: data0
SimpleAsyncTaskExecutor-10 processing partitionData = step22_master: data1
SimpleAsyncTaskExecutor-9 processing partitionData = step23_master: data1
SimpleAsyncTaskExecutor-8 processing partitionData = step23_master: data2
SimpleAsyncTaskExecutor-7 processing partitionData = step23_master: data0
SimpleAsyncTaskExecutor-5 processing partitionData = step21_master: data0
SimpleAsyncTaskExecutor-6 processing partitionData = step21_master: data1
main: step3
[main] INFO org.springframework.batch.core.job.SimpleStepHandler - Executing step: [step3]
[main] INFO org.springframework.batch.core.launch.support.SimpleJobLauncher - Job: [FlowJob: [name=job]] completed with the following parameters: [{}] and the following status: [COMPLETED]

step1 , step2 and step3 run in sequence where step2 is split in 3 sub steps running in parallel, where each sub step is a master step for 3 worker steps running in parallel.

Answer 2

I made it working now. The result I was not able to make it work is Thread deadlock. The spring batch is trying to insert/update metadata into database (HsqlDB). Each flow is in waiting status. Once I switch to different db, it works. Thank you Mahmoud for your input!!