Add usage queue worker autoscaling

crates/aether-runtime-state/src/redis/client.rs

@@ -1,7 +1,7 @@
 use crate::error::RedisResultExt;
 use crate::redis::RedisKeyspace;
 use crate::DataLayerError;
-use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
 use tracing::info;
@@ -9,6 +9,10 @@ use tracing::info;
 pub(crate) type RedisClient = redis::Client;
 pub(crate) type RedisManagedConnection = redis::aio::ConnectionManager;
 
+const DEFAULT_BLOCKING_STREAM_LANES_FALLBACK: usize = 4;
+const DEFAULT_BLOCKING_STREAM_LANES_CAP: usize = 16;
+const MAX_BLOCKING_STREAM_LANES_CAP: usize = 64;
+
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
 pub struct RedisClientConfig {
     pub url: String,
@@ -57,7 +61,21 @@ impl RedisClientFactory {
         &self,
         command_timeout_ms: Option<u64>,
     ) -> Result<RedisConnectionRouter, DataLayerError> {
-        RedisConnectionRouter::connect(self.connect_lazy()?, command_timeout_ms).await
+        self.connect_router_with_blocking_stream_lanes(command_timeout_ms, None)
+            .await
+    }
+
+    pub(crate) async fn connect_router_with_blocking_stream_lanes(
+        &self,
+        command_timeout_ms: Option<u64>,
+        blocking_stream_lanes: Option<usize>,
+    ) -> Result<RedisConnectionRouter, DataLayerError> {
+        RedisConnectionRouter::connect(
+            self.connect_lazy()?,
+            command_timeout_ms,
+            blocking_stream_lanes,
+        )
+        .await
     }
 }
 
@@ -84,7 +102,8 @@ impl RedisConnectionLane {
 pub(crate) struct RedisConnectionRouter {
     fast: RedisManagedConnection,
     stream: RedisManagedConnection,
-    blocking_stream: RedisManagedConnection,
+    blocking_stream: Arc<Vec<RedisManagedConnection>>,
+    blocking_stream_next: Arc<AtomicUsize>,
     admin: RedisManagedConnection,
     metrics: Arc<RedisConnectionMetrics>,
 }
@@ -93,6 +112,7 @@ impl std::fmt::Debug for RedisConnectionRouter {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("RedisConnectionRouter")
             .field("lanes", &["fast", "stream", "blocking_stream", "admin"])
+            .field("blocking_stream_lanes", &self.blocking_stream.len())
             .finish()
     }
 }
