Structured Concurrency for User Controlled Scheduling

[hymm]

Structured Concurrency for User Controlled Scheduling

Made this a discussion instead of an rfc as this doesn't feel very scoped and has a lot of open design questions.

Motivation

This is a discussion to add some support for structured concurrency in Bevy. The main new feature is to add the ability to run schedules from inside of a system from a scope. This allows for scopes to be nested and is a main feature of structured concurrency.

Systems are the base unit of concurrency inside of Bevy. By allowing running a schedule inside of a system, we have the basics of structured concurrency, where scopes nested inside other scopes must be completed before the outer scope is completed. Structured concurrency makes it easier to reason about how tasks are being run. This should allow building better abstractions for Application State and Stages and give power to users to create their own scheduling abstractions.

Also now that schedules are run from a system other systems can be scheduled around the schedule running system with before and after labeling.

Running schedules from scopes

For this discussion we define a schedule to be a group of systems that are topologically ordered with before and after labeling. They will need to be run from inside of an async scope inside of an async system.

#[derive(Schedule)]
const ScheduleA = "Schedule A";

fn setup(mut app: App) {
    // we configure the schedule through app
    app
        // this adds ScheduleA as a resource that can only be accessed by a system
        // that is part of BaseSchedule
        .add_resource_to_schedule(BaseSchedule, ScheduleA)
        // we might need a different api to add an async system
        .add_async_system_to_schedule(BaseSchedule, run_schedule_a)
        // we then add systems to the schedule
        .add_system_to_schedule(ScheduleA, system1);
}

async fn run_schedule_a(
    // ScheduleRes gets a resource that lives on the current schedule
    schedule: ScheduleRes<ScheduleA>,
    // The compute task pool is a resource shared by all schedules
    task_pool: ComputeTaskPool,
    // do not mutate world, only use this to pass to a schedule.
    // the executor is responsible for making sure that threads don't
    // write to the same data
    world: *mut World,
) {
    
    // awaiting the scope, waits for all the systems in the schedule to
    // complete. 
    task_pool.scope(|scope| {
        scope.spawn(schedule.run(world));
    }).await;
}

Examples for AppState and Stages

We can use the structured concurrency primitives to build things like stages and app state. These examples are not full designs, but rough code to suggest what can be accomplished with the new scheduling primitives.

AppState

App State is essentially a group of schedules with some sugar for lifecycle events. In this model a state is directly tied to the systems that run in that state.

#[derive(AppState)]
enum State1 {
    A,
    B,
}

fn setup(mut app: App) {
    app
        // initialize the state resources and systems on world
        // states belong to a schedule
        // in this case we add the state to the default schedule
        .add_state_to_schedule(DefaultSchedule, State1::A)
        // add systems to state
        .add_system_to_schedule(State1::A.on_enter(), system1)
        .add_system_to_schedule(State1::A.on_update(), system2)
        .add_system_to_schedule(State1::B.on_update(), system3)
        .add_system_to_schedule(State1::B.on_exit(), system4);
}

// this system would be added by .add_state
// It runs until no more states are queued. States would most likely
// be queued from inside the app state as queueing from outside might not
// be seen depending on execution order.
async fn execute_app_state(
    // do not operate directly on world, only use to pass to state for system params
    world: *mut World,
    task_pool: ComputeTaskPool,
    // queue would be something like an async channel
    queue: ScheduleRes<StateQueue<State1>>,
    current_state: Local<State1>,
) {
    // queue is asynchrous to allow pushing new states inside the scope
    while let Some(next_state) = queue.pop().await {
        if current_state != next_state {
            task_pool.scope(|scope| {
                scope.spawn(current_state.on_exit(world));
            }).await;
            
            current_state = next_state;
            
            task_pool.scope(|scope| {
                scope.spawn(current_state.on_enter(world));
            }).await;
        }
        
        task_pool.scope(|scope| {
            scope.spawn(current_state.on_update(world));
        }).await;
    }
}

Stages

Stages are a ordered list of schedules.

#[derive(Stages)]
enum StagesList {
    A,
    B,
    C,
}

fn setup_stages(mut app: App) {
    app
        .add_stages(StagesList)
        .add_system_to_schedule(Stages:A.on_update, system1)
}

async fn execute_stages(
    world: *mut World,
    task_pool: ComputeTaskPool,
    stages: ScheduleRes<StateQueue<State1>>
) {
    for stage in stages.iter() {
        task_pool.scope(|scope| {
            scope.spawn(stage.run(world));
        }).await
    }
}

New Features Needed

Running a schedule from a system would require a decent number of new features. The features required are listed below, but do not have a concrete design.

Async scopes and async systems

We need async scopes and async systems so we can release the thread the system is running on while running the schedule.

A common executor or a shared data access struct

