DigdagMultistageOperatorUtilityIdea.java

public class Multistage<P>
{
    public static class State<T>
    {
        private Object nextLambdaArgument;
        private boolean isNext;

        private State(Object nextLambdaArgument, boolean isNext)
        {
            this.nextLambdaArgument = nextLambdaArgument;
            this.isNext = isNext;
        }

        public static State repeat();
        {
            return new State(null, false);
        }

        public static <N> State<N> next(N nextStateValue)
        {
            return new State<N>(nextStateValue, true);
        }
    }

    private static class Stage
    {
        private final Class<?> stateClass;
        private final Function function;
    }

    private final List<Stage> stages;

    public static <N> Multistage<N> first(Class<N> nextStateClass, Supplier<Stage<N>> lambda)
    {
        return new Multistage<Void>().add(nextStateClass, (voidValue) -> lambda.get());
    }

    public <N> Multistage<N> add(Class<N> nextStateClass, Function<P, Stage<N>> lambda)
    {
        stages.add(new Stage(nextStateClass, lambda));
        return (Multistage) this;
    }

    public TaskResult last(Function<P, TaskResult> lambda)
    {
        stages.add(new Stage(TaskResult.class, lambda));
        return run();
    }

    private TaskResult run()
    {
        ...
    }
}


// How to use Multistage

public TaskResult runTask()
{
    final String someParam1 = params.get(...);
    final String someParam2 = params.get(...);

    return Multistage
        // First step generates a job id
        .first(String.class, () -> generateJobId())
        // next step submits a job, returns job info
        .add(JobInfo.class, (jobId) -> {
            if (jobExists(jobId)) {
                return State.repeat();
            }
            return State.next(submitJob(jobId));
        })
        // next step checks status of the job, returns job result
        .add(JobResult.class, (jobInfo) -> {
            if (isJobFinished(jobInfo)) {
                return State.next(getJobResult(jobInfo.getJobId()));
            }
            else {
                return State.repeat();
            }
        })
        // last stage creates TaskResult
        .last((jobResult) -> {
            return buildTaskResult(jobResult);
        });
}