Safe(r) minting of CapabilityRefs

timely/src/dataflow/channels/pullers/counter.rs

@@ -15,14 +15,36 @@ pub struct Counter<T: Ord+Clone+'static, D, P: Pull<BundleCore<T, D>>> {
     phantom: ::std::marker::PhantomData<D>,
 }
 
+/// A guard type that updates the change batch counts on drop
+pub struct ConsumedGuard<'a, T: Ord + Clone + 'static> {
+    consumed: &'a Rc<RefCell<ChangeBatch<T>>>,
+    time: Option<T>,
+    len: usize,
+}
+
+impl<'a, T:Ord+Clone+'static> Drop for ConsumedGuard<'a, T> {
+    fn drop(&mut self) {
+        self.consumed.borrow_mut().update(self.time.take().unwrap(), self.len as i64);
+    }
+}
+
 impl<T:Ord+Clone+'static, D: Container, P: Pull<BundleCore<T, D>>> Counter<T, D, P> {
     /// Retrieves the next timestamp and batch of data.
     #[inline]
     pub fn next(&mut self) -> Option<&mut BundleCore<T, D>> {
+        self.next_guarded().map(|(_guard, bundle)| bundle)
+    }
+
+    #[inline]
+    pub(crate) fn next_guarded(&mut self) -> Option<(ConsumedGuard<'_, T>, &mut BundleCore<T, D>)> {
         if let Some(message) = self.pullable.pull() {
             if message.data.len() > 0 {
-                self.consumed.borrow_mut().update(message.time.clone(), message.data.len() as i64);
-                Some(message)
+                let guard = ConsumedGuard {
+                    consumed: &self.consumed,
+                    time: Some(message.time.clone()),
+                    len: message.data.len(),
+                };
+                Some((guard, message))
             }
             else { None }
         }

timely/src/dataflow/operators/capability.rs

@@ -30,6 +30,7 @@ use crate::order::PartialOrder;
 use crate::progress::Timestamp;
 use crate::progress::ChangeBatch;
 use crate::scheduling::Activations;
+use crate::dataflow::channels::pullers::counter::ConsumedGuard;
 
 /// An internal trait expressing the capability to send messages with a given timestamp.
 pub trait CapabilityTrait<T: Timestamp> {
@@ -231,6 +232,8 @@ type CapabilityUpdates<T> = Rc<RefCell<Vec<Rc<RefCell<ChangeBatch<T>>>>>>;
 pub struct CapabilityRef<'cap, T: Timestamp+'cap> {
     time: &'cap T,
     internal: CapabilityUpdates<T>,
+    /// A drop guard that updates the consumed capability this CapabilityRef refers to on drop
+    _consumed_guard: ConsumedGuard<'cap, T>,
 }
 
 impl<'cap, T: Timestamp+'cap> CapabilityTrait<T> for CapabilityRef<'cap, T> {
@@ -244,10 +247,11 @@ impl<'cap, T: Timestamp+'cap> CapabilityTrait<T> for CapabilityRef<'cap, T> {
 impl<'cap, T: Timestamp + 'cap> CapabilityRef<'cap, T> {
     /// Creates a new capability reference at `time` while incrementing (and keeping a reference to)
     /// the provided [`ChangeBatch`].
-    pub(crate) fn new(time: &'cap T, internal: CapabilityUpdates<T>) -> Self {
+    pub(crate) fn new(time: &'cap T, internal: CapabilityUpdates<T>, guard: ConsumedGuard<'cap, T>) -> Self {
         CapabilityRef {
             time,
             internal,
+            _consumed_guard: guard,
         }
     }
 

timely/src/dataflow/operators/generic/handles.rs

@@ -49,13 +49,13 @@ impl<'a, T: Timestamp, D: Container, P: Pull<BundleCore<T, D>>> InputHandleCore<
     #[inline]
     pub fn next(&mut self) -> Option<(CapabilityRef<T>, RefOrMut<D>)> {
         let internal = &self.internal;
-        self.pull_counter.next().map(|bundle| {
+        self.pull_counter.next_guarded().map(|(guard, bundle)| {
             match bundle.as_ref_or_mut() {
                 RefOrMut::Ref(bundle) => {
-                    (CapabilityRef::new(&bundle.time, internal.clone()), RefOrMut::Ref(&bundle.data))
+                    (CapabilityRef::new(&bundle.time, internal.clone(), guard), RefOrMut::Ref(&bundle.data))
                 },
                 RefOrMut::Mut(bundle) => {
-                    (CapabilityRef::new(&bundle.time, internal.clone()), RefOrMut::Mut(&mut bundle.data))
+                    (CapabilityRef::new(&bundle.time, internal.clone(), guard), RefOrMut::Mut(&mut bundle.data))
                 },
             }
         })
@@ -81,22 +81,13 @@ impl<'a, T: Timestamp, D: Container, P: Pull<BundleCore<T, D>>> InputHandleCore<
     /// ```
     #[inline]
     pub fn for_each<F: FnMut(CapabilityRef<T>, RefOrMut<D>)>(&mut self, mut logic: F) {
-        // We inline `next()` so that we can use `self.logging` without cloning (and dropping) the logger.
-        let internal = &self.internal;
-        while let Some((cap, data)) = self.pull_counter.next().map(|bundle| {
-            match bundle.as_ref_or_mut() {
-                RefOrMut::Ref(bundle) => {
-                    (CapabilityRef::new(&bundle.time, internal.clone()), RefOrMut::Ref(&bundle.data))
-                },
-                RefOrMut::Mut(bundle) => {
-                    (CapabilityRef::new(&bundle.time, internal.clone()), RefOrMut::Mut(&mut bundle.data))
-                },
-            }
-        }) {
-            self.logging.as_mut().map(|l| l.log(crate::logging::GuardedMessageEvent { is_start: true }));
+        let mut logging = self.logging.take();
+        while let Some((cap, data)) = self.next() {
+            logging.as_mut().map(|l| l.log(crate::logging::GuardedMessageEvent { is_start: true }));
             logic(cap, data);
-            self.logging.as_mut().map(|l| l.log(crate::logging::GuardedMessageEvent { is_start: false }));
+            logging.as_mut().map(|l| l.log(crate::logging::GuardedMessageEvent { is_start: false }));
         }
+        self.logging = logging;
     }
 
 }