Skip to content

[dbsp] Avoid one data copy on remote exchange. #5942

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
blp wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[dbsp] Avoid one data copy on remote exchange. #5942

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions crates/dbsp/src/circuit/runtime.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ use core_affinity::{CoreId, get_core_ids};
use crossbeam::sync::{Parker, Unparker};
use enum_map::{Enum, EnumMap, enum_map};
use feldera_buffer_cache::ThreadType;
use feldera_storage::fbuf::FBuf;
use feldera_storage::fbuf::slab::{FBufSlabs, FBufSlabsStats, set_thread_slab_pool};
use feldera_types::config::{StorageCompression, StorageConfig, StorageOptions};
use feldera_types::memory_pressure::{
Expand Down Expand Up @@ -1342,7 +1343,7 @@ impl Consensus {
while !exchange.try_send_all_with_serializer(
Runtime::worker_index(),
repeat(local),
|local| vec![local as u8],
|local| FBuf::from_slice(&[local as u8]),
) {
if Runtime::kill_in_progress() {
return Err(SchedulerError::Killed);
Expand Down Expand Up @@ -1432,7 +1433,11 @@ where
while !exchange.try_send_all_with_serializer(
Runtime::worker_index(),
repeat(local.clone()),
|local| rmp_serde::to_vec(&local).unwrap(),
|local| {
let mut fbuf = FBuf::new();
rmp_serde::encode::write(&mut fbuf, &local).unwrap();
fbuf
},
) {
if Runtime::kill_in_progress() {
return Err(SchedulerError::Killed);
Expand Down
17 changes: 9 additions & 8 deletions crates/dbsp/src/operator/communication/exchange.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ use crate::{
};
use binrw::{BinRead, BinWrite};
use crossbeam_utils::CachePadded;
use feldera_storage::fbuf::FBuf;
use futures::{prelude::*, stream::FuturesUnordered};
use itertools::Itertools;
use rkyv::AlignedVec;
Expand Down Expand Up @@ -159,7 +160,7 @@ impl ExchangeServiceClient {
&self,
exchange_id: ExchangeId,
senders: Range<usize>,
data: Vec<Vec<Vec<u8>>>,
data: Vec<Vec<FBuf>>,
) -> std::io::Result<()> {
// We want to write each data buffer preceded by a header and followed
// by padding. To minimize the system calls required to do this, we
Expand Down Expand Up @@ -585,13 +586,13 @@ impl InnerExchange {

#[derive(Clone, Debug)]
pub enum Mailbox<T> {
Tx(Vec<u8>),
Tx(FBuf),
Rx(AlignedVec),
Plain(T),
}

impl<T> Mailbox<T> {
fn into_tx(self) -> Option<Vec<u8>> {
fn into_tx(self) -> Option<FBuf> {
match self {
Mailbox::Tx(bytes) => Some(bytes),
Mailbox::Rx(_) | Mailbox::Plain(_) => None,
Expand Down Expand Up @@ -792,7 +793,7 @@ where
mut serialize: F,
) -> bool
where
F: FnMut(T) -> Vec<u8> + Send + Sync,
F: FnMut(T) -> FBuf + Send + Sync,
{
self.try_send_all(
sender,
Expand Down Expand Up @@ -870,7 +871,7 @@ where
for host in runtime.layout().other_hosts() {
let receivers = &host.workers;
let mut serialized_bytes = 0;
let items: Vec<Vec<_>> = senders
let items: Vec<Vec<FBuf>> = senders
.clone()
.map(|sender| {
receivers
Expand Down Expand Up @@ -1134,7 +1135,7 @@ where
/// match location {
/// WorkerLocation::Local => vals.push(Mailbox::Plain(n)),
/// WorkerLocation::Remote => {
/// vals.push(Mailbox::Tx(to_bytes_dyn(&n).unwrap().into_vec()))
/// vals.push(Mailbox::Tx(to_bytes_dyn(&n).unwrap()))
/// }
/// }
/// }
Expand Down Expand Up @@ -1631,7 +1632,7 @@ mod tests {
if exchange.try_send_all_with_serializer(
Runtime::worker_index(),
repeat(round),
|round| to_bytes(&round).unwrap().into_vec(),
|round| to_bytes(&round).unwrap(),
) {
break;
}
Expand Down Expand Up @@ -1693,7 +1694,7 @@ mod tests {
match location {
WorkerLocation::Local => vals.push(Mailbox::Plain(n)),
WorkerLocation::Remote => {
vals.push(Mailbox::Tx(to_bytes_dyn(&n).unwrap().into_vec()))
vals.push(Mailbox::Tx(to_bytes_dyn(&n).unwrap()))
}
}
}
Expand Down
8 changes: 4 additions & 4 deletions crates/dbsp/src/operator/dynamic/balance/accumulate_trace_balanced.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1231,16 +1231,16 @@ where
fn serialize_with_flush<B, K, V, R>(
(batch, flush): (B, bool),
serializer_inner: &mut Option<SerializerInner>,
) -> Vec<u8>
) -> FBuf
where
B: BatchReader<Key = K, Val = V, Time = (), R = R>,
K: DataTrait + ?Sized,
V: DataTrait + ?Sized,
R: WeightTrait + ?Sized,
{
let mut vec = serialize_indexed_wset(&batch, serializer_inner.get_or_insert_default());
vec.push(flush as u8);
vec
let mut fbuf = serialize_indexed_wset(&batch, serializer_inner.get_or_insert_default());
fbuf.push(flush as u8);
fbuf
}

stream! {
Expand Down
4 changes: 2 additions & 2 deletions crates/dbsp/src/operator/dynamic/communication/shard.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -398,13 +398,13 @@ impl PairsSerializer {
.unwrap();
}

pub fn done(mut self, serializer: &mut SerializerInner) -> Vec<u8> {
pub fn done(mut self, serializer: &mut SerializerInner) -> FBuf {
serializer
.with(FBufSerializer::new(&mut self.fbuf), |s| {
s.serialize_value(&self.offsets)
})
.unwrap();
self.fbuf.into_vec()
self.fbuf
}
}

Expand Down
13 changes: 7 additions & 6 deletions crates/dbsp/src/operator/dynamic/time_series/waterline.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use dyn_clone::clone_box;
use feldera_storage::fbuf::FBuf;
use size_of::SizeOf;

use crate::circuit::checkpointer::Checkpoint;
Expand Down Expand Up @@ -94,9 +95,9 @@ where
WorkerLocation::Local => {
waterlines.push(Mailbox::Plain(clone_box(waterline.as_ref())))
}
WorkerLocation::Remote => {
waterlines.push(Mailbox::Tx(waterline.checkpoint().unwrap()))
}
WorkerLocation::Remote => waterlines.push(Mailbox::Tx(
FBuf::from_slice(&waterline.checkpoint().unwrap()),
)),
};
}
},
Expand Down Expand Up @@ -175,9 +176,9 @@ where
WorkerLocation::Local => {
waterlines.push(Mailbox::Plain(clone_box(waterline.as_ref())))
}
WorkerLocation::Remote => {
waterlines.push(Mailbox::Tx(waterline.checkpoint().unwrap()))
}
WorkerLocation::Remote => waterlines.push(Mailbox::Tx(
FBuf::from_slice(&waterline.checkpoint().unwrap()),
)),
};
}
},
Expand Down
9 changes: 3 additions & 6 deletions crates/dbsp/src/trace.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1454,22 +1454,19 @@ impl IndexedWSetSerializer {
});
}

pub fn done(mut self, serializer_inner: &mut SerializerInner) -> Vec<u8> {
pub fn done(mut self, serializer_inner: &mut SerializerInner) -> FBuf {
#[cfg(debug_assertions)]
debug_assert_eq!(self.state, State::Key);
self.offsets[0] = self.n_keys;
self.offsets[1] = self.n_values;
serializer_inner.with(FBufSerializer::new(&mut self.fbuf), |s| {
s.serialize_value(&self.offsets).unwrap()
});
self.fbuf.into_vec()
self.fbuf
}
}

pub fn serialize_indexed_wset<B, K, V, R>(
batch: &B,
serializer_inner: &mut SerializerInner,
) -> Vec<u8>
pub fn serialize_indexed_wset<B, K, V, R>(batch: &B, serializer_inner: &mut SerializerInner) -> FBuf
where
B: BatchReader<Key = K, Val = V, Time = (), R = R>,
K: DataTrait + ?Sized,
Expand Down
2 changes: 1 addition & 1 deletion crates/dbsp/src/trace/ord/fallback/indexed_wset.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -743,7 +743,7 @@ where
R: WeightTrait + ?Sized,
{
fn checkpoint(&self) -> Result<Vec<u8>, Error> {
Ok(serialize_indexed_wset(self, &mut SerializerInner::new()))
Ok(serialize_indexed_wset(self, &mut SerializerInner::new()).into_vec())
}

fn restore(&mut self, data: &[u8]) -> Result<(), Error> {
Expand Down
2 changes: 1 addition & 1 deletion crates/dbsp/src/trace/ord/vec/indexed_wset_batch.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -976,7 +976,7 @@ where
R: WeightTrait + ?Sized,
{
fn checkpoint(&self) -> Result<Vec<u8>, Error> {
Ok(serialize_indexed_wset(self, &mut SerializerInner::new()))
Ok(serialize_indexed_wset(self, &mut SerializerInner::new()).into_vec())
}

fn restore(&mut self, data: &[u8]) -> Result<(), Error> {
Expand Down
6 changes: 2 additions & 4 deletions crates/sqllib/src/string_interner.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -298,15 +298,13 @@ pub fn build_string_interner(
let mut locations = WorkerLocations::new();
match locations.next().unwrap() {
WorkerLocation::Local => outputs.push(Mailbox::Plain(spine.clone())),
WorkerLocation::Remote => {
outputs.push(Mailbox::Tx(to_bytes(&spine).unwrap().into_vec()))
}
WorkerLocation::Remote => outputs.push(Mailbox::Tx(to_bytes(&spine).unwrap())),
};
for location in locations {
match location {
WorkerLocation::Local => outputs.push(Mailbox::Plain(empty_by_id())),
WorkerLocation::Remote => {
outputs.push(Mailbox::Tx(to_bytes(&empty_by_id()).unwrap().into_vec()))
outputs.push(Mailbox::Tx(to_bytes(&empty_by_id()).unwrap()))
}
}
}
Expand Down
7 changes: 7 additions & 0 deletions crates/storage/src/fbuf.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -458,6 +458,13 @@ impl FBuf {
self.into_vec().into_boxed_slice()
}

/// Creates an `FBuf` by copying the slice.
pub fn from_slice(slice: &[u8]) -> Self {
let mut fbuf = FBuf::new();
fbuf.extend_from_slice(slice);
fbuf
}

/// Converts the vector into `Vec<u8>`.
///
/// This method reallocates and copies the underlying bytes. Any excess
Expand Down
Loading