Allow systems using Diagnostics to run in parallel

[MJohnson459]

Objective

I was trying to add some Diagnostics to have a better break down of performance but I noticed that the current implementation uses a ResMut which forces the functions to all run sequentially whereas before they could run in parallel. This created too great a performance penalty to be usable.

Solution

This PR reworks how the diagnostics work with a couple of breaking changes. The idea is to change how Diagnostics works by changing it to a SystemParam. This allows us to hold a Deferred buffer of measurements that can be applied later, avoiding the need for multiple mutable references to the hashmap. This means we can run systems that write diagnostic measurements in parallel.

Firstly, we rename the old Diagnostics to DiagnosticsStore. This clears up the original name for the new interface while allowing us to preserve more closely the original API.

Then we create a new Diagnostics struct which implements SystemParam and contains a deferred SystemBuffer. This can be used very similar to the old Diagnostics for writing new measurements.

fn system(diagnostics: ResMut<Diagnostics>) { diagnostics.new_measurement(ID, || 10.0)}
// changes to
fn system(mut diagnostics: Diagnostics) { diagnostics.new_measurement(ID, || 10.0)}

For reading the diagnostics, the user needs to change from Diagnostics to DiagnosticsStore but otherwise the function calls are the same.

Finally, we add a new method to the App for registering diagnostics. This replaces the old method of creating a startup system and adding it manually.

Testing it, this PR does indeed allow Diagnostic systems to be run in parallel.

Changelog

• Change Diagnostics to implement SystemParam which allows diagnostic systems to run in parallel.

Migration Guide

• Register Diagnostic's using the new app.register_diagnostic(Diagnostic::new(DIAGNOSTIC_ID, "diagnostic_name", 10));
• In systems for writing new measurements, change mut diagnostics: ResMut<Diagnostics> to mut diagnostics: Diagnostics to allow the systems to run in parallel.
• In systems for reading measurements, change diagnostics: Res<Diagnostics> to diagnostics: Res<DiagnosticsStore>.

[nicopap]

Would it be possible using the Deferred system parameter over a crossbeam channel?

[MJohnson459]

Would it be possible using the Deferred system parameter over a crossbeam channel?

That seems like a nicer option if it works, I can test it out. I'm not sure about is how that will work with the Diagnostics container which contains a HashMap vs individual Diagnostic instances.

[MJohnson459]

@nicopap I'm not that familiar with Deferred so let me check I'm understanding correctly. Is this what you were thinking of?

Rather than have a system require ResMut<Diagnostics> the idea would be create a new DiagnosticsParam that implements SystemParam and has a Deferred buffer internally. (Baiscally following this example #6817).

A dummy implementation would be something like this (in actuality I would assume the structs would be merged with existing ones).

#[derive(SystemParam)]
struct DiagnosticsParam<'s, F>
where
    F: FnOnce() -> f64,
{
    queue: Deferred<'s, DiagnosticsBuffer<F>>,
}

impl<'s, F> DiagnosticsParam<'s, F>
where
    F: FnOnce() -> f64,
{
    pub fn add_measurement(&mut self, id: DiagnosticId, value: F) {
        self.queue.0.insert(id, value)
    }
}

#[derive(Default)]
struct DiagnosticsBuffer<F: FnOnce() -> f64>(StableHashMap<DiagnosticId, F>);

impl<F> SystemBuffer for DiagnosticsBuffer<F>
where
    F: FnOnce() -> f64,
{
    fn apply(
        &mut self,
        system_meta: &bevy_ecs::system::SystemMeta,
        world: &mut bevy_ecs::world::World,
    ) {
        let mut diagnostics = world.resource_mut::<Diagnostics>();
        for (id, value) in self.0.into_iter() {
            diagnostics.add_measurement(id, value)
        }
    }
}

This doesn't like the function pointer that is part of the add_measurement as it isn't Default, and the lifetime is vague. If it was possible to get a immutable version of Diagnostics we could do the measurement up front as we can check if it's active. That might simplify it a bit. Otherwise I'm not really sure how to implement this.

[nicopap]

You can get the Diagnostics immutably. Add a diagnostics: Res<'w, Diagnostics> field to your DiagnosticsParam systemparam, check out the docs: https://docs.rs/bevy/latest/bevy/ecs/system/trait.SystemParam.html.

edit: To be clear, you need to drop the F generic parameter on DiagnosticsParam, and put it on the add_measurement method (or make value: impl FnOnce() -> f64). Afterward, the struct definition should look like DiagnosticsParam<'w, 's>

[MJohnson459]

Ok, I updated the PR to use a Deferred SystemParam and I verified it does run in parallel. The user experience has been hurt a little as now there are two structs they need to care about. The Diagnostics struct is still required to register a new diagnostic but when recording a diagnostic its DiagnosticsParam (naming tbd). If we rework how the Diagnostics registration works we might be able to bring it back down to one.

Once we are happy with the code the documentation will also need a brush up.

@nicopap anything else I'm missing?

[MJohnson459]

Final update for now. To try and fix the "usability" I did a little renaming and added a new trait helper method. Now the API looks more like it used to for the write case, and only differs for the read.

fn main() {
    App::new()
        ...
        .register_diagnostic(Diagnostic::new(SYSTEM_ITERATION_COUNT, "system_iteration_count", 10))
        .add_systems(Update, my_system)
        .run();
}

fn my_system(mut diagnostics: Diagnostics) {
    // Add a measurement of 10.0 for our diagnostic each time this system runs.
    diagnostics.add_measurement(SYSTEM_ITERATION_COUNT, || 10.0);
}

To get this I renamed Diagnostics to DiagnosticsStore and then added the RegisterDiagnostic trait to add a method to the App. This allows users to register the diagnostic like other bevy items rather than requiring a startup system while also hiding the name change of Diagnostics.

The DiagnosticsParam was then renamed to take the place of Diagnostics so the write system looks very similar to what it used to. The trade off here is that systems that read the Diagnostics data will need to change to read from DiagnosticsStore but as the methods are the same this shouldn't be too strenuous.

Note that I added the register_diagnostic function as a trait because I've used that pattern for third party plugins. In this case it would be equally feasible to add it directly to the App similar to register_type which a feature flag for bevy_diagnostics. Is there a preference one way or the other?

[nicopap]

Really nice work, just a small question on Diagnostic::add_measurement

[nicopap]

I'd prefer if this took a DiagnosticMeasurement over an Into<DiagnosticMeasurement>. Or even be made private. It's breaking, but using this method directly is footgunny, so better hint through the type that you aren't supposed to add the measurement through this API.

[MJohnson459]

I removed the Into<> and just use the DiagnosticMeasurement directly. I didn't make it private just because it seems like some people might be using the Diagnostic struct without Diagnostics and I erred on the side of reducing API changes. I'm happy to change it to private if preferred.

[nicopap]

I trust your judgment! I think It's fine to leave it public, especially with the new doc infos you added.

[nicopap]

Great work! I'd like to see this merged.

[alice-i-cecile]

Well documented, and a nice use of Deferred :)