updated consumed counts after capabilityrefs are dropped

CapabilityRefs are valid to exist as long as the data in the input are
not marked as consumed. This change makes sure that this is the case by
including an extra drop guard in the capability ref.

Signed-off-by: Petros Angelatos <petrosagg@gmail.com>

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

@@ -14,14 +14,36 @@ pub struct Counter<T: Ord+Clone+'static, D, P: Pull<Bundle<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, P: Pull<Bundle<T, D>>> Counter<T, D, P> {
     /// Retrieves the next timestamp and batch of data.
     #[inline]
     pub fn next(&mut self) -> Option<&mut Bundle<T, D>> {
+        self.next_guarded().map(|(_guard, bundle)| bundle)
+    }
+
+    #[inline]
+    pub(crate) fn next_guarded(&mut self) -> Option<(ConsumedGuard<'_, T>, &mut Bundle<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

@@ -43,13 +43,13 @@ impl<'a, T: Timestamp, D: Data, P: Pull<Bundle<T, D>>> InputHandle<T, D, P> {
     #[inline]
     pub fn next(&mut self) -> Option<(CapabilityRef<T>, RefOrMut<Vec<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))
                 },
             }
         })
@@ -75,22 +75,13 @@ impl<'a, T: Timestamp, D: Data, P: Pull<Bundle<T, D>>> InputHandle<T, D, P> {
     /// ```
     #[inline]
     pub fn for_each<F: FnMut(CapabilityRef<T>, RefOrMut<Vec<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;
     }
 
 }