From d29fcf99eebf7452a2dfb567d78228c533c6e813 Mon Sep 17 00:00:00 2001
From: Leonid Ryzhyk <lryzhyk@vmware.com>
Date: Wed, 28 Apr 2021 23:39:12 -0700
Subject: [PATCH] Lazy-allocate pusher memory.

The current implementation of channel pushers preallocates
`Message::default_length` entries per pusher.  For very large graphs this
adds up to a lot of memory being allocated even when the system is idle
with no outstanding messages.  This patch changes the allocation policy
to only allocate channel memory when there are messages to send and to
deallocate it at the end of a burst of messages (signaled by pushing
a `None` message), reducing the memory footprint to 0 in idle state at
the cost of some potential slow-down due to a larger number of allocations.

See #394.
---
 timely/src/dataflow/channels/mod.rs              |  5 -----
 timely/src/dataflow/channels/pushers/buffer.rs   |  5 ++++-
 timely/src/dataflow/channels/pushers/exchange.rs | 10 +++++++++-
 timely/src/dataflow/channels/pushers/tee.rs      |  4 +++-
 timely/src/dataflow/operators/input.rs           | 13 ++++++++++---
 5 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/timely/src/dataflow/channels/mod.rs b/timely/src/dataflow/channels/mod.rs
index 641f0a190b..acb3dd9de6 100644
--- a/timely/src/dataflow/channels/mod.rs
+++ b/timely/src/dataflow/channels/mod.rs
@@ -52,9 +52,4 @@ impl<T, D> Message<T, D> {
                 buffer.clear();
             }
         }
-
-        // TODO: Unclear we always want this here.
-        if buffer.capacity() != Self::default_length() {
-            *buffer = Vec::with_capacity(Self::default_length());
-        }
     }}
