Skip to main content

franka/robot/target_control/
slot.rs

1//! The lock-free mailbox between a commander and the realtime loop.
2
3use std::sync::atomic::{AtomicU64, Ordering};
4
5/// The latest target, `N` `f64`s, as a single-writer seqlock: the sequence number is odd
6/// while a write is in progress and changes with every write, so a reader that sees the
7/// same even number before and after loading the values has a consistent set. Reads never
8/// block and never spin unboundedly, which is what lets the realtime loop poll it every
9/// cycle; writes are wait-free too, but there must be only one writer at a time, which the
10/// handles serialise with a mutex on the user side.
11///
12/// The [`Default`] slot holds zeros with sequence 0, which [`TargetSlot::load`] reports as
13/// "never written" through its `false` return: the loop then keeps the target it has.
14#[derive(Debug)]
15pub struct TargetSlot<const N: usize> {
16    sequence: AtomicU64,
17    values: [AtomicU64; N],
18}
19
20impl<const N: usize> Default for TargetSlot<N> {
21    fn default() -> Self {
22        Self {
23            sequence: AtomicU64::new(0),
24            values: [const { AtomicU64::new(0) }; N],
25        }
26    }
27}
28
29impl<const N: usize> TargetSlot<N> {
30    /// A slot holding `initial`, already published (sequence 2).
31    pub fn new(initial: [f64; N]) -> Self {
32        let slot = Self::default();
33        slot.publish(initial);
34        slot
35    }
36
37    /// Publishes `target`. The single writer; wait-free.
38    pub fn publish(&self, target: [f64; N]) {
39        self.sequence.fetch_add(1, Ordering::SeqCst);
40        for (slot, value) in self.values.iter().zip(target) {
41            slot.store(value.to_bits(), Ordering::SeqCst);
42        }
43        self.sequence.fetch_add(1, Ordering::SeqCst);
44    }
45
46    /// Copies the latest consistent target into `into` and returns `true`; `false`, with
47    /// `into` untouched, if the writer was mid-update on every one of three tries or nothing
48    /// has been published yet.
49    pub fn load(&self, into: &mut [f64; N]) -> bool {
50        for _ in 0..3 {
51            let before = self.sequence.load(Ordering::SeqCst);
52            if before == 0 || before & 1 == 1 {
53                continue;
54            }
55            let candidate = self
56                .values
57                .each_ref()
58                .map(|slot| f64::from_bits(slot.load(Ordering::SeqCst)));
59            if self.sequence.load(Ordering::SeqCst) == before {
60                *into = candidate;
61                return true;
62            }
63        }
64        false
65    }
66
67    /// The number of completed writes, times two (the sequence number).
68    pub fn sequence(&self) -> u64 {
69        self.sequence.load(Ordering::SeqCst)
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use std::sync::mpsc::{sync_channel, TryRecvError};
77
78    #[test]
79    fn an_unwritten_slot_reports_nothing_and_leaves_the_output_alone() {
80        let slot = TargetSlot::<3>::default();
81        let mut into = [1.0, 2.0, 3.0];
82        assert!(!slot.load(&mut into));
83        assert_eq!(into, [1.0, 2.0, 3.0]);
84        assert_eq!(slot.sequence(), 0);
85    }
86
87    #[test]
88    fn publish_then_load_round_trips_every_bit_pattern() {
89        let slot = TargetSlot::<4>::new([0.0, -0.0, f64::MIN_POSITIVE, -1e300]);
90        let mut into = [f64::NAN; 4];
91        assert!(slot.load(&mut into));
92        assert_eq!(
93            into.map(f64::to_bits),
94            [0.0, -0.0, f64::MIN_POSITIVE, -1e300].map(f64::to_bits)
95        );
96        assert_eq!(slot.sequence(), 2);
97        slot.publish([1.0, 2.0, 3.0, 4.0]);
98        assert!(slot.load(&mut into));
99        assert_eq!(into, [1.0, 2.0, 3.0, 4.0]);
100        assert_eq!(slot.sequence(), 4);
101    }
102
103    #[test]
104    fn a_reader_never_sees_a_torn_triple_under_a_busy_writer() {
105        // Every published triple is `[k, 2k, 3k]`; a torn read would break that invariant.
106        const ROUNDS: u32 = 200;
107        const WRITES_PER_ROUND: u32 = 1_000;
108        let slot = TargetSlot::<3>::new([0.0; 3]);
109        let (done_tx, done_rx) = sync_channel(0);
110        let (resume_tx, resume_rx) = sync_channel(0);
111        let slot = &slot;
112        std::thread::scope(move |threads| {
113            threads.spawn(move || {
114                for round in 1..=ROUNDS {
115                    let last = round * WRITES_PER_ROUND;
116                    for k in last - WRITES_PER_ROUND + 1..=last {
117                        let k = f64::from(k);
118                        slot.publish([k, 2.0 * k, 3.0 * k]);
119                    }
120                    // Guarantee observable updates without assuming how the scheduler divides
121                    // time between the threads. A failed reader disconnects these channels.
122                    if done_tx.send(last).is_err() || resume_rx.recv().is_err() {
123                        return;
124                    }
125                }
126            });
127            let mut target = [0.0; 3];
128            for round in 1..=ROUNDS {
129                loop {
130                    if slot.load(&mut target) {
131                        assert_eq!(target[1], 2.0 * target[0], "torn read {target:?}");
132                        assert_eq!(target[2], 3.0 * target[0], "torn read {target:?}");
133                    }
134                    match done_rx.try_recv() {
135                        Ok(last) => {
136                            assert_eq!(last, round * WRITES_PER_ROUND);
137                            break;
138                        }
139                        Err(TryRecvError::Empty) => {}
140                        Err(TryRecvError::Disconnected) => panic!("writer stopped early"),
141                    }
142                }
143                // The paused writer gives every round a successful read, even on one CPU.
144                assert!(slot.load(&mut target));
145                let last = f64::from(round * WRITES_PER_ROUND);
146                assert_eq!(target, [last, 2.0 * last, 3.0 * last]);
147                resume_tx.send(()).unwrap();
148            }
149        });
150    }
151}