path: root/src/main.rs

//! continuum: a replicated Automerge key-value store synchronized over Mycelium.
//!
//! main.rs only wires things together: arguments, state construction, background
//! tasks, bootstrap and the HTTP server.

use std::collections::HashMap;
use std::future::IntoFuture;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
use std::time::Duration;

use anyhow::{Context, Result};
use automerge::AutoCommit;
use clap::Parser;
use tokio::sync::{broadcast, mpsc, watch, RwLock};
use tracing::{error, info};
use url::Url;

mod api;
mod auth;
mod bootstrap;
mod compaction;
mod firewall;
mod mycelium;
mod peer;
mod replica;
mod state;
mod store;
mod sync;
mod transport;

use auth::AuthState;
use firewall::{Firewall, FirewallTransport};
use mycelium::{MyceliumSetup, MyceliumSpawnConfig};
use replica::{Epoch, Replica};
use state::{AppState, NodeLifecycle, Timing};
use transport::Transport;

#[derive(Parser, Debug)]
#[command(version, about = "continuum messaging system", long_about = None)]
struct Args {
    #[arg(long)]
    debug_api: bool,

    #[arg(long, value_name = "HOSTNAME:PORT", default_value = "127.0.0.1:3000")]
    bind_addr: SocketAddr,

    #[arg(
        long,
        value_name = "SCHEME://HOSTNAME:PORT",
        default_value = None,
        conflicts_with_all = &["mycelium_key_file", "mycelium_network_key_file", "mycelium_network_name"],
    )]
    mycelium_rest_url: Option<Url>,

    #[arg(long)]
    pub mycelium_key_file: Option<PathBuf>,

    #[arg(long)]
    pub mycelium_network_key_file: Option<PathBuf>,

    #[arg(long)]
    pub mycelium_network_name: Option<String>,

    #[arg(long)]
    pub mycelium_tcp_listen_port: Option<u16>,

    #[arg(
        long,
        use_value_delimiter = true,
        default_value = "tls://127.0.0.1:9651"
    )]
    pub mycelium_peers: Vec<String>,

    #[arg(long, value_name = "STRING,[STRING...]", value_delimiter = ',')]
    seeds: Option<Vec<String>>,

    #[arg(long, value_name = "PATH")]
    htpasswd: String,

    /// Keepalive interval; dormancy, resync and probe timings derive from it.
    #[arg(long, value_name = "SECONDS", default_value_t = 60)]
    keepalive_secs: u64,

    /// Test hook: start on a non-genesis epoch of this generation (with a fixed dummy
    /// id), to check that nodes on different epochs refuse to sync.
    #[arg(long, hide = true, requires = "debug_api", value_name = "GENERATION")]
    debug_epoch: Option<u64>,
}

