Introduce a WindowWrapper to extend the lifetime of the window when using pipelined rendering

Author: Friz64

crates/bevy_window/src/raw_handle.rs

@@ -5,6 +5,38 @@ use raw_window_handle::{
     DisplayHandle, HandleError, HasDisplayHandle, HasWindowHandle, RawDisplayHandle,
     RawWindowHandle, WindowHandle,
 };
+use std::{any::Any, marker::PhantomData, ops::Deref, sync::Arc};
+
+/// A wrapper over a window.
+///
+/// This allows us to extend the lifetime of the window, so it doesn't get eagerly dropped while a
+/// pipelined renderer still has frames in flight that need to draw to it.
+///
+/// This is achieved by storing a shared reference to the window in the [`RawHandleWrapper`],
+/// which gets picked up by the renderer during extraction.
+#[derive(Debug)]
+pub struct WindowWrapper<W> {
+    reference: Arc<dyn Any + Send + Sync>,
+    ty: PhantomData<W>,
+}
+
+impl<W: Send + Sync + 'static> WindowWrapper<W> {
+    /// Creates a `WindowWrapper` from a window.
+    pub fn new(window: W) -> WindowWrapper<W> {
+        WindowWrapper {
+            reference: Arc::new(window),
+            ty: PhantomData,
+        }
+    }
+}
+
+impl<W: 'static> Deref for WindowWrapper<W> {
+    type Target = W;
+
+    fn deref(&self) -> &Self::Target {
+        self.reference.downcast_ref::<W>().unwrap()
+    }
+}
 
 /// A wrapper over [`RawWindowHandle`] and [`RawDisplayHandle`] that allows us to safely pass it across threads.
 ///
@@ -13,13 +45,25 @@ use raw_window_handle::{
 /// thread-safe.
 #[derive(Debug, Clone, Component)]
 pub struct RawHandleWrapper {
+    _window: Arc<dyn Any + Send + Sync>,
     /// Raw handle to a window.
     pub window_handle: RawWindowHandle,
     /// Raw handle to the display server.
     pub display_handle: RawDisplayHandle,
 }
 
 impl RawHandleWrapper {
+    /// Creates a `RawHandleWrapper` from a `WindowWrapper`.
+    pub fn new<W: HasWindowHandle + HasDisplayHandle + 'static>(
+        window: &WindowWrapper<W>,
+    ) -> Result<RawHandleWrapper, HandleError> {
+        Ok(RawHandleWrapper {
+            _window: window.reference.clone(),
+            window_handle: window.window_handle()?.as_raw(),
+            display_handle: window.display_handle()?.as_raw(),
+        })
+    }
+
     /// Returns a [`HasWindowHandle`] + [`HasDisplayHandle`] impl, which exposes [`WindowHandle`] and [`DisplayHandle`].
     ///
     /// # Safety

crates/bevy_winit/src/system.rs

@@ -11,7 +11,6 @@ use bevy_window::{
     RawHandleWrapper, Window, WindowClosed, WindowCreated, WindowMode, WindowResized,
 };
 
-use raw_window_handle::{HasDisplayHandle, HasWindowHandle};
 use winit::{
     dpi::{LogicalPosition, LogicalSize, PhysicalPosition, PhysicalSize},
     event_loop::EventLoopWindowTarget,
@@ -75,10 +74,7 @@ pub fn create_windows<F: QueryFilter + 'static>(
             .set_scale_factor(winit_window.scale_factor() as f32);
         commands
             .entity(entity)
-            .insert(RawHandleWrapper {
-                window_handle: winit_window.window_handle().unwrap().as_raw(),
-                display_handle: winit_window.display_handle().unwrap().as_raw(),
-            })
+            .insert(RawHandleWrapper::new(winit_window).unwrap())
             .insert(CachedWindow {
                 window: window.clone(),
             });

crates/bevy_winit/src/winit_windows.rs

@@ -3,7 +3,9 @@ use bevy_ecs::entity::Entity;
 
 use bevy_ecs::entity::EntityHashMap;
 use bevy_utils::{tracing::warn, HashMap};
-use bevy_window::{CursorGrabMode, Window, WindowMode, WindowPosition, WindowResolution};
+use bevy_window::{
+    CursorGrabMode, Window, WindowMode, WindowPosition, WindowResolution, WindowWrapper,
+};
 
 use winit::{
     dpi::{LogicalSize, PhysicalPosition},
@@ -20,7 +22,7 @@ use crate::{
 #[derive(Debug, Default)]
 pub struct WinitWindows {
     /// Stores [`winit`] windows by window identifier.
-    pub windows: HashMap<winit::window::WindowId, winit::window::Window>,
+    pub windows: HashMap<winit::window::WindowId, WindowWrapper<winit::window::Window>>,
     /// Maps entities to `winit` window identifiers.
     pub entity_to_winit: EntityHashMap<winit::window::WindowId>,
     /// Maps `winit` window identifiers to entities.
@@ -41,7 +43,7 @@ impl WinitWindows {
         adapters: &mut AccessKitAdapters,
         handlers: &mut WinitActionHandlers,
         accessibility_requested: &AccessibilityRequested,
-    ) -> &winit::window::Window {
+    ) -> &WindowWrapper<winit::window::Window> {
         let mut winit_window_builder = winit::window::WindowBuilder::new();
 
         // Due to a UIA limitation, winit windows need to be invisible for the
@@ -240,12 +242,12 @@ impl WinitWindows {
 
         self.windows
             .entry(winit_window.id())
-            .insert(winit_window)
+            .insert(WindowWrapper::new(winit_window))
             .into_mut()
     }
 
     /// Get the winit window that is associated with our entity.
-    pub fn get_window(&self, entity: Entity) -> Option<&winit::window::Window> {
+    pub fn get_window(&self, entity: Entity) -> Option<&WindowWrapper<winit::window::Window>> {
         self.entity_to_winit
             .get(&entity)
             .and_then(|winit_id| self.windows.get(winit_id))
@@ -261,7 +263,10 @@ impl WinitWindows {
     /// Remove a window from winit.
     ///
     /// This should mostly just be called when the window is closing.
-    pub fn remove_window(&mut self, entity: Entity) -> Option<winit::window::Window> {
+    pub fn remove_window(
+        &mut self,
+        entity: Entity,
+    ) -> Option<WindowWrapper<winit::window::Window>> {
         let winit_id = self.entity_to_winit.remove(&entity)?;
         self.winit_to_entity.remove(&winit_id);
         self.windows.remove(&winit_id)