Use UnsafeCell for PgLwLock and PgAtomic (#1696)

As explained in detail in #1223, due to the Postgres restart-on-crash
logic, `PgAtomic<T>::attach` and `PgLwLock<T>::attach` would be called
again after a crash. Using OnceCell prevented the restarted extension
from attaching new PgLwLocks, so Postgres failed to restart after a
crash. This addresses that by using UnsafeCell instead of OnceCell and
marking `attach` as unsafe.

pgrx/src/atomics.rs

@@ -7,30 +7,32 @@
 //LICENSE All rights reserved.
 //LICENSE
 //LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file.
-use once_cell::sync::OnceCell;
+#![deny(unsafe_op_in_unsafe_fn)]
+use std::cell::UnsafeCell;
 
 pub struct PgAtomic<T> {
-    inner: OnceCell<*mut T>,
+    inner: UnsafeCell<*mut T>,
 }
 
 impl<T> PgAtomic<T> {
     pub const fn new() -> Self {
-        Self { inner: OnceCell::new() }
+        Self { inner: UnsafeCell::new(std::ptr::null_mut()) }
     }
 }
 
 impl<T> PgAtomic<T>
 where
     T: atomic_traits::Atomic + Default,
 {
-    pub fn attach(&self, value: *mut T) {
-        self.inner.set(value).expect("This PgAtomic is not empty, can't re-attach");
+    /// SAFETY: Must only be called from inside the Postgres shared memory init hook
+    pub unsafe fn attach(&self, value: *mut T) {
+        unsafe {
+            *self.inner.get() = value;
+        }
     }
 
     pub fn get(&self) -> &T {
-        unsafe {
-            self.inner.get().expect("This PgAtomic as not been initialized").as_ref().unwrap()
-        }
+        unsafe { (*self.inner.get()).as_ref().expect("PgAtomic was not initialized") }
     }
 }
 

pgrx/src/lwlock.rs

@@ -11,6 +11,7 @@
 use crate::pg_sys;
 use core::ops::{Deref, DerefMut};
 use once_cell::sync::OnceCell;
+use std::cell::UnsafeCell;
 use std::fmt;
 use uuid::Uuid;
 
@@ -32,7 +33,7 @@ use uuid::Uuid;
 /// This lock can not be poisoned from Rust. Panic and Abort are handled by
 /// PostgreSQL cleanly.
 pub struct PgLwLock<T> {
-    inner: OnceCell<PgLwLockInner<T>>,
+    inner: UnsafeCell<PgLwLockInner<T>>,
     name: OnceCell<&'static str>,
 }
 
@@ -43,14 +44,13 @@ impl<T> PgLwLock<T> {
     /// Create an empty lock which can be created as a global with None as a
     /// sentinel value
     pub const fn new() -> Self {
-        PgLwLock { inner: OnceCell::new(), name: OnceCell::new() }
+        PgLwLock { inner: UnsafeCell::new(PgLwLockInner::empty()), name: OnceCell::new() }
     }
 
     /// Create a new lock for T by attaching a LWLock, which is looked up by name
     pub fn from_named(input_name: &'static str, value: *mut T) -> Self {
-        let inner = OnceCell::new();
         let name = OnceCell::new();
-        inner.set(PgLwLockInner::<T>::new(input_name, value)).unwrap();
+        let inner = UnsafeCell::new(PgLwLockInner::<T>::new(input_name, value));
         name.set(input_name).unwrap();
         PgLwLock { inner, name }
     }
@@ -69,19 +69,18 @@ impl<T> PgLwLock<T> {
 
     /// Obtain a shared lock (which comes with `&T` access)
     pub fn share(&self) -> PgLwLockShareGuard<T> {
-        self.inner.get().expect("Can't give out share, lock is in an empty state").share()
+        unsafe { self.inner.get().as_ref().unwrap().share() }
     }
 
     /// Obtain an exclusive lock (which comes with `&mut T` access)
     pub fn exclusive(&self) -> PgLwLockExclusiveGuard<T> {
-        self.inner.get().expect("Can't give out exclusive, lock is in an empty state").exclusive()
+        unsafe { self.inner.get().as_ref().unwrap().exclusive() }
     }
 
     /// Attach an empty PgLwLock lock to a LWLock, and wrap T
-    pub fn attach(&self, value: *mut T) {
-        self.inner
-            .set(PgLwLockInner::<T>::new(self.get_name(), value))
-            .expect("Can't attach, lock is not in an empty state");
+    /// SAFETY: Must only be called from inside the Postgres shared memory init hook
+    pub unsafe fn attach(&self, value: *mut T) {
+        *self.inner.get() = PgLwLockInner::<T>::new(self.get_name(), value);
     }
 }
 
@@ -107,7 +106,14 @@ impl<'a, T> PgLwLockInner<T> {
         }
     }
 
+    const fn empty() -> Self {
+        PgLwLockInner { lock_ptr: std::ptr::null_mut(), data: std::ptr::null_mut() }
+    }
+
     fn share(&self) -> PgLwLockShareGuard<T> {
+        if self.lock_ptr.is_null() {
+            panic!("PgLwLock is not initialized");
+        }
         unsafe {
             pg_sys::LWLockAcquire(self.lock_ptr, pg_sys::LWLockMode_LW_SHARED);
 
@@ -116,6 +122,9 @@ impl<'a, T> PgLwLockInner<T> {
     }
 
     fn exclusive(&self) -> PgLwLockExclusiveGuard<T> {
+        if self.lock_ptr.is_null() {
+            panic!("PgLwLock is not initialized");
+        }
         unsafe {
             pg_sys::LWLockAcquire(self.lock_ptr, pg_sys::LWLockMode_LW_EXCLUSIVE);
 

pgrx/src/shmem.rs

@@ -7,6 +7,7 @@
 //LICENSE All rights reserved.
 //LICENSE
 //LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file.
+#![deny(unsafe_op_in_unsafe_fn)]
 use crate::lwlock::*;
 use crate::{pg_sys, PgAtomic};
 use std::hash::Hash;
@@ -65,8 +66,8 @@ macro_rules! pg_shmem_init {
                     if let Some(i) = PREV_SHMEM_STARTUP_HOOK {
                         i();
                     }
+                    $thing.shmem_init();
                 }
-                $thing.shmem_init();
             }
         }
     };
@@ -87,8 +88,8 @@ macro_rules! pg_shmem_init {
                     if let Some(i) = PREV_SHMEM_REQUEST_HOOK {
                         i();
                     }
+                    $thing.pg_init();
                 }
-                $thing.pg_init();
             }
         }
 
