fix(ot): only request a flush for outstanding RCOT allocations

CHANGELOG.md

@@ -5,3 +5,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+
+### Fixed
+- `SharedRCOTSender` and `SharedRCOTReceiver` no longer report a pending flush for
+  allocations which have already been fulfilled, which could deadlock an instance
+  on the flush barrier.

crates/ot/src/rcot/shared/receiver.rs

@@ -18,6 +18,10 @@ use tokio::sync::Mutex;
 
 #[derive(Debug)]
 struct Buffer<U, V> {
+    /// Number of OTs which have been allocated but not yet setup.
+    ///
+    /// Reset to zero as soon as a flush fulfills the allocation, so that it
+    /// only ever counts *outstanding* work.
     count: usize,
     inputs: Vec<U>,
     macs: Vec<V>,
@@ -37,6 +41,14 @@ impl<U, V> Buffer<U, V> {
 struct State<U, V> {
     id_next: usize,
     alloc: usize,
+    /// Whether a flush is currently in progress.
+    ///
+    /// A flush is a collective operation: the barrier only releases once
+    /// *every* live instance has arrived. Instances decide independently
+    /// whether to flush, so once one of them commits to a flush all the
+    /// others must follow, even if the allocation which triggered it has
+    /// been fulfilled by the time they check.
+    flushing: bool,
     buffers: HashMap<usize, Buffer<U, V>>,
 }
 
@@ -45,6 +57,7 @@ impl<U, V> State<U, V> {
         Self {
             id_next: 0,
             alloc: 0,
+            flushing: false,
             buffers: HashMap::new(),
         }
     }
@@ -54,6 +67,26 @@ impl<U, V> State<U, V> {
         self.id_next += 1;
         id
     }
+
+    /// Returns `true` if any instance has an allocation which has not been
+    /// setup yet.
+    fn wants_flush(&self) -> bool {
+        self.buffers.values().any(|buffer| buffer.count > 0)
+    }
+
+    /// Returns `true` if this instance must participate in a flush, marking a
+    /// flush as in progress if one is not already.
+    fn enter_flush(&mut self) -> bool {
+        if !self.flushing {
+            if !self.wants_flush() {
+                return false;
+            }
+
+            self.flushing = true;
+        }
+
+        true
+    }
 }
 
 #[derive(Debug)]
@@ -202,11 +235,12 @@ where
     type Error = SharedRCOTReceiverError;
 
     fn wants_flush(&self) -> bool {
-        !self.state.lock().unwrap().buffers.is_empty()
+        let state = self.state.lock().unwrap();
+        state.flushing || state.wants_flush()
     }
 
     async fn flush(&mut self, ctx: &mut Context) -> Result<(), Self::Error> {
-        if !self.wants_flush() {
+        if !self.state.lock().unwrap().enter_flush() {
             return Ok(());
         }
 
@@ -216,11 +250,17 @@ where
             let mut inner = self.inner.lock().await;
 
             {
-                let state = self.state.lock().unwrap();
+                let mut state = self.state.lock().unwrap();
+                // Every instance is parked at the barrier until `proceed`
+                // below, so none of them can observe the flag in between and
+                // clearing it here also covers the error paths.
+                state.flushing = false;
                 for buffer in state.buffers.values() {
-                    inner
-                        .alloc(buffer.count)
-                        .map_err(SharedRCOTReceiverError::inner)?;
+                    if buffer.count > 0 {
+                        inner
+                            .alloc(buffer.count)
+                            .map_err(SharedRCOTReceiverError::inner)?;
+                    }
                 }
             }
 
@@ -234,10 +274,21 @@ where
             buffers.sort_by_key(|(id, _)| *id);
 
             for (_, buffer) in buffers {
+                if buffer.count == 0 {
+                    continue;
+                }
+
                 let output = inner
                     .try_recv_rcot(buffer.count)
                     .map_err(SharedRCOTReceiverError::inner)?;
 
+                // The allocation is fulfilled. It must be cleared here rather
+                // than when the instance picks the OTs up: until then the
+                // buffer is still present, and a non-zero count would make
+                // instances which are already done believe another flush is
+                // required.
+                buffer.count = 0;
+
                 // Optimization: avoid expensive copying of `choices` and
                 // `msgs` potentially containing millions of elements.
                 if output.choices.len() > buffer.inputs.len() {
@@ -255,17 +306,26 @@ where
 
         {
             let mut state = self.state.lock().unwrap();
-            if let Some(buffer) = state.buffers.remove(&self.id) {
+            if let Some(buffer) = state.buffers.get_mut(&self.id) {
+                let inputs = take(&mut buffer.inputs);
+                let macs = take(&mut buffer.macs);
+
+                // Only discard the buffer if it has no allocation outstanding,
+                // otherwise an allocation made during this flush would be lost.
+                if buffer.count == 0 {
+                    state.buffers.remove(&self.id);
+                }
+
                 // Optimization: avoid expensive copying of `inputs` and
                 // `macs` potentially containing millions of elements.
-                if buffer.inputs.len() > self.inputs.len() {
-                    let old_inputs = std::mem::replace(&mut self.inputs, buffer.inputs);
-                    let old_macs = std::mem::replace(&mut self.macs, buffer.macs);
+                if inputs.len() > self.inputs.len() {
+                    let old_inputs = std::mem::replace(&mut self.inputs, inputs);
+                    let old_macs = std::mem::replace(&mut self.macs, macs);
                     self.inputs.extend_from_slice(&old_inputs);
                     self.macs.extend_from_slice(&old_macs);
                 } else {
-                    self.inputs.extend_from_slice(&buffer.inputs);
-                    self.macs.extend_from_slice(&buffer.macs);
+                    self.inputs.extend_from_slice(&inputs);
+                    self.macs.extend_from_slice(&macs);
                 }
             }
         }
@@ -279,6 +339,50 @@ where
     }
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// A buffer which is still present but has no outstanding allocation must
+    /// not keep asking for flushes, otherwise instances which have already
+    /// finished are pulled back into a barrier nobody else will join.
+    #[test]
+    fn test_fulfilled_buffer_does_not_want_flush() {
+        let mut state = State::<(), ()>::new();
+        let id = state.register();
+
+        state.buffers.insert(id, Buffer::new(8));
+        assert!(state.wants_flush());
+
+        // The flush fulfilled the allocation, but the instance has not picked
+        // its OTs up yet, so the buffer is still in the map.
+        state.buffers.get_mut(&id).unwrap().count = 0;
+        assert!(!state.wants_flush());
+    }
+
+    /// Once one instance has committed to a flush, every other instance must
+    /// join it, even if the allocation which triggered it is already fulfilled
+    /// by the time they check.
+    #[test]
+    fn test_enter_flush_is_collective() {
+        let mut state = State::<(), ()>::new();
+        let first = state.register();
+        let second = state.register();
+
+        assert!(!state.enter_flush());
+
+        state.buffers.insert(first, Buffer::new(8));
+        assert!(state.enter_flush());
+
+        state.buffers.get_mut(&first).unwrap().count = 0;
+        assert!(state.enter_flush());
+
+        state.flushing = false;
+        state.buffers.insert(second, Buffer::new(0));
+        assert!(!state.enter_flush());
+    }
+}
+
 /// Error for [`SharedRCOTReceiver`].
 #[derive(Debug, thiserror::Error)]
 #[error(transparent)]

crates/ot/src/rcot/shared/sender.rs

@@ -18,6 +18,10 @@ use tokio::sync::Mutex;
 
 #[derive(Debug)]
 struct Buffer<U> {
+    /// Number of OTs which have been allocated but not yet setup.
+    ///
+    /// Reset to zero as soon as a flush fulfills the allocation, so that it
+    /// only ever counts *outstanding* work.
     count: usize,
     keys: Vec<U>,
 }
@@ -35,6 +39,14 @@ impl<U> Buffer<U> {
 struct State<U> {
     id_next: usize,
     alloc: usize,
+    /// Whether a flush is currently in progress.
+    ///
+    /// A flush is a collective operation: the barrier only releases once
+    /// *every* live instance has arrived. Instances decide independently
+    /// whether to flush, so once one of them commits to a flush all the
+    /// others must follow, even if the allocation which triggered it has
+    /// been fulfilled by the time they check.
+    flushing: bool,
     buffers: HashMap<usize, Buffer<U>>,
 }
 
@@ -43,6 +55,7 @@ impl<U> State<U> {
         Self {
             id_next: 0,
             alloc: 0,
+            flushing: false,
             buffers: HashMap::new(),
         }
     }
@@ -52,6 +65,26 @@ impl<U> State<U> {
         self.id_next += 1;
         id
     }
+
+    /// Returns `true` if any instance has an allocation which has not been
+    /// setup yet.
+    fn wants_flush(&self) -> bool {
+        self.buffers.values().any(|buffer| buffer.count > 0)
+    }
+
+    /// Returns `true` if this instance must participate in a flush, marking a
+    /// flush as in progress if one is not already.
+    fn enter_flush(&mut self) -> bool {
+        if !self.flushing {
+            if !self.wants_flush() {
+                return false;
+            }
+
+            self.flushing = true;
+        }
+
+        true
+    }
 }
 
 #[derive(Debug)]
@@ -199,11 +232,12 @@ where
     type Error = SharedRCOTSenderError;
 
     fn wants_flush(&self) -> bool {
-        !self.state.lock().unwrap().buffers.is_empty()
+        let state = self.state.lock().unwrap();
+        state.flushing || state.wants_flush()
     }
 
     async fn flush(&mut self, ctx: &mut Context) -> Result<(), Self::Error> {
-        if !self.wants_flush() {
+        if !self.state.lock().unwrap().enter_flush() {
             return Ok(());
         }
 
@@ -212,9 +246,15 @@ where
             let mut inner = self.inner.lock().await;
 
             {
-                let state = self.state.lock().unwrap();
+                let mut state = self.state.lock().unwrap();
+                // Every instance is parked at the barrier until `proceed`
+                // below, so none of them can observe the flag in between and
+                // clearing it here also covers the error paths.
+                state.flushing = false;
                 for Buffer { count, .. } in state.buffers.values() {
-                    inner.alloc(*count).map_err(SharedRCOTSenderError::inner)?;
+                    if *count > 0 {
+                        inner.alloc(*count).map_err(SharedRCOTSenderError::inner)?;
+                    }
                 }
             }
 
@@ -228,11 +268,22 @@ where
             buffers.sort_by_key(|(id, _)| *id);
 
             for (_, buffer) in buffers {
+                if buffer.count == 0 {
+                    continue;
+                }
+
                 let keys = inner
                     .try_send_rcot(buffer.count)
                     .map_err(SharedRCOTSenderError::inner)?
                     .keys;
 
+                // The allocation is fulfilled. It must be cleared here rather
+                // than when the instance picks the keys up: until then the
+                // buffer is still present, and a non-zero count would make
+                // instances which are already done believe another flush is
+                // required.
+                buffer.count = 0;
+
                 // Optimization: avoid expensive copying of `keys` potentially
                 // containing millions of elements.
                 if keys.len() > buffer.keys.len() {
@@ -247,7 +298,15 @@ where
 
         {
             let mut state = self.state.lock().unwrap();
-            if let Some(Buffer { keys, .. }) = state.buffers.remove(&self.id) {
+            if let Some(buffer) = state.buffers.get_mut(&self.id) {
+                let keys = take(&mut buffer.keys);
+
+                // Only discard the buffer if it has no allocation outstanding,
+                // otherwise an allocation made during this flush would be lost.
+                if buffer.count == 0 {
+                    state.buffers.remove(&self.id);
+                }
+
                 // Optimization: avoid expensive copying of `keys` potentially
                 // containing millions of elements.
                 if keys.len() > self.keys.len() {
@@ -268,6 +327,50 @@ where
     }
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// A buffer which is still present but has no outstanding allocation must
+    /// not keep asking for flushes, otherwise instances which have already
+    /// finished are pulled back into a barrier nobody else will join.
+    #[test]
+    fn test_fulfilled_buffer_does_not_want_flush() {
+        let mut state = State::<()>::new();
+        let id = state.register();
+
+        state.buffers.insert(id, Buffer::new(8));
+        assert!(state.wants_flush());
+
+        // The flush fulfilled the allocation, but the instance has not picked
+        // its keys up yet, so the buffer is still in the map.
+        state.buffers.get_mut(&id).unwrap().count = 0;
+        assert!(!state.wants_flush());
+    }
+
+    /// Once one instance has committed to a flush, every other instance must
+    /// join it, even if the allocation which triggered it is already fulfilled
+    /// by the time they check.
+    #[test]
+    fn test_enter_flush_is_collective() {
+        let mut state = State::<()>::new();
+        let first = state.register();
+        let second = state.register();
+
+        assert!(!state.enter_flush());
+
+        state.buffers.insert(first, Buffer::new(8));
+        assert!(state.enter_flush());
+
+        state.buffers.get_mut(&first).unwrap().count = 0;
+        assert!(state.enter_flush());
+
+        state.flushing = false;
+        state.buffers.insert(second, Buffer::new(0));
+        assert!(!state.enter_flush());
+    }
+}
+
 /// Error for [`SharedRCOTSender`].
 #[derive(Debug, thiserror::Error)]
 #[error(transparent)]