Memory leak in `EntitySpecializationTicks` and `SpecializedMaterialPipelineCache` when a material uses `NotShadowCaster`

[EmbersArc]

Bevy version and features

0.17.2

What you did

This came out of a longer debugging session I "documented" in a bluesky thread.

EDIT: This has been fixed in #21410. See below for a case that still leaks memory.

//! Example to reproduce memory leak issue.
//!
use bevy::{
    pbr::{ExtendedMaterial, MaterialExtension},
    prelude::*,
};
use bevy_render::render_resource::AsBindGroup;

fn main() {
    App::new()
        .add_plugins(DefaultPlugins)
        .add_plugins(MaterialPlugin::<MyExtendedMaterial1>::default())
        // Issue goes away when only a single material is present.
        .add_plugins(MaterialPlugin::<MyExtendedMaterial2>::default())
        .init_resource::<Handles>()
        .add_systems(Startup, setup)
        .add_systems(Update, update)
        .run();
}

#[derive(Resource)]
struct Handles(Handle<MyExtendedMaterial1>, Handle<Mesh>);

type MyExtendedMaterial1 = ExtendedMaterial<StandardMaterial, MyExtension1>;

#[derive(Default, Clone, AsBindGroup, Asset, TypePath)]
struct MyExtension1 {}

impl MaterialExtension for MyExtension1 {}

type MyExtendedMaterial2 = ExtendedMaterial<StandardMaterial, MyExtension2>;

#[derive(Default, Clone, AsBindGroup, Asset, TypePath)]
struct MyExtension2 {}

impl MaterialExtension for MyExtension2 {}

impl FromWorld for Handles {
    fn from_world(world: &mut World) -> Self {
        let material1_handle = world
            .resource_mut::<Assets<MyExtendedMaterial1>>()
            .add(MyExtendedMaterial1::default());
        let mesh_handle = world.resource_mut::<Assets<Mesh>>().add(Cuboid::default());
        Self(material1_handle, mesh_handle)
    }
}

fn update(
    mut commands: Commands,
    existing_meshes: Query<Entity, With<MeshMaterial3d<MyExtendedMaterial1>>>,
    handles: Res<Handles>,
) {
    dbg!(existing_meshes.count());

    const COUNT: usize = 100;

    for _ in 0..COUNT {
        commands.spawn((Mesh3d(handles.1.clone()), MeshMaterial3d(handles.0.clone())));
    }

    existing_meshes
        .iter()
        .take(COUNT)
        .for_each(|entity| commands.entity(entity).despawn());
}

fn setup(mut commands: Commands) {
    commands.spawn((
        Camera3d::default(),
        Transform::from_xyz(0.0, 7., 14.0).looking_at(Vec3::new(0., 1., 0.), Vec3::Y),
    ));
}

What went wrong

EntitySpecializationTicks and SpecializedMaterialPipelineCache grow with each set of spawned/despawned entities, eventually slowing the systems that use them to a crawl.

Running extract_entities_needs_specialization explicitly after early_sweep_material_instances fixes the leaks.

However

In the presence of entities with a NotShadowCaster component, the EntitySpecializationTicks keep growing:

//! Example to reproduce memory leak issue.
//!
use bevy::{
    light::NotShadowCaster,
    pbr::{ExtendedMaterial, MaterialExtension},
    prelude::*,
};
use bevy_render::render_resource::AsBindGroup;

fn main() {
    App::new()
        .add_plugins(DefaultPlugins)
        .add_plugins(MaterialPlugin::<MyExtendedMaterial1>::default())
        // Issue goes away when only a single material is present.
        .add_plugins(MaterialPlugin::<MyExtendedMaterial2>::default())
        .init_resource::<Handles>()
        .add_systems(Startup, setup)
        .add_systems(Update, update)
        .run();
}

#[derive(Resource)]
struct Handles(
    Handle<MyExtendedMaterial1>,
    Handle<MyExtendedMaterial2>,
    Handle<Mesh>,
);

type MyExtendedMaterial1 = ExtendedMaterial<StandardMaterial, MyExtension1>;

#[derive(Default, Clone, AsBindGroup, Asset, TypePath)]
struct MyExtension1 {}

impl MaterialExtension for MyExtension1 {}

type MyExtendedMaterial2 = ExtendedMaterial<StandardMaterial, MyExtension2>;

#[derive(Default, Clone, AsBindGroup, Asset, TypePath)]
struct MyExtension2 {}

impl MaterialExtension for MyExtension2 {}

impl FromWorld for Handles {
    fn from_world(world: &mut World) -> Self {
        let material1_handle = world
            .resource_mut::<Assets<MyExtendedMaterial1>>()
            .add(MyExtendedMaterial1::default());
        let material2_handle = world
            .resource_mut::<Assets<MyExtendedMaterial2>>()
            .add(MyExtendedMaterial2::default());
        let mesh_handle = world.resource_mut::<Assets<Mesh>>().add(Cuboid::default());
        Self(material1_handle, material2_handle, mesh_handle)
    }
}

fn update(
    mut commands: Commands,
    existing_meshes1: Query<Entity, With<MeshMaterial3d<MyExtendedMaterial1>>>,
    existing_meshes2: Query<Entity, With<MeshMaterial3d<MyExtendedMaterial2>>>,
    handles: Res<Handles>,
) {
    dbg!(existing_meshes1.count());
    dbg!(existing_meshes2.count());

    const COUNT: usize = 100;

    for _ in 0..COUNT {
        commands.spawn((Mesh3d(handles.2.clone()), MeshMaterial3d(handles.0.clone())));
        commands.spawn((
            Mesh3d(handles.2.clone()),
            MeshMaterial3d(handles.1.clone()),
            NotShadowCaster,
        ));
    }

    existing_meshes1
        .iter()
        .take(COUNT)
        .for_each(|entity| commands.entity(entity).despawn());
    existing_meshes2
        .iter()
        .take(COUNT)
        .for_each(|entity| commands.entity(entity).despawn());
}

fn setup(mut commands: Commands) {
    commands.spawn((
        Camera3d::default(),
        Transform::from_xyz(0.0, 7., 14.0).looking_at(Vec3::new(0., 1., 0.), Vec3::Y),
    ));
}

Additional context

Related: #21410

[alice-i-cecile]

Running extract_entities_needs_specialization explicitly after early_sweep_material_instances fixes the leaks.

I'd welcome a fix to add that system ordering constraint as a first step.

[EmbersArc]

With #21410 the first part looks like it's fixed. Now only the second part relating to NotShadowCaster components persists.