#[tokio::main]
async fn main() -> Result<()> {
    console_subscriber::init();

    let args = Args::parse();

    let mycelium_setup = {
        let mut spawn_config = MyceliumSpawnConfig {
            rest_url_opt: args.mycelium_rest_url,
            key_file_opt: args.mycelium_key_file,
            network_key_file: args.mycelium_network_key_file,
            peers: args.mycelium_peers,
            ..Default::default()
        };
        if let Some(net_name) = args.mycelium_network_name {
            spawn_config.network_name = net_name;
        }
        MyceliumSetup::bootstrap(spawn_config).await?
    };

    let mycelium_rest_url = mycelium_setup.rest_url.clone();

    info!(
        bind_addr = args.bind_addr.to_string(),
        mycelium_rest_url = mycelium_rest_url.to_string(),
        "Starting Automerge replication service"
    );

    let seeds: Vec<String> = args.seeds.unwrap_or_default();
    info!(seed_count = seeds.len(), "Loaded bootstrap configuration");

    let (sync_tx, sync_rx) = mpsc::unbounded_channel();
    let (change_tx, _) = watch::channel(0_u64);
    let (change_events, _) = broadcast::channel(256);

    let initial_epoch = match args.debug_epoch {
        Some(generation) => Epoch {
            generation,
            id: [0xee; 32],
        },
        None => Epoch::GENESIS,
    };
    let (epoch_tx, _) = watch::channel(initial_epoch);
    info!(epoch = %initial_epoch, "Local epoch");

    let firewall = Arc::new(Firewall::default());

    let mycelium_transport = Arc::new(mycelium::MyceliumTransport::connect(mycelium_setup).await?);
    info!(local_id = mycelium_transport.local_id(), "Transport ready");

    let state = AppState {
        doc: Arc::new(RwLock::new(Replica::new(initial_epoch, AutoCommit::new()))),
        epoch_tx,
        lifecycle: Arc::new(RwLock::new(NodeLifecycle::Joining {
            discovered_cluster_id: None,
        })),
        peer_states: Arc::new(RwLock::new(HashMap::new())),
        transport: Arc::new(FirewallTransport::new(
            mycelium_transport,
            Arc::clone(&firewall),
        )),
        firewall,
        sync_events: sync_tx,
        change_tx,
        initial_sync_done: Arc::new(AtomicBool::new(false)),
        change_events,
        timing: Timing::from_keepalive(Duration::from_secs(args.keepalive_secs)),
        compaction: Arc::new(compaction::Compaction::default()),
    };

    info!("Application state initialized");

    let consumer_state = state.clone();
    tokio::spawn(async move {
        let inner_state = consumer_state.clone();
        let _ = consumer_state
            .transport
            .start_consumer(Box::new(move |sender, msg| {
                let state_clone = inner_state.clone();
                Box::pin(async move { peer::handle_message(&state_clone, sender, msg).await })
            }))
            .await;
    });
    info!("Mycelium consumer task spawned");

    let sync_state = state.clone();
    tokio::spawn(async move {
        sync::run_sync_worker(sync_state, sync_rx).await;
    });
    info!("Sync worker task spawned");

    let peer_worker_state = state.clone();
    tokio::spawn(async move {
        peer::run_peer_worker(peer_worker_state).await;
    });
    info!("Peer worker task spawned");

    // Switches to committed epochs and catches up with newer ones. Starts before
    // bootstrapping: a node joining a compacted cluster can only finish its initial sync
    // after fetching the current epoch from a peer.
    tokio::spawn(compaction::run_worker(state.clone()));
    info!("Compaction worker task spawned");

    bootstrap::bootstrap(&state, &seeds)
        .await
        .context("Joining cluster")?;

    // Start the change publisher after bootstrapping: streams only exist from then on,
    // and it avoids diffing the initial sync.
    let publisher_state = state.clone();
    tokio::spawn(async move { sync::run_change_publisher(publisher_state).await });

    info!(
        cluster_id = ?state.cluster_id().await,
        ready = state.is_ready().await,
        "Bootstrap completed"
    );

    let auth_state = AuthState::new(&args.htpasswd);
    let app = api::router(state.clone(), auth_state, args.debug_api);

    let listener = tokio::net::TcpListener::bind(args.bind_addr)
        .await
        .with_context(|| format!("Binding local API to {}", args.bind_addr))?;

    info!(address = %listener.local_addr()?, "HTTP API listening");

    let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();

    let server_builder = axum::serve(listener, app).with_graceful_shutdown(async move {
        let _ = shutdown_rx.await;
        info!("Shutdown signal received. Draining API connections...")
    });

    let server_fut = server_builder.into_future();
    tokio::pin!(server_fut);

    tokio::select! {
        exit_status = state.transport.wait_for_exit() => {
            error!("Mycelium runtime process exited. Status: {:?}", exit_status);
            let _ = shutdown_tx.send(());
            let _ = server_fut.as_mut().await;
        }

        server_result = server_fut.as_mut() => {
            if let Err(e) = server_result {
                error!("Axum HTTP server encountered a runtime error: {:?}", e);
            } else {
                info!("Axum HTTP engine shutdown cleanly.");
            }
            std::mem::drop(state);
        }

        _ = tokio::signal::ctrl_c() => {
            info!("Received manual interrupt signal. Shutting down...");
            let _ = shutdown_tx.send(());
            let _ = server_fut.as_mut().await;
        }
    }

    info!("HTTP API server stopped");
    Ok(())
}

// === Footnotes ===============================================================
// [concurrent]: these comments are in place so we remember what to do if we go from async to concurrent.