[Ingress] Handle IngestRequest message

muhamadazmy · muhamadazmy · commit 8b8db7ffebbb · 2025-11-13T09:59:15.000+01:00
Summary:
Handle the incoming `IngestRequest` messages sent by the `ingress-core`
diff --git a/crates/worker/src/partition/leadership/leader_state.rs b/crates/worker/src/partition/leadership/leader_state.rs
@@ -21,8 +21,6 @@ use futures::future::OptionFuture;
 use futures::stream::FuturesUnordered;
 use futures::{FutureExt, StreamExt, stream};
 use metrics::counter;
-use restate_types::logs::Keys;
-use restate_wal_protocol::control::UpsertSchema;
 use tokio_stream::wrappers::{ReceiverStream, WatchStream};
 use tracing::{debug, trace};
 
@@ -35,12 +33,15 @@ use restate_types::identifiers::{
     WithPartitionKey,
 };
 use restate_types::invocation::client::{InvocationOutput, SubmittedInvocationNotification};
+use restate_types::logs::Keys;
+use restate_types::net::ingress::IngestRecord;
 use restate_types::net::partition_processor::{
     PartitionProcessorRpcError, PartitionProcessorRpcResponse,
 };
 use restate_types::time::MillisSinceEpoch;
 use restate_types::{SemanticRestateVersion, Version, Versioned};
 use restate_wal_protocol::Command;
+use restate_wal_protocol::control::UpsertSchema;
 use restate_wal_protocol::timer::TimerKeyValue;
 
 use crate::metric_definitions::{PARTITION_HANDLE_LEADER_ACTIONS, USAGE_LEADER_ACTION_COUNT};
@@ -380,14 +381,35 @@ impl LeaderState {
             Ok(commit_token) => {
                 self.awaiting_rpc_self_propose.push(SelfAppendFuture::new(
                     commit_token,
-                    success_response,
-                    reciprocal,
+                    |result: Result<(), PartitionProcessorRpcError>| {
+                        reciprocal.send(result.map(|_| success_response));
+                    },
                 ));
             }
             Err(e) => reciprocal.send(Err(PartitionProcessorRpcError::Internal(e.to_string()))),
         }
     }
 
+    pub async fn propose_many_with_callback<F>(
+        &mut self,
+        records: impl ExactSizeIterator<Item = IngestRecord>,
+        callback: F,
+    ) where
+        F: FnOnce(Result<(), PartitionProcessorRpcError>) + Send + Sync + 'static,
+    {
+        match self
+            .self_proposer
+            .propose_many_with_notification(records)
+            .await
+        {
+            Ok(commit_token) => {
+                self.awaiting_rpc_self_propose
+                    .push(SelfAppendFuture::new(commit_token, callback));
+            }
+            Err(e) => callback(Err(PartitionProcessorRpcError::Internal(e.to_string()))),
+        }
+    }
+
     pub fn handle_actions(
         &mut self,
         invoker_tx: &mut impl restate_invoker_api::InvokerHandle<InvokerStorageReader<PartitionStore>>,
@@ -594,42 +616,72 @@ impl LeaderState {
     }
 }
 
