Refactor frame scheduling

Combine the redraw state variables into one enum, and refactor to get
rid of the requirement that a VBlank must queue a subsequent redraw.
Also fix the bug where ongoing animations that produced no damage could
stall the redrawing.

4 files changed, 143 insertions, 44 deletions

diff --git a/src/backend/tty.rs b/src/backend/tty.rs
index 11595ef7..6f665432 100644
--- a/src/backend/tty.rs
+++ b/src/backend/tty.rs
@@ -1,5 +1,6 @@
 use std::cell::RefCell;
 use std::collections::{HashMap, HashSet};
+use std::mem;
 use std::os::fd::FromRawFd;
 use std::path::{Path, PathBuf};
 use std::rc::Rc;
@@ -38,7 +39,7 @@ use smithay_drm_extras::drm_scanner::{DrmScanEvent, DrmScanner};
 use smithay_drm_extras::edid::EdidInfo;
 
 use crate::config::Config;
-use crate::niri::{OutputRenderElements, State};
+use crate::niri::{OutputRenderElements, State, RedrawState};
 use crate::utils::get_monotonic_time;
 use crate::Niri;
 
@@ -703,34 +704,55 @@ impl Tty {
                 .plot(surface.sequence_delta_plot_name, delta);
         }
 
-        assert!(output_state.waiting_for_vblank);
-        assert!(output_state.estimated_vblank_timer.is_none());
-
-        output_state.waiting_for_vblank = false;
         output_state.frame_clock.presented(presentation_time);
         output_state.current_estimated_sequence = Some(meta.sequence);
-        niri.queue_redraw(output);
+
+        let redraw_needed = match mem::replace(&mut output_state.redraw_state, RedrawState::Idle) {
+            RedrawState::Idle => unreachable!(),
+            RedrawState::Queued(_) => unreachable!(),
+            RedrawState::WaitingForVBlank { redraw_needed } => redraw_needed,
+            RedrawState::WaitingForEstimatedVBlank(_) => unreachable!(),
+            RedrawState::WaitingForEstimatedVBlankAndQueued(_) => unreachable!(),
+        };
+
+        if redraw_needed || output_state.unfinished_animations_remain {
+            niri.queue_redraw(output);
+        } else {
+            niri.send_frame_callbacks(&output);
+        }
     }
 
