Reclaim memory where possible, dedup while iterating over remap_trace (#33367)

Shrinks heap allocations for some of the containers used in reclock.
Replaces a vec with a chained iterator and in-line dedup.

### Motivation

fixes MaterializeInc/database-issues#9587

Author: martykulma

Cargo.lock

@@ -9,6 +9,10 @@ publish = false
 [lints]
 workspace = true
 
+[[bench]]
+name = "reclock"
+harness = false
+
 [dependencies]
 ahash = { version = "0.8.12", default-features = false }
 bincode = "1.3.3"
@@ -29,8 +33,17 @@ tracing = "0.1.37"
 uuid = { version = "1.17.0", features = ["serde", "v4"] }
 workspace-hack = { version = "0.0.0", path = "../workspace-hack", optional = true }
 
+# This crate is used for capturing memory usage in a unit test. It is feature
+# flagged because it replaces the global memory allocator.  Unfortunately,
+# optional dependencies are not allowed in `dev-dependencies`, so it lives here.
+allocation-counter = { version = "0", optional = true }
+
+[dev-dependencies]
+criterion = { version = "0.6.0" }
+
 [features]
 default = ["mz-ore/default", "workspace-hack"]
+count-allocations = ["allocation-counter"]
 
 [package.metadata.cargo-udeps.ignore]
 normal = ["workspace-hack"]

src/timely-util/Cargo.toml

@@ -0,0 +1,128 @@
+// Copyright Materialize, Inc. and contributors. All rights reserved.
+//
+// Use of this software is governed by the Business Source License
+// included in the LICENSE file.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0.
+
+use std::sync::mpsc::Receiver;
+
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use differential_dataflow::{
+    ExchangeData,
+    input::{Input, InputSession},
+};
+
+use mz_ore::Overflowing;
+use mz_timely_util::capture::PusherCapture;
+use mz_timely_util::order::Partitioned;
+use mz_timely_util::reclock::reclock;
+use timely::communication::allocator::Thread;
+use timely::dataflow::Scope;
+use timely::dataflow::operators::capture::Event;
+use timely::dataflow::operators::unordered_input::UnorderedHandle;
+use timely::dataflow::operators::{ActivateCapability, Capture, UnorderedInput};
+use timely::progress::timestamp::Refines;
+use timely::progress::{Antichain, Timestamp};
+use timely::worker::Worker;
+
+type Diff = Overflowing<i64>;
+type FromTime = Partitioned<u64, u64>;
+type IntoTime = u64;
+type BindingHandle<FromTime> = InputSession<IntoTime, FromTime, Diff>;
+type DataHandle<D, FromTime> = (
+    UnorderedHandle<FromTime, (D, FromTime, Diff)>,
+    ActivateCapability<FromTime>,
+);
+type ReclockedStream<D> = Receiver<Event<IntoTime, Vec<(D, IntoTime, Diff)>>>;
+
+fn harness<FromTime, D, F, R>(as_of: Antichain<IntoTime>, test_logic: F) -> R
+where
+    FromTime: Timestamp + Refines<()>,
+    D: ExchangeData,
+    F: FnOnce(
+            &mut Worker<Thread>,
+            BindingHandle<FromTime>,
+            DataHandle<D, FromTime>,
+            ReclockedStream<D>,
+        ) -> R
+        + Send
+        + Sync
+        + 'static,
+    R: Send + 'static,
+{
+    timely::execute_directly(move |worker| {
+        let (bindings, data, data_cap, reclocked) = worker.dataflow::<(), _, _>(|scope| {
+            let (bindings, data_pusher, reclocked) =
+                scope.scoped::<IntoTime, _, _>("IntoScope", move |scope| {
+                    let (binding_handle, binding_collection) = scope.new_collection();
+                    let (data_pusher, reclocked_collection) = reclock(&binding_collection, as_of);
+                    let reclocked_capture = reclocked_collection.inner.capture();
+                    (binding_handle, data_pusher, reclocked_capture)
+                });
+
+            let (data, data_cap) = scope.scoped::<FromTime, _, _>("FromScope", move |scope| {
+                let ((handle, cap), data) = scope.new_unordered_input();
+                data.capture_into(PusherCapture(data_pusher));
+                (handle, cap)
+            });
+
+            (bindings, data, data_cap, reclocked)
+        });
+
+        test_logic(worker, bindings, (data, data_cap), reclocked)
+    })
+}
+
+fn step(worker: &mut Worker<Thread>) {
+    for _ in 0..4 {
+        worker.step();
+    }
+}
+
+fn bench_reclock_simple(c: &mut Criterion) {
+    let mut group = c.benchmark_group("bench_reclock_simple_group");
+
+    for n in [100_u64, 1_000_u64, 10_000_u64] {
+        group.bench_with_input(BenchmarkId::new("bench_reclock_simple", n), &n, |b, d| {
+            b.iter(|| {
+                let upto = *d;
+                let downgrade_interval = upto / 3;
+                let data_interval = upto / 10;
+                harness::<FromTime, u64, _, _>(
+                    Antichain::from_elem(0),
+                    move |worker, mut bindings, (mut data, mut data_cap), _| {
+                        for ts in 0..upto {
+                            if ts > 0 {
+                                bindings.update_at(
+                                    Partitioned::new_singleton(0, ts - 1),
+                                    ts,
+                                    Diff::MINUS_ONE,
+                                );
+                                if ts.is_multiple_of(downgrade_interval) {
+                                    data_cap.downgrade(&Partitioned::new_singleton(0, ts - 1));
+                                }
+                                if ts.is_multiple_of(data_interval) {
+                                    data.session(data_cap.clone()).give((
+                                        ts,
+                                        Partitioned::new_singleton(0, ts),
+                                        Diff::ONE,
+                                    ));
+                                }
+                            }
+                            bindings.update_at(Partitioned::new_singleton(0, ts), ts, Diff::ONE);
+                            bindings.advance_to(ts + 1);
+                            bindings.flush();
+                            step(worker);
+                        }
+                    },
+                )
+            })
+        });
+    }
+}
+
+criterion_group!(benches, bench_reclock_simple);
+criterion_main!(benches);

src/timely-util/benches/reclock.rs

@@ -248,7 +248,6 @@ where
         let mut source_frontier = MutableAntichain::new_bottom(FromTime::minimum());
 
         let mut binding_buffer = Vec::new();
-        let mut interesting_times = Vec::new();
 
         // Accumulation buffer for `remap_input` updates.
         use timely::progress::ChangeBatch;
@@ -365,13 +364,22 @@ where
                 // updates.
                 let mut min_time = IntoTime::minimum();
                 min_time.advance_by(remap_since.borrow());
-                interesting_times.push(min_time);
-                interesting_times.extend(remap_trace.iter().map(|(_, t, _)| t.clone()));
-                interesting_times.dedup();
-                for cur_time in interesting_times.drain(..) {
+                let mut prev_cur_time = None;
+                for cur_time in [min_time]
+                    .iter()
+                    .chain(remap_trace.iter().map(|(_, t, _)| t))
+                    .filter(|&v| {
+                        if prev_cur_time.is_some_and(|pv| pv == v) {
+                            false
+                        } else {
+                            prev_cur_time = Some(v);
+                            true
+                        }
+                    })
+                {
                     // 4.0. Load updates of `cur_time` from the trace into `cur_binding` to
                     //      construct the `[FromTime]` frontier that `cur_time` maps to.
-                    while let Some((t_from, _, diff)) = remap.next_if(|(_, t, _)| t == &cur_time) {
+                    while let Some((t_from, _, diff)) = remap.next_if(|(_, t, _)| t == cur_time) {
                         binding_buffer.push((t_from.clone(), *diff));
                     }
                     cur_binding.update_iter(binding_buffer.drain(..));
@@ -387,7 +395,7 @@ where
                     //      reclocked with the bindings until `prev_remap_upper`. For this reason
                     //      we only need to reconsider these updates when we start looking at new
                     //      bindings, i.e bindings that are beyond `prev_remap_upper`.
-                    if prev_remap_upper.less_equal(&cur_time) {
+                    if prev_remap_upper.less_equal(cur_time) {
                         deferred_source_updates.retain_mut(|batch| {
                             for (data, _, diff) in batch.extract(cur_binding) {
                                 session.give((data, cur_time.clone(), diff));
@@ -406,7 +414,7 @@ where
                         .iter()
                         .any(|t| !cur_binding.less_equal(t))
                     {
-                        reclocked_source_frontier.insert(cur_time);
+                        reclocked_source_frontier.insert(cur_time.clone());
                     }
                 }
 
@@ -419,6 +427,13 @@ where
                 consolidation::consolidate_updates(&mut remap_trace);
                 remap_trace
                     .sort_unstable_by(|(_, t1, _): &(_, IntoTime, _), (_, t2, _)| t1.cmp(t2));
+
+                // If using less than a quarter of the capacity, shrink the container. To avoid having
+                // to resize the container on a subsequent push, shrink to 2x the length, which is
+                // what push would grow it to.
+                if remap_trace.len() < remap_trace.capacity() / 4 {
+                    remap_trace.shrink_to(remap_trace.len() * 2);
+                }
             }
 
             // STEP 6. Tidy up deferred updates
@@ -436,6 +451,13 @@ where
                 let b = dsu.pop().unwrap();
                 dsu.push(a.merge_with(b));
             }
+
+            // If using less than a quarter of the capacity, shrink the container. To avoid having
+            // to resize the container on a subsequent push, shrink to 2x the length, which is
+            // what push would grow it to.
+            if deferred_source_updates.len() < deferred_source_updates.capacity() / 4 {
+                deferred_source_updates.shrink_to(deferred_source_updates.len() * 2);
+            }
         }
     });
 
@@ -1154,4 +1176,44 @@ mod test {
             },
         );
     }
+
+    #[cfg(feature = "count-allocations")]
+    #[mz_ore::test]
+    #[cfg_attr(miri, ignore)] // too slow
+    fn test_shrinking() {
+        let as_of = 1000_u64;
+
+        // Test that supplying a single big batch of unconsolidated bindings gets
+        // consolidated after a single worker step.
+
+        harness::<FromTime, u64, _, _>(
+            Antichain::from_elem(0),
+            move |worker, mut bindings, (_data, mut data_cap), _| {
+                let info1 = allocation_counter::measure(|| {
+                    step(worker);
+                    for ts in 0..as_of {
+                        if ts > 0 {
+                            bindings.update_at(
+                                Partitioned::new_singleton(0, ts - 1),
+                                ts,
+                                Diff::MINUS_ONE,
+                            );
+                        }
+                        bindings.update_at(Partitioned::new_singleton(0, ts), ts, Diff::ONE);
+                        bindings.advance_to(ts + 1);
+                        bindings.flush();
+                        step(worker);
+                    }
+                });
+                println!("info = {info1:?}");
+
+                let info2 = allocation_counter::measure(|| {
+                    data_cap.downgrade(&Partitioned::new_singleton(0, 900));
+                    step(worker);
+                });
+                println!("info = {info2:?}");
+                assert!(info1.bytes_current + info2.bytes_current < (info1.bytes_current / 4));
+            },
+        );
+    }
 }