perf(svc): forward messages shard replies without transcoding

fluxer_messages/src/router_impl.rs

@@ -21,6 +21,8 @@ impl RouterService for MessagesRouter {
     type Request = MessageRequest;
     type Response = MessageResponse;
 
+    const CACHES_RESPONSES: bool = false;
+
     fn service_name(&self) -> &str {
         "messages"
     }

fluxer_svc/src/router.rs

@@ -15,12 +15,15 @@ pub(crate) const SHARD_REQUEST_TIMEOUT: Duration = Duration::from_secs(5);
 const INFLIGHT_TTL: Duration = Duration::from_millis(200);
 const INFLIGHT_MAX_ENTRIES: u64 = 10_000;
 const MAX_ROUTER_REQUEST_BYTES: usize = 2 * 1024 * 1024;
+const LEGACY_SHARD_DECODE_ERROR: &[u8] = br#"{"error":"shard_request_decode_error"}"#;
 type InflightKey = (String, String);
 
 pub trait RouterService: Send + Sync + 'static {
     type Request: serde::Serialize + serde::de::DeserializeOwned + Send + Sync + 'static;
     type Response: serde::Serialize + serde::de::DeserializeOwned + Clone + Send + Sync + 'static;
 
+    const CACHES_RESPONSES: bool = true;
+
     fn service_name(&self) -> &str;
     fn route_key(request: &Self::Request) -> String;
     fn coalesce_key(_request: &Self::Request) -> Option<String> {
@@ -34,15 +37,19 @@ pub trait RouterService: Send + Sync + 'static {
     fn l1_invalidate(&self, _key: &str) {}
 }
 
+fn shard_subject<S: RouterService>(service: &S, ring: &HashRing, route_key: &str) -> String {
+    let shard_id = ring.owner(route_key);
+    format!("svc.{}.shard.{shard_id}", service.service_name())
+}
+
 async fn forward_to_shard<S: RouterService>(
     transport: &impl Transport,
     service: &S,
     ring: &HashRing,
     request: &S::Request,
     route_key: &str,
 ) -> anyhow::Result<Vec<u8>> {
-    let shard_id = ring.owner(route_key);
-    let shard_subject = format!("svc.{}.shard.{shard_id}", service.service_name());
+    let shard_subject = shard_subject(service, ring, route_key);
 
     let msgpack_payload = rmp_serde::to_vec_named(request)
         .map_err(|e| anyhow::anyhow!("failed to encode request as msgpack: {e}"))?;
@@ -52,6 +59,58 @@ async fn forward_to_shard<S: RouterService>(
         .await
 }
 
+async fn forward_to_shard_verbatim<S: RouterService>(
+    transport: &impl Transport,
+    service: &S,
+    ring: &HashRing,
+    request: &S::Request,
+    route_key: &str,
+    payload: &[u8],
+) -> anyhow::Result<Vec<u8>> {
+    let shard_subject = shard_subject(service, ring, route_key);
+    let response_bytes = transport
+        .request(&shard_subject, payload, SHARD_REQUEST_TIMEOUT)
+        .await?;
+    if response_bytes != LEGACY_SHARD_DECODE_ERROR {
+        return Ok(response_bytes);
+    }
+
+    warn!(
+        subject = shard_subject,
+        "shard rejected a pass-through request, retrying with the legacy msgpack encoding"
+    );
+    let legacy_bytes = forward_to_shard::<S>(transport, service, ring, request, route_key).await?;
+    match rmp_serde::from_slice::<S::Response>(&legacy_bytes) {
+        Ok(response) => serde_json::to_vec(&response).map_err(|err| {
+            anyhow::anyhow!("failed to encode legacy shard response as json: {err}")
+        }),
+        Err(err) => {
+            if serde_json::from_slice::<serde_json::Value>(&legacy_bytes).is_ok() {
+                Ok(legacy_bytes)
+            } else {
+                Err(anyhow::anyhow!(
+                    "failed to decode legacy shard response: {err}"
+                ))
+            }
+        }
+    }
+}
+
+async fn dispatch_to_shard<S: RouterService>(
+    transport: &impl Transport,
+    service: &S,
+    ring: &HashRing,
+    request: &S::Request,
+    route_key: &str,
+    payload: &[u8],
+) -> anyhow::Result<Vec<u8>> {
+    if S::CACHES_RESPONSES {
+        forward_to_shard::<S>(transport, service, ring, request, route_key).await
+    } else {
+        forward_to_shard_verbatim::<S>(transport, service, ring, request, route_key, payload).await
+    }
+}
+
 async fn handle_router_request<S, T>(
     msg: T::Message,
     transport: T,
@@ -117,27 +176,34 @@ async fn handle_router_request<S, T>(
         let forward_ring = ring.clone();
         let forward_route_key = route_key.clone();
         let forward_request = request.clone();
+        let forward_payload = if S::CACHES_RESPONSES {
+            Vec::new()
+        } else {
+            msg.payload().to_vec()
+        };
         let inflight_key = (route_key, coalesce_key);
         inflight
             .try_get_with(inflight_key, async move {
-                forward_to_shard::<S>(
+                dispatch_to_shard::<S>(
                     &forward_transport,
                     forward_service.as_ref(),
                     forward_ring.as_ref(),
                     &forward_request,
                     &forward_route_key,
+                    &forward_payload,
                 )
                 .await
             })
             .await
             .map_err(|err| anyhow::anyhow!("{err}"))
     } else {
-        forward_to_shard::<S>(
+        dispatch_to_shard::<S>(
             &transport,
             service.as_ref(),
             ring.as_ref(),
             &request,
             &route_key,
+            msg.payload(),
         )
         .await
     };
@@ -146,28 +212,36 @@ async fn handle_router_request<S, T>(
     metrics.record_request_duration(elapsed);
 
     match coalesce_result {
-        Ok(response_bytes) => match rmp_serde::from_slice::<S::Response>(&response_bytes) {
-            Ok(response) => {
-                service.l1_insert(&request, &response);
+        Ok(response_bytes) => {
+            if !S::CACHES_RESPONSES {
                 if msg.has_reply() {
-                    let json = serde_json::to_vec(&response).unwrap_or_default();
-                    let _ = reply_message(&msg, &transport, &json).await;
+                    let _ = reply_message(&msg, &transport, &response_bytes).await;
                 }
+                return;
             }
-            Err(err) => {
-                debug!(error = %err, "failed to decode shard response");
-                if msg.has_reply() {
-                    if serde_json::from_slice::<serde_json::Value>(&response_bytes).is_ok() {
-                        let _ = reply_message(&msg, &transport, &response_bytes).await;
-                        return;
+            match rmp_serde::from_slice::<S::Response>(&response_bytes) {
+                Ok(response) => {
+                    service.l1_insert(&request, &response);
+                    if msg.has_reply() {
+                        let json = serde_json::to_vec(&response).unwrap_or_default();
+                        let _ = reply_message(&msg, &transport, &json).await;
+                    }
+                }
+                Err(err) => {
+                    debug!(error = %err, "failed to decode shard response");
+                    if msg.has_reply() {
+                        if serde_json::from_slice::<serde_json::Value>(&response_bytes).is_ok() {
+                            let _ = reply_message(&msg, &transport, &response_bytes).await;
+                            return;
+                        }
+                        let error_response =
+                            serde_json::to_vec(&serde_json::json!({"error": "shard_decode_error"}))
+                                .unwrap_or_default();
+                        let _ = reply_message(&msg, &transport, &error_response).await;
                     }
-                    let error_response =
-                        serde_json::to_vec(&serde_json::json!({"error": "shard_decode_error"}))
-                            .unwrap_or_default();
-                    let _ = reply_message(&msg, &transport, &error_response).await;
                 }
             }
-        },
+        }
         Err(err) => {
             debug!(error = %err, "shard request failed (coalesced)");
             metrics.record_request_error();
@@ -601,6 +675,128 @@ mod tests {
         router_task.abort();
     }
 
+    struct PassThroughRouter;
+
+    impl RouterService for PassThroughRouter {
+        type Request = MockRequest;
+        type Response = MockResponse;
+
+        const CACHES_RESPONSES: bool = false;
+
+        fn service_name(&self) -> &str {
+            "passthrough-mock"
+        }
+
+        fn route_key(request: &MockRequest) -> String {
+            request.key.clone()
+        }
+
+        fn coalesce_key(request: &MockRequest) -> Option<String> {
+            Some(request.key.clone())
+        }
+    }
+
+    #[tokio::test]
+    async fn router_forwards_pass_through_requests_and_replies_verbatim() {
+        let transport = InMemoryTransport::new();
+        let mut shard_sub = transport
+            .subscribe("svc.passthrough-mock.shard.0")
+            .await
+            .unwrap();
+
+        let observed = Arc::new(Mutex::new(Vec::new()));
+        let shard_transport = transport.clone();
+        let shard_observed = observed.clone();
+        let shard_task = tokio::spawn(async move {
+            while let Some(msg) = shard_sub.next().await {
+                shard_observed.lock().unwrap().push(msg.payload().to_vec());
+                reply_message(&msg, &shard_transport, br#"{ "key" : "verbatim" }"#)
+                    .await
+                    .unwrap();
+            }
+        });
+
+        let router_config = test_config(4);
+        let router_transport = transport.clone();
+        let router_task = tokio::spawn(async move {
+            run_router(&router_config, PassThroughRouter, router_transport).await
+        });
+
+        tokio::time::sleep(Duration::from_millis(25)).await;
+
+        let request = serde_json::to_vec(&MockRequest {
+            key: "a".to_owned(),
+        })
+        .unwrap();
+
+        let response = transport
+            .request("svc.passthrough-mock", &request, Duration::from_secs(1))
+            .await
+            .unwrap();
+
+        assert_eq!(response, br#"{ "key" : "verbatim" }"#);
+        assert_eq!(observed.lock().unwrap().as_slice(), [request]);
+
+        shard_task.abort();
+        router_task.abort();
+    }
+
+    #[tokio::test]
+    async fn router_retries_pass_through_requests_that_legacy_shards_reject() {
+        let transport = InMemoryTransport::new();
+        let mut shard_sub = transport
+            .subscribe("svc.passthrough-mock.shard.0")
+            .await
+            .unwrap();
+
+        let shard_transport = transport.clone();
+        let shard_task = tokio::spawn(async move {
+            while let Some(msg) = shard_sub.next().await {
+                let response_bytes = match rmp_serde::from_slice::<MockRequest>(msg.payload()) {
+                    Ok(_) => rmp_serde::to_vec_named(&MockResponse {
+                        key: "legacy".to_owned(),
+                    })
+                    .unwrap(),
+                    Err(_) => serde_json::to_vec(
+                        &serde_json::json!({"error": "shard_request_decode_error"}),
+                    )
+                    .unwrap(),
+                };
+                reply_message(&msg, &shard_transport, &response_bytes)
+                    .await
+                    .unwrap();
+            }
+        });
+
+        let router_config = test_config(4);
+        let router_transport = transport.clone();
+        let router_task = tokio::spawn(async move {
+            run_router(&router_config, PassThroughRouter, router_transport).await
+        });
+
+        tokio::time::sleep(Duration::from_millis(25)).await;
+
+        let request = serde_json::to_vec(&MockRequest {
+            key: "a".to_owned(),
+        })
+        .unwrap();
+
+        let response = transport
+            .request("svc.passthrough-mock", &request, Duration::from_secs(1))
+            .await
+            .unwrap();
+
+        assert_eq!(
+            serde_json::from_slice::<MockResponse>(&response).unwrap(),
+            MockResponse {
+                key: "legacy".to_owned()
+            }
+        );
+
+        shard_task.abort();
+        router_task.abort();
+    }
+
     #[tokio::test]
     async fn router_sheds_requests_when_permits_are_exhausted() {
         let transport = InMemoryTransport::new();