telecom-platform/apps/packet-gateway/src/main.rs

use anyhow::{Context, Result};
use clap::Parser;
use redis::AsyncCommands;
use tokio::time::{interval, Duration};
use tracing::{info, warn, error, debug};

mod ebpf;
use ebpf::EbpfManager;

#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// Network interface to attach XDP program
    #[arg(short, long, default_value = "eth0")]
    interface: String,

    /// Redis connection URL
    #[arg(short, long, default_value = "redis://127.0.0.1/")]
    redis_url: String,

    /// Stats sync interval in seconds
    #[arg(short, long, default_value = "1")]
    sync_interval: u64,
}

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize tracing
    tracing_subscriber::fmt()
        .with_target(false)
        .compact()
        .init();

    dotenv::dotenv().ok();
    let args = Args::parse();

    info!("Starting Packet Gateway");
    info!("Interface: {}", args.interface);
    info!("Redis URL: {}", args.redis_url);

    // Initialize eBPF manager
    let mut ebpf_manager = EbpfManager::new(args.interface.clone()).await
        .context("Failed to initialize eBPF manager")?;
    
    // Load and attach XDP program
    ebpf_manager.attach()
        .context("Failed to attach eBPF program")?;
    
    // Connect to Redis
    let redis_client = redis::Client::open(args.redis_url.as_str())?;
    let mut redis_conn = redis_client.get_multiplexed_async_connection().await?;
    
    // Test connection
    let _: () = redis::cmd("PING").query_async(&mut redis_conn).await?;
    info!("Connected to Redis successfully");
    
    // Initial sync from Redis to eBPF maps
    info!("Performing initial sync from Redis to eBPF maps...");
    ebpf_manager.sync_from_redis(&mut redis_conn).await
        .context("Failed to sync from Redis to eBPF maps")?;
    info!("Initial sync completed");

    // Set up periodic synchronization loop
    let mut ticker = interval(Duration::from_secs(args.sync_interval));
    
    info!("Packet gateway running. Press Ctrl+C to exit.");
    
    // Main synchronization loop with graceful shutdown
    let mut counter = 0u64;
    loop {
        tokio::select! {
            _ = ticker.tick() => {
                // Use batch sync from eBPF to Redis
                if let Err(e) = ebpf_manager.sync_batch_to_redis(&mut redis_conn).await {
                    error!("Failed to batch sync eBPF maps to Redis: {}", e);
                    continue;
                }
                
                // Sync Redis data to eBPF maps (for credit updates, etc.)
                if let Err(e) = ebpf_manager.sync_from_redis(&mut redis_conn).await {
                    error!("Failed to sync from Redis to eBPF maps: {}", e);
                }
                
                counter += 1;
                
                // Log current stats (every 10 iterations to avoid spam)
                if counter % 10 == 0 {
                    if let Ok(stats) = ebpf_manager.get_packet_stats() {
                        let total_bytes: u64 = stats.iter().map(|s| s.bytes).sum();
                        info!("Current tracked IPs: {}, Total bytes processed: {}", stats.len(), total_bytes);
                    }
                    
                    if let Ok(credits) = ebpf_manager.get_credit_info() {
                        let low_credit_users = credits.iter().filter(|c| c.credit < 1000).count();
                        if low_credit_users > 0 {
                            warn!("{} users with low credit balance", low_credit_users);
                        }
                    }
                }
            }
            _ = tokio::signal::ctrl_c() => {
                info!("Received shutdown signal, cleaning up...");
                
                // Final sync before shutdown
                if let Err(e) = ebpf_manager.sync_batch_to_redis(&mut redis_conn).await {
                    error!("Failed to sync final stats to Redis: {}", e);
                }
                
                // Clean up eBPF maps
                if let Err(e) = ebpf_manager.cleanup_maps() {
                    error!("Failed to clean up eBPF maps: {}", e);
                }
                
                info!("Packet gateway shutdown complete");
                return Ok(());
            }
        }
}