diff --git a/lib/runtime/src/config/environment_names.rs b/lib/runtime/src/config/environment_names.rs index 2754472790f8..81a887054a28 100644 --- a/lib/runtime/src/config/environment_names.rs +++ b/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, diff --git a/lib/runtime/src/pipeline/network.rs b/lib/runtime/src/pipeline/network.rs index 33fa95891f9d..726b1a854c67 100644 --- a/lib/runtime/src/pipeline/network.rs +++ b/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 = OnceLock::new(); +static REQUEST_PLANE_PAYLOAD_CODEC: OnceLock = 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(&self, value: &T) -> Result> { + match self { + Self::Json => Ok(serde_json::to_vec(value)?), + Self::Msgpack => Ok(rmp_serde::to_vec_named(value)?), + } + } + + pub(crate) fn decode(&self, bytes: &[u8]) -> Result { + 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 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, @@ -407,10 +466,19 @@ pub struct Egress { #[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, + } #[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 = + 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 { #[serde(skip_serializing_if = "Option::is_none")] pub data: Option, diff --git a/lib/runtime/src/pipeline/network/egress/addressed_router.rs b/lib/runtime/src/pipeline/network/egress/addressed_router.rs index 9a7c9f8a7e6c..fcdcf3c0bf7e 100644 --- a/lib/runtime/src/pipeline/network/egress/addressed_router.rs +++ b/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( queue_start: Instant, tx_start: Instant, inflight_guard: InflightGuard, + payload_codec: RequestPlanePayloadCodec, ) -> ManyOut where U: Data + for<'de> Deserialize<'de> + MaybeError, @@ -76,7 +78,7 @@ where ); return Some(U::from_err(err)); } - match serde_json::from_slice::>(&res_bytes) { + match payload_codec.decode::>(&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( request: Option<&T>, ) -> Result 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> = 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( request_stream_provider: Option>, mut input_stream: crate::engine::DataStream, engine_ctx: Arc, + 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(), diff --git a/lib/runtime/src/pipeline/network/ingress/push_handler.rs b/lib/runtime/src/pipeline/network/ingress/push_handler.rs index 77703120f919..d796d93eb894 100644 --- a/lib/runtime/src/pipeline/network/ingress/push_handler.rs +++ b/lib/runtime/src/pipeline/network/ingress/push_handler.rs @@ -148,8 +148,12 @@ impl Ingress { /// classification (client-side disconnect vs. real failure), and the /// health-check notifier policy (notify only on non-error chunks and /// at clean stream end). - async fn pump_response_stream(&self, mut stream: ManyOut, publisher: &StreamSender) - where + async fn pump_response_stream( + &self, + mut stream: ManyOut, + publisher: &StreamSender, + payload_codec: RequestPlanePayloadCodec, + ) where U: Data + Serialize + MaybeError + std::fmt::Debug, { let context = stream.context(); @@ -167,8 +171,9 @@ impl Ingress { data: Some(resp), complete_final: false, }; - let resp_bytes = serde_json::to_vec(&resp_wrapper) - .expect("fatal error: invalid response object - this should never happen"); + let resp_bytes = payload_codec + .encode(&resp_wrapper) + .expect("fatal error: invalid request-plane response object"); if let Some(m) = self.metrics() { m.response_bytes.inc_by(resp_bytes.len() as u64); } @@ -210,8 +215,9 @@ impl Ingress { data: None, complete_final: true, }; - let resp_bytes = serde_json::to_vec(&resp_wrapper) - .expect("fatal error: invalid response object - this should never happen"); + let resp_bytes = payload_codec + .encode(&resp_wrapper) + .expect("fatal error: invalid request-plane response final object"); if let Some(m) = self.metrics() { m.response_bytes.inc_by(resp_bytes.len() as u64); } @@ -294,6 +300,7 @@ struct ParsedRequest { request: Req, response_connection_info: ConnectionInfo, frontend_send_ts_ns: Option, + payload_codec: RequestPlanePayloadCodec, } /// Per-shape strategy for turning a raw payload into a typed engine @@ -338,7 +345,19 @@ where "unary engine received a header-only envelope; expected a request payload", )) })?; - let request_t: T = serde_json::from_slice(&data)?; + let payload_codec = control_msg.payload_codec; + let request_t: T = payload_codec.decode(&data).map_err(|err| { + if let Some(m) = self.metrics() { + m.error_counter + .with_label_values(&[work_handler::error_types::DESERIALIZATION]) + .inc(); + } + PipelineError::DeserializationError(format!( + "Failed deserializing {} request payload: {}", + payload_codec.name(), + err + )) + })?; tracing::trace!( request_id = %control_msg.id, @@ -354,6 +373,7 @@ where request, response_connection_info: control_msg.connection_info, frontend_send_ts_ns: control_msg.frontend_send_ts_ns, + payload_codec, }) } } @@ -412,6 +432,7 @@ where control_msg.id.clone(), control_msg.metadata.clone(), ); + let payload_codec = control_msg.payload_codec; let context_arc: Arc = request_context.context(); // Open the request stream (upstream → worker) up front. The shared @@ -450,7 +471,7 @@ where if forwarder_ctx.is_killed() || forwarder_ctx.is_stopped() { break; } - match serde_json::from_slice::(&bytes) { + match payload_codec.decode::(&bytes) { Ok(item) => { if frame_tx.send(item).await.is_err() { tracing::debug!( @@ -462,6 +483,7 @@ where Err(e) => { tracing::error!( error = %e, + codec = payload_codec.name(), "failed to deserialize bidirectional request frame; killing context" ); forwarder_ctx.kill(); @@ -479,6 +501,7 @@ where request, response_connection_info: control_msg.connection_info, frontend_send_ts_ns: control_msg.frontend_send_ts_ns, + payload_codec, }) } } @@ -528,6 +551,7 @@ where request, response_connection_info, frontend_send_ts_ns, + payload_codec, } = self.parse_and_build_request(payload).await?; // Compute network transit time (T2 - T1) using cross-process wall-clock timestamps @@ -600,7 +624,8 @@ where } }; - self.pump_response_stream(stream, &publisher).await; + self.pump_response_stream(stream, &publisher, payload_codec) + .await; // Ensure the metrics guard is not dropped until the end of the function. // Drop fires "request completed" log via RAII.