BackgroundTasks: State pattern

Complex, multi-state tasks easily become a giant code-mess class. A good way to overcome this, is to decompose each execution steps to its own state. This wiki section describes core concepts and implementation details of using State pattern with Itmo.Dev.Platform.BackgroundTasks.

sample code

Idea

Main pattern implementation idea is described by couple of entities: task state, state handler.

Entities

Task state denotes a step in task execution it's object must be stored it task's execution metadata, to persist state changes between steps execution. Task state must store the data, required to execute the step. If one step is passing data to another, it has to be done through states.

State handlers are used to implement task execution steps. They receive handled state itself, and some context (usually task id and metadata). As states are stored inside execution metadata, they can be mutable, if you need to persist some step context in event of service being stopped (ex: when loading paginated data, you can store cursor in a state, and update it after each page processing. if application would be stopped, cursor will be restored after application restart).

Moves between the states should be encapsulated in state handlers.

State handler execution can return two kinds of results: Finished, FinishedWithResult. The first one is used to make a state move without returning any task result, so the task would process next state in the same run. Second kind is used to either suspend task execution (in case you need to wait on some operation callback from outside service), or end task execution (with success or failure).

Waiting states

As state move logic is encapsulated in state handlers, you have to ensure that this abstraction does not leak to outside service code.

Let's imagine that your task state machine is suspended, waiting on some outside callback. This callback is result of some operation, that you started in a previous step, and it has some id. When you receive callback, that this operation is finished, you have to match that operation id with an id of background task you need to proceed.

The naive way to approach this problem is to add a "match table" that will store pairs of operation ids and background task ids. That way, when you receive a callback, you query that table to determine which task you have to proceed.
Validation is when this approach comes short. In an event of failure, you probably do not want to proceed a task in an invalid state, so you need to check, whether task's state comes after the one that triggers on outside operation and suspends the task. The problem in that case, is that you probably moved your task's state machine to the state that comes after outside operation completion, and you must check the task state against it. This leads to discussed leakage itself, as state sequence encapsulated in state handlers, but you must know it here.

To solve that problem, you should use waiting states. This is a special kind of state, it does not explicitly denotes a step in your task execution, it is rather used to denote that the task is waiting for some specific outside operation.
Handler that starts this outside operation, must move state to this waiting state with suspended result. While waiting state handler must perform further state move to the next step of task execution.

When waiting state is added into a task state machine, you should use it to check against when proceeding task.
Implementation details are presented further in the article

Implementation

State machine

Define some type to be used as a base class for your states:

public abstract record TaskState;

Define an interface for state handler:

public record StateHandlerContext(BackgroundTaskId BackgroundTaskId, StateTaskMetadata Metadata);

public abstract record StateHandlerResult
{
    private StateHandlerResult() { }

    public sealed record Finished(TaskState State) : StateHandlerResult;

    public sealed record FinishedWithResult(
        TaskState State,
        BackgroundTaskExecutionResult<EmptyExecutionResult, StateTaskError> Result) : StateHandlerResult;
}

public interface IStateHandler<in TState> where TState : TaskState
{
    ValueTask<StateHandlerResult> HandleAsync(
        TState state,
        StateHandlerContext context,
        CancellationToken cancellationToken);
}

Starting state

To enforce abstraction, when implementing background task itself, you should always have a starting state, it's handler must perform a move operation to first step of your state machine.

public sealed record StartingState : TaskState;

public class StartingStateHandler : IStateHandler<StartingState>
{
    public ValueTask<StateHandlerResult> HandleAsync(
        StartingState state,
        StateHandlerContext context,
        CancellationToken cancellationToken)
    {
        var result = new StateHandlerResult.FinishedWithResult(
            new WaitingFirstState(Guid.NewGuid()),
            BackgroundTaskExecutionResult.Suspended.ForEmptyResult().ForError<StateTaskError>());

        return ValueTask.FromResult<StateHandlerResult>(result);
    }
}

In this example, starting state immediately puts a task to wait of operation associated with FirstState. For example purposes, an external operation id is randomly generated here. It is stored in a waiting state for this operation.

public sealed record WaitingFirstState(Guid OperationId) : TaskState;

Waiting states

Waiting state was introduced in a previous sample, as discussed before, is used to denote that the task is waiting for some operation completion (in this case, on operation associated with FirstState).

Let's implement some service code, to proceed task execution. It will run to process some callback of operation finish.