Systems running in different scopes need to coordinate so they don't have conflicting data accesses. This is currently done in the parallel executor of a stage. This logic would need to be lifted into an executor shared by all scopes or we could create a shared data access struct that each system would need to access before running. The shared struct would probably require a mutex or rwlock, but might lead to cleaner code. Would need to experiment. The schedules also need to coordinate running exclusive systems which would probably be done through access conflicts.

Shared Task Pool

The ComputeTaskPool needs to be shared by all scheduling systems without blocking other scheduling systems from running. The current implementation does not allow this.

Data owned by the schedule

Things like schedules cannot be owned by the world as the async scheduling needs to own them while it is running. We could possibly store this on the schedule somehow maybe inside of a sub world on the schedule?

Other changes

• rework schedule abstraction so it can be run inside a scope

Future Considerations

Allow scheduler systems to capture additional conflicts

async fn run_schedule_a(
    schedule: ScheduleRes<ScheduleA>,
    task_pool: ComputeTaskPool,
    world: *mut World,
    // marks access to component archetypes so that only systems that are part of the scope
    // can access. Other systems must wait until this system is done running
    // this does not access data, but just marks the data as unavailble for other systems
    access: QueryAccess<(&Player, &Transform)>
) {
    // we pass the access to the scope let systems spawned inside the scope to access the data
    task_pool.access_scope(access, |scope| {
        scope.spawn(schedule.run(world));
    }).await;
}

Cancellation

In this model schedules are always run to completion. In the future, it might make sense to add support for canceling schedules, but probably requires AsyncDrop support from rust.

Combinators for States

We'll want a system for schedules and app state that is compatible and logical with state combinators.

#[derive(AppState)]
enum State1 {
    A,
    B,
}

#[derive(AppState)]
enum State2 {
    C,
    D,
}

fn setup(mut app: App) {
    app
        .add_system_to_schedule(State1::A.and(State2::C), system1);
}

Drawbacks

• Not 100% sure, but this might end up adding more overhead to the executor.
• Before and after labeling doesn't work between schedules. This is natural as it is impossible to rely on dependencies that might not run.

Open Questions

• doing a fixed timestep would be relatively easy, but if we need to share the same timestep between different states that might be harder. Not sure how necessary that is might just as much sense to always nest the dependant states inside the fixed timestep.
• Giving access to users to *mut World is not ideal. We need this because we need to be able to give &mut World access from a schedule to multiple system task, and will prevent simultaneous access through the data access can run logic.
  • Does it work to pass a Pin<&mut World> instead of a *mut World? Pin only allows an &mut World through an unsafe function. Which we would use in the schedule run function or the system params. That feels like it would be safer.
  • can we keep the responsibility of passing &World and &mut World to the futures in the executor instead of a scheduling system?

Prior work for structured concurrency

• Structured Concurrency discussion for tokio
• Discussion comparing zio(scala) and trio(python)

[hymm]

To give some context to how I came to this design. I was poking around at the redesign for app states using side effect subgraphs. There were a lot of issues and structured concurrency cleaned things up in a very neat and tidy way.

1. An app state is a collection of states where each state when activated causes a set of systems to run. This collection is run mutually exclusively. i.e. two states from the same collection cannot be run at the same time.

   • In the side effect design this means we need to keep track of which systems from a specific state are running. We do this by making a final system that is dependant on all the systems finishing. The final system is then responsible for reading the state queue and starting new states running. This adds a bit of complexity in the executor as it needs to know to requeue the state producer until no new states are queued.
   • In structured concurrency we can just wait for the scheduling system to finish running.

2. The state producing and state consuming system want to be the same system, this leads to a dependancy cycle as we need to run the system both before and after a state is run.

   • In the subgraph design, we can break the dependancy with a soft edge that doesn't factor into calculating dependacy cycles.
   • Structured concurrency doesn't have this problem as there is just one system that loops internally.

3. State ordering has a higher priority than after/before ordering. If we have State1 with system_1 and State2 with system_2, we can add a system_2.after(system_1) dependancy, but run State2 before State1, so the after relation doesn't mean anything.

   • In both the subgraph design and structured concurrency design we can instead order around the state producing system. But this is an argument that systems that belong to a state don't actually belong to the schedule that spawns the state and so can be moved to their own independent schedule graph.

4. An addendum to 3. is that after/before labeling between systems in different states is ambiguous. Are you required to wait for the app state to be queued and then run or do you ignore the dependency because the system may never run.

   • In the first case the system should probably be part of the app state's systems.
   • In the second case you can probably order on the producing system instead, which you know will be guaranteed to be run.

5. Dealing with a dependency on a looping systems

   • For the side effect design it adds some complexity to the executor. It needs to not send finished to the dependants outside of the state until after the loop is finished, but send finished to systems inside of the loop each time it runs.
   • this is not a problem for stuctured concurrency as the schedule system can be treated like any other system.

Not sure if this was everything, but it's at least what I could reconstruct from my notes.