libs/kaptan-runtime: add lock-free IMU mailbox

libs/kaptan-runtime/src/mailbox.rs

@@ -0,0 +1,136 @@
+//! Lock-free single-producer, single-consumer latest-value mailboxes.
+
+use core::sync::atomic::{AtomicU32, Ordering};
+
+use kaptan_sensor::ImuSnapshot;
+use kaptan_time::Timestamp;
+
+/// A latest-value IMU mailbox.
+///
+/// One IMU worker publishes and one telemetry worker consumes. Publishing a
+/// newer sample intentionally replaces an unread older sample: telemetry is a
+/// best-effort observer and must never delay acquisition. The sequence counter
+/// is a seqlock; odd values indicate a writer and even values identify a
+/// stable snapshot.
+pub struct LatestImuMailbox {
+    sequence: AtomicU32,
+    timestamp_low: AtomicU32,
+    timestamp_high: AtomicU32,
+    acceleration: [AtomicU32; 3],
+    angular_rate: [AtomicU32; 3],
+    temperature: AtomicU32,
+}
+
+impl LatestImuMailbox {
+    /// Creates an empty mailbox.
+    pub const fn new() -> Self {
+        Self {
+            sequence: AtomicU32::new(0),
+            timestamp_low: AtomicU32::new(0),
+            timestamp_high: AtomicU32::new(0),
+            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.
+    pub fn publish(&self, sample: ImuSnapshot) {
+        let sequence = self.sequence.fetch_add(1, Ordering::Relaxed);
+        debug_assert_eq!(sequence & 1, 0, "SPSC producer must not reenter");
+
+        let timestamp = sample.timestamp.as_microseconds();
+        self.timestamp_low
+            .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);
+    }
+
+    /// Returns the newest stable sample different from `observed_sequence`.
+    pub fn take_newer(&self, observed_sequence: &mut u32) -> Option<ImuSnapshot> {
+        let before = self.sequence.load(Ordering::Acquire);
+        if before == *observed_sequence || before & 1 != 0 {
+            return None;
+        }
+
+        let timestamp = u64::from(self.timestamp_low.load(Ordering::Relaxed))
+            | (u64::from(self.timestamp_high.load(Ordering::Relaxed)) << 32);
+        let acceleration = load_vector(&self.acceleration);
+        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 {
+            return None;
+        }
+
+        *observed_sequence = before;
+        Some(ImuSnapshot::new(
+            Timestamp::from_microseconds(timestamp),
+            acceleration,
+            angular_rate,
+            temperature,
+        ))
+    }
+}
+
+impl Default for LatestImuMailbox {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+fn store_vector(destination: &[AtomicU32; 3], source: [i16; 3]) {
+    for (destination, source) in destination.iter().zip(source) {
+        destination.store(source as u16 as u32, Ordering::Relaxed);
+    }
+}
+
+fn load_vector(source: &[AtomicU32; 3]) -> [i16; 3] {
+    [
+        source[0].load(Ordering::Relaxed) as u16 as i16,
+        source[1].load(Ordering::Relaxed) as u16 as i16,
+        source[2].load(Ordering::Relaxed) as u16 as i16,
+    ]
+}
+
+#[cfg(test)]
+mod tests {
+    use super::LatestImuMailbox;
+    use kaptan_sensor::ImuSnapshot;
+    use kaptan_time::Timestamp;
+
+    #[test]
+    fn publishes_only_the_newest_unread_sample() {
+        let mailbox = LatestImuMailbox::new();
+        let mut observed = 0;
+        mailbox.publish(ImuSnapshot::new(
+            Timestamp::from_microseconds(1),
+            [1; 3],
+            [2; 3],
+            3,
+        ));
+        mailbox.publish(ImuSnapshot::new(
+            Timestamp::from_microseconds(2),
+            [4; 3],
+            [5; 3],
+            6,
+        ));
+
+        let sample = mailbox.take_newer(&mut observed).unwrap();
+        assert_eq!(sample.timestamp, Timestamp::from_microseconds(2));
+        assert_eq!(sample.acceleration_milligravity, [4; 3]);
+        assert!(mailbox.take_newer(&mut observed).is_none());
+    }
+}