refactor: extract runtime state backends

apps/aether-gateway/src/tunnel/embedded/local_relay.rs

@@ -134,16 +134,16 @@ fn map_request_admission_error(error: super::RequestAdmissionError) -> String {
             ..
         })
         | super::RequestAdmissionError::Distributed(
-            aether_runtime::DistributedConcurrencyError::Saturated { .. },
+            aether_runtime_state::RuntimeSemaphoreError::Saturated { .. },
         )
         | super::RequestAdmissionError::Distributed(
-            aether_runtime::DistributedConcurrencyError::Unavailable { .. },
+            aether_runtime_state::RuntimeSemaphoreError::Unavailable { .. },
         ) => "overloaded: hub relay overloaded".to_string(),
         super::RequestAdmissionError::Local(aether_runtime::ConcurrencyError::Closed {
             ..
         }) => "overloaded: hub relay gate closed".to_string(),
         super::RequestAdmissionError::Distributed(
-            aether_runtime::DistributedConcurrencyError::InvalidConfiguration(_),
+            aether_runtime_state::RuntimeSemaphoreError::InvalidConfiguration(_),
         ) => "overloaded: hub relay distributed gate invalid".to_string(),
     }
 }
@@ -174,10 +174,10 @@ pub async fn relay_request(
             ..
         }))
         | Err(super::RequestAdmissionError::Distributed(
-            aether_runtime::DistributedConcurrencyError::Saturated { .. },
+            aether_runtime_state::RuntimeSemaphoreError::Saturated { .. },
         ))
         | Err(super::RequestAdmissionError::Distributed(
-            aether_runtime::DistributedConcurrencyError::Unavailable { .. },
+            aether_runtime_state::RuntimeSemaphoreError::Unavailable { .. },
         )) => {
             return tunnel_error_response(
                 StatusCode::SERVICE_UNAVAILABLE,
@@ -195,7 +195,7 @@ pub async fn relay_request(
             );
         }
         Err(super::RequestAdmissionError::Distributed(
-            aether_runtime::DistributedConcurrencyError::InvalidConfiguration(_),
+            aether_runtime_state::RuntimeSemaphoreError::InvalidConfiguration(_),
         )) => {
             return tunnel_error_response(
                 StatusCode::SERVICE_UNAVAILABLE,

apps/aether-gateway/src/tunnel/embedded/mod.rs

@@ -8,10 +8,9 @@ use std::sync::Arc;
 
 use aether_runtime::{
     hold_admission_permit_until, prometheus_response, service_up_sample, AdmissionPermit,
-    ConcurrencyError, ConcurrencyGate, ConcurrencySnapshot, DistributedConcurrencyError,
-    DistributedConcurrencyGate, DistributedConcurrencySnapshot, MetricKind, MetricLabel,
-    MetricSample,
+    ConcurrencyError, ConcurrencyGate, ConcurrencySnapshot, MetricKind, MetricLabel, MetricSample,
 };
+use aether_runtime_state::{RuntimeSemaphore, RuntimeSemaphoreError, RuntimeSemaphoreSnapshot};
 use axum::extract::ws::WebSocketUpgrade;
 use axum::extract::State;
 use axum::response::{IntoResponse, Json};
@@ -33,13 +32,13 @@ pub struct AppState {
     pub max_streams: usize,
     data: Arc<GatewayDataState>,
     request_gate: Option<Arc<ConcurrencyGate>>,
-    distributed_request_gate: Option<Arc<DistributedConcurrencyGate>>,
+    distributed_request_gate: Option<Arc<RuntimeSemaphore>>,
 }
 
 #[derive(Debug)]
 enum RequestAdmissionError {
     Local(ConcurrencyError),
-    Distributed(DistributedConcurrencyError),
+    Distributed(RuntimeSemaphoreError),
 }
 
 impl AppState {
@@ -70,7 +69,7 @@ impl AppState {
         self
     }
 
-    pub fn with_distributed_request_gate(mut self, gate: DistributedConcurrencyGate) -> Self {
+    pub fn with_distributed_request_gate(mut self, gate: RuntimeSemaphore) -> Self {
         self.distributed_request_gate = Some(Arc::new(gate));
         self
     }
@@ -81,7 +80,7 @@ impl AppState {
 
     async fn distributed_request_concurrency_snapshot(
         &self,
-    ) -> Result<Option<DistributedConcurrencySnapshot>, DistributedConcurrencyError> {
+    ) -> Result<Option<RuntimeSemaphoreSnapshot>, RuntimeSemaphoreError> {
         match self.distributed_request_gate.as_ref() {
             Some(gate) => gate.snapshot().await.map(Some),
             None => Ok(None),
@@ -225,12 +224,10 @@ pub async fn ws_proxy(
     let request_permit = match state.try_acquire_request_permit().await {
         Ok(permit) => permit,
         Err(RequestAdmissionError::Local(ConcurrencyError::Saturated { .. }))
-        | Err(RequestAdmissionError::Distributed(DistributedConcurrencyError::Saturated {
-            ..
-        }))
-        | Err(RequestAdmissionError::Distributed(DistributedConcurrencyError::Unavailable {
-            ..
-        })) => return axum::http::StatusCode::SERVICE_UNAVAILABLE.into_response(),
+        | Err(RequestAdmissionError::Distributed(RuntimeSemaphoreError::Saturated { .. }))
+        | Err(RequestAdmissionError::Distributed(RuntimeSemaphoreError::Unavailable { .. })) => {
+            return axum::http::StatusCode::SERVICE_UNAVAILABLE.into_response()
+        }
         Err(RequestAdmissionError::Local(ConcurrencyError::Closed { gate })) => {
             warn!(
                 gate = gate,
@@ -238,9 +235,9 @@ pub async fn ws_proxy(
             );
             return axum::http::StatusCode::SERVICE_UNAVAILABLE.into_response();
         }
-        Err(RequestAdmissionError::Distributed(
-            DistributedConcurrencyError::InvalidConfiguration(message),
-        )) => {
+        Err(RequestAdmissionError::Distributed(RuntimeSemaphoreError::InvalidConfiguration(
+            message,
+        ))) => {
             warn!(
                 error = %message,
                 "standalone tunnel relay distributed request gate is invalid"

apps/aether-gateway/src/tunnel/mod.rs

@@ -14,6 +14,7 @@ use aether_data::repository::proxy_nodes::{
     ProxyNodeHeartbeatMutation, ProxyNodeTunnelStatusMutation, StoredProxyNode,
 };
 use aether_runtime::MetricSample;
+use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeState};
 use async_stream::stream;
 use axum::body::{Body, Bytes};
 use axum::extract::ws::WebSocketUpgrade;
@@ -102,6 +103,7 @@ pub(crate) struct TunnelAttachmentRecord {
 #[derive(Debug, Clone)]
 pub(crate) struct TunnelAttachmentDirectory {
     identity: Arc<TunnelInstanceIdentity>,
+    runtime_state: Arc<RuntimeState>,
 }
 
 impl TunnelAttachmentDirectory {
@@ -118,6 +120,7 @@ impl TunnelAttachmentDirectory {
                     .map(|value| value.clamp(15, 3600))
                     .unwrap_or(DEFAULT_ATTACHMENT_TTL_SECS),
             }),
+            runtime_state: Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())),
         }
     }
 
@@ -132,9 +135,15 @@ impl TunnelAttachmentDirectory {
                 relay_base_url: relay_base_url.map(Into::into),
                 attachment_ttl_secs,
             }),
+            runtime_state: Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())),
         }
     }
 
+    fn with_runtime_state(mut self, runtime_state: Arc<RuntimeState>) -> Self {
+        self.runtime_state = runtime_state;
+        self
+    }
+
     #[cfg(test)]
     pub(crate) fn for_tests(
         instance_id: &str,
@@ -257,7 +266,7 @@ impl TunnelAttachmentDirectory {
         data: &GatewayDataState,
         node_id: &str,
     ) -> Result<Option<TunnelAttachmentRecord>, String> {
-        match self.read_attachment_record_from_redis(data, node_id).await {
+        match self.read_attachment_record_from_runtime(node_id).await {
             Ok(Some(record)) => return Ok(Some(record)),
             Ok(None) => {}
             Err(error) => {
@@ -272,28 +281,18 @@ impl TunnelAttachmentDirectory {
             .await
     }
 
-    async fn read_attachment_record_from_redis(
+    async fn read_attachment_record_from_runtime(
         &self,
-        data: &GatewayDataState,
         node_id: &str,
     ) -> Result<Option<TunnelAttachmentRecord>, String> {
-        let Some(runner) = data.kv_runner() else {
-            return Ok(None);
-        };
-        let mut connection = runner
-            .client()
-            .get_multiplexed_async_connection()
+        let raw = self
+            .runtime_state
+            .kv_get(&tunnel_attachment_redis_key(node_id))
             .await
-            .map_err(|err| format!("attachment redis connect failed: {err}"))?;
-        let namespaced_key = runner.keyspace().key(&tunnel_attachment_redis_key(node_id));
-        let raw = redis::cmd("GET")
-            .arg(&namespaced_key)
-            .query_async::<Option<String>>(&mut connection)
-            .await
-            .map_err(|err| format!("attachment redis read failed: {err}"))?;
+            .map_err(|err| format!("attachment runtime read failed: {err}"))?;
         raw.map(|value| {
             serde_json::from_str::<TunnelAttachmentRecord>(&value)
-                .map_err(|err| format!("invalid redis tunnel attachment record: {err}"))
+                .map_err(|err| format!("invalid runtime tunnel attachment record: {err}"))
         })
         .transpose()
     }
@@ -323,21 +322,20 @@ impl TunnelAttachmentDirectory {
     ) -> Result<(), String> {
         let serialized = serde_json::to_string(record)
             .map_err(|err| format!("attachment serialization failed: {err}"))?;
-        if let Some(runner) = data.kv_runner() {
-            if let Err(error) = runner
-                .setex(
-                    &tunnel_attachment_redis_key(node_id),
-                    &serialized,
-                    Some(self.identity.attachment_ttl_secs),
-                )
-                .await
-            {
-                warn!(
-                    error = %error,
-                    node_id = %node_id,
-                    "failed to write tunnel attachment to redis; keeping system_config shadow only"
-                );
-            }
+        if let Err(error) = self
+            .runtime_state
+            .kv_set(
+                &tunnel_attachment_redis_key(node_id),
+                serialized.clone(),
+                Some(Duration::from_secs(self.identity.attachment_ttl_secs)),
+            )
+            .await
+        {
+            warn!(
+                error = %error,
+                node_id = %node_id,
+                "failed to write tunnel attachment to runtime state; keeping system_config shadow only"
+            );
         }
         let value = serde_json::to_value(record)
             .map_err(|err| format!("attachment serialization failed: {err}"))?;
@@ -352,14 +350,16 @@ impl TunnelAttachmentDirectory {
         data: &GatewayDataState,
         node_id: &str,
     ) -> Result<(), String> {
-        if let Some(runner) = data.kv_runner() {
-            if let Err(error) = runner.del(&tunnel_attachment_redis_key(node_id)).await {
-                warn!(
-                    error = %error,
-                    node_id = %node_id,
-                    "failed to delete tunnel attachment from redis; clearing system_config shadow anyway"
-                );
-            }
+        if let Err(error) = self
+            .runtime_state
+            .kv_delete(&tunnel_attachment_redis_key(node_id))
+            .await
+        {
+            warn!(
+                error = %error,
+                node_id = %node_id,
+                "failed to delete tunnel attachment from runtime state; clearing system_config shadow anyway"
+            );
         }
         data.delete_system_config_value(&tunnel_attachment_key(node_id))
             .await
@@ -396,6 +396,16 @@ impl EmbeddedTunnelState {
         Self::with_data_and_directory(data, TunnelAttachmentDirectory::from_environment())
     }
 
+    pub(crate) fn with_data_and_runtime_state(
+        data: Arc<GatewayDataState>,
+        runtime_state: Arc<RuntimeState>,
+    ) -> Self {
+        Self::with_data_and_directory(
+            data,
+            TunnelAttachmentDirectory::from_environment().with_runtime_state(runtime_state),
+        )
+    }
+
     pub(crate) fn with_data_and_identity(
         data: Arc<GatewayDataState>,
         instance_id: impl Into<String>,
@@ -408,6 +418,20 @@ impl EmbeddedTunnelState {
         )
     }
 
+    pub(crate) fn with_data_identity_and_runtime_state(
+        data: Arc<GatewayDataState>,
+        instance_id: impl Into<String>,
+        relay_base_url: Option<impl Into<String>>,
+        attachment_ttl_secs: u64,
+        runtime_state: Arc<RuntimeState>,
+    ) -> Self {
+        Self::with_data_and_directory(
+            data,
+            TunnelAttachmentDirectory::from_parts(instance_id, relay_base_url, attachment_ttl_secs)
+                .with_runtime_state(runtime_state),
+        )
+    }
+
     pub(crate) fn with_data_and_directory(
         data: Arc<GatewayDataState>,
         attachment_directory: TunnelAttachmentDirectory,