@@ -101,6 +121,7 @@ impl RedisConnectionRouter {
     pub(crate) async fn connect(
         client: RedisClient,
         command_timeout_ms: Option<u64>,
+        blocking_stream_lanes: Option<usize>,
     ) -> Result<Self, DataLayerError> {
         let fast = connect_lane(
             &client,
@@ -116,13 +137,9 @@ impl RedisConnectionRouter {
             command_timeout_ms,
         )
         .await?;
-        let blocking_stream = connect_lane(
-            &client,
-            connection_manager_config(command_timeout_ms),
-            RedisConnectionLane::BlockingStream,
-            command_timeout_ms,
-        )
-        .await?;
+        let blocking_stream =
+            connect_blocking_stream_lanes(&client, command_timeout_ms, blocking_stream_lanes)
+                .await?;
         let admin = connect_lane(
             &client,
             connection_manager_config(command_timeout_ms),
@@ -130,14 +147,17 @@ impl RedisConnectionRouter {
             command_timeout_ms,
         )
         .await?;
+        let blocking_stream_lanes = blocking_stream.len();
         info!(
             redis_lanes = "fast,stream,blocking_stream,admin",
+            redis_blocking_stream_lanes = blocking_stream_lanes,
             "runtime redis connection lanes initialized"
         );
         Ok(Self {
             fast,
             stream,
-            blocking_stream,
+            blocking_stream: Arc::new(blocking_stream),
+            blocking_stream_next: Arc::new(AtomicUsize::new(0)),
             admin,
             metrics: Arc::new(RedisConnectionMetrics::default()),
         })
@@ -147,7 +167,11 @@ impl RedisConnectionRouter {
         match lane {
             RedisConnectionLane::Fast => self.fast.clone(),
             RedisConnectionLane::Stream => self.stream.clone(),
-            RedisConnectionLane::BlockingStream => self.blocking_stream.clone(),
+            RedisConnectionLane::BlockingStream => {
+                let index = self.blocking_stream_next.fetch_add(1, Ordering::Relaxed)
+                    % self.blocking_stream.len();
+                self.blocking_stream[index].clone()
+            }
             RedisConnectionLane::Admin => self.admin.clone(),
         }
     }
@@ -228,6 +252,51 @@ fn connection_manager_config(
     config
 }
 
+async fn connect_blocking_stream_lanes(
+    client: &RedisClient,
+    command_timeout_ms: Option<u64>,
+    requested_lanes: Option<usize>,
+) -> Result<Vec<RedisManagedConnection>, DataLayerError> {
+    let lane_count = blocking_stream_lane_count(requested_lanes)?;
+    let mut lanes = Vec::with_capacity(lane_count);
+    for _ in 0..lane_count {
+        lanes.push(
+            connect_lane(
+                client,
+                connection_manager_config(command_timeout_ms),
+                RedisConnectionLane::BlockingStream,
+                command_timeout_ms,
+            )
+            .await?,
+        );
+    }
+    Ok(lanes)
+}
+
+fn blocking_stream_lane_count(requested_lanes: Option<usize>) -> Result<usize, DataLayerError> {
+    if matches!(requested_lanes, Some(0)) {
+        return Err(DataLayerError::InvalidConfiguration(
+            "runtime redis blocking_stream_lanes must be positive".to_string(),
+        ));
+    }
+
+    let default_lanes = default_blocking_stream_lane_count();
+    Ok(requested_lanes
+        .map(|lanes| lanes.max(default_lanes))
+        .unwrap_or(default_lanes)
+        .clamp(1, MAX_BLOCKING_STREAM_LANES_CAP))
+}
+
+fn default_blocking_stream_lane_count() -> usize {
+    std::thread::available_parallelism()
+        .map(|value| value.get())
+        .unwrap_or(DEFAULT_BLOCKING_STREAM_LANES_FALLBACK)
+        .clamp(
+            DEFAULT_BLOCKING_STREAM_LANES_FALLBACK,
+            DEFAULT_BLOCKING_STREAM_LANES_CAP,
+        )
+}
+
 async fn connect_lane(
     client: &RedisClient,
     config: redis::aio::ConnectionManagerConfig,
@@ -259,7 +328,10 @@ async fn connect_lane(
 
 #[cfg(test)]
 mod tests {
-    use super::{RedisClientConfig, RedisClientFactory};
+    use super::{
+        blocking_stream_lane_count, default_blocking_stream_lane_count, RedisClientConfig,
+        RedisClientFactory, MAX_BLOCKING_STREAM_LANES_CAP,
+    };
 
     #[test]
     fn factory_builds_lazy_client_from_valid_config() {
@@ -274,4 +346,28 @@ mod tests {
             .connect_lazy()
             .expect("lazy redis client should build");
     }
+
+    #[test]
+    fn blocking_stream_lane_count_uses_requested_as_floor() {
+        let default_lanes = default_blocking_stream_lane_count();
+
+        assert_eq!(
+            blocking_stream_lane_count(None).expect("default lanes"),
+            default_lanes
+        );
+        assert_eq!(
+            blocking_stream_lane_count(Some(1)).expect("requested below default"),
+            default_lanes
+        );
+        assert_eq!(
+            blocking_stream_lane_count(Some(default_lanes + 1)).expect("requested above default"),
+            default_lanes + 1
+        );
+        assert_eq!(
+            blocking_stream_lane_count(Some(MAX_BLOCKING_STREAM_LANES_CAP + 1))
+                .expect("requested above cap"),
+            MAX_BLOCKING_STREAM_LANES_CAP
+        );
+        assert!(blocking_stream_lane_count(Some(0)).is_err());
+    }
 }