chore: get rid of atomics, use volatile instead

Author: yannham

libdd-library-config/src/otel_process_ctx.rs

@@ -9,7 +9,7 @@
 //! Process context sharing implies concurrently writing to a memory area that another process
 //! might be actively reading. However, reading isn't done as direct memory accesses but goes
 //! through the OS, so the Rust definition of race conditions doesn't really apply. We also use
-//! atomics and fences, see MappingHeader's documentation. 
+//! atomics and fences, see MappingHeader's documentation.
 
 #[cfg(target_os = "linux")]
 #[cfg(target_has_atomic = "64")]

libdd-library-config/src/otel_thread_ctx.rs

@@ -8,51 +8,38 @@
 //! Since `rustc` doesn't currently support the TLSDESC dialect, we use a C shim to set and get the
 //! thread-local storage used for the context.
 //!
-//! As for [process context][crate::otel_process_ctx], we use atomics as an inter-process
-//! synchronization method; see the discussion there about assumptions and pitfalls.
+//! ## Synchronization
+//!
+//! Readers are constrained to the same thread as the writer and operate like async-signal
+//! handlers: the writer thread is always stopped while a reader runs. There is thus no
+//! cross-thread synchronization concerns. The only hazard is compiler reordering or not committing
+//! writes to memory, which is handled using volatile writes.
 
 #[cfg(target_os = "linux")]
 pub mod linux {
-    use std::sync::atomic::{AtomicPtr, AtomicU8, Ordering};
     use std::{
         ffi::c_void,
         mem,
         ops::{Deref, DerefMut},
         ptr::{self, NonNull},
-        sync::atomic::fence,
     };
 
     extern "C" {
         /// Return the address of the current thread's `custom_labels_current_set_v2` local.
         fn libdd_get_custom_labels_current_set_v2() -> *mut *mut c_void;
     }
 
-    /// Return a pointer to the TLS slot holding the context. The address calculation requires to
+    /// Return a pointer to the TLS slot holding the context. The address calculation requires a
     /// call to a C shim in order to use the TLSDESC dialect from Rust. Note that the returned
-    /// address is stable (per thread), so the result can be reused several time in a sequence.
+    /// address is stable (per thread), so the result can be reused several times in a sequence.
+    ///
+    /// Use `write_tls_slot` to make sure writes are visible to the reader. Since we're the only
+    /// writer, reads can be done directly.
     #[allow(clippy::missing_safety_doc)]
-    unsafe fn get_tls_ptr() -> *mut *mut c_void {
-        libdd_get_custom_labels_current_set_v2()
-    }
-
-    /// Call [get_tls_ptr], convert it to an atomic pointer, and provide proceed with the closure
-    /// `f`.
-    fn with_atomic_tls_ptr<T>(f: impl FnOnce(&AtomicPtr<ThreadContextRecord>) -> T) -> T {
-        const {
-            assert!(
-                mem::align_of::<AtomicPtr<ThreadContextRecord>>()
-                    == mem::align_of::<*mut ThreadContextRecord>()
-            )
-        }
-
-        // Safety: the const assertion above ensures the alignment is correct. The TLS slot is
-        // valid for writes during the lifetime of the program, and the scopped closure pattern
-        // makes sure the atomic doesn't escape, such that we don't mix accesses.
-        unsafe {
-            f(AtomicPtr::from_ptr(
-                get_tls_ptr() as *mut *mut ThreadContextRecord
-            ))
-        }
+    fn get_tls_ptr() -> *mut *mut ThreadContextRecord {
+        // Safety: this is just an FFI call, but there's no particular pre-condition to uphold: the
+        // TLS slot should always be accessible.
+        unsafe { libdd_get_custom_labels_current_set_v2().cast() }
     }
 
     /// Maximum size in bytes of the `attrs_data` field.
@@ -66,38 +53,19 @@ pub mod linux {
     ///
     /// # Synchronization
     ///
-    /// `valid` is the inter-process synchronization point between the writer and the reader, is
-    /// modified atomically and protected by appropriate fences.
+    /// Readers are async-signal handlers on the same thread; the writer is always stopped while a
+    /// reader runs. `valid` is written with `ptr::write_volatile` so the compiler cannot reorder
+    /// its store past the surrounding field writes:
     ///
-    /// - The writer sets `valid = 1` *after* all other fields are populated.
-    /// - The writer sets `valid = 0` *before* modifying fields during a reattach.
+    /// - The writer sets `valid = 1` *after* all other fields are populated (publish).
+    /// - The writer sets `valid = 0` *before* modifying fields in-place (modify).
     #[repr(C)]
     pub struct ThreadContextRecord {
         /// 128-bit trace identifier; all-zeroes means "no trace".
         pub trace_id: [u8; 16],
         /// 64-bit span identifier.
         pub span_id: [u8; 8],
-        /// Whether the record is ready/consistent.
-        ///
-        /// # Memory ordering
-        ///
-        /// We ensure all accesses to `valid` use `Ordering::SeqCst`. Indeed, the pattern is the
-        /// following: we set `valid` from 1 to 0, start modifying the record, and set it back to
-        /// 1. We want to avoid having a reader observing the modifications while still seeing
-        /// valid being set to 1. My current understanding is that doing with acquire-release
-        /// ordering is not possible, because **there's no such thing as an acquire-store**. So if
-        /// a reader sees `valid` set to 1, it doesn't establish any synchronization with the
-        /// `valid := 0` assignment, and thus there's no ordering guarantees of respective memory
-        /// operations.
-        ///
-        /// Using sequential consistency forces all stores and loads to `valid` to be totally
-        /// ordered, which makes the re-ordering impossible: if a reader sees `1`, then the read
-        /// happens-before the `valid := 0`, which happens-before the modifications. 
-        ///
-        /// Note that there's still an ABA problem: while a reader has seen `0` and started reading
-        /// the record, a writer could sneak in, and modify it before the reader finishes. However
-        /// this is currently a problem of the specification, not of this particular
-        /// implementation.
+        /// Whether the record is ready/consistent. Written with `ptr::write_volatile`.
         pub valid: u8,
         pub _reserved: u8,
         /// Number of populated bytes in `attrs_data`.
@@ -131,13 +99,12 @@ pub mod linux {
             record
         }
 
-        /// Set the `valid` field to 1 atomically.
-        pub fn set_valid_atomically(&mut self, value: u8, ordering: Ordering) {
-            // Safety: `AtomicU8` doesn't have alignment constraint (it's 1), and we hold a unique
-            // borrow to `self`, so `self.valid` is valid for both reads and writes. Finally, the
-            // atomic is materialized for just one store, so there's no interleaving with
-            // non-atomic accesses.
-            unsafe { AtomicU8::from_ptr(&mut self.valid as *mut u8).store(value, ordering) }
+        /// Write `value` to the `valid` field using a volatile store, preventing the compiler from
+        /// reordering this write past adjacent field writes.
+        fn write_valid(&mut self, value: u8) {
+            // Safety: we have an exclusive borrow of `self`, so `valid` is live and valid for
+            // writes.
+            unsafe { ptr::write_volatile(&mut self.valid as *mut u8, value) }
         }
 
         /// Encode `attributes` into `record.attrs_data` as packed key-value records. Existing data
@@ -213,7 +180,7 @@ pub mod linux {
         }
 
         /// Reconstruct a [ThreadContextRecord] from a raw pointer to `ThreadContextRecord` that is
-        /// either `null` or comes from [`Self::into_raw`]. Return `None` if `ptr` is null. 
+        /// either `null` or comes from [`Self::into_raw`]. Return `None` if `ptr` is null.
         ///
         /// # Safety
         ///
@@ -241,72 +208,59 @@ pub mod linux {
     }
 
     impl ThreadContext {
-        /// Atomically swap a new (or previously detached), already built thread context record.
-        /// The `valid` flag is set to `1`. Return the previously attached context, if any.
+        /// Publish a new (or previously detached) thread context record by writing its pointer
+        /// into the TLS slot. Sets `valid = 1` before publishing. Returns the previously attached
+        /// context, if any.
         pub fn attach(mut self) -> Option<ThreadContext> {
-            with_atomic_tls_ptr(|tls_ptr| {
-                unsafe {
-                    // Safety: `self` has ownership of the underlying record, which is alive and
-                    // valid for writes.
-                    //
-                    // We don't need any synchronization here, since the swap's ordering ensures
-                    // that a reader seeing the new pointer will see the record fully initialized
-                    // and `valid` set to 0 (and the current record can't be `self`, since `attach`
-                    // consumes self). Nobody could read the `valid = 0` before, since the record
-                    // isn't even shared yet.
-                    self.0.as_mut().valid = 1;
-                    // Safety: if there is a non-null pointer in here, it came from a previous call
-                    // to `ThreadContext::into_raw`
-                    ThreadContext::from_raw(tls_ptr.swap(self.into_raw(), Ordering::Release))
-                }
-            })
+            let slot = get_tls_ptr();
+            // Set `valid = 1` written before the TLS pointer is updated, so any reader that
+            // observes the new pointer also observes `valid = 1`.
+            self.write_valid(1);
+            Self::swap(slot, self.0.as_ptr())
+        }
+
+        /// Swap the current context with a pointer value. Return the previously attached context,
+        /// if any.
+        fn swap(
+            slot: *mut *mut ThreadContextRecord,
+            tgt: *mut ThreadContextRecord,
+        ) -> Option<ThreadContext> {
+            unsafe {
+                // Safety: TLS slot is always live and valid for reads and writes.
+                let prev = ThreadContext::from_raw(*slot);
+                ptr::write_volatile(slot, tgt);
+                // Safety: a non-null value in the slot came from a prior `into_raw` call.
+                prev
+            }
         }
 
-        /// Take a closure representing an update of the thread context record, and modify the
-        /// currently attached record in-place. Take care of setting `valid` to the proper values
-        /// and related synchronization.
+        /// Modify the currently attached record in-place. Sets `valid = 0` before the update and
+        /// `valid = 1` after, so a reader that fires between the two writes sees an inconsistent
+        /// record and skips it.
         ///
         /// # Error
         ///
-        /// Returns an error if there's currently no thread context record attached.
-        pub fn modify(f: impl FnOnce(&mut ThreadContextRecord)) -> anyhow::Result<()> {
-            with_atomic_tls_ptr(|tls_ptr| {
-                let current_ptr = tls_ptr.swap(ptr::null_mut(), Ordering::Relaxed);
-
-                // Safety: if the TLS slot is set to a non-null value, this value is a thread
-                // context record owned by the current thread.
-                if let Some(current) = unsafe { current_ptr.as_mut() } {
-                    current.set_valid_atomically(0, Ordering::Relaxed);
-                    fence(Ordering::SeqCst);
-                    f(current);
-                    fence(Ordering::SeqCst);
-                    current.set_valid_atomically(1, Ordering::Relaxed);
-                } else {
-                    return Err(anyhow::anyhow!("no context currently attached"));
-                };
-
-                tls_ptr.swap(current_ptr, Ordering::Relaxed);
-
+        /// Returns an error if there is no thread context record currently attached.
+        pub fn modify() -> anyhow::Result<()> {
+            let slot = get_tls_ptr();
+
+            // Safety: the tls slot is always valid for reads and writes.
+            if let Some(current) = unsafe { (*slot).as_mut() } {
+                current.write_valid(0);
+                todo!();
+                current.write_valid(1);
                 Ok(())
-            })
+            } else {
+                Err(anyhow::anyhow!("no context currently attached"))
+            }
         }
 
-        /// Detach the current record from the TLS slot.
-        ///
-        /// Sets the TLS pointer to `null`. If there was a record previously attached, sets `valid`
-        /// to `0` and return it.
+        /// Detach the current record from the TLS slot. Writes null to the slot, sets `valid = 0`
+        /// on the detached record, and returns it.
         pub fn detach() -> Option<ThreadContext> {
-            with_atomic_tls_ptr(|tls_ptr| {
-                let prev = tls_ptr.swap(ptr::null_mut(), Ordering::Relaxed);
-
-                // Safety: if the TLS slot is not null, it is pointing to a live
-                // `ThreadContextRecord` valid for writes.
-                unsafe {
-                    ThreadContext::from_raw(prev).map(|mut ctx| {
-                        ctx.0.as_mut().set_valid_atomically(0, Ordering::Relaxed);
-                        ctx
-                    })
-                }
+            Self::swap(get_tls_ptr(), ptr::null_mut()).map(|mut ctx| {
+                ctx.write_valid(0);
+                ctx
             })
         }
     }