From 2632a0c5f32fda09bdabf0e1e43d0040c66f32a1 Mon Sep 17 00:00:00 2001
From: Nova <technobaboo@gmail.com>
Date: Sun, 14 Sep 2025 00:40:07 -0700
Subject: [PATCH] fix(wayland): proper frame callback handling

---
 src/wayland/core/surface.rs | 22 +++++++++++++---------
 src/wayland/mod.rs          | 20 +++++++++++---------
 src/wayland/xdg/toplevel.rs |  4 ++--
 3 files changed, 26 insertions(+), 20 deletions(-)

diff --git a/src/wayland/core/surface.rs b/src/wayland/core/surface.rs
index 12fe004..b76dae1 100644
--- a/src/wayland/core/surface.rs
+++ b/src/wayland/core/surface.rs
@@ -57,7 +57,7 @@ pub struct SurfaceState {
 	pub geometry: Option<Geometry>,
 	pub min_size: Option<Vector2<u32>>,
 	pub max_size: Option<Vector2<u32>>,
-	frame_callbacks: Vec<Arc<Callback>>,
+	frame_callbacks: Registry<Callback>,
 }
 impl Default for SurfaceState {
 	fn default() -> Self {
@@ -67,7 +67,7 @@ impl Default for SurfaceState {
 			geometry: None,
 			min_size: None,
 			max_size: None,
-			frame_callbacks: Vec::new(),
+			frame_callbacks: Registry::new(),
 		}
 	}
 }
@@ -90,6 +90,7 @@ pub struct Surface {
 	pub role: OnceLock<SurfaceRole>,
 	pub panel_item: Mutex<Weak<PanelItem<XdgBackend>>>,
 	on_commit_handlers: Mutex<Vec<OnCommitCallback>>,
+	frame_callbacks: Registry<Callback>,
 	material: OnceLock<Handle<BevyMaterial>>,
 	pending_material_applications: Registry<ModelPart>,
 	presentation_feedback: Mutex<Vec<Arc<PresentationFeedback>>>,
@@ -122,6 +123,7 @@ impl Surface {
 			role: OnceLock::new(),
 			panel_item: Mutex::new(Weak::default()),
 			on_commit_handlers: Mutex::new(Vec::new()),
+			frame_callbacks: Registry::new(),
 			material: OnceLock::new(),
 			pending_material_applications: Registry::new(),
 			presentation_feedback: Mutex::default(),
@@ -156,7 +158,7 @@ impl Surface {
 	}
 
 	#[tracing::instrument(level = "debug", skip_all)]
-	pub fn pending_state(&self) -> parking_lot::MutexGuard<'_, DoubleBuffer<SurfaceState>> {
+	pub fn state_lock(&self) -> parking_lot::MutexGuard<'_, DoubleBuffer<SurfaceState>> {
 		self.state.lock()
 	}
 
@@ -232,8 +234,9 @@ impl Surface {
 	}
 	#[tracing::instrument(level = "debug", skip_all)]
 	pub fn frame_event(&self) {
-		for callback in self.state.lock().current.frame_callbacks.drain(..) {
-			let _ = self.message_sink.send(Message::Frame(callback));
+		let callbacks = self.frame_callbacks.take_valid_contents();
+		if !callbacks.is_empty() {
+			let _ = self.message_sink.send(Message::Frame(callbacks));
 		}
 	}
 	// pub fn size(&self) -> Option<Vector2<u32>> {
@@ -366,7 +369,7 @@ impl WlSurface for Surface {
 		callback_id: ObjectId,
 	) -> Result<()> {
 		let callback = client.insert(callback_id, Callback(callback_id));
-		self.state.lock().pending.frame_callbacks.push(callback);
+		self.state.lock().pending.frame_callbacks.add_raw(&callback);
 		Ok(())
 	}
 
@@ -399,8 +402,7 @@ impl WlSurface for Surface {
 	async fn commit(&self, _client: &mut Client, _sender_id: ObjectId) -> Result<()> {
 		// we want the upload to complete before we give the image to bevy
 		let buffer_option = self
-			.state
-			.lock()
+			.state_lock()
 			.pending
 			.buffer
 			.as_ref()
@@ -411,7 +413,9 @@ impl WlSurface for Surface {
 				.unwrap();
 		}
 		self.state.lock().apply();
-		self.state.lock().pending.frame_callbacks.clear();
+		for callback in self.current_state().frame_callbacks.get_valid_contents() {
+			self.frame_callbacks.add_raw(&callback)
+		}
 		let current_state = self.current_state();
 		let mut handlers = self.on_commit_handlers.lock();
 		handlers.retain(|f| (f)(self, &current_state));
diff --git a/src/wayland/mod.rs b/src/wayland/mod.rs
index 8459a05..a612aff 100644
--- a/src/wayland/mod.rs
+++ b/src/wayland/mod.rs
@@ -119,7 +119,7 @@ pub fn get_free_wayland_socket_path() -> Option<(PathBuf, FlockLock<File>)> {
 
 pub enum Message {
 	Disconnect,
-	Frame(Arc<Callback>),
+	Frame(Vec<Arc<Callback>>),
 	ReleaseBuffer(Arc<Buffer>),
 	CloseToplevel(Arc<Toplevel>),
 	ResizeToplevel {
@@ -215,17 +215,19 @@ impl WaylandClient {
 	) -> Result<bool, waynest::server::Error> {
 		match message {
 			Message::Disconnect => return Ok(true),
-			Message::Frame(callback) => {
+			Message::Frame(callbacks) => {
 				let now = rustix::time::clock_gettime(rustix::time::ClockId::Monotonic);
 				let now = Duration::new(now.tv_sec as u64, now.tv_nsec as u32);
 				let ms = (now.as_millis() % (u32::MAX as u128)) as u32;
-				callback.done(client, callback.0, ms).await?;
-				client
-					.get::<Display>(ObjectId::DISPLAY)
-					.unwrap()
-					.delete_id(client, ObjectId::DISPLAY, callback.0.as_raw())
-					.await?;
-				client.remove(callback.0);
+				for callback in callbacks {
+					callback.done(client, callback.0, ms).await?;
+					client
+						.get::<Display>(ObjectId::DISPLAY)
+						.unwrap()
+						.delete_id(client, ObjectId::DISPLAY, callback.0.as_raw())
+						.await?;
+					client.remove(callback.0);
+				}
 			}
 			Message::ReleaseBuffer(buffer) => {
 				buffer.release(client, buffer.id).await?;
diff --git a/src/wayland/xdg/toplevel.rs b/src/wayland/xdg/toplevel.rs
index 7f4a496..efcf478 100644
--- a/src/wayland/xdg/toplevel.rs
+++ b/src/wayland/xdg/toplevel.rs
@@ -246,7 +246,7 @@ impl XdgToplevel for Toplevel {
 		width: i32,
 		height: i32,
 	) -> Result<()> {
-		self.wl_surface().pending_state().pending.max_size = if width == 0 && height == 0 {
+		self.wl_surface().state_lock().pending.max_size = if width == 0 && height == 0 {
 			None
 		} else {
 			Some([width as u32, height as u32].into())
@@ -261,7 +261,7 @@ impl XdgToplevel for Toplevel {
 		width: i32,
 		height: i32,
 	) -> Result<()> {
-		self.xdg_surface.wl_surface.pending_state().pending.min_size = if width == 0 && height == 0
+		self.xdg_surface.wl_surface.state_lock().pending.min_size = if width == 0 && height == 0
 		{
 			None
 		} else {