rusticl/mesa: make PipeScreen refcounted

Reviewed-by: Seán de Búrca <sdeburca@fastmail.net>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/37377>

src/gallium/frontends/rusticl/core/device.rs

@@ -34,7 +34,6 @@ use std::mem::transmute;
 use std::num::NonZeroU64;
 use std::ops::Deref;
 use std::os::raw::*;
-use std::sync::Arc;
 use std::sync::Mutex;
 use std::sync::MutexGuard;
 
@@ -1132,7 +1131,7 @@ impl DeviceBase {
         self.reusable_ctx.lock().unwrap()
     }
 
-    pub fn screen(&self) -> &Arc<PipeScreen> {
+    pub fn screen(&self) -> &PipeScreen {
         &self.screen
     }
 

src/gallium/frontends/rusticl/mesa/pipe/context.rs

@@ -13,7 +13,6 @@ use mesa_rust_util::has_required_feature;
 use std::os::raw::*;
 use std::ptr;
 use std::ptr::*;
-use std::sync::Arc;
 
 #[repr(u32)]
 #[derive(Clone, Copy, Debug, PartialEq)]
@@ -31,7 +30,6 @@ impl From<PipeContextPrio> for u32 {
 
 pub struct PipeContext {
     pipe: NonNull<pipe_context>,
-    screen: Arc<PipeScreen>,
     pub prio: PipeContextPrio,
 }
 
@@ -52,11 +50,11 @@ impl From<RWFlags> for pipe_map_flags {
 }
 
 impl PipeContext {
-    pub fn new(prio: PipeContextPrio, screen: &Arc<PipeScreen>) -> Option<Self> {
+    pub fn new(prio: PipeContextPrio, screen: &PipeScreen) -> Option<Self> {
+        let screen = screen.to_owned();
         let context = screen.create_context(prio);
         let s = Self {
             pipe: NonNull::new(context)?,
-            screen: Arc::clone(screen),
             prio: prio,
         };
 
@@ -65,6 +63,10 @@ impl PipeContext {
             return None;
         }
 
+        // As the raw type already stores the pointer we can safely leak it here and turn it back
+        // into the owned wrapper on `drop`.
+        let ptr = screen.into_raw();
+        assert_eq!(ptr, s.screen().pipe());
         Some(s)
     }
 
@@ -72,6 +74,11 @@ impl PipeContext {
         self.pipe
     }
 
+    pub fn screen(&self) -> &PipeScreen {
+        // SAFETY: self.pipe() is a valid pointer.
+        PipeScreen::from_raw(&unsafe { self.pipe().as_ref() }.screen)
+    }
+
     pub fn buffer_subdata(
         &self,
         res: &PipeResourceOwned,
@@ -637,7 +644,7 @@ impl PipeContext {
             let mut fence = ptr::null_mut();
             self.pipe.as_ref().flush.unwrap()(self.pipe.as_ptr(), &mut fence, 0);
             // TODO: handle properly
-            PipeFence::new(fence, &self.screen).unwrap()
+            PipeFence::new(fence, self.screen()).unwrap()
         }
     }
 
@@ -650,7 +657,7 @@ impl PipeContext {
                 fence_fd.fd,
                 fence_type,
             );
-            PipeFence::new(fence, &self.screen)
+            PipeFence::new(fence, self.screen())
         }
     }
 
@@ -695,9 +702,15 @@ impl PipeContext {
 impl Drop for PipeContext {
     fn drop(&mut self) {
         self.flush().wait();
+        let screen = self.screen().pipe();
         unsafe {
             self.pipe.as_ref().destroy.unwrap()(self.pipe.as_ptr());
         }
+
+        // In new we check that screen is identical to the pointer we retrieved from into_raw. We
+        // convert the pointer back to an owned reference so we release our reference to prevent
+        // leaking memory.
+        unsafe { PipeScreenOwned::from_raw(screen) };
     }
 }
 

src/gallium/frontends/rusticl/mesa/pipe/fence.rs

@@ -7,7 +7,7 @@ use libc_rust_gen::close;
 use mesa_rust_gen::*;
 use mesa_rust_util::ptr::ThreadSafeCPtr;
 
-use std::{ffi::c_int, sync::Arc};
+use std::ffi::c_int;
 
 pub struct FenceFd {
     pub fd: i32,
@@ -23,16 +23,16 @@ impl Drop for FenceFd {
 
 pub struct PipeFence {
     fence: ThreadSafeCPtr<pipe_fence_handle>,
-    screen: Arc<PipeScreen>,
+    screen: PipeScreenOwned,
 }
 
 unsafe impl Send for PipeFence {}
 
 impl PipeFence {
-    pub fn new(fence: *mut pipe_fence_handle, screen: &Arc<PipeScreen>) -> Option<Self> {
+    pub fn new(fence: *mut pipe_fence_handle, screen: &PipeScreen) -> Option<Self> {
         Some(Self {
             fence: unsafe { ThreadSafeCPtr::new(fence)? },
-            screen: Arc::clone(screen),
+            screen: screen.to_owned(),
         })
     }
 

src/gallium/frontends/rusticl/mesa/pipe/screen.rs

@@ -11,9 +11,12 @@ use crate::util::disk_cache::*;
 use mesa_rust_gen::*;
 use mesa_rust_util::has_required_feature;
 use mesa_rust_util::ptr::ThreadSafeCPtr;
+use mesa_rust_util::static_assert;
 
+use std::borrow::Borrow;
 use std::ffi::c_int;
 use std::ffi::CStr;
+use std::mem;
 use std::num::NonZeroU64;
 use std::ops::Deref;
 use std::os::raw::c_schar;
@@ -21,20 +24,24 @@ use std::os::raw::c_uchar;
 use std::os::raw::c_void;
 use std::ptr;
 use std::ptr::NonNull;
-use std::sync::Arc;
+use std::sync::atomic::AtomicI32;
+use std::sync::atomic::Ordering;
 
 #[derive(PartialEq)]
 pub struct PipeScreenWithLdev {
     ldev: PipeLoaderDevice,
-    screen: Arc<PipeScreen>,
+    screen: PipeScreenOwned,
 }
 
 #[derive(PartialEq)]
 #[repr(transparent)]
-pub struct PipeScreen {
+pub struct PipeScreenOwned {
     screen: ThreadSafeCPtr<pipe_screen>,
 }
 
+#[repr(transparent)]
+pub struct PipeScreen(pipe_screen);
+
 pub const UUID_SIZE: usize = PIPE_UUID_SIZE as usize;
 const LUID_SIZE: usize = PIPE_LUID_SIZE as usize;
 
@@ -86,13 +93,18 @@ impl PipeScreenWithLdev {
             return None;
         }
 
-        Some(Self {
-            ldev,
+        let screen = Self {
+            ldev: ldev,
             // SAFETY: `pipe_screen` is considered a thread-safe type
-            screen: Arc::new(PipeScreen {
+            screen: PipeScreenOwned {
                 screen: unsafe { ThreadSafeCPtr::new(screen)? },
-            }),
-        })
+            },
+        };
+
+        // We use SeqCst here as refcnt might be accessed behind a mutex.
+        screen.refcnt().store(1, Ordering::SeqCst);
+
+        Some(screen)
     }
 
     pub fn driver_name(&self) -> &CStr {
@@ -105,24 +117,79 @@ impl PipeScreenWithLdev {
 }
 
 impl Deref for PipeScreenWithLdev {
-    type Target = Arc<PipeScreen>;
+    type Target = PipeScreenOwned;
 
     fn deref(&self) -> &Self::Target {
         &self.screen
     }
 }
 
-impl PipeScreen {
-    fn screen(&self) -> &pipe_screen {
-        // SAFETY: We own the pointer, so it's valid for every caller of this function as we are
-        //         responsible of freeing it.
-        unsafe { self.screen.as_ref() }
+impl Borrow<PipeScreen> for PipeScreenOwned {
+    fn borrow(&self) -> &PipeScreen {
+        // SAFETY: PipeScreen is transparent over pipe_screen, so we can convert a &pipe_screen to
+        //         &PipeScreen.
+        unsafe { mem::transmute(self.screen) }
+    }
+}
+
+impl Deref for PipeScreenOwned {
+    type Target = PipeScreen;
+
+    fn deref(&self) -> &Self::Target {
+        self.borrow()
+    }
+}
+
+impl Drop for PipeScreenOwned {
+    fn drop(&mut self) {
+        if self.refcnt().fetch_sub(1, Ordering::SeqCst) == 1 {
+            unsafe { self.screen().destroy.unwrap()(self.pipe()) }
+        }
+    }
+}
+
+impl PipeScreenOwned {
+    /// Turns a raw pointer into an owned reference.
+    ///
+    /// # Safety
+    ///
+    /// `screen` must be equivalent to a pointer retrieved via [PipeScreenOwned::into_raw].
+    /// This function does not increase reference count; use with a pointer not accounted
+    /// for in the reference count could lead to undefined behavior.
+    pub(super) unsafe fn from_raw<'s>(screen: *mut pipe_screen) -> Self {
+        // SAFETY: PipeScreenOwned is transparent over *mut pipe_screen
+        unsafe { mem::transmute(screen) }
     }
 
-    fn pipe(&self) -> *mut pipe_screen {
+    /// Turns self into a raw pointer leaking the reference count.
+    pub(super) fn into_raw(self) -> *mut pipe_screen {
+        // SAFETY: PipeScreenOwned is transparent over *mut pipe_screen
+        unsafe { mem::transmute(self) }
+    }
+}
+
+impl PipeScreen {
+    fn screen(&self) -> &pipe_screen {
+        &self.0
+    }
+
+    pub(super) fn pipe(&self) -> *mut pipe_screen {
         // screen methods are all considered thread safe, so we can just pass the mut pointer
         // around.
-        self.screen.as_ptr()
+        ((&self.0) as *const pipe_screen).cast_mut()
+    }
+
+    fn refcnt(&self) -> &AtomicI32 {
+        static_assert!(mem::align_of::<i32>() >= mem::align_of::<AtomicI32>());
+
+        let refcnt: *const _ = &self.screen().refcnt;
+
+        // SAFETY: refcnt is supposed to be atomically accessed
+        unsafe { AtomicI32::from_ptr(refcnt.cast_mut()) }
+    }
+
+    pub(super) fn from_raw<'s>(screen: &'s *mut pipe_screen) -> &'s Self {
+        unsafe { mem::transmute(*screen) }
     }
 
     pub fn caps(&self) -> &pipe_caps {
@@ -452,7 +519,7 @@ impl PipeScreen {
         }
     }
 
-    pub fn create_semaphore(self: &Arc<PipeScreen>) -> Option<PipeFence> {
+    pub fn create_semaphore(&self) -> Option<PipeFence> {
         let fence = unsafe { self.screen().semaphore_create.unwrap()(self.pipe()) };
         PipeFence::new(fence, self)
     }
@@ -495,9 +562,19 @@ impl PipeScreen {
     }
 }
 
-impl Drop for PipeScreen {
-    fn drop(&mut self) {
-        unsafe { self.screen().destroy.unwrap()(self.pipe()) }
+impl ToOwned for PipeScreen {
+    type Owned = PipeScreenOwned;
+
+    fn to_owned(&self) -> Self::Owned {
+        let refcnt = self.refcnt().fetch_add(1, Ordering::SeqCst);
+
+        // refcnt is not supposed to be 0 at any given point in time.
+        assert!(refcnt > 0, "Reference count underflow detected!");
+
+        PipeScreenOwned {
+            // SAFETY: self.pipe() is a valid pointer.
+            screen: unsafe { ThreadSafeCPtr::new(self.pipe()).unwrap() },
+        }
     }
 }