[Bifrost] BackgroundAppender serializes payload to ensure correct size estimation

Part of #4130, #4132

Author: AhmedSoliman

crates/bifrost/src/background_appender.rs

@@ -10,6 +10,7 @@
 
 use std::num::NonZeroUsize;
 
+use bytes::BytesMut;
 use futures::FutureExt;
 use pin_project::pin_project;
 use restate_types::logs::Record;
@@ -70,6 +71,7 @@ where
         Ok(AppenderHandle {
             inner_handle: Some(handle),
             sender: Some(LogSender {
+                arena: BytesMut::default(),
                 tx,
                 record_size_limit,
                 _phantom: std::marker::PhantomData,
@@ -213,8 +215,8 @@ impl<T> AppenderHandle<T> {
     }
 
     /// If you need an owned LogSender, clone this.
-    pub fn sender(&self) -> &LogSender<T> {
-        self.sender.as_ref().unwrap()
+    pub fn sender(&mut self) -> &mut LogSender<T> {
+        self.sender.as_mut().unwrap()
     }
 
     /// Waits for the underlying appender task to finish.
@@ -231,13 +233,24 @@ impl<T> AppenderHandle<T> {
     }
 }
 
-#[derive(Clone)]
 pub struct LogSender<T> {
+    arena: BytesMut,
     tx: tokio::sync::mpsc::Sender<AppendOperation>,
     record_size_limit: NonZeroUsize,
     _phantom: std::marker::PhantomData<T>,
 }
 
+impl<T> Clone for LogSender<T> {
+    fn clone(&self) -> Self {
+        Self {
+            arena: BytesMut::default(),
+            tx: self.tx.clone(),
+            record_size_limit: self.record_size_limit,
+            _phantom: std::marker::PhantomData,
+        }
+    }
+}
+
 impl<T: StorageEncode> LogSender<T> {
     fn check_record_size<E>(&self, record: &Record) -> Result<(), EnqueueError<E>> {
         let record_size = record.estimated_encode_size();
@@ -252,7 +265,7 @@ impl<T: StorageEncode> LogSender<T> {
 
     /// Attempt to enqueue a record to the appender. Returns immediately if the
     /// appender is pushing back or if the appender is draining or drained.
-    pub fn try_enqueue<A>(&self, record: A) -> Result<(), EnqueueError<A>>
+    pub fn try_enqueue<A>(&mut self, record: A) -> Result<(), EnqueueError<A>>
     where
         A: Into<InputRecord<T>>,
     {
@@ -262,15 +275,15 @@ impl<T: StorageEncode> LogSender<T> {
             Err(mpsc::error::TrySendError::Closed(_)) => return Err(EnqueueError::Closed(record)),
         };
 
-        let record = record.into().into_record();
+        let record = record.into().into_record().ensure_encoded(&mut self.arena);
         self.check_record_size(&record)?;
         permit.send(AppendOperation::Enqueue(record));
         Ok(())
     }
 
     /// Enqueues an append and returns a commit token
     pub fn try_enqueue_with_notification<A>(
-        &self,
+        &mut self,
         record: A,
     ) -> Result<CommitToken, EnqueueError<A>>
     where
@@ -283,22 +296,22 @@ impl<T: StorageEncode> LogSender<T> {
         };
 
         let (tx, rx) = oneshot::channel();
-        let record = record.into().into_record();
+        let record = record.into().into_record().ensure_encoded(&mut self.arena);
         self.check_record_size(&record)?;
         permit.send(AppendOperation::EnqueueWithNotification(record, tx));
         Ok(CommitToken { rx })
     }
 
     /// Waits for capacity on the channel and returns an error if the appender is
     /// draining or drained.
-    pub async fn enqueue<A>(&self, record: A) -> Result<(), EnqueueError<A>>
+    pub async fn enqueue<A>(&mut self, record: A) -> Result<(), EnqueueError<A>>
     where
         A: Into<InputRecord<T>>,
     {
         let Ok(permit) = self.tx.reserve().await else {
             return Err(EnqueueError::Closed(record));
         };
-        let record = record.into().into_record();
+        let record = record.into().into_record().ensure_encoded(&mut self.arena);
         self.check_record_size(&record)?;
         permit.send(AppendOperation::Enqueue(record));
 
@@ -310,7 +323,7 @@ impl<T: StorageEncode> LogSender<T> {
     ///
     /// Attempts to enqueue all records in the iterator. This will immediately return if there is
     /// no capacity in the channel to enqueue _all_ records.
-    pub fn try_enqueue_many<I, A>(&self, records: I) -> Result<(), EnqueueError<I>>
+    pub fn try_enqueue_many<I, A>(&mut self, records: I) -> Result<(), EnqueueError<I>>
     where
         I: Iterator<Item = A> + ExactSizeIterator,
         A: Into<InputRecord<T>>,
@@ -322,7 +335,7 @@ impl<T: StorageEncode> LogSender<T> {
         };
 
         for (permit, record) in std::iter::zip(permits, records) {
-            let record = record.into().into_record();
+            let record = record.into().into_record().ensure_encoded(&mut self.arena);
             self.check_record_size(&record)?;
             permit.send(AppendOperation::Enqueue(record));
         }
@@ -334,7 +347,7 @@ impl<T: StorageEncode> LogSender<T> {
     ///
     /// The method is cancel safe in the sense that if enqueue_many is used in a `tokio::select!`,
     /// no records are enqueued if another branch completed.
-    pub async fn enqueue_many<I, A>(&self, records: I) -> Result<(), EnqueueError<I>>
+    pub async fn enqueue_many<I, A>(&mut self, records: I) -> Result<(), EnqueueError<I>>
     where
         I: Iterator<Item = A> + ExactSizeIterator,
         A: Into<InputRecord<T>>,
@@ -344,7 +357,7 @@ impl<T: StorageEncode> LogSender<T> {
         };
 
         for (permit, record) in std::iter::zip(permits, records) {
-            let record = record.into().into_record();
+            let record = record.into().into_record().ensure_encoded(&mut self.arena);
             self.check_record_size(&record)?;
             permit.send(AppendOperation::Enqueue(record));
         }
@@ -355,7 +368,7 @@ impl<T: StorageEncode> LogSender<T> {
     /// Enqueues a record and returns a [`CommitToken`] future that's resolved when the record is
     /// committed.
     pub async fn enqueue_with_notification<A>(
-        &self,
+        &mut self,
         record: A,
     ) -> Result<CommitToken, EnqueueError<A>>
     where
@@ -366,7 +379,7 @@ impl<T: StorageEncode> LogSender<T> {
         };
 
         let (tx, rx) = oneshot::channel();
-        let record = record.into().into_record();
+        let record = record.into().into_record().ensure_encoded(&mut self.arena);
         self.check_record_size(&record)?;
         permit.send(AppendOperation::EnqueueWithNotification(record, tx));
 

crates/bifrost/src/bifrost.rs

@@ -141,28 +141,6 @@ impl Bifrost {
             .await
     }
 
-    /// Appends a batch of records to a log. The log id must exist, otherwise the
-    /// operation fails with [`Error::UnknownLogId`]. The returned Lsn is the Lsn of the last
-    /// record in this batch. This will only return after all records have been stored.
-    #[instrument(
-        level="trace",
-        skip(self, batch),
-        fields(
-            otel.name = "Bifrost: append_batch",
-        )
-    )]
-    pub async fn append_batch<T: StorageEncode>(
-        &self,
-        log_id: LogId,
-        error_recovery_strategy: ErrorRecoveryStrategy,
-        batch: Vec<impl Into<InputRecord<T>>>,
-    ) -> Result<Lsn> {
-        self.inner.fail_if_shutting_down()?;
-        self.inner
-            .append_batch(log_id, error_recovery_strategy, batch)
-            .await
-    }
-
     /// Read the next record from the LSN provided. The `from` indicates the LSN where we will
     /// start reading from. This means that the record returned will have a LSN that is equal
     /// or greater than `from`. If no records are committed yet at this LSN, this read operation
@@ -343,17 +321,6 @@ impl BifrostInner {
             .await
     }
 
-    pub async fn append_batch<T: StorageEncode>(
-        self: &Arc<Self>,
-        log_id: LogId,
-        error_recovery_strategy: ErrorRecoveryStrategy,
-        batch: Vec<impl Into<InputRecord<T>>>,
-    ) -> Result<Lsn> {
-        Appender::new(log_id, error_recovery_strategy, Arc::clone(self))
-            .append_batch(batch)
-            .await
-    }
-
     #[cfg(any(test, feature = "test-util"))]
     pub async fn read(
         self: &Arc<Self>,
@@ -1357,11 +1324,11 @@ mod tests {
         let background_appender: crate::BackgroundAppender<String> = bifrost
             .create_background_appender(LogId::new(0), ErrorRecoveryStrategy::Wait, 10, 10)?;
 
-        let handle = background_appender.start("test-appender")?;
+        let mut handle = background_appender.start("test-appender")?;
         let sender = handle.sender();
 
-        // Any record will have an estimated size of ~2KB due to PolyBytes::Typed constant estimate
-        let payload = "test".to_string();
+        // A string with 100 bytes
+        let payload = String::from_utf8(vec![b't'; 100]).unwrap();
 
         // try_enqueue should fail with RecordTooLarge
         let result = sender.try_enqueue(payload.clone());
@@ -1413,7 +1380,7 @@ mod tests {
         let background_appender: crate::BackgroundAppender<String> = bifrost
             .create_background_appender(LogId::new(0), ErrorRecoveryStrategy::Wait, 10, 10)?;
 
-        let handle = background_appender.start("test-appender")?;
+        let mut handle = background_appender.start("test-appender")?;
         let sender = handle.sender();
 
         // With a 10KB limit, the ~2KB estimated record should succeed

crates/bifrost/src/providers/local_loglet/record_format.rs

@@ -124,6 +124,7 @@ fn write_legacy_payload(record: &Record, serde_buffer: &mut BytesMut) -> BytesMu
     // encoding the user payload.
     let body = match record.body() {
         PolyBytes::Bytes(raw_bytes) => raw_bytes.clone(),
+        PolyBytes::Both(_, raw_bytes) => raw_bytes.clone(),
         PolyBytes::Typed(encodeable) => {
             StorageCodec::encode_and_split(encodeable.deref(), serde_buffer)
                 .expect("record serde is infallible")
@@ -183,6 +184,7 @@ fn write_record(record: &Record, buf: &mut BytesMut) -> BytesMut {
     // serialize payload
     match record.body() {
         PolyBytes::Bytes(raw_bytes) => buf.put_slice(raw_bytes),
+        PolyBytes::Both(_, raw_bytes) => buf.put_slice(raw_bytes),
         PolyBytes::Typed(encodeable) => {
             StorageCodec::encode(encodeable.deref(), buf).expect("record serde is infallible")
         }

crates/bifrost/src/providers/replicated_loglet/sequencer.rs

@@ -15,7 +15,6 @@ use std::sync::{
     atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering},
 };
 
-use bytes::BytesMut;
 use crossbeam_utils::CachePadded;
 use tokio::sync::Semaphore;
 use tokio_util::{sync::CancellationToken, task::TaskTracker};
@@ -284,31 +283,13 @@ impl<T: TransportConnect> Sequencer<T> {
                 .fetch_add(len, std::sync::atomic::Ordering::Relaxed),
         );
 
-        // Add to cache before we commit those records, we also do this to try and cache the
-        // records before we transform/serialize their payloads.
-        //
         // The records being in cache does not mean they are committed, all readers must respect
         // the result of find_tail() or the global_known_tail.
         self.record_cache
             .extend(*self.sequencer_shared_state.loglet_id(), offset, &payloads);
 
         let (loglet_commit, commit_resolver) = LogletCommit::deferred();
 
-        // Why not in bifrost's appender? because we don't need serialization in all loglet providers.
-        // In-memory loglet won't require this setup and it would be a waste of cpu-cycles to serialize.
-        // estimate total to allocate one big allocation for all payloads. This is not particularly
-        // useful with flexbuffers, but will be useful when Store/Record are representative as
-        // protobuf or flatbuffers.
-        let estimated_bufsize: usize = payloads
-            .iter()
-            .map(|record| record.estimated_encode_size())
-            .sum();
-        let mut buf = BytesMut::with_capacity(estimated_bufsize);
-        let payloads = payloads
-            .iter()
-            .map(|record| record.to_encoded(&mut buf))
-            .collect();
-
         let appender = SequencerAppender::new(
             Arc::clone(&self.sequencer_shared_state),
             self.networking.clone(),

crates/types/src/logs/record.rs

@@ -46,14 +46,15 @@ impl Record {
         size_of::<Keys>() + size_of::<NanosSinceEpoch>() + self.body.estimated_encode_size()
     }
 
-    pub fn to_encoded(&self, buf: &mut BytesMut) -> Self {
-        let keys = self.keys.clone();
+    pub fn ensure_encoded(self, buf: &mut BytesMut) -> Self {
+        let keys = self.keys;
         let created_at = self.created_at;
-        let body = match &self.body {
-            PolyBytes::Bytes(bytes) => PolyBytes::Bytes(bytes.clone()),
+        let body = match self.body {
+            PolyBytes::Bytes(bytes) => PolyBytes::Bytes(bytes),
+            PolyBytes::Both(v, bytes) => PolyBytes::Both(v, bytes),
             PolyBytes::Typed(typed) => {
-                StorageCodec::encode(&**typed, buf).expect("serde is infallible");
-                PolyBytes::Bytes(buf.split().freeze())
+                StorageCodec::encode(&*typed, buf).expect("serde is infallible");
+                PolyBytes::Both(typed, buf.split().freeze())
             }
         };
         Self {
@@ -83,7 +84,7 @@ impl Record {
                 let mut buf = std::io::Cursor::new(slice);
                 StorageCodec::decode(&mut buf)?
             }
-            PolyBytes::Typed(value) => {
+            PolyBytes::Both(value, _) | PolyBytes::Typed(value) => {
                 let target_arc: Arc<T> = value.downcast_arc().map_err(|_| {
                 StorageDecodeError::DecodeValue(
                     anyhow::anyhow!(
@@ -113,7 +114,7 @@ impl Record {
                 let mut buf = std::io::Cursor::new(slice);
                 Arc::new(StorageCodec::decode(&mut buf)?)
             }
-            PolyBytes::Typed(value) => {
+            PolyBytes::Typed(value) | PolyBytes::Both(value, _) => {
                 value.downcast_arc().map_err(|_| {
                 StorageDecodeError::DecodeValue(
                     anyhow::anyhow!(

crates/types/src/logs/record_cache.rs

@@ -8,12 +8,16 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+use std::sync::Arc;
+
 use moka::{
     ops::compute::Op,
     policy::EvictionPolicy,
     sync::{Cache, CacheBuilder},
 };
 
+use crate::storage::PolyBytes;
+
 use super::{LogletId, LogletOffset, Record, SequenceNumber};
 
 /// Unique record key across different loglets.
@@ -64,14 +68,28 @@ impl RecordCache {
                 let Some(existing) = existing else {
                     return Op::Put(record.clone());
                 };
-                match (
-                    existing.value().body().is_encoded(),
-                    record.body().is_encoded(),
-                ) {
-                    // both are encoded, we don't want to replace the existing value.
-                    (true, true) | (false, false) | (false, true) => Op::Nop,
-                    // replace the existing value if the new one is deserialized.
-                    (true, false) => Op::Put(record.clone()),
+                match (existing.value().body(), record.body()) {
+                    (PolyBytes::Bytes(_), PolyBytes::Bytes(_)) => Op::Nop,
+                    (PolyBytes::Bytes(_), PolyBytes::Typed(_)) => Op::Put(record.clone()),
+                    (PolyBytes::Bytes(_), PolyBytes::Both(typed, _)) => {
+                        // we only need to cache the typed value, let's repackage it.
+                        Op::Put(Record::from_parts(
+                            record.created_at(),
+                            record.keys().clone(),
+                            PolyBytes::Typed(Arc::clone(typed)),
+                        ))
+                    }
+                    // Shouldn't happen (we only cache Typed or Bytes), but let's handle it anyway.
+                    (PolyBytes::Both(typed, _), _) =>
+                    // repackge the existing value into Typed only
+                    {
+                        Op::Put(Record::from_parts(
+                            existing.value().created_at(),
+                            existing.value().keys().clone(),
+                            PolyBytes::Typed(Arc::clone(typed)),
+                        ))
+                    }
+                    (PolyBytes::Typed(_), _) => Op::Nop,
                 }
             });
     }