rostrum 14.0.1

extern crate rostrum;

extern crate anyhow;
#[macro_use]
extern crate log;

use anyhow::Result;
use std::{collections::HashSet, sync::Arc};

use rostrum::{
    app::App,
    cache::TransactionCache,
    chain::Chain,
    chaindef::Block,
    config::Config,
    daemon::Daemon,
    discovery::{announcer::PeerAnnouncer, discoverer::PeerDiscoverer},
    doslimit::{ConnectionLimits, GlobalLimits},
    index::Index,
    metrics::Metrics,
    query::Query,
    rpc::daemon::server::Server,
    signal::SignalNotification,
    signal::{NetworkNotifier, Waiter},
    store::{is_compatible_version, DBContents, DBStore},
};

/// Ensure index get flushed to disk on exit.
/// In case a thread holds ownership of Arc<DBStore>, the thread
/// may not release its reference count in time for drop to get called.
struct FlushRAII {
    store: Arc<DBStore>,
}

impl Drop for FlushRAII {
    fn drop(&mut self) {
        info!("Flushing indexes to disk...");
        if let Err(e) = self.store.flush() {
            warn!("Flushing indexes failed: {e}");
            return;
        }
        info!("Flushing indexes done.")
    }
}

async fn run_server(config: Arc<Config>) -> Result<()> {
    let network_notifier = Arc::new(NetworkNotifier::new());
    let signal = Waiter::start(network_notifier.clone()).await;
    let metrics = Arc::new(
        Metrics::new(config.monitoring_addr, config.metrics).expect("Failed to start metrics"),
    );

    let daemon = Arc::new(Daemon::new(&config, network_notifier.clone(), &metrics).await?);

    // Perform initial indexing.
    if config.reindex {
        info!("Configuration flag `reindex` enabled. Running a full reindex.");
        DBStore::destroy(&config.db_path);
    }

    let compatible = is_compatible_version(&config.db_path, &metrics);

    if !compatible {
        info!("Incompatible database. Running full reindex.");
        DBStore::destroy(&config.db_path);
    }
    let store = Arc::new(
        DBStore::open(
            DBContents::ConfirmedIndex,
            &config.db_path,
            &metrics,
            !compatible || config.reindex, // force new flag if not compatible
        )
        .unwrap(),
    );
    let _flush_raii = FlushRAII {
        store: Arc::clone(&store),
    };
    let genesis: Block = daemon
        .get_genesis()
        .await
        .expect("Failed to get genesis block");
    let genesis_hash = genesis.block_hash();
    let index = Index::load(
        store,
        Chain::new(genesis),
        daemon.clone(),
        &metrics,
        &config,
    )
    .await?;

    let discoverer = PeerDiscoverer::new(&config, &metrics).await;
    let announcer = Arc::new(PeerAnnouncer::new(genesis_hash, &config).await);

    let app = App::new(
        index,
        daemon.clone(),
        &config,
        discoverer.clone(),
        announcer.clone(),
    )?;
    let mempool = Arc::new({
        let mut mempool_path = config.db_path.clone();
        mempool_path.push("mempool");
        rostrum::mempool::Tracker::new(&mempool_path, &metrics)
    });

    let tx_cache = TransactionCache::new(config.tx_cache_size as u64, &metrics);
    let query = Query::new(
        app.clone(),
        &metrics,
        tx_cache,
        mempool.clone(),
        config.network_type,
    )?;
    let connection_limits = ConnectionLimits::new(
        config.rpc_timeout,
        config.scripthash_subscription_limit,
        config.scripthash_alias_bytes_limit,
    );
    let global_limits = Arc::new(GlobalLimits::new(
        config.rpc_max_connections,
        config.rpc_max_connections_shared_prefix,
        &metrics,
    ));

    app.first_update().await?;

    // Electrum RPC server
    // We start electrum server AFTER first update, because many applications
    // just assume that the server index is up-to-date
    let server = Server::start(
        config.clone(),
        query.clone(),
        metrics.clone(),
        connection_limits,
        global_limits,
        network_notifier,
    )
    .await;

    // Now that our RPC is up; we can start announcing ourselves to other peers
    announcer
        .start_announce_thread(discoverer, &config, &metrics)
        .await;

    loop {
        match signal.wait(config.wait_duration).await {
            Ok(sig) => {
                match sig {
                    SignalNotification::NewTxs(tx_invs) => {
                        let txs_changed = mempool.update_from_txids(tx_invs, query.tx()).await;
                        server
                            .notify_scripthash_subscriptions(&[], txs_changed)
                            .await;
                        continue;
                    }
                    SignalNotification::NewBlock(blockhash) => {
                        trace!("NewBlock notification");

                        if app.index().chain().contains(&blockhash).await {
                            // we have already seen this block
                            continue;
                        }

                        // Block not seen. Proceed to pull full update.
                    }
                    SignalNotification::Timeout => {
                        daemon.p2p_keepalive().await;
                        // Proceed to pull full update
                    }
                    SignalNotification::NewFullTx(tx) => {
                        let txid = tx.txid();
                        match mempool.update_from_tx(tx, query.tx()).await {
                            Ok(_) => {
                                server
                                    .notify_scripthash_subscriptions(&[], HashSet::from([txid]))
                                    .await;
                            }
                            Err(e) => {
                                trace!("Update from tx broadcast failed: {}", e);
                            }
                        }
                        continue;
                    }
                }
            }
            Err(err) => {
                info!("Stopping server: {}", err);
                break;
            }
        }

        // New block, or full pull (via timeout). Do a full update.
        let txs_changed = query.update_mempool().await?;
        let (headers_changed, new_tip) = app.update_blocks().await?;

        server
            .notify_scripthash_subscriptions(&headers_changed, txs_changed)
            .await;
        if let Some(header) = new_tip {
            server.notify_subscriptions_chaintip(header).await;
        }
    }
    Ok(())
}

#[tokio::main]
async fn main() -> Result<()> {
    let config = Arc::new(Config::from_args());

    // initialize ssl
    let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();

    if let Err(e) = run_server(config).await {
        info!(
            "Server stopped: {}",
            e.chain().map(|x| x.to_string()).collect::<String>()
        );
        if e.to_string().contains("Shutdown triggered")
            || e.to_string().contains("Interrupted by signal")
        {
            Ok(())
        } else {
            Err(e)
        }
    } else {
        Ok(())
    }
}