fix(svc): shed overload instead of buffering router requests (#2171)

config/env/development.env

@@ -43,7 +43,6 @@ FLUXER_SVC_NATS_URL=nats://nats:4222
 FLUXER_SVC_SHARD_COUNT=1
 FLUXER_SVC_CACHE_TTL_MS=30000
 FLUXER_SVC_CACHE_HARD_TTL_MS=600000
-FLUXER_SVC_MAX_CONCURRENT_REQUESTS=64
 
 FLUXER_S3_ENDPOINT=http://127.0.0.1:8333
 FLUXER_S3_PUBLIC_ENDPOINT=http://localhost:8088

deploy/self-hosting/docker-compose.yml

@@ -353,6 +353,7 @@ services:
     image: ${FLUXER_REGISTRY:-ghcr.io/${FLUXER_REGISTRY_OWNER:-fluxerapp}}/fluxer-snowflakes:${FLUXER_IMAGE_TAG:-v1}
     environment:
       <<: *fluxer-env
+      FLUXER_SVC_NAME: snowflakes
       FLUXER_SVC_MODE: router
     depends_on:
       nats: {condition: service_started}
@@ -362,6 +363,7 @@ services:
     image: ${FLUXER_REGISTRY:-ghcr.io/${FLUXER_REGISTRY_OWNER:-fluxerapp}}/fluxer-snowflakes:${FLUXER_IMAGE_TAG:-v1}
     environment:
       <<: *fluxer-env
+      FLUXER_SVC_NAME: snowflakes
       FLUXER_SVC_MODE: shard
       FLUXER_SVC_SHARD_ID: "0"
     depends_on:
@@ -417,6 +419,7 @@ services:
     image: ${FLUXER_REGISTRY:-ghcr.io/${FLUXER_REGISTRY_OWNER:-fluxerapp}}/fluxer-messages:${FLUXER_IMAGE_TAG:-v1}
     environment:
       <<: *fluxer-env
+      FLUXER_SVC_NAME: messages
       FLUXER_SVC_MODE: router
     depends_on:
       nats: {condition: service_started}
@@ -426,6 +429,7 @@ services:
     image: ${FLUXER_REGISTRY:-ghcr.io/${FLUXER_REGISTRY_OWNER:-fluxerapp}}/fluxer-messages:${FLUXER_IMAGE_TAG:-v1}
     environment:
       <<: *fluxer-env
+      FLUXER_SVC_NAME: messages
       FLUXER_SVC_MODE: shard
       FLUXER_SVC_SHARD_ID: "0"
       FLUXER_POSTGRES_MAX_CONNECTIONS: "20"

fluxer_svc/src/config.rs

@@ -4,6 +4,10 @@ use std::env;
 use std::net::{IpAddr, SocketAddr};
 use std::time::Duration;
 
+const DEFAULT_MAX_CONCURRENT_REQUESTS: usize = 64;
+const MESSAGES_MAX_CONCURRENT_REQUESTS: usize = 192;
+const SNOWFLAKES_MAX_CONCURRENT_REQUESTS: usize = 320;
+
 #[derive(Clone, Debug)]
 pub struct ServiceConfig {
     pub service_name: String,
@@ -141,6 +145,12 @@ impl ServiceConfig {
             .unwrap_or(20)
             .max(1);
 
+        let max_concurrent_requests = optional_from(&get, "FLUXER_SVC_MAX_CONCURRENT_REQUESTS")
+            .map(|v| v.parse::<usize>())
+            .transpose()?
+            .unwrap_or_else(|| default_max_concurrent_requests(&service_name))
+            .max(1);
+
         Ok(Self {
             service_name,
             mode,
@@ -155,11 +165,7 @@ impl ServiceConfig {
                 .unwrap_or(100_000),
             cache_ttl: Duration::from_millis(cache_ttl_ms),
             cache_hard_ttl: Duration::from_millis(cache_hard_ttl_ms),
-            max_concurrent_requests: optional_from(&get, "FLUXER_SVC_MAX_CONCURRENT_REQUESTS")
+            max_concurrent_requests,
-                .map(|v| v.parse::<usize>())
-                .transpose()?
-                .unwrap_or(64)
-                .max(1),
             scylla_hosts,
             scylla_keyspace: optional_from(&get, "FLUXER_CASSANDRA_KEYSPACE")
                 .unwrap_or_else(|| "fluxer".to_owned()),
@@ -184,6 +190,14 @@ impl ServiceConfig {
     }
 }
 
+fn default_max_concurrent_requests(service_name: &str) -> usize {
+    match service_name {
+        "messages" => MESSAGES_MAX_CONCURRENT_REQUESTS,
+        "snowflakes" => SNOWFLAKES_MAX_CONCURRENT_REQUESTS,
+        _ => DEFAULT_MAX_CONCURRENT_REQUESTS,
+    }
+}
+
 pub fn optional_env(name: &str) -> Option<String> {
     optional_from(&|key| env::var(key).ok(), name)
 }
@@ -364,6 +378,32 @@ mod tests {
         assert_eq!("fluxer_kv_dev", cfg.postgres_kv_table);
     }
 
+    #[test]
+    fn raises_concurrency_defaults_for_hot_path_services() {
+        assert_eq!(
+            MESSAGES_MAX_CONCURRENT_REQUESTS,
+            config_from_pairs(&[("FLUXER_SVC_NAME", "messages")]).max_concurrent_requests
+        );
+        assert_eq!(
+            SNOWFLAKES_MAX_CONCURRENT_REQUESTS,
+            config_from_pairs(&[("FLUXER_SVC_NAME", "snowflakes")]).max_concurrent_requests
+        );
+        assert_eq!(
+            DEFAULT_MAX_CONCURRENT_REQUESTS,
+            config_from_pairs(&[("FLUXER_SVC_NAME", "users")]).max_concurrent_requests
+        );
+    }
+
+    #[test]
+    fn concurrency_env_override_wins_over_service_default() {
+        let cfg = config_from_pairs(&[
+            ("FLUXER_SVC_NAME", "messages"),
+            ("FLUXER_SVC_MAX_CONCURRENT_REQUESTS", "32"),
+        ]);
+
+        assert_eq!(32, cfg.max_concurrent_requests);
+    }
+
     #[test]
     fn reads_legacy_cassandra_backend_aliases() {
         let cfg = config_from_pairs(&[("FLUXER_DATABASE_BACKEND", "scylla")]);

fluxer_svc/src/router.rs

@@ -7,7 +7,7 @@ use crate::transport::{Transport, TransportMessage, TransportSubscriber, reply_m
 use moka::future::Cache;
 use std::sync::Arc;
 use std::time::Duration;
-use tokio::sync::Semaphore;
+use tokio::sync::{Semaphore, TryAcquireError};
 use tokio::task::JoinSet;
 use tracing::{debug, info, warn};
 
@@ -246,9 +246,21 @@ where
                     }
                 };
 
-                let permit = match req_permits.clone().acquire_owned().await {
+                let permit = match req_permits.clone().try_acquire_owned() {
                     Ok(permit) => permit,
-                    Err(_) => return anyhow::Ok(()),
+                    Err(TryAcquireError::NoPermits) => {
+                        debug!("shedding router request, no permits available");
+                        req_metrics.record_request();
+                        req_metrics.record_request_error();
+                        if msg.has_reply() {
+                            let error_response =
+                                serde_json::to_vec(&serde_json::json!({"error": "overloaded"}))
+                                    .unwrap_or_default();
+                            let _ = reply_message(&msg, &req_transport, &error_response).await;
+                        }
+                        continue;
+                    }
+                    Err(TryAcquireError::Closed) => return anyhow::Ok(()),
                 };
                 let transport = req_transport.clone();
                 let service = req_service.clone();
@@ -589,6 +601,68 @@ mod tests {
         router_task.abort();
     }
 
+    #[tokio::test]
+    async fn router_sheds_requests_when_permits_are_exhausted() {
+        let transport = InMemoryTransport::new();
+        let mut shard_sub = transport.subscribe("svc.mock.shard.0").await.unwrap();
+
+        let (forwarded_tx, forwarded_rx) = tokio::sync::oneshot::channel();
+        let shard_task = tokio::spawn(async move {
+            let _msg = shard_sub.next().await.unwrap();
+            let _ = forwarded_tx.send(());
+            std::future::pending::<()>().await;
+        });
+
+        let router_config = test_config(1);
+        let router_transport = transport.clone();
+        let router_task =
+            tokio::spawn(
+                async move { run_router(&router_config, MockRouter, router_transport).await },
+            );
+
+        tokio::time::sleep(Duration::from_millis(25)).await;
+
+        let request_a = serde_json::to_vec(&MockRequest {
+            key: "a".to_owned(),
+        })
+        .unwrap();
+        let request_b = serde_json::to_vec(&MockRequest {
+            key: "b".to_owned(),
+        })
+        .unwrap();
+
+        let client_a = {
+            let transport = transport.clone();
+            tokio::spawn(async move {
+                transport
+                    .request("svc.mock", &request_a, Duration::from_secs(10))
+                    .await
+            })
+        };
+
+        tokio::time::timeout(Duration::from_millis(250), forwarded_rx)
+            .await
+            .expect("router should forward the first request and hold the only permit")
+            .unwrap();
+
+        let response_b = tokio::time::timeout(
+            Duration::from_millis(250),
+            transport.request("svc.mock", &request_b, Duration::from_secs(1)),
+        )
+        .await
+        .expect("shed reply should not wait behind the in-flight request")
+        .unwrap();
+
+        assert_eq!(
+            serde_json::from_slice::<serde_json::Value>(&response_b).unwrap(),
+            serde_json::json!({"error": "overloaded"})
+        );
+
+        client_a.abort();
+        shard_task.abort();
+        router_task.abort();
+    }
+
     fn test_config(max_concurrent_requests: usize) -> ServiceConfig {
         ServiceConfig {
             service_name: "mock".to_owned(),

fluxer_svc/src/shard.rs

@@ -5,7 +5,7 @@ use crate::metrics::{ServiceMetrics, now_ms};
 use crate::transport::{Transport, TransportMessage, TransportSubscriber, reply_message};
 use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, Ordering};
-use tokio::sync::Semaphore;
+use tokio::sync::{Semaphore, TryAcquireError};
 use tokio::task::JoinSet;
 use tracing::{debug, info, warn};
 
@@ -98,11 +98,18 @@ where
                             );
                             continue;
                         }
-                        let raw_payload = msg.payload().to_vec();
+                        let permit = match shard_permits.clone().try_acquire_owned() {
-                        let permit = match shard_permits.clone().acquire_owned().await {
                             Ok(permit) => permit,
-                            Err(_) => return anyhow::Ok(()),
+                            Err(TryAcquireError::NoPermits) => {
+                                debug!("shedding shard request, no permits available");
+                                shard_metrics.record_request();
+                                shard_metrics.record_request_error();
+                                reply_shard_error(&msg, &transport, "overloaded").await;
+                                continue;
+                            }
+                            Err(TryAcquireError::Closed) => return anyhow::Ok(()),
                         };
+                        let raw_payload = msg.payload().to_vec();
 
                         tokio::spawn(async move {
                             let _permit = permit;
@@ -210,3 +217,121 @@ async fn reply_shard_error(msg: &impl TransportMessage, transport: &impl Transpo
         debug!(error = %err, "failed to send shard error reply");
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::config::{DatabaseBackend, Mode, ServiceConfig};
+    use crate::transport::InMemoryTransport;
+    use serde::{Deserialize, Serialize};
+    use std::time::Duration;
+    use tokio::sync::Notify;
+
+    #[derive(Serialize, Deserialize)]
+    struct MockRequest {
+        key: String,
+    }
+
+    #[derive(Serialize, Deserialize)]
+    struct MockResponse {
+        key: String,
+    }
+
+    struct BlockingShard {
+        started: Arc<Notify>,
+    }
+
+    impl ShardService for BlockingShard {
+        type Request = MockRequest;
+        type Response = MockResponse;
+
+        fn service_name(&self) -> &str {
+            "mock"
+        }
+
+        async fn handle(&self, _request: MockRequest) -> anyhow::Result<MockResponse> {
+            self.started.notify_one();
+            std::future::pending().await
+        }
+    }
+
+    #[tokio::test]
+    async fn shard_sheds_requests_when_permits_are_exhausted() {
+        let transport = InMemoryTransport::new();
+        let started = Arc::new(Notify::new());
+        let shard = BlockingShard {
+            started: started.clone(),
+        };
+        let config = test_config(1);
+        let shard_transport = transport.clone();
+        let shard_task =
+            tokio::spawn(async move { run_shard(&config, shard, shard_transport).await });
+
+        tokio::time::sleep(Duration::from_millis(25)).await;
+
+        let payload = rmp_serde::to_vec_named(&MockRequest {
+            key: "a".to_owned(),
+        })
+        .unwrap();
+
+        let client = {
+            let transport = transport.clone();
+            let payload = payload.clone();
+            tokio::spawn(async move {
+                transport
+                    .request("svc.mock.shard.0", &payload, Duration::from_secs(10))
+                    .await
+            })
+        };
+
+        tokio::time::timeout(Duration::from_millis(250), started.notified())
+            .await
+            .expect("shard should start the first request and hold the only permit");
+
+        let response = tokio::time::timeout(
+            Duration::from_millis(250),
+            transport.request("svc.mock.shard.0", &payload, Duration::from_secs(1)),
+        )
+        .await
+        .expect("shed reply should not wait behind the in-flight request")
+        .unwrap();
+
+        assert_eq!(
+            serde_json::from_slice::<serde_json::Value>(&response).unwrap(),
+            serde_json::json!({"error": "overloaded"})
+        );
+
+        client.abort();
+        shard_task.abort();
+    }
+
+    fn test_config(max_concurrent_requests: usize) -> ServiceConfig {
+        ServiceConfig {
+            service_name: "mock".to_owned(),
+            mode: Mode::Shard,
+            database_backend: DatabaseBackend::Postgres,
+            shard_id: 0,
+            shard_count: 1,
+            listen_addr: "127.0.0.1:0".parse().unwrap(),
+            nats_url: "memory".to_owned(),
+            cache_max_entries: 100,
+            cache_ttl: Duration::from_secs(30),
+            cache_hard_ttl: Duration::from_secs(600),
+            max_concurrent_requests,
+            scylla_hosts: Vec::new(),
+            scylla_keyspace: "fluxer".to_owned(),
+            scylla_username: None,
+            scylla_password: None,
+            postgres_url: None,
+            postgres_host: "127.0.0.1".to_owned(),
+            postgres_port: 5432,
+            postgres_database: "fluxer".to_owned(),
+            postgres_username: "fluxer".to_owned(),
+            postgres_password: Some("fluxer".to_owned()),
+            postgres_ssl: false,
+            postgres_ssl_ca: None,
+            postgres_max_connections: 1,
+            postgres_kv_table: "fluxer_kv".to_owned(),
+        }
+    }
+}

fluxer_svc/src/transport.rs

@@ -8,6 +8,7 @@ use std::time::Duration;
 use tokio::sync::Notify;
 use tracing::{info, warn};
 
+const NATS_SUBSCRIPTION_CAPACITY: usize = 8_192;
 const SLOW_CONSUMER_LOG_INTERVAL_MS: u64 = 1_000;
 const SLOW_CONSUMER_LOG_NEVER_MS: u64 = u64::MAX;
 
@@ -111,7 +112,10 @@ impl NatsTransport {
             }
         });
 
-        let client = options.connect(url).await?;
+        let client = options
+            .subscription_capacity(NATS_SUBSCRIPTION_CAPACITY)
+            .connect(url)
+            .await?;
         info!(url, "connected to NATS");
 
         Ok(Self {