perf(runtime): add request-plane msgpack payload codec

lib/runtime/src/config/environment_names.rs

@@ -513,6 +513,13 @@ pub mod router {
     pub const DYN_ROUTER_POLICY_CONFIG: &str = "DYN_ROUTER_POLICY_CONFIG";
 }
 
+/// Request plane transport environment variables
+pub mod request_plane {
+    /// Request plane payload codec selection: "json" or "msgpack".
+    /// JSON is the compatibility default.
+    pub const DYN_REQUEST_PLANE_CODEC: &str = "DYN_REQUEST_PLANE_CODEC";
+}
+
 /// TCP response stream server (CallHome listener) environment variables
 pub mod tcp_response_stream {
     /// Port for the TCP response stream server.
@@ -744,6 +751,7 @@ mod tests {
             router::DYN_ROUTER_QUEUE_THRESHOLD,
             router::DYN_ROUTER_QUEUE_POLICY,
             router::DYN_ROUTER_POLICY_CONFIG,
+            request_plane::DYN_REQUEST_PLANE_CODEC,
             // TCP Response Stream
             tcp_response_stream::DYN_TCP_RESPONSE_STREAM_PORT,
             tcp_response_stream::DYN_TCP_RESPONSE_STREAM_HOST,

lib/runtime/src/pipeline/network.rs

@@ -25,7 +25,7 @@ use derive_builder::Builder;
 use futures::StreamExt;
 // io::Cursor, TryStreamExt
 use super::{AsyncEngine, AsyncEngineContext, AsyncEngineContextProvider, ResponseStream};
-use serde::{Deserialize, Serialize};
+use serde::{Deserialize, Serialize, de::DeserializeOwned};
 
 use super::{
     AsyncTransportEngine, Context, Data, Error, ManyIn, ManyOut, PipelineError, PipelineIO,
@@ -41,6 +41,7 @@ use prometheus::{CounterVec, Histogram, IntCounter, IntCounterVec, IntGauge};
 pub(crate) const DEFAULT_TCP_MAX_MESSAGE_SIZE: usize = 32 * 1024 * 1024;
 
 static TCP_MAX_MESSAGE_SIZE: OnceLock<usize> = OnceLock::new();
+static REQUEST_PLANE_PAYLOAD_CODEC: OnceLock<RequestPlanePayloadCodec> = OnceLock::new();
 
 /// Read the configured TCP max message size once and share it across client,
 /// server, and zero-copy decoder code paths.
@@ -53,6 +54,62 @@ pub(crate) fn get_tcp_max_message_size() -> usize {
     })
 }
 
+#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
+#[serde(rename_all = "snake_case")]
+pub(crate) enum RequestPlanePayloadCodec {
+    // TODO(jthomson04): Migrate the default to Msgpack after the 1.3 release.
+    #[default]
+    Json,
+    Msgpack,
+}
+
+impl RequestPlanePayloadCodec {
+    pub(crate) fn configured() -> Self {
+        *REQUEST_PLANE_PAYLOAD_CODEC.get_or_init(Self::from_env)
+    }
+
+    fn from_env() -> Self {
+        match std::env::var(
+            crate::config::environment_names::request_plane::DYN_REQUEST_PLANE_CODEC,
+        )
+        .as_deref()
+        {
+            Err(_) | Ok("") | Ok("json") => Self::Json,
+            Ok("msgpack") => Self::Msgpack,
+            Ok(other) => {
+                tracing::warn!(
+                    env_var =
+                        crate::config::environment_names::request_plane::DYN_REQUEST_PLANE_CODEC,
+                    value = other,
+                    "invalid request plane payload codec, defaulting to json"
+                );
+                Self::Json
+            }
+        }
+    }
+
+    pub(crate) fn name(&self) -> &'static str {
+        match self {
+            Self::Json => "json",
+            Self::Msgpack => "msgpack",
+        }
+    }
+
+    pub(crate) fn encode<T: Serialize>(&self, value: &T) -> Result<Vec<u8>> {
+        match self {
+            Self::Json => Ok(serde_json::to_vec(value)?),
+            Self::Msgpack => Ok(rmp_serde::to_vec_named(value)?),
+        }
+    }
+
+    pub(crate) fn decode<T: DeserializeOwned>(&self, bytes: &[u8]) -> Result<T> {
+        match self {
+            Self::Json => Ok(serde_json::from_slice(bytes)?),
+            Self::Msgpack => Ok(rmp_serde::from_slice(bytes)?),
+        }
+    }
+}
+
 pub trait Codable: PipelineIO + Serialize + for<'de> Deserialize<'de> {}
 impl<T: PipelineIO + Serialize + for<'de> Deserialize<'de>> Codable for T {}
 
