cloudlysis/aggregate/src/main.rs

use aggregate::config::Settings;
use aggregate::gateway::server::GrpcCommandServer;
use aggregate::http_server;
use aggregate::observability::Observability;
use aggregate::runtime::RuntimeExecutor;
use aggregate::server::AdminServer;
use aggregate::storage::StorageClient;
use aggregate::stream::StreamClient;
use aggregate::swarm::TenantPlacementKvClient;
use aggregate::{aggregate::AggregateHandler, placement::TenantPlacementManager};
use futures::StreamExt;
use std::sync::Arc;
use std::time::Duration;

#[tokio::main]
async fn main() {
    match std::env::args().nth(1).as_deref() {
        Some("-h") | Some("--help") => {
            print_help();
            return;
        }
        Some("serve") | None => serve().await,
        Some(other) => {
            eprintln!("Unknown command: {}", other);
            print_help();
        }
    }
}

async fn serve() {
    let settings = load_settings();

    let observability = Observability::default();
    let health_checker = aggregate::server::HealthChecker::new();
    let admin = Arc::new(AdminServer::new(
        observability,
        health_checker,
        settings.shard_id.clone(),
    ));

    spawn_health_probe(admin.clone(), settings.clone());
    spawn_placement_watcher(admin.placement_manager(), settings.clone());

    let storage = StorageClient::open(settings.storage_path.clone()).unwrap();
    let stream = StreamClient::new(settings.nats_url.clone()).await.unwrap();
    let _ = stream.setup_stream().await;
    let executor = RuntimeExecutor::new();

    let handler = AggregateHandler::new(
        storage,
        stream,
        executor,
        settings.decide_program.clone(),
        settings.apply_program.clone(),
    )
    .with_snapshot_threshold(settings.snapshot_threshold)
    .with_max_retries(settings.max_retries);

    let grpc_addr: std::net::SocketAddr = settings.grpc_addr.parse().unwrap();
    let grpc_service = GrpcCommandServer::new(
        handler,
        admin.placement_manager(),
        admin.observability(),
        settings.multi_tenant_enabled,
        settings
            .default_tenant_id
            .as_ref()
            .map(aggregate::types::TenantId::new),
    )
    .service();

    let addr = std::env::var("AGGREGATE_HTTP_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".to_string());
    let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();

    let (shutdown_tx, _) = tokio::sync::broadcast::channel::<()>(1);
    let mut http_shutdown = shutdown_tx.subscribe();
    let mut grpc_shutdown = shutdown_tx.subscribe();

    let http_task = tokio::spawn(async move {
        http_server::serve(listener, admin, async move {
            let _ = http_shutdown.recv().await;
        })
        .await;
    });

    let grpc_task = tokio::spawn(async move {
        tonic::transport::Server::builder()
            .add_service(grpc_service)
            .serve_with_shutdown(grpc_addr, async move {
                let _ = grpc_shutdown.recv().await;
            })
            .await
            .unwrap();
    });

    let _ = tokio::signal::ctrl_c().await;
    let _ = shutdown_tx.send(());

    let _ = tokio::join!(http_task, grpc_task);
}

fn print_help() {
    println!(
        "aggregate\n\nUSAGE:\n    aggregate [COMMAND]\n\nCOMMANDS:\n    serve    Start the HTTP server (default)\n\nOPTIONS:\n    -h, --help    Print help\n"
    );
}

fn load_settings() -> Settings {
    if let Ok(path) = std::env::var("AGGREGATE_CONFIG_PATH") {
        if let Ok(settings) = Settings::load_from_file_with_env_overrides(path) {
            return settings;
        }
    }

    Settings::from_env().unwrap_or_default()
}

fn spawn_health_probe(admin: Arc<AdminServer>, settings: Settings) {
    tokio::spawn(async move {
        loop {
            let storage_ok = StorageClient::open(settings.storage_path.clone()).is_ok();
            admin.health_checker().set_storage_healthy(storage_ok);

            let stream_ok = tokio::time::timeout(Duration::from_secs(1), async {
                let stream = StreamClient::new(settings.nats_url.clone()).await?;
                let _ = stream.setup_stream().await;
                Ok::<_, aggregate::types::AggregateError>(())
            })
            .await
            .is_ok_and(|r| r.is_ok());

            admin.health_checker().set_stream_healthy(stream_ok);

            tokio::time::sleep(Duration::from_secs(5)).await;
        }
    });
}

fn spawn_placement_watcher(placement: Arc<TenantPlacementManager>, settings: Settings) {
    tokio::spawn(async move {
        loop {
            let client = TenantPlacementKvClient::connect(
                settings.nats_url.clone(),
                settings.placement_bucket.clone(),
            )
            .await;

            let client = match client {
                Ok(c) => c,
                Err(_) => {
                    tokio::time::sleep(Duration::from_secs(1)).await;
                    continue;
                }
            };

            if let Ok(Some(value)) = client.get_json(&settings.placement_key).await {
                apply_placement_value(&placement, &settings.shard_id, value).await;
            }

            let watch = client.watch_json(&settings.placement_key).await;
            let mut stream = match watch {
                Ok(s) => s,
                Err(_) => {
                    tokio::time::sleep(Duration::from_secs(1)).await;
                    continue;
                }
            };

            while let Some(update) = stream.next().await {
                if let Ok(value) = update {
                    apply_placement_value(&placement, &settings.shard_id, value).await;
                }
            }

            tokio::time::sleep(Duration::from_secs(1)).await;
        }
    });
}

async fn apply_placement_value(
    placement: &TenantPlacementManager,
    shard_id: &str,
    value: serde_json::Value,
) {
    if let Some(map) = value.as_object() {
        let placement_map = map
            .iter()
            .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string())))
            .collect::<std::collections::HashMap<_, _>>();
        placement
            .apply_placement_map(shard_id, &placement_map)
            .await;
        return;
    }

    if let Some(map) = value.get("placement").and_then(|v| v.as_object()) {
        let placement_map = map
            .iter()
            .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string())))
            .collect::<std::collections::HashMap<_, _>>();
        placement
            .apply_placement_map(shard_id, &placement_map)
            .await;
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn binary_exists() {
        assert!(std::env::current_exe().is_ok());
    }
}