envelope-v2: Move ingestion to use V2

Switch the code paths that uses IngestionClient to start ingesting
v2 Envelope instead of v1

Author: muhamadazmy

crates/admin/src/rest_api/invocations.rs

@@ -12,6 +12,7 @@ use axum::Json;
 use axum::extract::{Path, Query, State};
 use axum::http::StatusCode;
 use futures::future;
+use serde::Deserialize;
 use tracing::warn;
 
 use restate_admin_rest_model::invocations::{
@@ -28,12 +29,12 @@ use restate_types::invocation::client::{
 };
 use restate_types::invocation::{InvocationTermination, PurgeInvocationRequest, TerminationFlavor};
 use restate_types::journal_v2::EntryIndex;
-use restate_wal_protocol::{Command, Envelope};
-use serde::Deserialize;
+use restate_types::logs::{BodyWithKeys, Keys};
+use restate_wal_protocol::v2;
+use restate_wal_protocol::v2::commands;
 
 use super::error::*;
 use crate::generate_meta_api_error;
-use crate::rest_api::create_envelope_header;
 use crate::state::AdminServiceState;
 
 #[derive(Debug, Default, Deserialize, utoipa::ToSchema)]
@@ -88,30 +89,43 @@ where
         .parse::<InvocationId>()
         .map_err(|e| MetaApiError::InvalidField("invocation_id", e.to_string()))?;
 
-    let cmd = match mode.unwrap_or_default() {
-        DeletionMode::Cancel => Command::TerminateInvocation(InvocationTermination {
-            invocation_id,
-            flavor: TerminationFlavor::Cancel,
-            response_sink: None,
-        }),
-        DeletionMode::Kill => Command::TerminateInvocation(InvocationTermination {
-            invocation_id,
-            flavor: TerminationFlavor::Kill,
-            response_sink: None,
-        }),
-        DeletionMode::Purge => Command::PurgeInvocation(PurgeInvocationRequest {
-            invocation_id,
-            response_sink: None,
-        }),
+    let envelope = match mode.unwrap_or_default() {
+        DeletionMode::Cancel => v2::Envelope::new(
+            v2::Dedup::None,
+            commands::TerminateInvocationCommand::from(InvocationTermination {
+                invocation_id,
+                flavor: TerminationFlavor::Cancel,
+                response_sink: None,
+            }),
+        )
+        .into_raw(),
+        DeletionMode::Kill => v2::Envelope::new(
+            v2::Dedup::None,
+            commands::TerminateInvocationCommand::from(InvocationTermination {
+                invocation_id,
+                flavor: TerminationFlavor::Kill,
+                response_sink: None,
+            }),
+        )
+        .into_raw(),
+        DeletionMode::Purge => v2::Envelope::new(
+            v2::Dedup::None,
+            commands::PurgeInvocationCommand::from(PurgeInvocationRequest {
+                invocation_id,
+                response_sink: None,
+            }),
+        )
+        .into_raw(),
     };
 
     let partition_key = invocation_id.partition_key();
 
-    let envelope = Envelope::new(create_envelope_header(partition_key), cmd);
-
     let result = state
         .ingestion_client
-        .ingest(partition_key, envelope)
+        .ingest(
+            partition_key,
+            BodyWithKeys::new(envelope, Keys::Single(partition_key)),
+        )
         .await
         .map_err(|err| {
             warn!("Could not ingest invocation termination command: {err}");

crates/admin/src/rest_api/mod.rs

@@ -29,10 +29,8 @@ use utoipa::OpenApi;
 use utoipa_axum::{router::OpenApiRouter, routes};
 
 use restate_core::network::TransportConnect;
-use restate_types::identifiers::PartitionKey;
 use restate_types::invocation::client::InvocationClient;
 use restate_types::schema::registry::{DiscoveryClient, MetadataService, TelemetryClient};
-use restate_wal_protocol::{Destination, Header, Source};
 
 use crate::state::AdminServiceState;
 
@@ -187,16 +185,6 @@ where
         .with_state(state)
 }
 
-fn create_envelope_header(partition_key: PartitionKey) -> Header {
-    Header {
-        source: Source::ControlPlane {},
-        dest: Destination::Processor {
-            partition_key,
-            dedup: None,
-        },
-    }
-}
-
 /// # Error description response
 ///
 /// Error details of the response

crates/admin/src/rest_api/services.rs

@@ -8,12 +8,11 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-use tracing::{debug, warn};
-
 use axum::Json;
 use axum::extract::{Path, State};
 use bytes::Bytes;
 use http::StatusCode;
+use tracing::{debug, warn};
 
 use restate_admin_rest_model::services::ListServicesResponse;
 use restate_admin_rest_model::services::*;
@@ -22,13 +21,14 @@ use restate_core::network::TransportConnect;
 use restate_errors::warn_it;
 use restate_types::config::Configuration;
 use restate_types::identifiers::{ServiceId, WithPartitionKey};
+use restate_types::logs::{BodyWithKeys, Keys};
 use restate_types::schema::registry::MetadataService;
 use restate_types::schema::service::ServiceMetadata;
 use restate_types::state_mut::ExternalStateMutation;
 use restate_types::{Scope, schema};
-use restate_wal_protocol::{Command, Envelope};
+use restate_wal_protocol::v2;
+use restate_wal_protocol::v2::commands;
 
-use super::create_envelope_header;
 use super::error::*;
 use crate::state::AdminServiceState;
 
@@ -250,14 +250,17 @@ where
         state: new_state,
     };
 
-    let envelope = Envelope::new(
-        create_envelope_header(partition_key),
-        Command::PatchState(patch_state),
+    let envelop = v2::Envelope::new(
+        v2::Dedup::None,
+        commands::PatchStateCommand::from(patch_state),
     );
 
     let result = state
         .ingestion_client
-        .ingest(partition_key, envelope)
+        .ingest(
+            partition_key,
+            BodyWithKeys::new(envelop.into_raw(), Keys::Single(partition_key)),
+        )
         .await
         .map_err(|err| {
             warn!("Could not ingest state patching command: {err}");

crates/admin/src/service.rs

@@ -13,8 +13,6 @@ use std::time::Duration;
 
 use axum::error_handling::HandleErrorLayer;
 use http::{Request, Response, StatusCode};
-use restate_ingestion_client::IngestionClient;
-use restate_wal_protocol::Envelope;
 use tower::ServiceBuilder;
 use tower_http::classify::ServerErrorsFailureClass;
 use tower_http::compression::CompressionLayer;
@@ -24,6 +22,7 @@ use tracing::{Span, debug, info, info_span};
 use restate_admin_rest_model::version::AdminApiVersion;
 use restate_core::network::{TransportConnect, net_util};
 use restate_core::{MetadataWriter, TaskCenter};
+use restate_ingestion_client::IngestionClient;
 use restate_limiter::rule_book::RuleBookObserver;
 use restate_metadata_store::MetadataStoreClient;
 use restate_service_client::HttpClient;
@@ -36,6 +35,7 @@ use restate_types::net::address::AdminPort;
 use restate_types::net::listener::Listeners;
 use restate_types::schema::registry::SchemaRegistry;
 use restate_util_time::DurationExt;
+use restate_wal_protocol::v2::{Envelope, Raw};
 
 use crate::rest_api::{MAX_ADMIN_API_VERSION, MIN_ADMIN_API_VERSION};
 use crate::schema_registry_integration::{MetadataService, TelemetryClient};
@@ -47,7 +47,7 @@ pub struct BuildError(#[from] restate_service_client::BuildError);
 
 pub struct AdminService<Metadata, Discovery, Telemetry, Invocations, Transport> {
     listeners: Listeners<AdminPort>,
-    ingestion_client: IngestionClient<Transport, Envelope>,
+    ingestion_client: IngestionClient<Transport, Envelope<Raw>>,
     schema_registry: SchemaRegistry<Metadata, Discovery, Telemetry>,
     serdes_client: SerdesClient,
     invocation_client: Invocations,
@@ -65,7 +65,7 @@ where
     pub fn new(
         listeners: Listeners<AdminPort>,
         metadata_writer: MetadataWriter,
-        ingestion_client: IngestionClient<Transport, Envelope>,
+        ingestion_client: IngestionClient<Transport, Envelope<Raw>>,
         invocation_client: Invocations,
         serdes_client: SerdesClient,
         service_discovery: ServiceDiscovery,

crates/admin/src/state.rs

@@ -8,22 +8,23 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+use std::sync::Arc;
+
 use restate_core::network::TransportConnect;
 use restate_ingestion_client::IngestionClient;
 use restate_limiter::rule_book::RuleBookObserver;
 use restate_metadata_store::MetadataStoreClient;
 use restate_service_protocol_v4::serdes::SerdesClient;
 use restate_storage_query_datafusion::context::QueryContext;
 use restate_types::schema::registry::SchemaRegistry;
-use restate_wal_protocol::Envelope;
-use std::sync::Arc;
+use restate_wal_protocol::v2::{Envelope, Raw};
 
 #[derive(Clone, derive_builder::Builder)]
 pub struct AdminServiceState<Metadata, Discovery, Telemetry, Invocations, Transport> {
     pub schema_registry: SchemaRegistry<Metadata, Discovery, Telemetry>,
     pub serdes_client: SerdesClient,
     pub invocation_client: Invocations,
-    pub ingestion_client: IngestionClient<Transport, Envelope>,
+    pub ingestion_client: IngestionClient<Transport, Envelope<Raw>>,
     /// Used by handlers that mutate cluster-global metadata-store keys
     /// directly (e.g. the rule book) via `read_modify_write`.
     pub metadata_store_client: MetadataStoreClient,
@@ -41,7 +42,7 @@ where
         schema_registry: SchemaRegistry<Metadata, Discovery, Telemetry>,
         serdes_client: SerdesClient,
         invocation_client: Invocations,
-        ingestion_client: IngestionClient<Transport, Envelope>,
+        ingestion_client: IngestionClient<Transport, Envelope<Raw>>,
         metadata_store_client: MetadataStoreClient,
         query_context: Option<QueryContext>,
         rule_book_observer: Option<Arc<dyn RuleBookObserver>>,

crates/ingestion-client/src/client.rs

@@ -21,7 +21,7 @@ use restate_core::{
 use restate_types::{
     identifiers::PartitionKey,
     live::Live,
-    logs::{HasRecordKeys, Keys},
+    logs::{BodyWithKeys, HasRecordKeys, Keys},
     net::ingest::IngestRecord,
     partitions::{FindPartition, PartitionTable, PartitionTableError},
     storage::{StorageCodec, StorageEncode},
@@ -271,6 +271,16 @@ impl InputRecord<String> {
     }
 }
 
+impl<T> From<BodyWithKeys<T>> for InputRecord<T>
+where
+    T: StorageEncode,
+{
+    fn from(value: BodyWithKeys<T>) -> Self {
+        let (keys, record) = value.split();
+        Self { keys, record }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::{num::NonZeroUsize, time::Duration};

crates/ingress-kafka/src/builder.rs

@@ -18,21 +18,20 @@ use opentelemetry::trace::{Span, SpanContext, TraceContextExt};
 use opentelemetry_sdk::propagation::TraceContextPropagator;
 use rdkafka::Message;
 use rdkafka::message::BorrowedMessage;
-use tracing::{info_span, trace};
-
 use rdkafka::message::Headers;
+use tracing::{info_span, trace};
 
-use restate_storage_api::deduplication_table::DedupInformation;
 use restate_types::Scope;
-use restate_types::identifiers::{InvocationId, WithPartitionKey, partitioner};
+use restate_types::identifiers::{InvocationId, partitioner};
 use restate_types::invocation::{Header, InvocationTarget, ServiceInvocation, SpanRelation};
 use restate_types::limit_key::LimitKey;
 use restate_types::live::Live;
 use restate_types::schema::Schema;
 use restate_types::schema::invocation_target::{DeploymentStatus, InvocationTargetResolver};
 use restate_types::schema::subscriptions::{EventInvocationTargetTemplate, Sink, Subscription};
+use restate_types::sharding::{PartitionKey, WithPartitionKey};
 use restate_util_string::{ReString, RestateString, RestrictedValueError};
-use restate_wal_protocol::{Command, Destination, Envelope, Source};
+use restate_wal_protocol::v2::{Dedup, Envelope, commands};
 
 use crate::Error;
 
@@ -62,7 +61,7 @@ impl EnvelopeBuilder {
         producer_id: u128,
         consumer_group_id: &str,
         msg: BorrowedMessage<'_>,
-    ) -> Result<Envelope, Error> {
+    ) -> Result<(PartitionKey, Envelope<commands::InvokeCommand>), Error> {
         // Prepare ingress span
         let ingress_span = info_span!(
             "kafka_ingress_consume",
@@ -106,7 +105,11 @@ impl EnvelopeBuilder {
             (None, LimitKey::None)
         };
 
-        let dedup = DedupInformation::producer(producer_id, msg.offset() as u64);
+        let dedup = Dedup::Arbitrary {
+            prefix: None,
+            producer_id: producer_id.into(),
+            seq: msg.offset() as u64,
+        };
 
         let invocation = InvocationBuilder::create(
             &self.subscription,
@@ -130,23 +133,11 @@ impl EnvelopeBuilder {
             cause,
         })?;
 
-        Ok(self.wrap_service_invocation_in_envelope(invocation, dedup))
-    }
-
-    fn wrap_service_invocation_in_envelope(
-        &self,
-        service_invocation: Box<ServiceInvocation>,
-        dedup_information: DedupInformation,
-    ) -> Envelope {
-        let header = restate_wal_protocol::Header {
-            source: Source::Ingress {},
-            dest: Destination::Processor {
-                partition_key: service_invocation.partition_key(),
-                dedup: Some(dedup_information),
-            },
-        };
-
-        Envelope::new(header, Command::Invoke(service_invocation))
+        let partition_key = invocation.partition_key();
+        Ok((
+            partition_key,
+            Envelope::new(dedup, commands::InvokeCommand::from(invocation)),
+        ))
     }
 
     fn generate_events_attributes(msg: &impl Message, subscription_id: &str) -> Vec<Header> {
@@ -225,7 +216,7 @@ impl InvocationBuilder {
         topic: &str,
         partition: i32,
         offset: i64,
-    ) -> Result<Box<ServiceInvocation>, anyhow::Error> {
+    ) -> Result<ServiceInvocation, anyhow::Error> {
         let Sink::Invocation {
             event_invocation_target_template,
         } = subscription.sink();
@@ -325,11 +316,11 @@ impl InvocationBuilder {
         );
 
         // Finally generate service invocation
-        let mut service_invocation = Box::new(ServiceInvocation::initialize(
+        let mut service_invocation = ServiceInvocation::initialize(
             invocation_id,
             invocation_target,
             restate_types::invocation::Source::Subscription(subscription.id()),
-        ));
+        );
         service_invocation.with_related_span(SpanRelation::parent(ingress_span_context));
         service_invocation.argument = payload;
         service_invocation.headers = headers;