From f0bbbc44597dbb2842a010223f9b990bd4d09974 Mon Sep 17 00:00:00 2001 From: nutcas3 Date: Sat, 25 Apr 2026 00:50:42 +0300 Subject: [PATCH] feat: Refactor credit management to use multiplexed Redis connections and explicit AsyncCommands trait for improved connection pooling - Update all Redis operations to use get_multiplexed_async_connection instead of get_async_connection for better connection pooling - Replace context error handling with map_err for RedisConnection/RedisOperation errors - Add explicit redis::AsyncCommands trait qualification for all Redis operations (get/set/del/incr/decr/expire) - Add ChargingSession import to types --- .../src/charging/credit_management.rs | 84 +++++++++---------- 1 file changed, 42 insertions(+), 42 deletions(-) diff --git a/apps/charging-engine/src/charging/credit_management.rs b/apps/charging-engine/src/charging/credit_management.rs index f2988db..03b1b2b 100644 --- a/apps/charging-engine/src/charging/credit_management.rs +++ b/apps/charging-engine/src/charging/credit_management.rs @@ -1,45 +1,45 @@ use redis::AsyncCommands; use tracing::{info, warn, error, debug}; -use super::types::{SubscriberAccount, UsageEvent}; +use super::types::{SubscriberAccount, UsageEvent, ChargingSession}; use crate::errors::{ChargingError, ChargingResult, ErrorContext}; -impl super::ChargingEngine { +impl crate::charging::ChargingEngine { pub async fn get_balance(&self, ip: &str) -> ChargingResult { - let mut conn = self.redis_client.get_async_connection().await - .context("Failed to get Redis connection")?; + let mut conn = self.redis_client.get_multiplexed_async_connection().await + .map_err(|e| crate::errors::ChargingError::RedisConnection(e.to_string()))?; let key = format!("credit:{}", ip); - let balance: u64 = conn.get(&key).await - .context("Failed to get balance from Redis")?; + let balance: u64 = redis::AsyncCommands::get(&mut conn, &key).await + .map_err(|e| crate::errors::ChargingError::RedisOperation(e.to_string()))?; debug!("Retrieved balance {} for IP: {}", balance, ip); Ok(balance) } pub async fn add_credit(&self, ip: &str, bytes_to_add: u64) -> ChargingResult { - let mut conn = self.redis_client.get_async_connection().await - .context("Failed to get Redis connection")?; + let mut conn = self.redis_client.get_multiplexed_async_connection().await + .map_err(|e| crate::errors::ChargingError::RedisConnection(e.to_string()))?; let key = format!("credit:{}", ip); - let new_balance: u64 = conn.incr(&key, bytes_to_add).await - .context("Failed to add credit in Redis")?; + let new_balance: u64 = redis::AsyncCommands::incr(&mut conn, &key, bytes_to_add).await + .map_err(|e| crate::errors::ChargingError::RedisOperation(e.to_string()))?; // Set expiration if this is a new key - let _: () = conn.expire(&key, 86400).await.unwrap_or(()); + let _: () = redis::AsyncCommands::expire(&mut conn, &key, 86400).await.unwrap_or(()); info!("Added {} bytes credit to IP: {}, new balance: {}", bytes_to_add, ip, new_balance); Ok(new_balance) } pub async fn deduct_credit(&self, ip: &str, bytes_used: u64) -> ChargingResult { - let mut conn = self.redis_client.get_async_connection().await - .context("Failed to get Redis connection")?; + let mut conn = self.redis_client.get_multiplexed_async_connection().await + .map_err(|e| crate::errors::ChargingError::RedisConnection(e.to_string()))?; let key = format!("credit:{}", ip); // Check current balance first - let current_balance: u64 = conn.get(&key).await.unwrap_or(0); + let current_balance: u64 = redis::AsyncCommands::get(&mut conn, &key).await.unwrap_or(0); if current_balance < bytes_used { return Err(ChargingError::InsufficientCredit { @@ -48,8 +48,8 @@ impl super::ChargingEngine { }); } - let new_balance: u64 = conn.decr(&key, bytes_used).await - .context("Failed to deduct credit in Redis")?; + let new_balance: u64 = redis::AsyncCommands::decr(&mut conn, &key, bytes_used).await + .map_err(|e| crate::errors::ChargingError::RedisOperation(e.to_string()))?; info!("Deducted {} bytes credit from IP: {}, new balance: {}", bytes_used, ip, new_balance); Ok(new_balance) @@ -61,76 +61,76 @@ impl super::ChargingEngine { } pub async fn block_user(&self, ip: &str) -> ChargingResult<()> { - let mut conn = self.redis_client.get_async_connection().await - .context("Failed to get Redis connection")?; + let mut conn = self.redis_client.get_multiplexed_async_connection().await + .map_err(|e| crate::errors::ChargingError::RedisConnection(e.to_string()))?; let key = format!("block:{}", ip); - let _: () = conn.set(&key, "blocked").await - .context("Failed to block user in Redis")?; + let _: () = redis::AsyncCommands::set(&mut conn, &key, "blocked").await + .map_err(|e| crate::errors::ChargingError::RedisOperation(e.to_string()))?; // Set expiration for block (24 hours) - let _: () = conn.expire(&key, 86400).await.unwrap_or(()); + let _: () = redis::AsyncCommands::expire(&mut conn, &key, 86400).await.unwrap_or(()); warn!("User IP: {} has been blocked", ip); Ok(()) } pub async fn unblock_user(&self, ip: &str) -> ChargingResult<()> { - let mut conn = self.redis_client.get_async_connection().await - .context("Failed to get Redis connection")?; + let mut conn = self.redis_client.get_multiplexed_async_connection().await + .map_err(|e| crate::errors::ChargingError::RedisConnection(e.to_string()))?; let key = format!("block:{}", ip); - let _: () = conn.del(&key).await - .context("Failed to unblock user in Redis")?; + let _: () = redis::AsyncCommands::del(&mut conn, &key).await + .map_err(|e| crate::errors::ChargingError::RedisOperation(e.to_string()))?; info!("User IP: {} has been unblocked", ip); Ok(()) } pub async fn is_user_blocked(&self, ip: &str) -> ChargingResult { - let mut conn = self.redis_client.get_async_connection().await - .context("Failed to get Redis connection")?; + let mut conn = self.redis_client.get_multiplexed_async_connection().await + .map_err(|e| crate::errors::ChargingError::RedisConnection(e.to_string()))?; let key = format!("block:{}", ip); - let blocked: Option = conn.get(&key).await - .context("Failed to check block status in Redis")?; + let blocked: Option = redis::AsyncCommands::get(&mut conn, &key).await + .map_err(|e| crate::errors::ChargingError::RedisOperation(e.to_string()))?; Ok(blocked.is_some()) } pub async fn get_subscriber_account(&self, imsi: &str) -> ChargingResult> { - let mut conn = self.redis_client.get_async_connection().await - .context("Failed to get Redis connection")?; + let mut conn = self.redis_client.get_multiplexed_async_connection().await + .map_err(|e| crate::errors::ChargingError::RedisConnection(e.to_string()))?; let key = format!("account:{}", imsi); - let account: Option = conn.get(&key).await - .context("Failed to get subscriber account from Redis")?; + let account: Option = redis::AsyncCommands::get(&mut conn, &key).await + .map_err(|e| crate::errors::ChargingError::RedisOperation(e.to_string()))?; Ok(account) } pub async fn update_subscriber_account(&self, account: &SubscriberAccount) -> ChargingResult<()> { - let mut conn = self.redis_client.get_async_connection().await - .context("Failed to get Redis connection")?; + let mut conn = self.redis_client.get_multiplexed_async_connection().await + .map_err(|e| crate::errors::ChargingError::RedisConnection(e.to_string()))?; let key = format!("account:{}", account.imsi); - let _: () = conn.set(&key, account).await - .context("Failed to update subscriber account in Redis")?; + let _: () = redis::AsyncCommands::set(&mut conn, &key, account).await + .map_err(|e| crate::errors::ChargingError::RedisOperation(e.to_string()))?; debug!("Updated subscriber account for IMSI: {}", account.imsi); Ok(()) } pub async fn record_usage_event(&self, event: &UsageEvent) -> ChargingResult<()> { - let mut conn = self.redis_client.get_async_connection().await - .context("Failed to get Redis connection")?; + let mut conn = self.redis_client.get_multiplexed_async_connection().await + .map_err(|e| crate::errors::ChargingError::RedisConnection(e.to_string()))?; let key = format!("usage:{}:{}", event.imsi, event.session_id); - let _: () = conn.set(&key, event).await - .context("Failed to record usage event in Redis")?; + let _: () = redis::AsyncCommands::set(&mut conn, &key, event).await + .map_err(|e| crate::errors::ChargingError::RedisOperation(e.to_string()))?; // Set expiration for usage events (7 days) - let _: () = conn.expire(&key, 604800).await.unwrap_or(()); + let _: () = redis::AsyncCommands::expire(&mut conn, &key, 604800).await.unwrap_or(()); info!("Recorded usage event for IMSI: {}, cost: ${:.4}", event.imsi, event.cost); Ok(())