Handle dmabuf blocker separately in toplevel pre-commit

Will be needed for transactions.

3 files changed, 107 insertions, 43 deletions

diff --git a/src/handlers/compositor.rs b/src/handlers/compositor.rs
index 583e88bd..140d00e0 100644
--- a/src/handlers/compositor.rs
+++ b/src/handlers/compositor.rs
@@ -8,8 +8,9 @@ use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
 use smithay::reexports::wayland_server::{Client, Resource};
 use smithay::wayland::buffer::BufferHandler;
 use smithay::wayland::compositor::{
-    add_blocker, add_pre_commit_hook, get_parent, is_sync_subsurface, with_states,
-    BufferAssignment, CompositorClientState, CompositorHandler, CompositorState, SurfaceAttributes,
+    add_blocker, add_pre_commit_hook, get_parent, is_sync_subsurface, remove_pre_commit_hook,
+    with_states, BufferAssignment, CompositorClientState, CompositorHandler, CompositorState,
+    SurfaceAttributes,
 };
 use smithay::wayland::dmabuf::get_dmabuf;
 use smithay::wayland::shell::xdg::XdgToplevelSurfaceData;
@@ -46,40 +47,7 @@ impl CompositorHandler for State {
     }
 
     fn new_surface(&mut self, surface: &WlSurface) {
-        add_pre_commit_hook::<Self, _>(surface, move |state, _dh, surface| {
-            let maybe_dmabuf = with_states(surface, |surface_data| {
-                surface_data
-                    .cached_state
-                    .get::<SurfaceAttributes>()
-                    .pending()
-                    .buffer
-                    .as_ref()
-                    .and_then(|assignment| match assignment {
-                        BufferAssignment::NewBuffer(buffer) => get_dmabuf(buffer).cloned().ok(),
-                        _ => None,
-                    })
-            });
-            if let Some(dmabuf) = maybe_dmabuf {
-                if let Ok((blocker, source)) = dmabuf.generate_blocker(Interest::READ) {
-                    if let Some(client) = surface.client() {
-                        let res =
-                            state
-                                .niri
-                                .event_loop
-                                .insert_source(source, move |_, _, state| {
-                                    let display_handle = state.niri.display_handle.clone();
-                                    state
-                                        .client_compositor_state(&client)
-                                        .blocker_cleared(state, &display_handle);
-                                    Ok(())
-                                });
-                        if res.is_ok() {
-                            add_blocker(surface, blocker);
-                        }
-                    }
-                }
-            }
-        });
+        self.add_default_dmabuf_pre_commit_hook(surface);
     }
 
     fn commit(&mut self, surface: &WlSurface) {
@@ -157,6 +125,8 @@ impl CompositorHandler for State {
                         })
                         .map(|(mapped, _)| mapped.window.clone());
 
+                    // The mapped pre-commit hook deals with dma-bufs on its own.
+                    self.remove_default_dmabuf_pre_commit_hook(toplevel.wl_surface());
                     let hook = add_mapped_toplevel_pre_commit_hook(toplevel);
                     let mapped = Mapped::new(window, rules, hook);
                     let window = mapped.window.clone();
@@ -246,6 +216,7 @@ impl CompositorHandler for State {
                         });
 
                     self.niri.layout.remove_window(&window);
+                    self.add_default_dmabuf_pre_commit_hook(surface);
 
                     if was_active {
                         self.maybe_warp_cursor_to_focus();
@@ -356,6 +327,8 @@ impl CompositorHandler for State {
         self.niri
             .root_surface
             .retain(|k, v| k != surface && v != surface);
+
+        self.niri.dmabuf_pre_commit_hook.remove(surface);
     }
 }
 
