riftenlabs-indexer/src/signal.rs

// Copyright (C) 2024-2026 Whiterun LLC
//
// This software is licensed under the GNU Affero General Public License (AGPL), version 3.0 or later.
// A copy of the license can be found in the LICENSE file or at https://www.gnu.org/licenses/agpl-3.0.html

use rocket::fairing::{Fairing, Info, Kind};
use rocket::http::{ContentType, Status};
use rocket::outcome::Outcome;
use rocket::{Orbit, Request, Response, Rocket, State};
use std::io::Cursor;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;

use crate::rpc::err::ApiErrorCode;
use crate::IbdState;

static SHUTDOWN_FLAG: AtomicBool = AtomicBool::new(false);

/// Returns true if a shutdown has been requested.
pub fn shutdown_requested() -> bool {
    SHUTDOWN_FLAG.load(Ordering::Relaxed)
}

/// Requests a graceful shutdown of all background threads.
pub fn request_shutdown() {
    SHUTDOWN_FLAG.store(true, Ordering::SeqCst);
}

/// Rocket fairing that calls `request_shutdown()` when Rocket shuts down.
pub struct ShutdownFairing;

#[rocket::async_trait]
impl Fairing for ShutdownFairing {
    fn info(&self) -> Info {
        Info {
            name: "Shutdown Fairing",
            kind: Kind::Shutdown,
        }
    }

    async fn on_shutdown(&self, _rocket: &Rocket<Orbit>) {
        log::info!("Shutdown signal received, requesting graceful shutdown...");
        request_shutdown();
    }
}

/// Rocket fairing that returns 503 Service Unavailable during initial block download.
pub struct IbdCheckFairing;

#[rocket::async_trait]
impl Fairing for IbdCheckFairing {
    fn info(&self) -> Info {
        Info {
            name: "IBD Check Fairing",
            kind: Kind::Response,
        }
    }

    async fn on_response<'r>(&self, request: &'r Request<'_>, response: &mut Response<'r>) {
        // Get IbdState from managed state
        let ibd_state = request.guard::<&State<Arc<IbdState>>>().await;

        if let Outcome::Success(state) = ibd_state {
            // If serve_during_ibd is true, allow all requests
            if state.serve_during_ibd {
                return;
            }

            // If initial sync is complete, allow all requests
            if state.initial_sync_complete.load(Ordering::Relaxed) {
                return;
            }

            // Initial sync is in progress and we should block requests
            let current = state.current_height.load(Ordering::Relaxed);
            let target = state.target_height.load(Ordering::Relaxed);

            let message = if target > 0 {
                format!("Initial sync in progress (block {} of {})", current, target)
            } else {
                "Initial sync in progress".to_string()
            };

            let body = serde_json::json!({
                "error": {
                    "code": ApiErrorCode::IndexingInProgress.to_string(),
                    "message": message
                }
            });

            response.set_status(Status::ServiceUnavailable);
            response.set_header(ContentType::JSON);
            response.set_sized_body(None, Cursor::new(body.to_string()));
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rocket::http::Status;
    use rocket::local::blocking::Client;
    use rocket::{get, routes};
    use serde_json::Value;
    use std::sync::atomic::AtomicU64;

    #[get("/test")]
    fn test_endpoint() -> &'static str {
        "ok"
    }

    #[test]
    fn test_ibd_in_progress_returns_503() {
        let ibd_state = Arc::new(IbdState {
            initial_sync_complete: AtomicBool::new(false),
            current_height: AtomicU64::new(850000),
            target_height: AtomicU64::new(880000),
            serve_during_ibd: false,
        });

        let rocket = rocket::build()
            .attach(IbdCheckFairing)
            .manage(ibd_state)
            .mount("/", routes![test_endpoint]);

        let client = Client::tracked(rocket).expect("valid rocket instance");
        let response = client.get("/test").dispatch();

        assert_eq!(response.status(), Status::ServiceUnavailable);

        let body: Value = serde_json::from_str(&response.into_string().unwrap()).unwrap();
        assert_eq!(body["error"]["code"], "INDEXING_IN_PROGRESS");
        assert!(body["error"]["message"]
            .as_str()
            .unwrap()
            .contains("850000 of 880000"));
    }

    #[test]
    fn test_ibd_complete_returns_normal_response() {
        let ibd_state = Arc::new(IbdState {
            initial_sync_complete: AtomicBool::new(true),
            current_height: AtomicU64::new(880000),
            target_height: AtomicU64::new(880000),
            serve_during_ibd: false,
        });

        let rocket = rocket::build()
            .attach(IbdCheckFairing)
            .manage(ibd_state)
            .mount("/", routes![test_endpoint]);

        let client = Client::tracked(rocket).expect("valid rocket instance");
        let response = client.get("/test").dispatch();

        assert_eq!(response.status(), Status::Ok);
        assert_eq!(response.into_string().unwrap(), "ok");
    }
}