@@ -88,6 +145,8 @@ pub(crate) struct RequestControlMessage {
     pub(crate) id: String,
     pub(crate) request_type: RequestType,
     pub(crate) response_type: ResponseType,
+    #[serde(default)]
+    pub(crate) payload_codec: RequestPlanePayloadCodec,
     pub(crate) connection_info: ConnectionInfo,
     #[serde(default, skip_serializing_if = "std::collections::BTreeMap::is_empty")]
     pub(crate) metadata: std::collections::BTreeMap<String, String>,
@@ -407,10 +466,19 @@ pub struct Egress<Req: PipelineIO, Resp: PipelineIO> {
 #[cfg(test)]
 mod tests {
     use super::{
-        DEFAULT_SEND_BUFFER_COUNT, RequestControlMessage, RequestType, ResponseType, StreamOptions,
+        DEFAULT_SEND_BUFFER_COUNT, NetworkStreamWrapper, RequestControlMessage,
+        RequestPlanePayloadCodec, RequestType, ResponseType, StreamOptions,
     };
     use crate::engine::AsyncEngineContextProvider;
     use crate::pipeline::Context;
+    use serde::{Deserialize, Serialize};
+
+    #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)]
+    struct TestPayload {
+        id: u64,
+        text: String,
+        tokens: Vec<u32>,
+    }
 
     #[test]
     fn stream_options_send_buffer_count_defaults_to_64() {
@@ -458,11 +526,53 @@ mod tests {
         assert_eq!(message.id, "request-123");
         assert!(matches!(message.request_type, RequestType::SingleIn));
         assert!(matches!(message.response_type, ResponseType::ManyOut));
+        assert_eq!(message.payload_codec, RequestPlanePayloadCodec::Json);
         assert_eq!(message.connection_info.transport, "tcp");
         assert_eq!(message.connection_info.info, "{}");
         assert!(message.metadata.is_empty());
         assert!(message.frontend_send_ts_ns.is_none());
     }
+
+    #[test]
+    fn request_control_message_decodes_msgpack_payload_codec() {
+        let json = r#"{
+            "id": "request-123",
+            "request_type": "single_in",
+            "response_type": "many_out",
+            "payload_codec": "msgpack",
+            "connection_info": {
+                "transport": "tcp",
+                "info": "{}"
+            }
+        }"#;
+
+        let message: RequestControlMessage =
+            serde_json::from_str(json).expect("control message should deserialize");
+
+        assert_eq!(message.payload_codec, RequestPlanePayloadCodec::Msgpack);
+    }
+
+    #[test]
+    fn request_plane_payload_codec_round_trips_response_wrapper_json_and_msgpack() {
+        let wrapper = NetworkStreamWrapper {
+            data: Some(TestPayload {
+                id: 42,
+                text: "line\nquote\"slash\\unicode 中".to_string(),
+                tokens: vec![1, 2, 3, 65535],
+            }),
+            complete_final: false,
+        };
+
+        for codec in [
+            RequestPlanePayloadCodec::Json,
+            RequestPlanePayloadCodec::Msgpack,
+        ] {
+            let encoded = codec.encode(&wrapper).expect("wrapper should encode");
+            let decoded: NetworkStreamWrapper<TestPayload> =
+                codec.decode(&encoded).expect("wrapper should decode");
+            assert_eq!(decoded, wrapper);
+        }
+    }
 }
 
 #[async_trait]
@@ -624,7 +734,7 @@ pub trait PushWorkHandler: Send + Sync {
 /// can be due to network issues that only the egress component can detect.
 */
 /// TODO: Detect end-of-stream using Server-Sent Events (SSE). This will be removed.
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Serialize, Deserialize, Debug, PartialEq, Eq)]
 pub struct NetworkStreamWrapper<U> {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub data: Option<U>,

lib/runtime/src/pipeline/network/egress/addressed_router.rs

@@ -23,6 +23,7 @@ use crate::pipeline::network::NetworkStreamWrapper;
 use crate::pipeline::network::PendingConnections;
 use crate::pipeline::network::RegisteredStream;
 use crate::pipeline::network::RequestControlMessage;
+use crate::pipeline::network::RequestPlanePayloadCodec;
 use crate::pipeline::network::RequestType;
 use crate::pipeline::network::ResponseType;
 use crate::pipeline::network::StreamOptions;
@@ -54,6 +55,7 @@ fn decode_response_stream<U>(
     queue_start: Instant,
     tx_start: Instant,
     inflight_guard: InflightGuard,
+    payload_codec: RequestPlanePayloadCodec,
 ) -> ManyOut<U>
 where
     U: Data + for<'de> Deserialize<'de> + MaybeError,
@@ -76,7 +78,7 @@ where
                 );
                 return Some(U::from_err(err));
             }