diff --git a/timely/src/dataflow/channels/pushers/buffer.rs b/timely/src/dataflow/channels/pushers/buffer.rs
index 288d5ff884..6fdebdd0fa 100644
--- a/timely/src/dataflow/channels/pushers/buffer.rs
+++ b/timely/src/dataflow/channels/pushers/buffer.rs
@@ -23,7 +23,7 @@ impl<T, D, P: Push<Bundle<T, D>>> Buffer<T, D, P> where T: Eq+Clone {
     pub fn new(pusher: P) -> Buffer<T, D, P> {
         Buffer {
             time: None,
-            buffer: Vec::with_capacity(Message::<T, D>::default_length()),
+            buffer: Vec::new(),
             pusher,
         }
     }
@@ -65,6 +65,9 @@ impl<T, D, P: Push<Bundle<T, D>>> Buffer<T, D, P> where T: Eq+Clone {
 
     // internal method for use by `Session`.
     fn give(&mut self, data: D) {
+        if self.buffer.capacity() == 0 {
+            self.buffer = Vec::with_capacity(Message::<T, D>::default_length());
+        }
         self.buffer.push(data);
         // assert!(self.buffer.capacity() == Message::<O::Data>::default_length());
         if self.buffer.len() == self.buffer.capacity() {
diff --git a/timely/src/dataflow/channels/pushers/exchange.rs b/timely/src/dataflow/channels/pushers/exchange.rs
index 5b37a84386..08dc45b3d7 100644
--- a/timely/src/dataflow/channels/pushers/exchange.rs
+++ b/timely/src/dataflow/channels/pushers/exchange.rs
@@ -18,7 +18,7 @@ impl<T: Clone, D, P: Push<Bundle<T, D>>, H: FnMut(&T, &D)->u64>  Exchange<T, D,
     pub fn new(pushers: Vec<P>, key: H) -> Exchange<T, D, P, H> {
         let mut buffers = vec![];
         for _ in 0..pushers.len() {
-            buffers.push(Vec::with_capacity(Message::<T, D>::default_length()));
+            buffers.push(Vec::new());
         }
         Exchange {
             pushers,
@@ -64,6 +64,9 @@ impl<T: Eq+Data, D: Data, P: Push<Bundle<T, D>>, H: FnMut(&T, &D)->u64> Push<Bun
                 for datum in data.drain(..) {
                     let index = (((self.hash_func)(time, &datum)) & mask) as usize;
 
+                    if self.buffers[index].capacity() == 0 {
+                        self.buffers[index] = Vec::with_capacity(Message::<T, D>::default_length());
+                    }
                     self.buffers[index].push(datum);
                     if self.buffers[index].len() == self.buffers[index].capacity() {
                         self.flush(index);
@@ -82,6 +85,9 @@ impl<T: Eq+Data, D: Data, P: Push<Bundle<T, D>>, H: FnMut(&T, &D)->u64> Push<Bun
             else {
                 for datum in data.drain(..) {
                     let index = (((self.hash_func)(time, &datum)) % self.pushers.len() as u64) as usize;
+                    if self.buffers[index].capacity() == 0 {
+                        self.buffers[index] = Vec::with_capacity(Message::<T, D>::default_length());
+                    }
                     self.buffers[index].push(datum);
                     if self.buffers[index].len() == self.buffers[index].capacity() {
                         self.flush(index);
@@ -95,6 +101,8 @@ impl<T: Eq+Data, D: Data, P: Push<Bundle<T, D>>, H: FnMut(&T, &D)->u64> Push<Bun
             for index in 0..self.pushers.len() {
                 self.flush(index);
                 self.pushers[index].push(&mut None);
+                // Free up buffer space.
+                self.buffers[index] = Vec::new();
             }
         }
     }
diff --git a/timely/src/dataflow/channels/pushers/tee.rs b/timely/src/dataflow/channels/pushers/tee.rs
index f5e3915eaf..f31496ff01 100644
--- a/timely/src/dataflow/channels/pushers/tee.rs
+++ b/timely/src/dataflow/channels/pushers/tee.rs
@@ -31,6 +31,8 @@ impl<T: Data, D: Data> Push<Bundle<T, D>> for Tee<T, D> {
             for index in 1..pushers.len() {
                 pushers[index-1].push(&mut None);
             }
+            // Free up buffer space.
+            self.buffer = Vec::new();
         }
         if pushers.len() > 0 {
             let last = pushers.len() - 1;
@@ -44,7 +46,7 @@ impl<T, D> Tee<T, D> {
     pub fn new() -> (Tee<T, D>, TeeHelper<T, D>) {
         let shared = Rc::new(RefCell::new(Vec::new()));
         let port = Tee {
-            buffer: Vec::with_capacity(Message::<T, D>::default_length()),
+            buffer: Vec::new(),
             shared: shared.clone(),
         };
 
diff --git a/timely/src/dataflow/operators/input.rs b/timely/src/dataflow/operators/input.rs
index 11d02896d0..e59a52c28e 100644
--- a/timely/src/dataflow/operators/input.rs
+++ b/timely/src/dataflow/operators/input.rs
@@ -212,8 +212,8 @@ impl<T:Timestamp, D: Data> Handle<T, D> {
             activate: Vec::new(),
             progress: Vec::new(),
             pushers: Vec::new(),
-            buffer1: Vec::with_capacity(Message::<T, D>::default_length()),
-            buffer2: Vec::with_capacity(Message::<T, D>::default_length()),
+            buffer1: Vec::new(),
+            buffer2: Vec::new(),
             now_at: T::minimum(),
         }
     }
@@ -288,7 +288,11 @@ impl<T:Timestamp, D: Data> Handle<T, D> {
 
     // closes the current epoch, flushing if needed, shutting if needed, and updating the frontier.
     fn close_epoch(&mut self) {
-        if !self.buffer1.is_empty() { self.flush(); }
+        if !self.buffer1.is_empty() {
+            self.flush();
+            self.buffer1 = Vec::new();
+            self.buffer2 = Vec::new();
+        }
         for pusher in self.pushers.iter_mut() {
             pusher.done();
         }
@@ -305,6 +309,9 @@ impl<T:Timestamp, D: Data> Handle<T, D> {
     /// Sends one record into the corresponding timely dataflow `Stream`, at the current epoch.
     pub fn send(&mut self, data: D) {
         // assert!(self.buffer1.capacity() == Message::<T, D>::default_length());
+        if self.buffer1.capacity() == 0 {
+            self.buffer1 = Vec::with_capacity(Message::<T, D>::default_length());
+        }
         self.buffer1.push(data);
         if self.buffer1.len() == self.buffer1.capacity() {
             self.flush();