-    fn on_estimated_vblank_timer(&self, niri: &mut Niri, output: &Output) {
+    fn on_estimated_vblank_timer(&self, niri: &mut Niri, output: Output) {
         let span = tracy_client::span!("Tty::on_estimated_vblank_timer");
 
         let name = output.name();
         span.emit_text(&name);
 
-        let Some(output_state) = niri.output_state.get_mut(output) else {
+        let Some(output_state) = niri.output_state.get_mut(&output) else {
             error!("missing output state for {name}");
             return;
         };
 
-        assert!(!output_state.waiting_for_vblank);
-        let token = output_state.estimated_vblank_timer.take();
-        assert!(token.is_some());
+        match mem::replace(&mut output_state.redraw_state, RedrawState::Idle) {
+            RedrawState::Idle => unreachable!(),
+            RedrawState::Queued(_) => unreachable!(),
+            RedrawState::WaitingForVBlank { .. } => unreachable!(),
+            RedrawState::WaitingForEstimatedVBlank(_) => (),
+            // The timer fired just in front of a redraw.
+            RedrawState::WaitingForEstimatedVBlankAndQueued((_, idle)) => {
+                output_state.redraw_state = RedrawState::Queued(idle);
+                return;
+            }
+        }
 
         if let Some(sequence) = output_state.current_estimated_sequence.as_mut() {
             *sequence = sequence.wrapping_add(1);
 
-            niri.send_frame_callbacks(output);
+            if output_state.unfinished_animations_remain {
+                niri.queue_redraw(output);
+            } else {
+                niri.send_frame_callbacks(&output);
+            }
         }
     }
 
@@ -793,10 +815,18 @@ impl Tty {
                     match drm_compositor.queue_frame(data) {
                         Ok(()) => {
                             let output_state = niri.output_state.get_mut(output).unwrap();
-                            output_state.waiting_for_vblank = true;
-                            if let Some(token) = output_state.estimated_vblank_timer.take() {
-                                niri.event_loop.remove(token);
-                            }
+                            let new_state = RedrawState::WaitingForVBlank {
+                                redraw_needed: false,
+                            };
+                            match mem::replace(&mut output_state.redraw_state, new_state) {
+                                RedrawState::Idle => unreachable!(),
+                                RedrawState::Queued(_) => (),
+                                RedrawState::WaitingForVBlank { .. } => unreachable!(),
+                                RedrawState::WaitingForEstimatedVBlank(_) => unreachable!(),
+                                RedrawState::WaitingForEstimatedVBlankAndQueued((token, _)) => {
+                                    niri.event_loop.remove(token);
+                                }
+                            };
 
                             return Some(&surface.dmabuf_feedback);
                         }
@@ -904,8 +934,15 @@ fn queue_estimated_vblank_timer(
     target_presentation_time: Duration,
 ) {
     let output_state = niri.output_state.get_mut(&output).unwrap();
-    if output_state.estimated_vblank_timer.is_some() {
-        return;
+    match mem::take(&mut output_state.redraw_state) {
+        RedrawState::Idle => unreachable!(),
+        RedrawState::Queued(_) => (),
+        RedrawState::WaitingForVBlank { .. } => unreachable!(),
+        RedrawState::WaitingForEstimatedVBlank(token)
+        | RedrawState::WaitingForEstimatedVBlankAndQueued((token, _)) => {
+            output_state.redraw_state = RedrawState::WaitingForEstimatedVBlank(token);
+            return;
+        }
     }
 
     let now = get_monotonic_time();
@@ -915,9 +952,9 @@ fn queue_estimated_vblank_timer(
         .insert_source(timer, move |_, _, data| {
             data.backend
                 .tty()
-                .on_estimated_vblank_timer(&mut data.niri, &output);
+                .on_estimated_vblank_timer(&mut data.niri, output.clone());
             TimeoutAction::Drop
         })
         .unwrap();
-    output_state.estimated_vblank_timer = Some(token);
+    output_state.redraw_state = RedrawState::WaitingForEstimatedVBlank(token);
 }
diff --git a/src/backend/winit.rs b/src/backend/winit.rs
index e0fbf272..7639b568 100644
--- a/src/backend/winit.rs
+++ b/src/backend/winit.rs
@@ -1,5 +1,6 @@
 use std::cell::RefCell;
 use std::collections::HashMap;
+use std::mem;
 use std::rc::Rc;
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
@@ -18,7 +19,7 @@ use smithay::utils::Transform;
 use smithay::wayland::dmabuf::DmabufFeedback;
 
 use crate::config::Config;
-use crate::niri::{OutputRenderElements, State};
+use crate::niri::{OutputRenderElements, RedrawState, State};
 use crate::utils::get_monotonic_time;
 use crate::Niri;
 
@@ -182,7 +183,18 @@ impl Winit {
                 0,
                 wp_presentation_feedback::Kind::empty(),
             );
+        }
+
+        let output_state = niri.output_state.get_mut(output).unwrap();
+        match mem::replace(&mut output_state.redraw_state, RedrawState::Idle) {
+            RedrawState::Idle => unreachable!(),
+            RedrawState::Queued(_) => (),
+            RedrawState::WaitingForVBlank { .. } => unreachable!(),
+            RedrawState::WaitingForEstimatedVBlank(_) => unreachable!(),
+            RedrawState::WaitingForEstimatedVBlankAndQueued(_) => unreachable!(),
+        }
 
+        if output_state.unfinished_animations_remain {
             self.backend.window().request_redraw();
         }
 
diff --git a/src/layout.rs b/src/layout.rs
index 4a272f9d..f878ed9e 100644
--- a/src/layout.rs
+++ b/src/layout.rs
@@ -1328,6 +1328,11 @@ impl<W: LayoutElement> Monitor<W> {
         }
     }
 
+    pub fn are_animations_ongoing(&self) -> bool {
+        self.workspace_idx_anim.is_some()
+            || self.workspaces.iter().any(|ws| ws.are_animations_ongoing())
+    }
+
     pub fn update_config(&mut self, config: &Config) {
         for ws in &mut self.workspaces {
             ws.update_config(config);
@@ -1464,6 +1469,10 @@ impl<W: LayoutElement> Workspace<W> {
         }
     }
 
+    pub fn are_animations_ongoing(&self) -> bool {
+        self.view_offset_anim.is_some()
+    }
+
     pub fn update_config(&mut self, config: &Config) {
         let c = &config.focus_ring;
         self.focus_ring.is_off = c.off;
diff --git a/src/niri.rs b/src/niri.rs
index 48761702..b2fe0c41 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -6,7 +6,7 @@ use std::process::Command;
 use std::rc::Rc;
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
-use std::{env, thread};
+use std::{env, mem, thread};
 
 use _server_decoration::server::org_kde_kwin_server_decoration_manager::Mode as KdeDecorationsMode;
 use anyhow::Context;
@@ -132,16 +132,10 @@ pub struct Niri {
 
 pub struct OutputState {
     pub global: GlobalId,
-    // Set if there's a redraw queued on the event loop. Reset in redraw() which means that you
-    // cannot queue more than one redraw at once.
-    pub queued_redraw: Option<Idle<'static>>,
-    // Set to `true` when the output was redrawn and is waiting for a VBlank. Upon VBlank a redraw
-    // will always be queued, so you cannot queue a redraw while waiting for a VBlank.
-    pub waiting_for_vblank: bool,
-    // When we did a redraw which did not result in a KMS frame getting queued, this is a timer
-    // that will fire at the estimated VBlank time.
-    pub estimated_vblank_timer: Option<RegistrationToken>,
     pub frame_clock: FrameClock,
+    pub redraw_state: RedrawState,
+    // After the last redraw, some ongoing animations still remain.
+    pub unfinished_animations_remain: bool,
     /// Estimated sequence currently displayed on this output.
     ///
     /// When a frame is presented on this output, this becomes the real sequence from the VBlank
@@ -153,6 +147,21 @@ pub struct OutputState {
     pub current_estimated_sequence: Option<u32>,
 }
 
+#[derive(Default)]
+pub enum RedrawState {
+    /// The compositor is idle.
+    #[default]
+    Idle,
+    /// A redraw is queued.
+    Queued(Idle<'static>),
+    /// We submitted a frame to the KMS and waiting for it to be presented.
+    WaitingForVBlank { redraw_needed: bool },
+    /// We did not submit anything to KMS and made a timer to fire at the estimated VBlank.
+    WaitingForEstimatedVBlank(RegistrationToken),
+    /// A redraw is queued on top of the above.
+    WaitingForEstimatedVBlankAndQueued((RegistrationToken, Idle<'static>)),
+}
+
 // Not related to the one in Smithay.
 //
 // This state keeps track of when a surface last received a frame callback.
@@ -733,9 +742,8 @@ impl Niri {
 
         let state = OutputState {
             global,
-            queued_redraw: None,
-            waiting_for_vblank: false,
-            estimated_vblank_timer: None,
+            redraw_state: RedrawState::Idle,
+            unfinished_animations_remain: false,
             frame_clock: FrameClock::new(refresh_interval),
             current_estimated_sequence: None,
         };
@@ -748,9 +756,16 @@ impl Niri {
         self.global_space.unmap_output(output);
         // FIXME: reposition outputs so they are adjacent.
 
-        let mut state = self.output_state.remove(output).unwrap();
-        if let Some(idle) = state.queued_redraw.take() {
-            idle.cancel();
+        let state = self.output_state.remove(output).unwrap();
+        match state.redraw_state {
+            RedrawState::Idle => (),
+            RedrawState::Queued(idle) => idle.cancel(),
+            RedrawState::WaitingForVBlank { .. } => (),
+            RedrawState::WaitingForEstimatedVBlank(token) => self.event_loop.remove(token),
+            RedrawState::WaitingForEstimatedVBlankAndQueued((token, idle)) => {
+                self.event_loop.remove(token);
+                idle.cancel();
+            }
         }
 
         // Disable the output global and remove some time later to give the clients some time to
@@ -917,15 +932,34 @@ impl Niri {
     /// Schedules an immediate redraw if one is not already scheduled.
     pub fn queue_redraw(&mut self, output: Output) {
         let state = self.output_state.get_mut(&output).unwrap();
+        let token = match mem::take(&mut state.redraw_state) {
+            RedrawState::Idle => None,
+            RedrawState::WaitingForEstimatedVBlank(token) => Some(token),
+
+            // A redraw is already queued, put it back and do nothing.
+            value @ (RedrawState::Queued(_)
+            | RedrawState::WaitingForEstimatedVBlankAndQueued(_)) => {
+                state.redraw_state = value;
+                return;
+            }
 
-        if state.queued_redraw.is_some() || state.waiting_for_vblank {
-            return;
-        }
+            // We're waiting for VBlank, request a redraw afterwards.
+            RedrawState::WaitingForVBlank { .. } => {
+                state.redraw_state = RedrawState::WaitingForVBlank {
+                    redraw_needed: true,
+                };
+                return;
+            }
+        };
 
         let idle = self.event_loop.insert_idle(move |state| {
             state.niri.redraw(&mut state.backend, &output);
         });
-        state.queued_redraw = Some(idle);
+
+        state.redraw_state = match token {
+            Some(token) => RedrawState::WaitingForEstimatedVBlankAndQueued((token, idle)),
+            None => RedrawState::Queued(idle),
+        };
     }
 
     pub fn pointer_element(
@@ -1070,14 +1104,21 @@ impl Niri {
         };
 
         let state = self.output_state.get_mut(output).unwrap();
-        let presentation_time = state.frame_clock.next_presentation_time();
+        assert!(matches!(
+            state.redraw_state,
+            RedrawState::Queued(_) | RedrawState::WaitingForEstimatedVBlankAndQueued(_)
+        ));
 
-        assert!(state.queued_redraw.take().is_some());
-        assert!(!state.waiting_for_vblank);
+        let presentation_time = state.frame_clock.next_presentation_time();
 
         // Update from the config and advance the animations.
         self.monitor_set.update_config(&self.config.borrow());
         self.monitor_set.advance_animations(presentation_time);
+        state.unfinished_animations_remain = self
+            .monitor_set
+            .monitor_for_output(output)
+            .unwrap()
+            .are_animations_ongoing();
 
         // Render the elements.
         let elements = self.render(renderer, output, true);