@@ -103,8 +104,8 @@ macro_rules! pg_shmem_init {
                     if let Some(i) = PREV_SHMEM_STARTUP_HOOK {
                         i();
                     }
+                    $thing.shmem_init();
                 }
-                $thing.shmem_init();
             }
         }
     };
@@ -118,7 +119,8 @@ pub trait PgSharedMemoryInitialization {
 
     /// Automatically called by the `pg_shmem_init!()` macro, when Postgres is initializing its
     /// shared memory system
-    fn shmem_init(&'static self);
+    /// SAFETY: Must only be called from inside the Postgres shared memory init hook
+    unsafe fn shmem_init(&'static self);
 }
 
 impl<T> PgSharedMemoryInitialization for PgLwLock<T>
@@ -129,8 +131,11 @@ where
         PgSharedMem::pg_init_locked(self);
     }
 
-    fn shmem_init(&'static self) {
-        PgSharedMem::shmem_init_locked(self);
+    /// SAFETY: Must only be called from inside the Postgres shared memory init hook
+    unsafe fn shmem_init(&'static self) {
+        unsafe {
+            PgSharedMem::shmem_init_locked(self);
+        }
     }
 }
 
@@ -142,8 +147,11 @@ where
         PgSharedMem::pg_init_atomic(self);
     }
 
-    fn shmem_init(&'static self) {
-        PgSharedMem::shmem_init_atomic(self);
+    /// SAFETY: Must only be called from inside the Postgres shared memory init hook
+    unsafe fn shmem_init(&'static self) {
+        unsafe {
+            PgSharedMem::shmem_init_atomic(self);
+        }
     }
 }
 
@@ -168,7 +176,8 @@ impl PgSharedMem {
     }
 
     /// Must be run from the shared memory init hook, use for types which are guarded by a `LWLock`
-    pub fn shmem_init_locked<T: Default + PGRXSharedMemory>(lock: &PgLwLock<T>) {
+    /// SAFETY: Must only be called from inside the Postgres shared memory init hook
+    pub unsafe fn shmem_init_locked<T: Default + PGRXSharedMemory>(lock: &PgLwLock<T>) {
         let mut found = false;
         unsafe {
             let shm_name = alloc::ffi::CString::new(lock.get_name()).expect("CString::new failed");
@@ -188,7 +197,8 @@ impl PgSharedMem {
     }
 
     /// Must be run from the shared memory init hook, use for rust atomics behind `PgAtomic`
-    pub fn shmem_init_atomic<T: atomic_traits::Atomic + Default>(atomic: &PgAtomic<T>) {
+    /// SAFETY: Must only be called from inside the Postgres shared memory init hook
+    pub unsafe fn shmem_init_atomic<T: atomic_traits::Atomic + Default>(atomic: &PgAtomic<T>) {
         unsafe {
             let shm_name = alloc::ffi::CString::new(Uuid::new_v4().to_string())
                 .expect("CString::new() failed");