+trait CallbackInner: Send + Sync + 'static {
+    fn call(self: Box<Self>, result: Result<(), PartitionProcessorRpcError>);
+}
+
+impl<F> CallbackInner for F
+where
+    F: FnOnce(Result<(), PartitionProcessorRpcError>) + Send + Sync + 'static,
+{
+    fn call(self: Box<Self>, result: Result<(), PartitionProcessorRpcError>) {
+        self(result)
+    }
+}
+
+struct Callback {
+    inner: Box<dyn CallbackInner>,
+}
+
+impl Callback {
+    fn call(self, result: Result<(), PartitionProcessorRpcError>) {
+        self.inner.call(result);
+    }
+}
+
+impl<I> From<I> for Callback
+where
+    I: CallbackInner,
+{
+    fn from(value: I) -> Self {
+        Self {
+            inner: Box::new(value),
+        }
+    }
+}
+
 struct SelfAppendFuture {
     commit_token: CommitToken,
-    response: Option<(PartitionProcessorRpcResponse, RpcReciprocal)>,
+    callback: Option<Callback>,
 }
 
 impl SelfAppendFuture {
-    fn new(
-        commit_token: CommitToken,
-        success_response: PartitionProcessorRpcResponse,
-        response_reciprocal: RpcReciprocal,
-    ) -> Self {
+    fn new(commit_token: CommitToken, callback: impl Into<Callback>) -> Self {
         Self {
             commit_token,
-            response: Some((success_response, response_reciprocal)),
+            callback: Some(callback.into()),
         }
     }
 
     fn fail_with_internal(&mut self) {
-        if let Some((_, reciprocal)) = self.response.take() {
-            reciprocal.send(Err(PartitionProcessorRpcError::Internal(
+        if let Some(callback) = self.callback.take() {
+            callback.call(Err(PartitionProcessorRpcError::Internal(
                 "error when proposing to bifrost".to_string(),
             )));
         }
     }
 
     fn fail_with_lost_leadership(&mut self, this_partition_id: PartitionId) {
-        if let Some((_, reciprocal)) = self.response.take() {
-            reciprocal.send(Err(PartitionProcessorRpcError::LostLeadership(
+        if let Some(callback) = self.callback.take() {
+            callback.call(Err(PartitionProcessorRpcError::LostLeadership(
                 this_partition_id,
             )));
         }
     }
 
     fn succeed_with_appended(&mut self) {
-        if let Some((success_response, reciprocal)) = self.response.take() {
-            reciprocal.send(Ok(success_response));
+        if let Some(callback) = self.callback.take() {
+            callback.call(Ok(()))
         }
     }
 }
diff --git a/crates/worker/src/partition/leadership/mod.rs b/crates/worker/src/partition/leadership/mod.rs
@@ -46,14 +46,15 @@ use restate_types::errors::GenericError;
 use restate_types::identifiers::{InvocationId, PartitionKey, PartitionProcessorRpcRequestId};
 use restate_types::identifiers::{LeaderEpoch, PartitionLeaderEpoch};
 use restate_types::message::MessageIndex;
+use restate_types::net::ingress::IngestRecord;
 use restate_types::net::partition_processor::{
     PartitionProcessorRpcError, PartitionProcessorRpcResponse,
 };
 use restate_types::partitions::Partition;
 use restate_types::partitions::state::PartitionReplicaSetStates;
 use restate_types::retries::with_jitter;
 use restate_types::schema::Schema;
-use restate_types::storage::StorageEncodeError;
+use restate_types::storage::{StorageDecodeError, StorageEncodeError};
 use restate_wal_protocol::Command;
 use restate_wal_protocol::control::{AnnounceLeader, PartitionDurability};
 use restate_wal_protocol::timer::TimerKeyValue;
@@ -82,7 +83,9 @@ pub(crate) enum Error {
     #[error("failed writing to bifrost: {0}")]
     Bifrost(#[from] restate_bifrost::Error),
     #[error("failed serializing payload: {0}")]
-    Codec(#[from] StorageEncodeError),
+    Encode(#[from] StorageEncodeError),
+    #[error("failed deserializing payload: {0}")]
+    Decode(#[from] StorageDecodeError),
     #[error(transparent)]
     Shutdown(#[from] ShutdownError),
     #[error("error when self proposing")]
@@ -612,6 +615,26 @@ impl<I> LeadershipState<I> {
             }
         }
     }
+
+    /// propose to this partition
+    pub async fn propose_many_with_callback<F>(
+        &mut self,
+        records: impl ExactSizeIterator<Item = IngestRecord>,
+        callback: F,
+    ) where
+        F: FnOnce(Result<(), PartitionProcessorRpcError>) + Send + Sync + 'static,
+    {
+        match &mut self.state {
+            State::Follower | State::Candidate { .. } => callback(Err(
+                PartitionProcessorRpcError::NotLeader(self.partition.partition_id),
+            )),
+            State::Leader(leader_state) => {
+                leader_state
+                    .propose_many_with_callback(records, callback)
+                    .await;
+            }
+        }
+    }
 }
 #[derive(Debug, derive_more::From)]
 struct TimerReader(PartitionStore);
diff --git a/crates/worker/src/partition/leadership/self_proposer.rs b/crates/worker/src/partition/leadership/self_proposer.rs
@@ -14,7 +14,9 @@ use futures::never::Never;
 
 use restate_bifrost::{Bifrost, CommitToken, ErrorRecoveryStrategy};
 use restate_storage_api::deduplication_table::{DedupInformation, EpochSequenceNumber};
-use restate_types::{identifiers::PartitionKey, logs::LogId};
+use restate_types::{
+    identifiers::PartitionKey, logs::LogId, net::ingress::IngestRecord, storage::StorageCodec,
+};
 use restate_wal_protocol::{Command, Destination, Envelope, Header, Source};
 
 use crate::partition::leadership::Error;
@@ -100,6 +102,44 @@ impl SelfProposer {
         Ok(commit_token)
     }
 
+    pub async fn propose_many_with_notification(
+        &mut self,
+        records: impl ExactSizeIterator<Item = IngestRecord>,
+    ) -> Result<CommitToken, Error> where {
+        let sender = self.bifrost_appender.sender();
+
+        // This is ideally should be implemented
+        // by using `sender.enqueue_many`
+        // but since we have no guarantee over the
+        // underlying channel size. a `reserve_many()` might
+        // block forever waiting for n permits that will
+        // never be available.
+        //
+        // sender
+        //     .enqueue_many(records)
+        //     .await
+        //     .map_err(|_| Error::SelfProposer)?;
+        //
+        // so instead we do this.
+
+        for mut record in records {
+            // todo: unfortunately we need to decode tha pyaload first although
+            // the appended will need to encode it eventually. Maybe if there
+            // is a way to pass the raw encoded data directly to the appender
+            let envelope = StorageCodec::decode(&mut record.record)?;
+
+            sender
+                .enqueue(Arc::new(envelope))
+                .await
+                .map_err(|_| Error::SelfProposer)?;
+        }
+
+        sender
+            .notify_committed()
+            .await
+            .map_err(|_| Error::SelfProposer)
+    }
+
     fn create_header(&mut self, partition_key: PartitionKey) -> Header {
         let esn = self.epoch_sequence_number;
         self.epoch_sequence_number = self.epoch_sequence_number.next();
diff --git a/crates/worker/src/partition/mod.rs b/crates/worker/src/partition/mod.rs
@@ -31,7 +31,7 @@ use tracing::{Span, debug, error, info, instrument, trace, warn};
 
 use restate_bifrost::loglet::FindTailOptions;
 use restate_bifrost::{Bifrost, LogEntry, MaybeRecord};
-use restate_core::network::{Oneshot, Reciprocal, ServiceMessage, Verdict};
+use restate_core::network::{Incoming, Oneshot, Reciprocal, Rpc, ServiceMessage, Verdict};
 use restate_core::{Metadata, ShutdownError, cancellation_watcher, my_node_id};
 use restate_partition_store::{PartitionStore, PartitionStoreTransaction};
 use restate_storage_api::deduplication_table::{
@@ -47,6 +47,7 @@ use restate_types::config::Configuration;
 use restate_types::identifiers::LeaderEpoch;
 use restate_types::logs::{KeyFilter, Lsn, Record, SequenceNumber};
 use restate_types::net::RpcRequest;
+use restate_types::net::ingress::{IngestResponse, ReceivedIngestRequest};
 use restate_types::net::partition_processor::{
     PartitionLeaderService, PartitionProcessorRpcError, PartitionProcessorRpcRequest,
     PartitionProcessorRpcResponse,
@@ -463,15 +464,7 @@ where
                     self.status.effective_mode = self.leadership_state.effective_mode();
                 }
                 Some(msg) = self.network_leader_svc_rx.recv() => {
-                    match msg {
-                        ServiceMessage::Rpc(msg) if msg.msg_type() == PartitionProcessorRpcRequest::TYPE => {
-                            let msg = msg.into_typed::<PartitionProcessorRpcRequest>();
-                            // note: split() decodes the payload
-                            let (response_tx, body) = msg.split();
-                            self.on_rpc(response_tx, body, &mut partition_store, live_schemas.live_load()).await;
-                        }
-                        msg => { msg.fail(Verdict::MessageUnrecognized); }
-                    }
+                    self.on_rpc(msg, &mut partition_store, live_schemas.live_load()).await;
                 }
                 _ = status_update_timer.tick() => {
                     if durable_lsn_watch.has_changed().map_err(|e| ProcessorError::Other(e.into()))? {
@@ -599,7 +592,7 @@ where
         Ok(())
     }
 
-    async fn on_rpc(
+    async fn on_pp_rpc_request(
         &mut self,
         response_tx: Reciprocal<
             Oneshot<Result<PartitionProcessorRpcResponse, PartitionProcessorRpcError>>,
@@ -615,6 +608,57 @@ where
         )
         .await;
     }
+
+    async fn on_rpc(
+        &mut self,
+        msg: ServiceMessage<PartitionLeaderService>,
+        partition_store: &mut PartitionStore,
+        schemas: &Schema,
+    ) {
+        match msg {
+            ServiceMessage::Rpc(msg) if msg.msg_type() == PartitionProcessorRpcRequest::TYPE => {
+                let msg = msg.into_typed::<PartitionProcessorRpcRequest>();
+                // note: split() decodes the payload
+                let (response_tx, body) = msg.split();
+                self.on_pp_rpc_request(response_tx, body, partition_store, schemas)
+                    .await;
+            }
+            ServiceMessage::Rpc(msg) if msg.msg_type() == ReceivedIngestRequest::TYPE => {
+                self.on_pp_ingest_request(msg.into_typed()).await;
+            }
+            msg => {
+                msg.fail(Verdict::MessageUnrecognized);
+            }
+        }
+    }
+
+    async fn on_pp_ingest_request(&mut self, msg: Incoming<Rpc<ReceivedIngestRequest>>) {
+        let (reciprocal, request) = msg.split();
+        self.leadership_state
+            .propose_many_with_callback(
+                request.records.into_iter(),
+                |result: Result<(), PartitionProcessorRpcError>| match result {
+                    Ok(_) => reciprocal.send(IngestResponse::Ack),
+                    Err(err) => match err {
+                        PartitionProcessorRpcError::NotLeader(id)
+                        | PartitionProcessorRpcError::LostLeadership(id) => {
+                            reciprocal.send(IngestResponse::NotLeader { of: id })
+                        }
+                        PartitionProcessorRpcError::Starting => {
+                            reciprocal.send(IngestResponse::Starting)
+                        }
+                        PartitionProcessorRpcError::Stopping => {
+                            reciprocal.send(IngestResponse::Stopping)
+                        }
+                        PartitionProcessorRpcError::Internal(msg) => {
+                            reciprocal.send(IngestResponse::Internal { msg })
+                        }
+                    },
+                },
+            )
+            .await;
+    }
+
     async fn maybe_advance<'a>(
         &mut self,
         maybe_record: LogEntry,