@@ -371,3 +344,57 @@ impl ShmHandler for State {
 
 delegate_compositor!(State);
 delegate_shm!(State);
+
+impl State {
+    pub fn add_default_dmabuf_pre_commit_hook(&mut self, surface: &WlSurface) {
+        let hook = add_pre_commit_hook::<Self, _>(surface, move |state, _dh, surface| {
+            let maybe_dmabuf = with_states(surface, |surface_data| {
+                surface_data
+                    .cached_state
+                    .get::<SurfaceAttributes>()
+                    .pending()
+                    .buffer
+                    .as_ref()
+                    .and_then(|assignment| match assignment {
+                        BufferAssignment::NewBuffer(buffer) => get_dmabuf(buffer).cloned().ok(),
+                        _ => None,
+                    })
+            });
+            if let Some(dmabuf) = maybe_dmabuf {
+                if let Ok((blocker, source)) = dmabuf.generate_blocker(Interest::READ) {
+                    if let Some(client) = surface.client() {
+                        let res =
+                            state
+                                .niri
+                                .event_loop
+                                .insert_source(source, move |_, _, state| {
+                                    let display_handle = state.niri.display_handle.clone();
+                                    state
+                                        .client_compositor_state(&client)
+                                        .blocker_cleared(state, &display_handle);
+                                    Ok(())
+                                });
+                        if res.is_ok() {
+                            add_blocker(surface, blocker);
+                            trace!("added default dmabuf blocker");
+                        }
+                    }
+                }
+            }
+        });
+
+        let s = surface.clone();
+        if let Some(prev) = self.niri.dmabuf_pre_commit_hook.insert(s, hook) {
+            error!("tried to add dmabuf pre-commit hook when there was already one");
+            remove_pre_commit_hook(surface, prev);
+        }
+    }
+
+    pub fn remove_default_dmabuf_pre_commit_hook(&mut self, surface: &WlSurface) {
+        if let Some(hook) = self.niri.dmabuf_pre_commit_hook.remove(surface) {
+            remove_pre_commit_hook(surface, hook);
+        } else {
+            error!("tried to remove dmabuf pre-commit hook but there was none");
+        }
+    }
+}
diff --git a/src/handlers/xdg_shell.rs b/src/handlers/xdg_shell.rs
index 4dc353b5..5f84da3b 100644
--- a/src/handlers/xdg_shell.rs
+++ b/src/handlers/xdg_shell.rs
@@ -1,5 +1,6 @@
 use std::cell::Cell;
 
+use calloop::Interest;
 use smithay::desktop::{
     find_popup_root_surface, get_popup_toplevel_coords, layer_map_for_output, utils, LayerSurface,
     PopupKeyboardGrab, PopupKind, PopupManager, PopupPointerGrab, PopupUngrabStrategy, Window,
@@ -17,8 +18,10 @@ use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
 use smithay::reexports::wayland_server::{self, Resource, WEnum};
 use smithay::utils::{Logical, Rectangle, Serial};
 use smithay::wayland::compositor::{
-    add_pre_commit_hook, with_states, BufferAssignment, HookId, SurfaceAttributes,
+    add_blocker, add_pre_commit_hook, with_states, BufferAssignment, CompositorHandler as _,
+    HookId, SurfaceAttributes,
 };
+use smithay::wayland::dmabuf::get_dmabuf;
 use smithay::wayland::input_method::InputMethodSeat;
 use smithay::wayland::shell::kde::decoration::{KdeDecorationHandler, KdeDecorationState};
 use smithay::wayland::shell::wlr_layer::{self, Layer};
@@ -491,6 +494,7 @@ impl XdgShellHandler for State {
         let was_active = active_window == Some(&window);
 
         self.niri.layout.remove_window(&window);
+        self.add_default_dmabuf_pre_commit_hook(surface.wl_surface());
 
         if was_active {
             self.maybe_warp_cursor_to_focus();
@@ -1005,10 +1009,17 @@ pub fn add_mapped_toplevel_pre_commit_hook(toplevel: &ToplevelSurface) -> HookId
             return;
         };
 
-        let (got_unmapped, commit_serial) = with_states(surface, |states| {
-            let got_unmapped = {
+        let (got_unmapped, dmabuf, commit_serial) = with_states(surface, |states| {
+            let (got_unmapped, dmabuf) = {
                 let mut guard = states.cached_state.get::<SurfaceAttributes>();
-                matches!(guard.pending().buffer, Some(BufferAssignment::Removed))
+                match guard.pending().buffer.as_ref() {
+                    Some(BufferAssignment::NewBuffer(buffer)) => {
+                        let dmabuf = get_dmabuf(buffer).cloned().ok();
+                        (false, dmabuf)
+                    }
+                    Some(BufferAssignment::Removed) => (true, None),
+                    None => (false, None),
+                }
             };
 
             let role = states
@@ -1018,9 +1029,12 @@ pub fn add_mapped_toplevel_pre_commit_hook(toplevel: &ToplevelSurface) -> HookId
                 .lock()
                 .unwrap();
 
-            (got_unmapped, role.configure_serial)
+            (got_unmapped, dmabuf, role.configure_serial)
         });
 
+        let mut dmabuf_blocker =
+            dmabuf.and_then(|dmabuf| dmabuf.generate_blocker(Interest::READ).ok());
+
         let animate = if let Some(serial) = commit_serial {
             mapped.should_animate_commit(serial)
         } else {
@@ -1028,6 +1042,25 @@ pub fn add_mapped_toplevel_pre_commit_hook(toplevel: &ToplevelSurface) -> HookId
             false
         };
 
+        if let Some((blocker, source)) = dmabuf_blocker.take() {
+            if let Some(client) = surface.client() {
+                let res = state
+                    .niri
+                    .event_loop
+                    .insert_source(source, move |_, _, state| {
+                        let display_handle = state.niri.display_handle.clone();
+                        state
+                            .client_compositor_state(&client)
+                            .blocker_cleared(state, &display_handle);
+                        Ok(())
+                    });
+                if res.is_ok() {
+                    add_blocker(surface, blocker);
+                    trace!("added toplevel dmabuf blocker");
+                }
+            }
+        }
+
         let window = mapped.window.clone();
         if got_unmapped {
             state.backend.with_primary_renderer(|renderer| {
diff --git a/src/niri.rs b/src/niri.rs
index ac07f113..7a79c0f4 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -67,8 +67,8 @@ use smithay::utils::{
     Transform, SERIAL_COUNTER,
 };
 use smithay::wayland::compositor::{
-    with_states, with_surface_tree_downward, CompositorClientState, CompositorState, SurfaceData,
-    TraversalAction,
+    with_states, with_surface_tree_downward, CompositorClientState, CompositorState, HookId,
+    SurfaceData, TraversalAction,
 };
 use smithay::wayland::cursor_shape::CursorShapeManagerState;
 use smithay::wayland::dmabuf::DmabufState;
@@ -189,6 +189,9 @@ pub struct Niri {
     // normal get_parent() is cleared out.
     pub root_surface: HashMap<WlSurface, WlSurface>,
 
+    // Dmabuf readiness pre-commit hook for a surface.
+    pub dmabuf_pre_commit_hook: HashMap<WlSurface, HookId>,
+
     pub output_state: HashMap<Output, OutputState>,
     pub output_by_name: HashMap<String, Output>,
 
@@ -1732,6 +1735,7 @@ impl Niri {
             output_by_name: HashMap::new(),
             unmapped_windows: HashMap::new(),
             root_surface: HashMap::new(),
+            dmabuf_pre_commit_hook: HashMap::new(),
             monitors_active: true,
 
             devices: HashSet::new(),