refactor(wayland): continue implementing better buffer release

Signed-off-by: Schmarni <marnistromer@gmail.com>

src/wayland/core/buffer.rs

@@ -16,24 +16,22 @@ use waynest::{
 
 #[derive(Debug)]
 pub struct BufferUsage {
-	pub buffer: Weak<Buffer>,
+	pub buffer: Arc<Buffer>,
 	message_sink: MessageSink,
 }
 impl BufferUsage {
 	pub fn new(client: &Client, buffer: &Arc<Buffer>) -> Arc<Self> {
 		Arc::new(Self {
-			buffer: Arc::downgrade(buffer),
+			buffer: buffer.clone(),
 			message_sink: client.message_sink(),
 		})
 	}
 }
 impl Drop for BufferUsage {
 	fn drop(&mut self) {
-		if let Some(buffer) = self.buffer.upgrade() {
-			self.message_sink
-				.send(Message::ReleaseBuffer(buffer))
-				.unwrap();
-		}
+		self.message_sink
+			.send(Message::ReleaseBuffer(self.buffer.clone()))
+			.unwrap();
 	}
 }
 
@@ -42,7 +40,6 @@ pub enum BufferBacking {
 	Shm(ShmBufferBacking),
 	Dmabuf(DmabufBacking),
 }
-impl BufferBacking {}
 
 #[derive(Debug, Dispatcher)]
 pub struct Buffer {
@@ -83,6 +80,9 @@ impl Buffer {
 			BufferBacking::Dmabuf(backing) => backing.size(),
 		}
 	}
+	pub fn uses_buffer_usage(&self) -> bool {
+		matches!(self.backing, BufferBacking::Dmabuf(_))
+	}
 }
 
 impl WlBuffer for Buffer {

src/wayland/core/surface.rs

@@ -38,7 +38,7 @@ pub enum SurfaceRole {
 #[derive(Debug, Clone)]
 pub struct BufferState {
 	pub buffer: Arc<Buffer>,
-	pub usage: Arc<BufferUsage>,
+	pub usage: Option<Arc<BufferUsage>>,
 }
 
 #[derive(Debug, Clone)]
@@ -232,7 +232,9 @@ impl WlSurface for Surface {
 		self.state.lock().pending.buffer = buffer.and_then(|b| {
 			let buffer = client.get::<Buffer>(b)?;
 			Some(BufferState {
-				usage: BufferUsage::new(client, &buffer),
+				usage: buffer
+					.uses_buffer_usage()
+					.then(|| BufferUsage::new(client, &buffer)),
 				buffer,
 			})
 		});

src/wayland/dmabuf/buffer_backing.rs

@@ -1,8 +1,5 @@
 use super::buffer_params::BufferParams;
-use crate::{
-	core::registry::Registry,
-	wayland::{MessageSink, RENDER_DEVICE, core::buffer::BufferUsage},
-};
+use crate::wayland::{MessageSink, RENDER_DEVICE};
 use bevy::{
 	asset::{Assets, Handle},
 	image::Image,
@@ -18,9 +15,6 @@ use std::sync::{Arc, OnceLock};
 use tracing::info;
 use waynest::server::protocol::stable::linux_dmabuf_v1::zwp_linux_buffer_params_v1::Flags;
 
-// clear this out on end frame, add to it whenever a commit with the surface with the buffer is present
-const IN_USE_DMABUFS: Registry<BufferUsage> = Registry::new();
-
 /// Parameters for a shared memory buffer
 pub struct DmabufBacking {
 	message_sink: Option<MessageSink>,
@@ -88,9 +82,7 @@ impl DmabufBacking {
 		&self,
 		dmatexes: &ImportedDmatexs,
 		images: &mut Assets<Image>,
-		usage: Arc<BufferUsage>,
 	) -> Option<Handle<Image>> {
-		IN_USE_DMABUFS.add_raw(usage);
 		info!("updating dmabuf tex");
 		self.pending_imported_dmatex
 			.lock()

src/wayland/mod.rs

@@ -16,13 +16,17 @@ use crate::{
 use bevy::app::{App, Plugin, Update};
 use bevy::ecs::schedule::IntoScheduleConfigs;
 use bevy::ecs::system::{Res, ResMut};
+use bevy::prelude::{Deref, DerefMut};
 use bevy::render::renderer::RenderDevice;
 use bevy::render::{Render, RenderApp};
 use bevy::{asset::Assets, ecs::resource::Resource, image::Image};
 use bevy_dmabuf::import::ImportedDmatexs;
+use bevy_mod_xr::session::XrRenderSet;
 use cluFlock::{FlockLock, ToFlock};
+use core::buffer::BufferUsage;
 use core::{buffer::Buffer, callback::Callback, display::Display, surface::WL_SURFACE_REGISTRY};
 use mint::Vector2;
+use rustc_hash::FxHashMap;
 use std::fs::File;
 use std::{
 	fs::{self, OpenOptions},
@@ -285,17 +289,41 @@ impl Plugin for WaylandPlugin {
 	fn build(&self, app: &mut App) {
 		app.add_systems(PreFrameWait, early_frame);
 		app.add_systems(Update, update_graphics);
-		app.sub_app_mut(RenderApp).add_systems(
-			Render,
-			init_render_device.run_if(|| RENDER_DEVICE.get().is_none()),
-		);
+		app.sub_app_mut(RenderApp)
+			.init_resource::<UsedBuffers>()
+			.add_systems(
+				Render,
+				init_render_device.run_if(|| RENDER_DEVICE.get().is_none()),
+			);
 	}
 	fn finish(&self, app: &mut App) {
 		app.sub_app_mut(RenderApp)
-			.add_systems(Render, setup_vulkano_context);
+			.add_systems(Render, setup_vulkano_context)
+			.add_systems(Render, push_used_buffers.in_set(XrRenderSet::PreRender))
+			.add_systems(
+				Render,
+				release_unneeded_buffers.in_set(XrRenderSet::PostRender),
+			);
 	}
 }
 
+#[derive(Resource, Default, Deref, DerefMut)]
+struct UsedBuffers(FxHashMap<ObjectId, Arc<BufferUsage>>);
+
+fn push_used_buffers(mut buffers: ResMut<UsedBuffers>) {
+	for surface in WL_SURFACE_REGISTRY.get_valid_contents() {
+		if let Some(buffer) = surface.current_state().buffer.and_then(|b| b.usage).clone() {
+			buffers.insert(buffer.buffer.id, buffer);
+		}
+	}
+}
+
+fn release_unneeded_buffers(mut buffers: ResMut<UsedBuffers>) {
+	dbg!(buffers.len());
+	buffers.retain(|_, v| Arc::downgrade(v).strong_count() > 1);
+	dbg!(buffers.len());
+}
+
 fn init_render_device(dev: Res<RenderDevice>) {
 	_ = RENDER_DEVICE.set(dev.clone());
 }