Wait for lock surfaces before locking

Fixes the red flash when locking.

1 files changed, 13 insertions, 6 deletions

diff --git a/src/handlers/compositor.rs b/src/handlers/compositor.rs
index 527f88b7..43a405ea 100644
--- a/src/handlers/compositor.rs
+++ b/src/handlers/compositor.rs
@@ -21,7 +21,7 @@ use smithay::{delegate_compositor, delegate_shm};
 use super::xdg_shell::add_mapped_toplevel_pre_commit_hook;
 use crate::handlers::XDG_ACTIVATION_TOKEN_TIMEOUT;
 use crate::layout::{ActivateWindow, AddWindowTarget};
-use crate::niri::{ClientState, State};
+use crate::niri::{ClientState, LockState, State};
 use crate::utils::transaction::Transaction;
 use crate::utils::{is_mapped, send_scale_transform};
 use crate::window::{InitialConfigureState, Mapped, ResolvedWindowRules, Unmapped};
@@ -408,16 +408,23 @@ impl CompositorHandler for State {
         }
 
         // This might be a lock surface.
-        if self.niri.is_locked() {
-            for (output, state) in &self.niri.output_state {
-                if let Some(lock_surface) = &state.lock_surface {
-                    if lock_surface.wl_surface() == &root_surface {
+        for (output, state) in &self.niri.output_state {
+            if let Some(lock_surface) = &state.lock_surface {
+                if lock_surface.wl_surface() == &root_surface {
+                    if matches!(self.niri.lock_state, LockState::WaitingForSurfaces { .. }) {
+                        self.niri.maybe_continue_to_locking();
+                    } else {
                         self.niri.queue_redraw(&output.clone());
-                        return;
                     }
+
+                    return;
                 }
             }
         }
+
+        // This message can trigger for lock surfaces that had a commit right after we unlocked
+        // the session, but that's ok, we don't need to handle them.
+        trace!("commit on an unrecognized surface: {surface:?}, root: {root_surface:?}");
     }
 
     fn destroyed(&mut self, surface: &WlSurface) {