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
This commit is contained in:
nutcas3
2026-04-25 00:50:42 +03:00
parent 0fc60e553d
commit f0bbbc4459
@@ -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<u64> {
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<u64> {
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<u64> {
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<bool> {
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<String> = conn.get(&key).await
.context("Failed to check block status in Redis")?;
let blocked: Option<String> = 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<Option<SubscriberAccount>> {
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<SubscriberAccount> = conn.get(&key).await
.context("Failed to get subscriber account from Redis")?;
let account: Option<SubscriberAccount> = 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(())