From 996638546b020b335327c5d73c3f6b93a809b357 Mon Sep 17 00:00:00 2001 From: Dawid Nowak Date: Tue, 12 May 2026 16:19:50 +0100 Subject: [PATCH] Changing Redis connection management Signed-off-by: Dawid Nowak --- Cargo.lock | 12 ++++++++ Cargo.toml | 1 + crates/contextforge-gateway-rs-lib/Cargo.toml | 1 + .../contextforge-gateway-rs-lib/src/common.rs | 16 ++++++++++ .../session_store/redis_session_store.rs | 28 ++++++++++-------- crates/contextforge-gateway-rs-lib/src/lib.rs | 24 ++++++++++----- .../src/tests/gateway_end_to_end.rs | 6 ++-- .../user_config_store/redis_config_store.rs | 29 ++++++++++--------- crates/contextforge-gateway-rs/src/logging.rs | 12 ++++++-- crates/contextforge-gateway-rs/src/main.rs | 9 ++---- 10 files changed, 96 insertions(+), 42 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4e42894a..ac0780d8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -351,6 +351,15 @@ dependencies = [ "url", ] +[[package]] +name = "backon" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef" +dependencies = [ + "fastrand", +] + [[package]] name = "base16ct" version = "0.2.0" @@ -2983,11 +2992,14 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72d32a1ac9123f0d84fda64bfc02a271d9868483162dd2d9099b5c362ece064c" dependencies = [ + "arc-swap", "arcstr", "async-lock", + "backon", "bytes", "cfg-if", "combine", + "futures-channel", "futures-util", "itoa", "num-bigint", diff --git a/Cargo.toml b/Cargo.toml index d80bd89c..a0adfc7b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -55,6 +55,7 @@ redis = { version = "1.2.1", features = [ "default", "tokio-rustls-comp", "tls-rustls", + "connection-manager" ] } clap = { version = "4.5.60", features = ["derive", "env"] } thiserror = "2.0.18" diff --git a/crates/contextforge-gateway-rs-lib/Cargo.toml b/crates/contextforge-gateway-rs-lib/Cargo.toml index c0ae6dcb..4a41b69e 100644 --- a/crates/contextforge-gateway-rs-lib/Cargo.toml +++ b/crates/contextforge-gateway-rs-lib/Cargo.toml @@ -63,6 +63,7 @@ mockito = "1.7.2" axum-test = "20.0.0" test-log = "0.2.20" axum-server = { version = "0.8.0", features = ["tls-rustls"] } +futures.workspace = true [lints] workspace = true diff --git a/crates/contextforge-gateway-rs-lib/src/common.rs b/crates/contextforge-gateway-rs-lib/src/common.rs index de6841ec..869d009e 100644 --- a/crates/contextforge-gateway-rs-lib/src/common.rs +++ b/crates/contextforge-gateway-rs-lib/src/common.rs @@ -90,6 +90,16 @@ pub enum RedisConnectionMode { Mtls, } +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)] +#[derive(Default)] +pub enum LogRotation { + Minutely, + #[default] + Hourly, + Daily, + Never, +} + #[derive(Debug, Clone, Parser, Default)] #[command(name = "contextforge-gateway-rs")] #[command(about = "Minimal, fast and experimental Gateway/Dataplane for ContextForge")] @@ -150,6 +160,12 @@ pub struct Config { #[arg(long, env = "CONTEXTFORGE_GATEWAY_RS_REDIS_TLS_REDIS_CLIENT_CERTIFICATE")] pub redis_tls_client_certificate: Option, + + #[arg(long, env = "CONTEXTFORGE_GATEWAY_LOG_NAME")] + pub log_name: Option, + + #[arg(long, env = "CONTEXTFORGE_GATEWAY_LOG_ROTATION")] + pub log_rotation: Option, } #[derive(Error, Debug)] diff --git a/crates/contextforge-gateway-rs-lib/src/gateway/session_store/redis_session_store.rs b/crates/contextforge-gateway-rs-lib/src/gateway/session_store/redis_session_store.rs index d6c65a1b..a6572763 100644 --- a/crates/contextforge-gateway-rs-lib/src/gateway/session_store/redis_session_store.rs +++ b/crates/contextforge-gateway-rs-lib/src/gateway/session_store/redis_session_store.rs @@ -2,7 +2,11 @@ use std::sync::Arc; use async_trait::async_trait; use lru_time_cache::LruCache; -use redis::{AsyncCommands, RedisError, cmd}; +use redis::{ + AsyncCommands, RedisError, + aio::{ConnectionManager, ConnectionManagerConfig}, + cmd, +}; use tokio::sync::Mutex; use super::{SessionMapping, SessionStoreError, UserSession, UserSessionStore}; @@ -13,19 +17,23 @@ use crate::{ #[derive(Clone)] pub struct RedisUserSessionStore { - redis_client: RedisClient, cache: Arc>>, + connection: ConnectionManager, } + impl RedisUserSessionStore { #[expect(dead_code, reason = "Redis-backed user sessions are implemented but not wired by default")] - pub fn new(redis_client: RedisClient) -> Self { - Self { - redis_client, + pub async fn new(redis_client: &RedisClient) -> crate::Result { + Ok(Self { cache: Arc::new(Mutex::new(LruCache::with_expiry_duration_and_capacity( LRU_CACHE_EXPIRY_DURATION, LRU_CACHE_ENTRIES, ))), - } + connection: redis_client + .get_connection_manager_with_config(ConnectionManagerConfig::default()) + .await + .map_err(|_| SessionStoreError::InvalidConnection)?, + }) } } @@ -44,9 +52,7 @@ impl UserSessionStore for RedisUserSessionStore { return Err(SessionStoreError::DataEncoding); }; - let Ok(mut connection) = self.redis_client.get_multiplexed_async_connection().await else { - return Err(SessionStoreError::InvalidConnection); - }; + let mut connection = self.connection.clone(); let maybe_user_session: Result>, RedisError> = cmd("GET").arg(key).take().query_async(&mut connection).await; @@ -77,9 +83,7 @@ impl UserSessionStore for RedisUserSessionStore { return Err(SessionStoreError::DataEncoding); }; - let Ok(mut connection) = self.redis_client.get_multiplexed_async_connection().await else { - return Err(SessionStoreError::InvalidConnection); - }; + let mut connection = self.connection.clone(); if connection.set::<&[u8], &[u8], String>(&key, &encoded).await.is_ok() { self.cache.lock().await.insert(session_key.clone(), mapping.clone()); diff --git a/crates/contextforge-gateway-rs-lib/src/lib.rs b/crates/contextforge-gateway-rs-lib/src/lib.rs index 4a1e382b..c8cde43e 100644 --- a/crates/contextforge-gateway-rs-lib/src/lib.rs +++ b/crates/contextforge-gateway-rs-lib/src/lib.rs @@ -27,9 +27,8 @@ use layers::session_id::SessionId; use tower_http::cors::{Any, CorsLayer}; use transports::{DownstreamTls, Tcp}; use typed_builder::TypedBuilder; -pub use user_config_store::RedisUserConfigStore; -pub use crate::common::Config; +pub use crate::common::{Config, LogRotation}; pub type Error = Box; pub type Result = std::result::Result; @@ -42,22 +41,33 @@ use crate::{ claims_id::claims_layer, session_id::SessionIdLayer, user_config_store::user_config_store_layer, virtual_host_id::virtual_host_id_layer, }, - user_config_store::UserConfigStore, + user_config_store::{RedisUserConfigStore, UserConfigStore}, }; +#[derive(Clone)] +pub enum UserConfigStoreType { + Redis, + Test(Arc), +} + #[derive(Clone, TypedBuilder)] #[builder(field_defaults(setter(prefix = "with_")))] pub struct Gateway { config: Config, session_manager: Arc, - user_config_store: Arc, + user_config_store_type: UserConfigStoreType, } impl Gateway { pub async fn run_gateway(self) -> Result<()> { let config = &self.config; let session_manager = self.session_manager; - let user_config_store = self.user_config_store; + + let user_config_store = match self.user_config_store_type { + UserConfigStoreType::Redis => Arc::new(get_config_store(config).await?), + UserConfigStoreType::Test(store) => store, + }; + let user_config_store = user_config_store as Arc; let user_session_store = LocalUserSessionStore::new(); @@ -137,7 +147,7 @@ impl Gateway { } } -pub fn get_config_store(config: &Config) -> Result { +pub async fn get_config_store(config: &Config) -> Result { let redis_config = RedisConfig::try_from(config)?; - Ok(RedisUserConfigStore::new(RedisClient::try_from(redis_config)?)) + RedisUserConfigStore::new(&RedisClient::try_from(redis_config)?).await } diff --git a/crates/contextforge-gateway-rs-lib/src/tests/gateway_end_to_end.rs b/crates/contextforge-gateway-rs-lib/src/tests/gateway_end_to_end.rs index eae834fd..7ad77489 100644 --- a/crates/contextforge-gateway-rs-lib/src/tests/gateway_end_to_end.rs +++ b/crates/contextforge-gateway-rs-lib/src/tests/gateway_end_to_end.rs @@ -158,8 +158,9 @@ async fn create_gateway_with_four_counters(user: &str, config: Config) -> crate: let gateway = Gateway::builder() .with_config(config.clone()) - .with_user_config_store(Arc::new(mocked_user_config_store)) + //.with_user_config_store(Arc::new(mocked_user_config_store)) .with_session_manager(Arc::new(LocalSessionManager::default())) + .with_user_config_store_type(crate::UserConfigStoreType::Test(Arc::new(mocked_user_config_store))) .build(); let gateway = async move { @@ -225,8 +226,9 @@ async fn create_tls_gateway_with_four_tls_counters(user: &str, config: Config) - let gateway = Gateway::builder() .with_config(config.clone()) - .with_user_config_store(Arc::new(mocked_user_config_store)) + //.with_user_config_store(Arc::new(mocked_user_config_store)) .with_session_manager(Arc::new(LocalSessionManager::default())) + .with_user_config_store_type(crate::UserConfigStoreType::Test(Arc::new(mocked_user_config_store))) .build(); let gateway = async move { diff --git a/crates/contextforge-gateway-rs-lib/src/user_config_store/redis_config_store.rs b/crates/contextforge-gateway-rs-lib/src/user_config_store/redis_config_store.rs index b02da341..84898298 100644 --- a/crates/contextforge-gateway-rs-lib/src/user_config_store/redis_config_store.rs +++ b/crates/contextforge-gateway-rs-lib/src/user_config_store/redis_config_store.rs @@ -3,7 +3,12 @@ use std::sync::Arc; use async_trait::async_trait; use contextforge_gateway_rs_apis::user_store::UserConfig; use lru_time_cache::LruCache; -use redis::{AsyncCommands, RedisError, cmd}; +use redis::{ + AsyncCommands, RedisError, + aio::{ConnectionManager, ConnectionManagerConfig}, + cmd, +}; + use tokio::sync::Mutex; use super::{ConfigStoreError, UserConfigStore}; @@ -15,18 +20,21 @@ use crate::{ #[derive(Clone)] pub struct RedisUserConfigStore { - redis_client: RedisClient, + connection: ConnectionManager, cache: Arc>>, } impl RedisUserConfigStore { - pub fn new(redis_client: RedisClient) -> Self { - Self { - redis_client, + pub async fn new(redis_client: &RedisClient) -> crate::Result { + Ok(Self { + connection: redis_client + .get_connection_manager_with_config(ConnectionManagerConfig::default()) + .await + .map_err(|_| ConfigStoreError::InvalidConnection)?, cache: Arc::new(Mutex::new(LruCache::with_expiry_duration_and_capacity( LRU_CACHE_EXPIRY_DURATION, LRU_CACHE_ENTRIES, ))), - } + }) } } @@ -45,10 +53,7 @@ impl UserConfigStore for RedisUserConfigStore { return Err(ConfigStoreError::DataEncoding); }; - let Ok(mut connection) = self.redis_client.get_multiplexed_async_connection().await else { - return Err(ConfigStoreError::InvalidConnection); - }; - + let mut connection = self.connection.clone(); let maybe_user_config: Result>, RedisError> = cmd("GET").arg(key).take().query_async(&mut connection).await; @@ -74,9 +79,7 @@ impl UserConfigStore for RedisUserConfigStore { return Err(ConfigStoreError::DataEncoding); }; - let Ok(mut connection) = self.redis_client.get_multiplexed_async_connection().await else { - return Err(ConfigStoreError::InvalidConnection); - }; + let mut connection = self.connection.clone(); if connection.set::<&[u8], &[u8], String>(&key, &encoded).await.is_ok() { self.cache.lock().await.insert(user_key.key.to_owned(), config.clone()); diff --git a/crates/contextforge-gateway-rs/src/logging.rs b/crates/contextforge-gateway-rs/src/logging.rs index 71f5e7f2..cfe7523b 100644 --- a/crates/contextforge-gateway-rs/src/logging.rs +++ b/crates/contextforge-gateway-rs/src/logging.rs @@ -1,4 +1,4 @@ -use contextforge_gateway_rs_lib::Config; +use contextforge_gateway_rs_lib::{Config, LogRotation}; use opentelemetry::trace::TracerProvider; use opentelemetry_otlp::WithExportConfig; use opentelemetry_sdk::trace::{RandomIdGenerator, Sampler}; @@ -20,7 +20,15 @@ const CONTROLLER_NAME: &str = "CONTEXTFORGE-GATEWAY-RS"; pub fn init_tracing_logging(configuration: &Config) -> Guard { let registry = Registry::default(); - let file_appender = tracing_appender::rolling::minutely(".", "contextforge-gateway-rs.log"); + let log_name = configuration.log_name.clone().unwrap_or("contextforge-gateway-rs.log".to_owned()); + + let file_appender = match configuration.log_rotation.clone().unwrap_or_default() { + LogRotation::Minutely => tracing_appender::rolling::minutely(".", log_name), + LogRotation::Hourly => tracing_appender::rolling::hourly(".", log_name), + LogRotation::Daily => tracing_appender::rolling::daily(".", log_name), + LogRotation::Never => tracing_appender::rolling::never(".", log_name), + }; + let (non_blocking_appender, guard) = tracing_appender::non_blocking(file_appender); let file_filter = tracing_subscriber::EnvFilter::new(std::env::var("RUST_FILE_LOG").unwrap_or_else(|_| "debug".to_owned())); diff --git a/crates/contextforge-gateway-rs/src/main.rs b/crates/contextforge-gateway-rs/src/main.rs index e57ee98e..1792ea11 100644 --- a/crates/contextforge-gateway-rs/src/main.rs +++ b/crates/contextforge-gateway-rs/src/main.rs @@ -4,7 +4,7 @@ mod runtime; use std::sync::Arc; use clap::Parser; -use contextforge_gateway_rs_lib::{Config, Gateway, get_config_store}; +use contextforge_gateway_rs_lib::{Config, Gateway, UserConfigStoreType}; use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; use rustls::crypto; use tikv_jemallocator::Jemalloc; @@ -22,14 +22,11 @@ fn main() -> Result<(), Box> { let runtime = runtime::Runtime::from(&config); - let user_config_store = get_config_store(&config)?; let gateway = Gateway::builder() .with_config(config) - .with_user_config_store(Arc::new(user_config_store)) + .with_user_config_store_type(UserConfigStoreType::Redis) .with_session_manager(Arc::new(LocalSessionManager::default())) .build(); - _ = runtime.execute(gateway); - - Ok(()) + runtime.execute(gateway) }