Restore ingestion pipelining with epoch-pinned ordering

crates/bifrost/src/background_appender.rs

@@ -357,6 +357,26 @@ impl<T: StorageEncode> LogSender<T> {
         Ok(())
     }
 
+    /// Waits for capacity on the channel and returns an error if the appender is
+    /// draining or drained.
+    ///
+    /// Unlike [`enqueue`](Self::enqueue), this does not check the record size and
+    /// accepts a record of any size.
+    ///
+    /// Callers have to ensure that record is not larger than the network message size limit.
+    pub async fn enqueue_unchecked<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().ensure_encoded(&mut self.arena);
+        permit.send(AppendOperation::Enqueue(record));
+
+        Ok(())
+    }
+
     /// Attempt to enqueue a record to the appender. Returns immediately if the
     /// appender is pushing back or if the appender is draining or drained.
     ///

crates/bifrost/src/error.rs

@@ -67,6 +67,18 @@ pub enum EnqueueError<T> {
     },
 }
 
+impl<T> EnqueueError<T> {
+    pub fn drop_payload(self) -> EnqueueError<()> {
+        match self {
+            Self::Closed(_) => EnqueueError::Closed(()),
+            Self::Full(_) => EnqueueError::Full(()),
+            Self::RecordTooLarge { record_size, limit } => {
+                EnqueueError::RecordTooLarge { record_size, limit }
+            }
+        }
+    }
+}
+
 #[derive(Clone, Debug, thiserror::Error)]
 pub enum AdminError {
     #[error("log {0} is permanently sealed")]

crates/bifrost/src/lib.rs

@@ -29,7 +29,7 @@ pub use appender::Appender;
 pub use background_appender::{AppenderHandle, BackgroundAppender, CommitToken, LogSender};
 pub use bifrost::{Bifrost, ErrorRecoveryStrategy};
 pub use bifrost_admin::{BifrostAdmin, MaybeSealedSegment};
-pub use error::{Error, Result};
+pub use error::{EnqueueError, Error, Result};
 pub use read_stream::LogReadStream;
 pub use read_stream_registry::ActiveReadStreamRegistry;
 pub use record::{InputRecord, LogEntry, MaybeRecord, RecordKind};

crates/core/src/partitions.rs

@@ -118,4 +118,8 @@ impl PartitionRouting {
             .expect("partition replica set states are never dropped while in use");
         *state
     }
+
+    pub fn partition_replica_set_state(&self) -> &PartitionReplicaSetStates {
+        &self.partition_replica_set_states
+    }
 }

crates/core/src/task_center/task_kind.rs

@@ -158,6 +158,8 @@ pub enum TaskKind {
     LogStoreWriter,
     // - Datafusion
     DfScanner,
+    #[strum(props(runtime = "default"))]
+    IngestionSession,
 }
 
 impl TaskKind {

crates/ingestion-client/Cargo.toml

@@ -7,6 +7,9 @@ rust-version.workspace = true
 license.workspace = true
 publish = false
 
+[features]
+test-util = []
+
 [dependencies]
 
 restate-core = { workspace = true }
@@ -15,6 +18,7 @@ restate-workspace-hack = { workspace = true }
 
 bytes = { workspace = true }
 dashmap = { workspace = true }
+derive_builder = { workspace = true }
 futures = { workspace = true }
 pin-project = { workspace = true }
 thiserror = { workspace = true }

crates/ingestion-client/src/chunks_size.rs

@@ -30,6 +30,10 @@ impl<F, S: Stream> ChunksSize<F, S>
 where
     F: Fn(&S::Item) -> usize,
 {
+    pub fn new(stream: S, max_size: usize, size_fn: F) -> Self {
+        Self::with_buffered(stream, max_size, size_fn, Vec::default())
+    }
+
     /// Creates a chunker, pre-seeding the accumulation buffer with `buffered` items so they lead the
     /// next emitted chunk. Pass an empty `Vec` to start fresh. Seeding is used to carry a previously
     /// over-pulled item across a recreated `ChunksSize` (e.g. after a reconnect) without losing it or