-            match serde_json::from_slice::<NetworkStreamWrapper<U>>(&res_bytes) {
+            match payload_codec.decode::<NetworkStreamWrapper<U>>(&res_bytes) {
                 Ok(item) => {
                     is_complete_final = item.complete_final;
                     if let Some(data) = item.data {
@@ -90,8 +92,13 @@ where
                     }
                 }
                 Err(err) => {
-                    let json_str = String::from_utf8_lossy(&res_bytes);
-                    tracing::warn!(%err, %json_str, "Failed deserializing JSON to response");
+                    let response_bytes_len = res_bytes.len();
+                    tracing::warn!(
+                        %err,
+                        codec = payload_codec.name(),
+                        response_bytes_len,
+                        "failed deserializing request-plane response"
+                    );
                     Some(U::from_err(DynamoError::msg(err.to_string())))
                 }
             }
@@ -143,8 +150,9 @@ fn build_request_envelope<T>(
     request: Option<&T>,
 ) -> Result<bytes::Bytes, Error>
 where
-    T: serde::Serialize + ?Sized,
+    T: serde::Serialize,
 {
+    let payload_codec = RequestPlanePayloadCodec::configured();
     let request_id = context.id();
     let request_type = if send_conn_info.is_some() {
         RequestType::ManyIn
@@ -155,6 +163,7 @@ where
         id: request_id.to_string(),
         request_type,
         response_type: ResponseType::ManyOut,
+        payload_codec,
         connection_info: recv_conn_info,
         metadata: context.metadata().clone(),
         frontend_send_ts_ns: None,
@@ -163,7 +172,7 @@ where
 
     let ctrl = serialize_control_message(&control_message)?;
     let data: Option<Vec<u8>> = match request {
-        Some(req) => Some(serde_json::to_vec(req)?),
+        Some(req) => Some(payload_codec.encode(req)?),
         None => None,
     };
 
@@ -200,6 +209,7 @@ async fn spawn_request_stream_forwarder<T>(
     request_stream_provider: Option<StreamProvider<StreamSender>>,
     mut input_stream: crate::engine::DataStream<T>,
     engine_ctx: Arc<dyn crate::engine::AsyncEngineContext>,
+    payload_codec: RequestPlanePayloadCodec,
 ) -> Result<(), Error>
 where
     T: serde::Serialize + Send + 'static,
@@ -244,7 +254,7 @@ where
                     None => break,
                 },
             };
-            let bytes = match serde_json::to_vec(&item) {
+            let bytes = match payload_codec.encode(&item) {
                 Ok(b) => b,
                 Err(e) => {
                     // Stream-side framing failure: the engine sees a
@@ -253,6 +263,7 @@ where
                     // dropping frames.
                     tracing::error!(
                         error = %e,
+                        codec = payload_codec.name(),
                         "failed to serialize bidirectional request frame; killing context"
                     );
                     engine_ctx.kill();
@@ -470,6 +481,7 @@ impl AddressedPushRouter {
         let inflight_guard = InflightGuard::new();
 
         let enable_request_stream = input_stream.is_some();
+        let payload_codec = RequestPlanePayloadCodec::configured();
 
         // Hold the `RegisteredStream` as their RAII cleanup stays armed while held,
         // which simplifies the cancellation of registration on error. Each side is
@@ -543,8 +555,13 @@ impl AddressedPushRouter {
                 let (_conn_info, provider) = r.into_parts();
                 provider
             });
-            spawn_request_stream_forwarder(request_stream_provider, stream, engine_ctx.clone())
-                .await?;
+            spawn_request_stream_forwarder(
+                request_stream_provider,
+                stream,
+                engine_ctx.clone(),
+                payload_codec,
+            )
+            .await?;
         }
 
         let _nvtx_wait = dynamo_nvtx_range!("transport.tcp.wait_backend");
@@ -592,6 +609,7 @@ impl AddressedPushRouter {
             queue_start,
             tx_start,
             inflight_guard,
+            payload_codec,
         ))
     }
 
@@ -750,8 +768,8 @@ where
 #[cfg(test)]
 mod tests {
     use super::{
-        CONTROL_MESSAGE_MAX_BYTES, ConnectionInfo, RequestControlMessage, RequestType,
-        ResponseType, serialize_control_message,
+        CONTROL_MESSAGE_MAX_BYTES, ConnectionInfo, RequestControlMessage, RequestPlanePayloadCodec,
+        RequestType, ResponseType, serialize_control_message,
     };
     use std::collections::BTreeMap;
 
@@ -760,6 +778,7 @@ mod tests {
             id: "request-123".to_string(),
             request_type: RequestType::SingleIn,
             response_type: ResponseType::ManyOut,
+            payload_codec: RequestPlanePayloadCodec::Json,
             connection_info: ConnectionInfo {
                 transport: "tcp".to_string(),
                 info: "{}".to_string(),