How to borrow WINIT_WINDOWS without panicking

[TeamDman]

I am on a side quest trying to port some code from another project that used Bevy 0.16.0.
The code is responsible for setting the icons for individual windows.

Here, we see an example where the icon for the window itself is unset.

The code for Bevy 0.16.0 used windows: NonSend<WinitWindows>, to get access to the WindowWrapper<Window> corresponding to the Entity of a window.

I defined a custom component for making it easy to set the icon for a window.

#[derive(Debug, Component, Reflect)]
pub struct WindowIcon {
    pub image: Handle<Image>,
}

https://github.com/TeamDman/youre-muted-btw/blob/4783e1efe184083b053882cc066c9932588551c5/crates/window_icon_plugin/src/lib.rs

In Bevy 0.17.0-rc.1, the NonSend<WinitWindows> has been replaced with a thread_local! WINIT_WINDOWS

#19575

I have since rewritten my logic for setting the custom window icon

https://github.com/TeamDman/teamy-mft/blob/2a073c6101a1f6ac8060aa60ce127a64661a5127/src/engine/window_icon_plugin.rs

use bevy::app::MainScheduleOrder;
use bevy::ecs::schedule::ExecutorKind;
use bevy::ecs::schedule::ScheduleLabel;
use bevy::ecs::system::NonSendMarker;
use bevy::prelude::*;
use bevy_winit::WINIT_WINDOWS;
use std::any::type_name;

pub struct WindowIconPlugin;

impl Plugin for WindowIconPlugin {
    fn build(&self, app: &mut App) {
        app.register_type::<WindowIcon>();
        app.add_message::<AddWindowIconWithRetry>();
        app.add_systems(Update, emit_window_icon_requests);

        // Set up a custom schedule to run the window icon update as an exclusive system.
        // We want to avoid `RefCell already borrowed` issues.
        let mut custom_update_schedule = Schedule::new(UpdateWindowIconsSchedule);
        custom_update_schedule.set_executor_kind(ExecutorKind::SingleThreaded);
        app.add_schedule(custom_update_schedule);

        let mut main_schedule_order = app.world_mut().resource_mut::<MainScheduleOrder>();
        main_schedule_order.insert_after(Last, UpdateWindowIconsSchedule);

        app.add_systems(UpdateWindowIconsSchedule, set_window_icon);
    }
}

#[derive(ScheduleLabel, Debug, Hash, PartialEq, Eq, Clone)]
struct UpdateWindowIconsSchedule;

#[derive(Debug, Component, Reflect)]
pub struct WindowIcon {
    pub image: Handle<Image>,
}

#[derive(Debug, Clone, Reflect, Message)]
pub struct AddWindowIconWithRetry {
    pub image: Handle<Image>,
    pub window: Entity,
}

impl WindowIcon {
    pub fn new(image: Handle<Image>) -> Self {
        Self { image }
    }
}

// Sends an AddWindowIconWithRetry message for any entity that just received WindowIcon
fn emit_window_icon_requests(
    mut q: Query<(Entity, &WindowIcon), Added<WindowIcon>>,
    mut writer: MessageWriter<AddWindowIconWithRetry>,
) {
    for (entity, icon) in &mut q {
        debug!("Detected new WindowIcon on {:?}", entity);
        writer.write(AddWindowIconWithRetry {
            image: icon.image.clone(),
            window: entity,
        });
    }
}

