libs/kaptan-runtime: generalize latest-value mailbox

libs/kaptan-runtime/src/mailbox.rs

@@ -1,118 +1,157 @@
 //! Lock-free single-producer, single-consumer latest-value mailboxes.
 
-use core::sync::atomic::{AtomicU32, Ordering};
+use core::{
+    marker::PhantomData,
+    sync::atomic::{AtomicU32, Ordering},
+};
 
-use kaptan_sensor::ImuSnapshot;
+use kaptan_sensor::{imu::ImuSnapshot, range::RangeSnapshot};
 use kaptan_time::Timestamp;
 
-/// A latest-value IMU mailbox.
+/// Fixed-width snapshot representation for [`LatestMailbox`].
 ///
-/// One IMU worker publishes and one output worker consumes. Publishing a newer
+/// The representation consists only of atomic words. Implementations must
-/// sample intentionally replaces an unread older sample: an output observer is
+/// encode and decode every field without loss so an unread older observation
-/// best-effort and must never delay acquisition. The sequence counter is a
+/// can be safely replaced by a newer one.
-/// seqlock; odd values indicate a writer and even values identify a stable
+pub trait MailboxValue<const WORDS: usize>: Copy {
-/// snapshot.
+    /// Converts a snapshot into its atomic-word representation.
-pub struct LatestImuMailbox {
+    fn into_words(self) -> [u32; WORDS];
-    sequence: AtomicU32,
+
-    timestamp_low: AtomicU32,
+    /// Reconstructs a snapshot from its atomic-word representation.
-    timestamp_high: AtomicU32,
+    fn from_words(words: [u32; WORDS]) -> Self;
-    acceleration: [AtomicU32; 3],
-    angular_rate: [AtomicU32; 3],
-    temperature: AtomicU32,
 }
 
-impl LatestImuMailbox {
+/// Lock-free latest-value SPSC mailbox.
+///
+/// Publishing a newer value intentionally replaces an unread older value: the
+/// output path is best-effort and must never delay acquisition. The sequence
+/// counter is a seqlock; odd values indicate a writer and even values identify
+/// a stable snapshot. Exactly one producer and one consumer may use an
+/// instance concurrently. Sequence operations are sequentially consistent so
+/// a reader cannot validate an old even sequence around concurrently written
+/// atomic payload words; the payload words themselves need only be relaxed.
+pub struct LatestMailbox<Value, const WORDS: usize> {
+    sequence: AtomicU32,
+    words: [AtomicU32; WORDS],
+    value: PhantomData<Value>,
+}
+
+impl<Value: MailboxValue<WORDS>, const WORDS: usize> LatestMailbox<Value, WORDS> {
     /// Creates an empty mailbox.
     pub const fn new() -> Self {
         Self {
             sequence: AtomicU32::new(0),
-            timestamp_low: AtomicU32::new(0),
+            words: [const { AtomicU32::new(0) }; WORDS],
-            timestamp_high: AtomicU32::new(0),
+            value: PhantomData,
-            acceleration: [const { AtomicU32::new(0) }; 3],
-            angular_rate: [const { AtomicU32::new(0) }; 3],
-            temperature: AtomicU32::new(0),
         }
     }
 
-    /// Publishes a sample, replacing any unread older value.
+    /// Publishes a snapshot, replacing any unread older snapshot.
-    pub fn publish(&self, sample: ImuSnapshot) {
+    pub fn publish(&self, value: Value) {
-        let sequence = self.sequence.fetch_add(1, Ordering::Relaxed);
+        let sequence = self.sequence.fetch_add(1, Ordering::SeqCst);
         debug_assert_eq!(sequence & 1, 0, "SPSC producer must not reenter");
 
-        let timestamp = sample.timestamp.as_microseconds();
+        for (destination, word) in self.words.iter().zip(value.into_words()) {
-        self.timestamp_low
+            destination.store(word, Ordering::Relaxed);
-            .store(timestamp as u32, Ordering::Relaxed);
+        }
-        self.timestamp_high
-            .store((timestamp >> 32) as u32, Ordering::Relaxed);
-        store_vector(&self.acceleration, sample.acceleration_milligravity);
-        store_vector(
-            &self.angular_rate,
-            sample.angular_rate_milliradians_per_second,
-        );
-        self.temperature.store(
-            sample.temperature_centidegrees_celsius as u16 as u32,
-            Ordering::Relaxed,
-        );
 
         self.sequence
-            .store(sequence.wrapping_add(2), Ordering::Release);
+            .store(sequence.wrapping_add(2), Ordering::SeqCst);
     }
 
-    /// Returns the newest stable sample different from `observed_sequence`.
+    /// Returns the newest stable snapshot different from `observed_sequence`.
-    pub fn take_newer(&self, observed_sequence: &mut u32) -> Option<ImuSnapshot> {
+    pub fn take_newer(&self, observed_sequence: &mut u32) -> Option<Value> {
-        let before = self.sequence.load(Ordering::Acquire);
+        let before = self.sequence.load(Ordering::SeqCst);
         if before == *observed_sequence || before & 1 != 0 {
             return None;
         }
 
-        let timestamp = u64::from(self.timestamp_low.load(Ordering::Relaxed))
+        let mut words = [0; WORDS];
-            | (u64::from(self.timestamp_high.load(Ordering::Relaxed)) << 32);
+        for (destination, source) in words.iter_mut().zip(&self.words) {
-        let acceleration = load_vector(&self.acceleration);
+            *destination = source.load(Ordering::Relaxed);
-        let angular_rate = load_vector(&self.angular_rate);
+        }
-        let temperature = self.temperature.load(Ordering::Relaxed) as u16 as i16;
 
-        if self.sequence.load(Ordering::Acquire) != before {
+        if self.sequence.load(Ordering::SeqCst) != before {
             return None;
         }
 
         *observed_sequence = before;
-        Some(ImuSnapshot::new(
+        Some(Value::from_words(words))
-            Timestamp::from_microseconds(timestamp),
-            acceleration,
-            angular_rate,
-            temperature,
-        ))
     }
 }
 
-impl Default for LatestImuMailbox {
+impl<Value: MailboxValue<WORDS>, const WORDS: usize> Default for LatestMailbox<Value, WORDS> {
     fn default() -> Self {
         Self::new()
     }
 }
 
-fn store_vector(destination: &[AtomicU32; 3], source: [i16; 3]) {
+impl MailboxValue<9> for ImuSnapshot {
-    for (destination, source) in destination.iter().zip(source) {
+    fn into_words(self) -> [u32; 9] {
-        destination.store(source as u16 as u32, Ordering::Relaxed);
+        let timestamp = self.timestamp.as_microseconds();
+        [
+            timestamp as u32,
+            (timestamp >> 32) as u32,
+            self.acceleration_milligravity[0] as u16 as u32,
+            self.acceleration_milligravity[1] as u16 as u32,
+            self.acceleration_milligravity[2] as u16 as u32,
+            self.angular_rate_milliradians_per_second[0] as u16 as u32,
+            self.angular_rate_milliradians_per_second[1] as u16 as u32,
+            self.angular_rate_milliradians_per_second[2] as u16 as u32,
+            self.temperature_centidegrees_celsius as u16 as u32,
+        ]
+    }
+
+    fn from_words(words: [u32; 9]) -> Self {
+        ImuSnapshot::new(
+            timestamp_from_words(words[0], words[1]),
+            [
+                words[2] as u16 as i16,
+                words[3] as u16 as i16,
+                words[4] as u16 as i16,
+            ],
+            [
+                words[5] as u16 as i16,
+                words[6] as u16 as i16,
+                words[7] as u16 as i16,
+            ],
+            words[8] as u16 as i16,
+        )
     }
 }
 
-fn load_vector(source: &[AtomicU32; 3]) -> [i16; 3] {
+impl MailboxValue<3> for RangeSnapshot {
-    [
+    fn into_words(self) -> [u32; 3] {
-        source[0].load(Ordering::Relaxed) as u16 as i16,
+        let timestamp = self.timestamp.as_microseconds();
-        source[1].load(Ordering::Relaxed) as u16 as i16,
+        [
-        source[2].load(Ordering::Relaxed) as u16 as i16,
+            timestamp as u32,
-    ]
+            (timestamp >> 32) as u32,
+            self.distance_millimetres,
+        ]
+    }
+
+    fn from_words(words: [u32; 3]) -> Self {
+        RangeSnapshot::new(timestamp_from_words(words[0], words[1]), words[2])
+    }
+}
+
+/// Latest-value mailbox for IMU telemetry snapshots.
+pub type LatestImuMailbox = LatestMailbox<ImuSnapshot, 9>;
+/// Latest-value mailbox for range telemetry snapshots.
+pub type LatestRangeMailbox = LatestMailbox<RangeSnapshot, 3>;
+
+const fn timestamp_from_words(low: u32, high: u32) -> Timestamp {
+    Timestamp::from_microseconds((low as u64) | ((high as u64) << 32))
 }
 
 #[cfg(test)]
 mod tests {
-    use super::LatestImuMailbox;
+    use super::{LatestImuMailbox, LatestRangeMailbox};
-    use kaptan_sensor::ImuSnapshot;
+    use kaptan_sensor::{imu::ImuSnapshot, range::RangeSnapshot};
     use kaptan_time::Timestamp;
 
     #[test]
-    fn publishes_only_the_newest_unread_sample() {
+    fn publishes_only_the_newest_unread_imu_sample() {
         let mailbox = LatestImuMailbox::new();
         let mut observed = 0;
         mailbox.publish(ImuSnapshot::new(
@@ -133,4 +172,18 @@ mod tests {
         assert_eq!(sample.acceleration_milligravity, [4; 3]);
         assert!(mailbox.take_newer(&mut observed).is_none());
     }
+
+    #[test]
+    fn publishes_only_the_newest_unread_range_sample() {
+        let mailbox = LatestRangeMailbox::new();
+        let mut observed = 0;
+        mailbox.publish(RangeSnapshot::new(Timestamp::from_microseconds(1), 100));
+        mailbox.publish(RangeSnapshot::new(Timestamp::from_microseconds(2), 200));
+
+        assert_eq!(
+            mailbox.take_newer(&mut observed),
+            Some(RangeSnapshot::new(Timestamp::from_microseconds(2), 200))
+        );
+        assert!(mailbox.take_newer(&mut observed).is_none());
+    }
 }