fix(wayland): proper shm double buffering

src/wayland/core/buffer.rs

@@ -1,10 +1,12 @@
-use std::sync::Arc;
+use std::sync::{Arc, atomic::AtomicBool};
 
 #[cfg(feature = "dmabuf")]
 use crate::wayland::dmabuf::buffer_backing::DmabufBacking;
 use crate::{
 	core::registry::Registry,
-	wayland::{GraphicsInfo, core::shm_buffer_backing::ShmBufferBacking},
+	wayland::{
+		GraphicsInfo, MessageSink, core::shm_buffer_backing::ShmBufferBacking, util::ClientExt,
+	},
 };
 use bevy::{
 	asset::{Assets, Handle},
@@ -36,11 +38,21 @@ impl BufferBacking {
 pub struct Buffer {
 	pub id: ObjectId,
 	backing: BufferBacking,
+	pub message_sink: MessageSink,
+	pub rendered: AtomicBool,
 }
 
 impl Buffer {
 	pub fn new(client: &mut Client, id: ObjectId, backing: BufferBacking) -> Arc<Self> {
-		let buffer = client.insert(id, Self { id, backing });
+		let buffer = client.insert(
+			id,
+			Self {
+				id,
+				backing,
+				message_sink: client.message_sink(),
+				rendered: AtomicBool::new(false),
+			},
+		);
 		BUFFER_REGISTRY.add_raw(&buffer);
 		buffer
 	}
@@ -71,6 +83,14 @@ impl Buffer {
 		self.backing.can_release_after_update()
 	}
 
+	pub fn is_transparent(&self) -> bool {
+		match &self.backing {
+			BufferBacking::Shm(backing) => backing.is_transparent(),
+			#[cfg(feature = "dmabuf")]
+			BufferBacking::Dmabuf(backing) => backing.is_transparent(),
+		}
+	}
+
 	pub fn size(&self) -> Vector2<usize> {
 		match &self.backing {
 			BufferBacking::Shm(backing) => backing.size(),

src/wayland/core/shm_buffer_backing.rs

@@ -5,7 +5,8 @@ use bevy::{
 	render::render_resource::{Extent3d, TextureDimension, TextureFormat},
 };
 use mint::Vector2;
-use std::sync::{Arc, OnceLock};
+use parking_lot::Mutex;
+use std::sync::Arc;
 use waynest::server::protocol::core::wayland::wl_shm::Format;
 
 /// Parameters for a shared memory buffer
@@ -16,7 +17,7 @@ pub struct ShmBufferBacking {
 	stride: usize,
 	size: Vector2<usize>,
 	format: Format,
-	image: OnceLock<Handle<Image>>,
+	image: Mutex<Handle<Image>>,
 }
 
 impl ShmBufferBacking {
@@ -33,25 +34,24 @@ impl ShmBufferBacking {
 			stride,
 			size,
 			format,
-			image: OnceLock::new(),
+			image: Mutex::new(Handle::default()),
 		}
 	}
 
 	pub fn update_tex(&self, images: &mut Assets<Image>) -> Option<Handle<Image>> {
-		let image = self.image.get_or_init(|| {
-			let image = Image::new_fill(
-				Extent3d {
-					width: self.size.x as u32,
-					height: self.size.y as u32,
-					depth_or_array_layers: 1,
-				},
-				TextureDimension::D2,
-				&[255, 0, 255, 255],
-				TextureFormat::Rgba8UnormSrgb,
-				RenderAssetUsages::all(),
-			);
-			images.add(image)
-		});
+		let mut image_handle = self.image.lock();
+		images.remove(image_handle.id());
+		let mut image = Image::new_fill(
+			Extent3d {
+				width: self.size.x as u32,
+				height: self.size.y as u32,
+				depth_or_array_layers: 1,
+			},
+			TextureDimension::D2,
+			&[255, 0, 255, 255],
+			TextureFormat::Rgba8UnormSrgb,
+			RenderAssetUsages::all(),
+		);
 
 		let src_data_lock = self.pool.data_lock();
 		let mut src_cursor = self.offset;
@@ -64,7 +64,7 @@ impl ShmBufferBacking {
 			return None;
 		}
 
-		let dst_data = images.get_mut(image).unwrap().data.get_or_insert_with(|| {
+		let dst_data = image.data.get_or_insert_with(|| {
 			let length = self.size.x * self.size.y * 4;
 			vec![255; length]
 		});
@@ -93,7 +93,16 @@ impl ShmBufferBacking {
 			src_cursor += self.stride - (self.size.x * 4);
 		}
 
-		Some(image.clone())
+		*image_handle = images.add(image);
+		Some(image_handle.clone())
+	}
+
+	pub fn is_transparent(&self) -> bool {
+		match self.format {
+			Format::Xrgb8888 => false,
+			Format::Argb8888 => true,
+			_ => true,
+		}
 	}
 
 	pub fn size(&self) -> Vector2<usize> {

src/wayland/core/surface.rs

@@ -12,10 +12,11 @@ use crate::{
 use bevy::{
 	asset::{Assets, Handle},
 	image::Image,
+	render::alpha::AlphaMode,
 };
 use mint::Vector2;
 use parking_lot::Mutex;
-use std::sync::{Arc, OnceLock};
+use std::sync::{Arc, OnceLock, atomic::Ordering};
 use waynest::{
 	server::{
 		Client, Dispatcher, Result,
@@ -139,16 +140,14 @@ impl Surface {
 		});
 
 		if let Some(new_tex) = buffer.update_tex(images) {
-			materials
-				.get_mut(material)
-				.unwrap()
-				.base_color_texture
-				.replace(new_tex);
-
-			// For SHM buffers, we can release immediately after copying to GPU
-			if buffer.can_release_after_update() {
-				let _ = self.message_sink.send(Message::ReleaseBuffer(buffer));
-			}
+			buffer.rendered.store(true, Ordering::Relaxed);
+			let material = materials.get_mut(material).unwrap();
+			material.base_color_texture.replace(new_tex);
+			material.alpha_mode = if buffer.is_transparent() {
+				AlphaMode::Premultiplied
+			} else {
+				AlphaMode::Opaque
+			};
 		}
 
 		self.apply_surface_materials();

src/wayland/mod.rs

@@ -24,6 +24,7 @@ use core::{
 #[cfg(feature = "dmabuf")]
 use dmabuf::buffer_params::BufferParams;
 use mint::Vector2;
+use std::sync::atomic::Ordering;
 use std::{
 	fs::{self, OpenOptions},
 	io::{self, ErrorKind},
@@ -285,6 +286,14 @@ impl Wayland {
 
 	pub fn early_frame(graphics_info: &mut GraphicsInfo) {
 		for buffer in BUFFER_REGISTRY.get_valid_contents() {
+			if buffer.can_release_after_update() {
+				if buffer.rendered.load(Ordering::Relaxed) {
+					let _ = buffer
+						.message_sink
+						.send(Message::ReleaseBuffer(buffer.clone()));
+				}
+				buffer.rendered.store(false, Ordering::Relaxed);
+			}
 			buffer.init_tex(graphics_info);
 		}
 		for surface in WL_SURFACE_REGISTRY.get_valid_contents() {