From 2e251220d13e61acb6e5dd4de706bc65d48b64df Mon Sep 17 00:00:00 2001 From: Frank McSherry Date: Sat, 15 Aug 2026 18:29:12 -0400 Subject: [PATCH 1/2] Share pipeline resources through worker allocator --- communication/src/allocator/generic.rs | 11 +- communication/src/allocator/mod.rs | 6 + communication/src/allocator/process.rs | 4 + communication/src/allocator/thread.rs | 173 +++++++++++++++--- .../src/allocator/zero_copy/allocator.rs | 6 + .../allocator/zero_copy/allocator_process.rs | 8 + 6 files changed, 186 insertions(+), 22 deletions(-) diff --git a/communication/src/allocator/generic.rs b/communication/src/allocator/generic.rs index a3f68b12e..07a324520 100644 --- a/communication/src/allocator/generic.rs +++ b/communication/src/allocator/generic.rs @@ -98,7 +98,16 @@ impl Allocator { (crate::allocator::thread::ThreadPusher, crate::allocator::thread::ThreadPuller) { - crate::allocator::thread::Thread::new_from(identifier, Rc::clone(self.events())) + let recycler = match self { + Allocator::Thread(t) => t.recycler(), + Allocator::Process(p) => p.recycler(), + Allocator::Tcp(z) => z.recycler(), + }; + crate::allocator::thread::Thread::new_from_with_recycler( + identifier, + Rc::clone(self.events()), + recycler, + ) } } diff --git a/communication/src/allocator/mod.rs b/communication/src/allocator/mod.rs index 8c7cf6893..5113b906b 100644 --- a/communication/src/allocator/mod.rs +++ b/communication/src/allocator/mod.rs @@ -232,4 +232,10 @@ impl Process { Process::Bytes(pb) => pb.await_events(duration), } } + pub(crate) fn recycler(&self) -> thread::RecyclerHandle { + match self { + Process::Typed(p) => p.recycler(), + Process::Bytes(pb) => pb.recycler(), + } + } } diff --git a/communication/src/allocator/process.rs b/communication/src/allocator/process.rs index 5080d040c..266059484 100644 --- a/communication/src/allocator/process.rs +++ b/communication/src/allocator/process.rs @@ -70,6 +70,10 @@ pub struct Process { impl Process { /// Access the wrapped inner allocator. pub fn inner(&mut self) -> &mut Thread { &mut self.inner } + + pub(crate) fn recycler(&self) -> crate::allocator::thread::RecyclerHandle { + self.inner.recycler() + } } impl PeerBuilder for Process { diff --git a/communication/src/allocator/thread.rs b/communication/src/allocator/thread.rs index 9857ed5ba..891de3512 100644 --- a/communication/src/allocator/thread.rs +++ b/communication/src/allocator/thread.rs @@ -3,7 +3,8 @@ use std::rc::Rc; use std::cell::RefCell; use std::time::Duration; -use std::collections::VecDeque; +use std::any::{Any, TypeId}; +use std::collections::{HashMap, VecDeque}; use crate::allocator::{Allocate, AllocateBuilder}; use crate::allocator::counters::Pusher as CountPusher; @@ -24,13 +25,19 @@ impl AllocateBuilder for ThreadBuilder { pub struct Thread { /// Shared counts of messages in channels. events: Rc>>, + /// Reusable values shared by all pipeline channels allocated by this worker. + recycler: RecyclerHandle, } impl Allocate for Thread { fn index(&self) -> usize { 0 } fn peers(&self) -> usize { 1 } fn allocate(&mut self, identifier: usize) -> (Vec>>, Box>) { - let (pusher, puller) = Thread::new_from(identifier, Rc::clone(&self.events)); + let (pusher, puller) = Thread::new_from_with_recycler( + identifier, + Rc::clone(&self.events), + Rc::clone(&self.recycler), + ); (vec![Box::new(pusher)], Box::new(puller)) } fn events(&self) -> &Rc>> { @@ -58,49 +65,173 @@ impl Thread { pub fn new_from(identifier: usize, events: Rc>>) -> (ThreadPusher, ThreadPuller) { - let shared = Rc::new(RefCell::new((VecDeque::::new(), VecDeque::::new()))); - let pusher = Pusher { target: Rc::clone(&shared) }; + Self::new_from_with_recycler(identifier, events, Default::default()) + } + + pub(crate) fn new_from_with_recycler( + identifier: usize, + events: Rc>>, + recycler: RecyclerHandle, + ) -> (ThreadPusher, ThreadPuller) { + let shared = Rc::new(RefCell::new(VecDeque::::new())); + let pusher = Pusher { + target: Rc::clone(&shared), + recycler: Rc::clone(&recycler), + }; let pusher = CountPusher::new(pusher, identifier, Rc::clone(&events)); - let puller = Puller { source: shared, current: None }; + let puller = Puller { source: shared, current: None, recycler }; let puller = CountPuller::new(puller, identifier, events); (pusher, puller) } + + pub(crate) fn recycler(&self) -> RecyclerHandle { + Rc::clone(&self.recycler) + } +} + +/// Maximum number of reusable values retained for each concrete channel type. +const PER_TYPE_LIMIT: usize = 2; + +pub(crate) type RecyclerHandle = Rc>; + +/// Worker-owned storage shared by pipeline channels of the same concrete type. +#[derive(Default)] +pub(crate) struct Recycler { + pools: HashMap>, +} + +impl Recycler { + fn recycle(&mut self, value: T) -> Result<(), T> { + let pool = self + .pools + .entry(TypeId::of::()) + .or_insert_with(|| Box::new(Vec::::new())) + .downcast_mut::>() + .expect("TypeId must identify the recycler pool type"); + if pool.len() < PER_TYPE_LIMIT { + pool.push(value); + Ok(()) + } else { + Err(value) + } + } + + fn acquire(&mut self) -> Option { + self.pools + .get_mut(&TypeId::of::()) + .and_then(|pool| pool.downcast_mut::>()) + .and_then(Vec::pop) + } } /// The push half of an intra-thread channel. pub struct Pusher { - target: Rc, VecDeque)>>, + target: Rc>>, + recycler: RecyclerHandle, } -impl Push for Pusher { +impl Push for Pusher { #[inline] fn push(&mut self, element: &mut Option) { - let mut borrow = self.target.borrow_mut(); + let sent = element.is_some(); if let Some(element) = element.take() { - borrow.0.push_back(element); + self.target.borrow_mut().push_back(element); + } + // `Push::done` calls `push` with `None`; do not remove and immediately + // drop a pooled value merely to communicate that control signal. + if sent { + *element = self.recycler.borrow_mut().acquire(); } - *element = borrow.1.pop_front(); } } /// The pull half of an intra-thread channel. -pub struct Puller { +pub struct Puller { current: Option, - source: Rc, VecDeque)>>, + source: Rc>>, + recycler: RecyclerHandle, } -impl Pull for Puller { +impl Pull for Puller { #[inline] fn pull(&mut self) -> &mut Option { - let mut borrow = self.source.borrow_mut(); - // if let Some(element) = self.current.take() { - // // TODO : Arbitrary constant. - // if borrow.1.len() < 16 { - // borrow.1.push_back(element); - // } - // } - self.current = borrow.0.pop_front(); + if let Some(element) = self.current.take() { + let _ = self.recycler.borrow_mut().recycle(element); + } + self.current = self.source.borrow_mut().pop_front(); &mut self.current } } + +impl Drop for Puller { + fn drop(&mut self) { + if let Some(element) = self.current.take() { + let _ = self.recycler.borrow_mut().recycle(element); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn channels_of_the_same_type_share_returned_values() { + let allocator = Thread::default(); + let events = Rc::clone(&allocator.events); + let recycler = allocator.recycler(); + let (mut first_push, mut first_pull) = Thread::new_from_with_recycler::>( + 0, + Rc::clone(&events), + Rc::clone(&recycler), + ); + let (mut second_push, _second_pull) = + Thread::new_from_with_recycler::>(1, events, recycler); + + let value = Vec::with_capacity(1024); + let allocation = value.as_ptr(); + let mut slot = Some(value); + first_push.push(&mut slot); + assert!(slot.is_none()); + + first_pull.pull().as_mut().unwrap().clear(); + assert!(first_pull.pull().is_none()); + + let mut second = Some(Vec::new()); + second_push.push(&mut second); + let returned = second.expect("second channel should acquire the returned value"); + assert_eq!(returned.as_ptr(), allocation); + assert_eq!(returned.capacity(), 1024); + } + + #[test] + fn recycler_bounds_each_concrete_type() { + let mut recycler = Recycler::default(); + assert!(recycler.recycle::>(Vec::new()).is_ok()); + assert!(recycler.recycle::>(Vec::new()).is_ok()); + assert!(recycler.recycle::>(Vec::new()).is_err()); + + assert!(recycler.recycle::>(Vec::new()).is_ok()); + assert!(recycler.recycle::>(Vec::new()).is_ok()); + assert!(recycler.recycle::>(Vec::new()).is_err()); + } + + #[test] + fn done_does_not_discard_a_recycled_value() { + let recycler = RecyclerHandle::default(); + recycler + .borrow_mut() + .recycle::>(Vec::with_capacity(1024)) + .unwrap(); + let (mut push, _pull) = Thread::new_from_with_recycler::>( + 0, + Rc::new(RefCell::new(Vec::new())), + Rc::clone(&recycler), + ); + + push.done(); + + assert_eq!(recycler.borrow_mut().acquire::>().unwrap().capacity(), 1024); + } +} diff --git a/communication/src/allocator/zero_copy/allocator.rs b/communication/src/allocator/zero_copy/allocator.rs index 59ce805a1..05e7f2aed 100644 --- a/communication/src/allocator/zero_copy/allocator.rs +++ b/communication/src/allocator/zero_copy/allocator.rs @@ -149,6 +149,12 @@ pub struct TcpAllocator { to_local: HashMap>>>, // to worker-local typed pullers. } +impl TcpAllocator { + pub(crate) fn recycler(&self) -> crate::allocator::thread::RecyclerHandle { + self.inner.recycler() + } +} + impl Allocate for TcpAllocator { fn index(&self) -> usize { self.index } fn peers(&self) -> usize { self.peers } diff --git a/communication/src/allocator/zero_copy/allocator_process.rs b/communication/src/allocator/zero_copy/allocator_process.rs index 72a1e94b8..0f60dbcd1 100644 --- a/communication/src/allocator/zero_copy/allocator_process.rs +++ b/communication/src/allocator/zero_copy/allocator_process.rs @@ -100,6 +100,7 @@ impl ProcessBuilder { sends, recvs, to_local: HashMap::new(), + recycler: Default::default(), } } } @@ -130,6 +131,13 @@ pub struct ProcessAllocator { sends: Vec>>>, // sends[x] -> goes to thread x. recvs: Vec, // recvs[x] <- from thread x. to_local: HashMap>>>, // to worker-local typed pullers. + recycler: crate::allocator::thread::RecyclerHandle, +} + +impl ProcessAllocator { + pub(crate) fn recycler(&self) -> crate::allocator::thread::RecyclerHandle { + Rc::clone(&self.recycler) + } } impl Allocate for ProcessAllocator { From cd0d537eefeec18713fbf67659133dd10ce45ef9 Mon Sep 17 00:00:00 2001 From: Frank McSherry Date: Sun, 16 Aug 2026 14:37:06 -0400 Subject: [PATCH 2/2] Cache typed pipeline recycler pools --- communication/src/allocator/thread.rs | 70 +++++++++++++-------------- 1 file changed, 35 insertions(+), 35 deletions(-) diff --git a/communication/src/allocator/thread.rs b/communication/src/allocator/thread.rs index 891de3512..5076b3abc 100644 --- a/communication/src/allocator/thread.rs +++ b/communication/src/allocator/thread.rs @@ -74,12 +74,13 @@ impl Thread { recycler: RecyclerHandle, ) -> (ThreadPusher, ThreadPuller) { let shared = Rc::new(RefCell::new(VecDeque::::new())); + let pool = recycler.borrow_mut().pool::(); let pusher = Pusher { target: Rc::clone(&shared), - recycler: Rc::clone(&recycler), + pool: Rc::clone(&pool), }; let pusher = CountPusher::new(pusher, identifier, Rc::clone(&events)); - let puller = Puller { source: shared, current: None, recycler }; + let puller = Puller { source: shared, current: None, pool }; let puller = CountPuller::new(puller, identifier, events); (pusher, puller) } @@ -101,26 +102,25 @@ pub(crate) struct Recycler { } impl Recycler { - fn recycle(&mut self, value: T) -> Result<(), T> { - let pool = self + fn pool(&mut self) -> Pool { + Rc::clone(self .pools .entry(TypeId::of::()) - .or_insert_with(|| Box::new(Vec::::new())) - .downcast_mut::>() - .expect("TypeId must identify the recycler pool type"); - if pool.len() < PER_TYPE_LIMIT { - pool.push(value); - Ok(()) - } else { - Err(value) - } + .or_insert_with(|| Box::new(Pool::::default())) + .downcast_ref::>() + .expect("TypeId must identify the recycler pool type")) } +} + +type Pool = Rc>>; - fn acquire(&mut self) -> Option { - self.pools - .get_mut(&TypeId::of::()) - .and_then(|pool| pool.downcast_mut::>()) - .and_then(Vec::pop) +fn recycle(pool: &Pool, value: T) -> Result<(), T> { + let mut pool = pool.borrow_mut(); + if pool.len() < PER_TYPE_LIMIT { + pool.push(value); + Ok(()) + } else { + Err(value) } } @@ -128,7 +128,7 @@ impl Recycler { /// The push half of an intra-thread channel. pub struct Pusher { target: Rc>>, - recycler: RecyclerHandle, + pool: Pool, } impl Push for Pusher { @@ -141,7 +141,7 @@ impl Push for Pusher { // `Push::done` calls `push` with `None`; do not remove and immediately // drop a pooled value merely to communicate that control signal. if sent { - *element = self.recycler.borrow_mut().acquire(); + *element = self.pool.borrow_mut().pop(); } } } @@ -150,14 +150,14 @@ impl Push for Pusher { pub struct Puller { current: Option, source: Rc>>, - recycler: RecyclerHandle, + pool: Pool, } impl Pull for Puller { #[inline] fn pull(&mut self) -> &mut Option { if let Some(element) = self.current.take() { - let _ = self.recycler.borrow_mut().recycle(element); + let _ = recycle(&self.pool, element); } self.current = self.source.borrow_mut().pop_front(); &mut self.current @@ -167,7 +167,7 @@ impl Pull for Puller { impl Drop for Puller { fn drop(&mut self) { if let Some(element) = self.current.take() { - let _ = self.recycler.borrow_mut().recycle(element); + let _ = recycle(&self.pool, element); } } } @@ -208,22 +208,22 @@ mod tests { #[test] fn recycler_bounds_each_concrete_type() { let mut recycler = Recycler::default(); - assert!(recycler.recycle::>(Vec::new()).is_ok()); - assert!(recycler.recycle::>(Vec::new()).is_ok()); - assert!(recycler.recycle::>(Vec::new()).is_err()); - - assert!(recycler.recycle::>(Vec::new()).is_ok()); - assert!(recycler.recycle::>(Vec::new()).is_ok()); - assert!(recycler.recycle::>(Vec::new()).is_err()); + let bytes = recycler.pool::>(); + assert!(recycle(&bytes, Vec::new()).is_ok()); + assert!(recycle(&bytes, Vec::new()).is_ok()); + assert!(recycle(&bytes, Vec::new()).is_err()); + + let words = recycler.pool::>(); + assert!(recycle(&words, Vec::new()).is_ok()); + assert!(recycle(&words, Vec::new()).is_ok()); + assert!(recycle(&words, Vec::new()).is_err()); } #[test] fn done_does_not_discard_a_recycled_value() { let recycler = RecyclerHandle::default(); - recycler - .borrow_mut() - .recycle::>(Vec::with_capacity(1024)) - .unwrap(); + let pool = recycler.borrow_mut().pool::>(); + recycle(&pool, Vec::with_capacity(1024)).unwrap(); let (mut push, _pull) = Thread::new_from_with_recycler::>( 0, Rc::new(RefCell::new(Vec::new())), @@ -232,6 +232,6 @@ mod tests { push.done(); - assert_eq!(recycler.borrow_mut().acquire::>().unwrap().capacity(), 1024); + assert_eq!(pool.borrow_mut().pop().unwrap().capacity(), 1024); } }