proto-server.rs

use tokio::signal;
use tokio::sync::watch;
use tokio::task;

/// State machine that represents the states the Mina Indexer can be in
#[derive(Debug, Clone)]
enum State {
    Ready(Ready),
    Init(Init),
    Running(Running),
    Shutdown(Shutdown),
}

#[derive(Debug, Clone)]
struct Ready;

#[derive(Debug, Clone)]
struct Init;

#[derive(Debug, Clone)]
struct Running;

#[derive(Debug, Clone)]
struct Shutdown;

impl Ready {
    async fn init(self) -> State {
        println!("Transitioning from Ready to Init");
        State::Init(Init)
    }
}

impl Init {
    async fn run(self) -> State {
        println!("Transitioning from Init to Running");
        State::Running(Running)
    }
}

impl Running {
    async fn shutdown(self) -> State {
        println!("Transitioning from Running to Shutdown");
        State::Shutdown(Shutdown)
    }
}

impl Shutdown {
    async fn complete(self) {
        println!("Shutdown complete");
    }
}

struct MinaIndexer {
    state: State,
    shutdown_trigger: watch::Receiver<()>,
}

impl MinaIndexer {
    fn new(shutdown_trigger: watch::Receiver<()>) -> Self {
        MinaIndexer {
            state: State::Ready(Ready),
            shutdown_trigger,
        }
    }

    async fn run(&mut self) {
        loop {
            match &mut self.state {
                State::Ready(state) => {
                    self.state = state.clone().init().await;
                }
                State::Init(state) => {
                    self.state = state.clone().run().await;
                }
                State::Running(state) => {
                    tokio::select! {
                        _ = self.shutdown_trigger.changed() => {
                            self.state = state.clone().shutdown().await;
                        }
                    }
                }
                State::Shutdown(state) => {
                    state.clone().complete().await;
                    break;
                }
            }
        }
    }
}

struct ShutdownMonitor {
    shutdown_trigger: watch::Sender<()>,
}

impl ShutdownMonitor {
    fn new() -> (Self, watch::Receiver<()>) {
        let (shutdown_trigger, shutdown_signal) = watch::channel(());
        (Self { shutdown_trigger }, shutdown_signal)
    }

    async fn monitor_shutdown(self) {
        signal::ctrl_c().await.expect("Failed to listen for ctrl_c signal");
        self.shutdown_trigger.send(()).expect("Failed to send shutdown signal");
    }
}

#[tokio::main]
async fn main() {
    let (shutdown_monitor, shutdown_signal) = ShutdownMonitor::new();
    let mut mina_indexer = MinaIndexer::new(shutdown_signal);

    // Spawn a task to run the state machine
    let sm_handle = task::spawn(async move {
        mina_indexer.run().await;
    });

    // Spawn a task to monitor shutdown signal
    let monitor_handle = task::spawn(async move {
        shutdown_monitor.monitor_shutdown().await;
    });

    // Wait for both tasks to complete
    let _ = tokio::join!(sm_handle, monitor_handle);

    println!("Cleaning up rocksdb")
}