public async Task ProceedFirstStateAsync(Guid operationId, CancellationToken cancellationToken)
{
    var query = BackgroundTaskQuery.Build(builder => builder
        .WithName(StateBackgroundTask.Name)
        .WithState(BackgroundTaskState.Suspended)
        .WithExecutionMetadata(new StateTaskExecutionMetadata { State = new WaitingFirstState(operationId) }));

    await _runner
        .ProceedBackgroundTask
        .WithQuery(query)
        .WithoutExecutionMetadataModification()
        .ProceedAsync(cancellationToken);
}

In the example above, we create a query for background task with the name we need, in a suspended state. Additionally, we add an execution metadata to include only background tasks in WaitingFistState with received operation id.

Keep in mind a difference between BackgroundTask's state and state stored in execution metadata. First is a state of a task as a whole (whether it is completed, suspended, or failed), and the latter is the current execution state (your state machine).

That way, if there would be no such task, we can handle a result returned by ProceedAsync method, and perform some actions accordingly.

When task is proceeded, it is scheduled for execution, and will resume in WaitingFirstState state, so we have to implement a handler, that will move task's state to the next step.

public class WaitingFirstStateHandler : IStateHandler<WaitingFirstState>
{
    public ValueTask<StateHandlerResult> HandleAsync(
        WaitingFirstState state,
        StateHandlerContext context,
        CancellationToken cancellationToken)
    {
        var result = new StateHandlerResult.Finished(new FirstState());
        return ValueTask.FromResult<StateHandlerResult>(result);
    }
}

In our case, after WaitingFirstState - comes FirstState. Note that StateHandlerResult.Finished is used here, as the task does not have to wait for anything, it can process FirstState in the current execution.

Background task handler

When using State pattern, your background task will become an engine, running this state machine. It has to 1) execute handler for current state, 2) process handler result

Handling current state

To handle current state you need to determine which handler is designated for it. Simplest, and probably the best, way to do it, is to match state types with switch expression. You would have to inject IServiceProvider into a background task to create the handlers themselves.

You can register your handler in DI directly, or you can use ActivatorUtilities to create them (as used in further example).

private ValueTask<StateHandlerResult> HandleAsync(
    TaskState taskState,
    StateHandlerContext context,
    CancellationToken cancellationToken)
{
    return taskState switch
    {
        StartingState state => ActivatorUtilities
            .CreateInstance<StartingStateHandler>(_serviceProvider)
            .HandleAsync(state, context, cancellationToken),

        WaitingFirstState state => ActivatorUtilities
            .CreateInstance<WaitingFirstStateHandler>(_serviceProvider)
            .HandleAsync(state, context, cancellationToken),

        FirstState state => ActivatorUtilities
            .CreateInstance<FirstStateHandler>(_serviceProvider)
            .HandleAsync(state, context, cancellationToken),

        CompletedState state => ActivatorUtilities
            .CreateInstance<CompletedStateHandler>(_serviceProvider)
            .HandleAsync(state, context, cancellationToken),

        _ => throw new ArgumentOutOfRangeException(nameof(taskState), taskState, "Could not resolve state handler"),
    };
}

To simplify type inference of handlers and etc, you can execute HandleAsync method directly in switch expression.

Updating state machine

After current state handler execution, you must update task context, to reflect your state machine changes. So your's task ExecuteAsync method would look like so:

public async Task<BackgroundTaskExecutionResult<EmptyExecutionResult, StateTaskError>> ExecuteAsync(
    BackgroundTaskExecutionContext<StateTaskMetadata, StateTaskExecutionMetadata> executionContext,
    CancellationToken cancellationToken)
{
    executionContext.ExecutionMetadata.State ??= new StartingState();

    var context = new StateHandlerContext(executionContext.Id, executionContext.Metadata);

    while (true)
    {
        var result = await HandleAsync(executionContext.ExecutionMetadata.State, context, cancellationToken);

        if (result is StateHandlerResult.Finished finished)
        {
            executionContext.ExecutionMetadata.State = finished.State;
            continue;
        }

        if (result is StateHandlerResult.FinishedWithResult finishedWithResult)
        {
            executionContext.ExecutionMetadata.State = finishedWithResult.State;
            return finishedWithResult.Result;
        }

        return BackgroundTaskExecutionResult
            .Failure
            .ForEmptyResult()
            .WithError(new StateTaskError($"Invalid handler result = {result}"));
    }
}

Note that first line of ExecuteAsync initialises StartingState when needed (would be executed on task's first run)

In the example above, we process handler result. We always update current state in execution metadata, to the state in result. If the state is FinishedWithResult, we return the BackgroundTaskExecutionResult from it, finishing current task execution.

If retuned handler result does not match either Finished or FinishedWithResult cases - something went really wrong, so we return failure result from the task.

Summary

This task implementation will continuously