fix tests

Signed-off-by: Moritz Hoffmann <[email protected]>

Author: antiguru

src/timely-util/src/builder_async.rs

@@ -211,25 +211,34 @@ pub enum Event<T: Timestamp, C, D> {
     Progress(Antichain<T>),
 }
 
+/// Shared part of an async output handle
 struct AsyncOutputHandleInner<T: Timestamp, CB: ContainerBuilder> {
+    /// Handle to write to the output stream.
     output: Output<T, CB::Container>,
+    /// Current capability held by this output handle.
     capability: Option<Capability<T>>,
+    /// Container builder to accumulate data before sending at `capability`.
     builder: CB,
 }
 
 impl<T: Timestamp, CB: ContainerBuilder> AsyncOutputHandleInner<T, CB> {
+    /// Write all pending data to the output stream.
     fn flush(&mut self) {
         while let Some(container) = self.builder.finish() {
             self.output
                 .give(self.capability.as_ref().expect("must exist"), container);
         }
     }
+
+    /// Cease this output handle, flushing all pending data and releasing its capability.
     fn cease(&mut self) {
         self.flush();
         let _ = self.output.activate();
         self.capability = None;
     }
 
+    /// Provides data at the time specified by the capability. Flushes automatically when the
+    /// capability time changes.
     fn give<D>(&mut self, cap: &Capability<T>, data: D)
     where
         CB: PushInto<D>,
@@ -238,7 +247,7 @@ impl<T: Timestamp, CB: ContainerBuilder> AsyncOutputHandleInner<T, CB> {
             && cap.time() != capability.time()
         {
             self.flush();
-            self.capability = Some(cap.clone());
+            self.capability = None;
         }
         if self.capability.is_none() {
             self.capability = Some(cap.clone());

src/timely-util/src/containers.rs

@@ -19,7 +19,6 @@ pub mod stack;
 
 pub(crate) use alloc::alloc_aligned_zeroed;
 pub use alloc::{enable_columnar_lgalloc, set_enable_columnar_lgalloc};
-pub use provided_builder::ProvidedBuilder;
 
 mod alloc {
     use mz_ore::region::Region;
@@ -50,38 +49,3 @@ mod alloc {
         ENABLE_COLUMNAR_LGALLOC.set(enabled);
     }
 }
-
-mod provided_builder {
-    use timely::Container;
-    use timely::container::ContainerBuilder;
-
-    /// A container builder that doesn't support pushing elements, and is only suitable for pushing
-    /// whole containers at Timely sessions. See [`give_container`] for more information.
-    ///
-    ///  [`give_container`]: timely::dataflow::channels::pushers::buffer::Session::give_container
-    pub struct ProvidedBuilder<C> {
-        _marker: std::marker::PhantomData<C>,
-    }
-
-    impl<C> Default for ProvidedBuilder<C> {
-        fn default() -> Self {
-            Self {
-                _marker: std::marker::PhantomData,
-            }
-        }
-    }
-
-    impl<C: Container + Clone + 'static> ContainerBuilder for ProvidedBuilder<C> {
-        type Container = C;
-
-        #[inline(always)]
-        fn extract(&mut self) -> Option<&mut Self::Container> {
-            None
-        }
-
-        #[inline(always)]
-        fn finish(&mut self) -> Option<&mut Self::Container> {
-            None
-        }
-    }
-}

src/timely-util/src/reclock.rs

@@ -613,9 +613,12 @@ mod test {
     /// test provides a test logic closure which accepts four arguments:
     ///
     /// * A reference to the worker that allows the test to step the computation
-    /// * A `BindingHandle` that allows the test to manipulate the remap bindings
-    /// * A `DataHandle` that allows the test to submit the data to be reclocked
-    /// * A `ReclockedStream` that allows observing the result of the reclocking process
+    /// * A [`BindingHandle`] that allows the test to manipulate the remap bindings
+    /// * A [`DataHandle`] that allows the test to submit the data to be reclocked
+    /// * A [`ReclockedStream`] that allows observing the result of the reclocking process
+    ///
+    /// Note that the `DataHandle` contains a capability that should be dropped or downgraded before
+    /// calling [`step`] to process data at the time.
     fn harness<FromTime, D, F, R>(as_of: Antichain<IntoTime>, test_logic: F) -> R
     where
         FromTime: Timestamp + Refines<()>,
@@ -674,6 +677,7 @@ mod test {
                 data.activate()
                     .session(&data_cap)
                     .give(('a', Partitioned::minimum(), Diff::ONE));
+                drop(data_cap);
                 step(worker);
                 let extracted = reclocked.extract();
                 let expected = vec![(0, vec![('a', 0, Diff::ONE)])];
@@ -966,6 +970,7 @@ mod test {
                     Partitioned::new_singleton(0, 50),
                     Diff::ONE,
                 ));
+                drop(data_cap);
                 step(worker);
                 assert_eq!(
                     reclocked.try_recv(),
@@ -1024,6 +1029,7 @@ mod test {
                 data.activate()
                     .session(&data_cap)
                     .give_iterator(source_updates1.iter().cloned());
+                drop(data_cap);
                 step(worker);
                 reclocked.extract()
             },
@@ -1054,6 +1060,7 @@ mod test {
                 data.activate()
                     .session(&data_cap)
                     .give_iterator(source_updates.iter().cloned());
+                drop(data_cap);
                 step(worker);
                 reclocked.extract()
             },

test/sqllogictest/introspection/relations.slt

@@ -186,9 +186,10 @@ Differential␠Logging␠Demux  Consolidate␠Differential(BatcherCapacity)  all
 Differential␠Logging␠Demux  Consolidate␠Differential(BatcherRecords)  alloc::vec::Vec<((usize,␠()),␠mz_repr::timestamp::Timestamp,␠mz_ore::overflowing::Overflowing<i64>)>
 Differential␠Logging␠Demux  Consolidate␠Differential(BatcherSize)  alloc::vec::Vec<((usize,␠()),␠mz_repr::timestamp::Timestamp,␠mz_ore::overflowing::Overflowing<i64>)>
 Differential␠Logging␠Demux  Consolidate␠Differential(Sharing)  alloc::vec::Vec<((usize,␠()),␠mz_repr::timestamp::Timestamp,␠mz_ore::overflowing::Overflowing<i64>)>
+FlatMapReachability  Consolidate␠Timely(Reachability)  mz_timely_util::columnar::Column<(((bool,␠usize,␠usize,␠usize,␠mz_repr::timestamp::Timestamp),␠()),␠mz_repr::timestamp::Timestamp,␠mz_ore::overflowing::Overflowing<i64>)>
 Replay␠compute␠logs  Compute␠Logging␠Demux  mz_timely_util::columnar::Column<(core::time::Duration,␠mz_compute::logging::compute::ComputeEvent)>
 Replay␠differential␠logs  Differential␠Logging␠Demux  alloc::vec::Vec<(core::time::Duration,␠differential_dataflow::logging::DifferentialEvent)>
-Replay␠reachability␠logs  Consolidate␠Timely(Reachability)  mz_timely_util::columnar::Column<(((bool,␠usize,␠usize,␠usize,␠mz_repr::timestamp::Timestamp),␠()),␠mz_repr::timestamp::Timestamp,␠mz_ore::overflowing::Overflowing<i64>)>
+Replay␠reachability␠logs  FlatMapReachability  mz_timely_util::columnar::Column<(core::time::Duration,␠(usize,␠alloc::vec::Vec<(usize,␠usize,␠bool,␠mz_repr::timestamp::Timestamp,␠mz_ore::overflowing::Overflowing<i64>)>))>
 Replay␠timely␠logs  Timely␠Logging␠Demux  alloc::vec::Vec<(core::time::Duration,␠timely::logging::TimelyEvent)>
 Timely␠Logging␠Demux  Consolidate␠Timely(Addresses)  alloc::vec::Vec<((usize,␠alloc::vec::Vec<usize>),␠mz_repr::timestamp::Timestamp,␠mz_ore::overflowing::Overflowing<i64>)>
 Timely␠Logging␠Demux  Consolidate␠Timely(BatchesReceived)  alloc::vec::Vec<((mz_compute::logging::timely::MessageDatum,␠()),␠mz_repr::timestamp::Timestamp,␠mz_ore::overflowing::Overflowing<i64>)>
@@ -241,6 +242,7 @@ GROUP BY type;
 1  alloc::vec::Vec<alloc::vec::Vec<differential_dataflow::containers::TimelyStack<((usize,␠alloc::vec::Vec<usize>),␠mz_repr::timestamp::Timestamp,␠mz_ore::overflowing::Overflowing<i64>)>>>
 1  mz_timely_util::columnar::Column<(((bool,␠usize,␠usize,␠usize,␠mz_repr::timestamp::Timestamp),␠()),␠mz_repr::timestamp::Timestamp,␠mz_ore::overflowing::Overflowing<i64>)>
 1  mz_timely_util::columnar::Column<((mz_compute::logging::timely::ChannelDatum,␠()),␠mz_repr::timestamp::Timestamp,␠mz_ore::overflowing::Overflowing<i64>)>
+1  mz_timely_util::columnar::Column<(core::time::Duration,␠(usize,␠alloc::vec::Vec<(usize,␠usize,␠bool,␠mz_repr::timestamp::Timestamp,␠mz_ore::overflowing::Overflowing<i64>)>))>
 1  mz_timely_util::columnar::Column<(core::time::Duration,␠mz_compute::logging::compute::ComputeEvent)>
 32  alloc::vec::Vec<alloc::rc::Rc<differential_dataflow::trace::implementations::ord_neu::val_batch::OrdValBatch<mz_compute::row_spine::spines::RowRowLayout<((mz_repr::row::Row,␠mz_repr::row::Row),␠mz_repr::timestamp::Timestamp,␠mz_ore::overflowing::Overflowing<i64>)>>>>
 32  mz_timely_util::columnar::Column<((mz_repr::row::Row,␠mz_repr::row::Row),␠mz_repr::timestamp::Timestamp,␠mz_ore::overflowing::Overflowing<i64>)>