fn set_window_icon(
    mut events: ParamSet<(
        MessageReader<AddWindowIconWithRetry>,
        MessageWriter<AddWindowIconWithRetry>,
    )>,
    assets: Res<Assets<Image>>,
    _main_thread_marker: NonSendMarker,
) {
    if events.p0().is_empty() {
        return;
    }
    // WINIT_WINDOWS.with_borrow_mut(|windows| { info!("Ahoy!") });
    let outgoing = WINIT_WINDOWS.with_borrow_mut(|windows| {
        let mut outgoing = Vec::new();
        for event in events.p0().read() {
            debug!(?event.window, ?windows, "Handling {}", type_name::<AddWindowIconWithRetry>());

            // Identify window
            let target_window = windows.get_window(event.window);
            let Some(window) = target_window else {
                warn!(
                    ?windows,
                    "Window {:?} does not exist, retrying later...",
                    event.window
                );
                outgoing.push(event.clone());
                continue;
            };

            // Fetch the image asset
            let Some(image) = assets.get(&event.image) else {
                error!(
                    "Image handle {:?} not found in assets, the window will not have our custom icon",
                    event.image
                );
                continue;
            };

            // Acquire pixel data from the image
            let Some(image_data) = image.data.clone() else {
                error!(
                    "Image handle {:?} has no data, the window will not have our custom icon",
                    event.image
                );
                continue;
            };

            // Convert between formats
            let icon = match winit::window::Icon::from_rgba(
                image_data,
                image.texture_descriptor.size.width,
                image.texture_descriptor.size.height,
            ) {
                Ok(icon) => icon,
                Err(e) => {
                    error!("Failed to construct window icon: {:?}", e);
                    continue;
                }
            };

            // Set the window icon
            info!(image_size = ?image.size(), "Setting window icon");
            window.set_window_icon(Some(icon));
        }
        outgoing
    });

    for event in outgoing {
        events.p1().write(event);
    }
}

However, this panics.

The application panicked (crashed).
Message:  RefCell already borrowed
Location: G:\Programming\Caches\RUSTUP_HOME\toolchains\nightly-x86_64-pc-windows-msvc\lib/rustlib/src/rust\library\std\src\thread\local.rs:559

Simplifying the system, it still panicks.

fn set_window_icon(
    mut events: ParamSet<(
        MessageReader<AddWindowIconWithRetry>,
        MessageWriter<AddWindowIconWithRetry>,
    )>,
    assets: Res<Assets<Image>>,
    _main_thread_marker: NonSendMarker,
) {
    if events.p0().is_empty() {
        return;
    }
    WINIT_WINDOWS.with_borrow_mut(|windows| { info!("Ahoy!") });
}

unless I disable the system all together, which doesn't really help lol

//        app.add_systems(UpdateWindowIconsSchedule, set_window_icon); // disabled

The migration guide mentions

bevy/release-content/migration-guides/replace_non_send_resources.md

Line 48 in 2451112

If a borrow is attempted while the data is borrowed elsewhere, the method will panic.

which makes sense.

Because it is a thread-local variable, I've included _main_thread_marker: NonSendMarker, in the system to ensure we get the correct WINIT_WINDOWS.

Because it is a RefCell, to try and guarantee non-contentious access, I've tried adding my own schedule which to my understanding means that nothing else will run at the same time.

I've tried commenting out the set_executor_kind in my custom schedule

// custom_update_schedule.set_executor_kind(ExecutorKind::SingleThreaded);

I've tried running in Last

-        app.add_systems(UpdateWindowIconsSchedule, set_window_icon);
+        app.add_systems(Last, set_window_icon);

From what I can tell, the Bevy logic is running in Last

bevy/crates/bevy_winit/src/lib.rs

Lines 138 to 149 in 2451112

	.add_systems(
	Last,
	(
	// `exit_on_all_closed` only checks if windows exist but doesn't access data,
	// so we don't need to care about its ordering relative to `changed_windows`
	changed_windows.ambiguous_with(exit_on_all_closed),
	changed_cursor_options,
	despawn_windows,
	check_keyboard_focus_lost,
	)
	.chain(),
	);

and it immediately borrows WINIT_WINDOWS

bevy/crates/bevy_winit/src/system.rs

Lines 569 to 581 in 2451112

	pub(crate) fn changed_cursor_options(
	mut changed_windows: Query<
	(
	Entity,
	&Window,
	&mut CursorOptions,
	&mut CachedCursorOptions,
	),
	Changed<CursorOptions>,
	>,
	_non_send_marker: NonSendMarker,
	) {
	WINIT_WINDOWS.with_borrow(|winit_windows| {

...

And it is at this point that I realized the problem and thus the fix

-    let outgoing = WINIT_WINDOWS.with_borrow_mut(|windows| {
+    let outgoing = WINIT_WINDOWS.with_borrow(|windows| {

I don't need a mutable borrow at all!

All is well now :D

[TeamDman]

The real problem was that I was using with_borrow_mut when with_borrow is sufficient.