Monorepo consolidation: workspace, shared types, transport plans, docker/swam assets

aggregate/src/main.rs

@@ -0,0 +1,213 @@
+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());
+    }
+}