iota_rest_kv/
main.rs

1// Copyright (c) 2025 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4use std::{fs, path::PathBuf};
5
6use anyhow::Result;
7use clap::Parser;
8use kv_store_client::KvStoreConfig;
9use serde::{Deserialize, Serialize};
10use server::Server;
11use tokio_util::sync::CancellationToken;
12use tracing::Level;
13use tracing_subscriber::FmtSubscriber;
14
15mod errors;
16mod extractors;
17mod kv_store_client;
18mod routes;
19mod server;
20mod types;
21
22/// The main CLI application.
23#[derive(Parser, Clone, Debug)]
24#[clap(
25    name = "KV Store REST API",
26    about = "A HTTP server exposing key-value data of the IOTA network through a REST API."
27)]
28struct Cli {
29    #[clap(long, default_value = "INFO", env = "LOG_LEVEL")]
30    log_level: Level,
31    /// The yaml config file path.
32    #[clap(short, long)]
33    config: PathBuf,
34}
35
36#[derive(Serialize, Deserialize, Clone, Debug)]
37#[serde(rename_all = "kebab-case")]
38pub struct RestApiConfig {
39    #[serde(flatten)]
40    pub kv_store_config: KvStoreConfig,
41    pub server_address: std::net::SocketAddr,
42}
43
44#[tokio::main]
45async fn main() -> Result<()> {
46    let cli = Cli::parse();
47
48    init_tracing(cli.log_level);
49
50    let raw_config = fs::read_to_string(cli.config).expect("failed to read config file");
51    let config = serde_yaml::from_str::<RestApiConfig>(&raw_config)?;
52
53    let token = CancellationToken::new();
54
55    shutdown_signal_listener(token.clone());
56
57    let server = Server::new(config, token).await?;
58    server.serve().await
59}
60
61/// Initialize the tracing with custom subscribers.
62fn init_tracing(log_level: Level) {
63    let subscriber = FmtSubscriber::builder().with_max_level(log_level).finish();
64    tracing::subscriber::set_global_default(subscriber).expect("setting default subscriber failed");
65}
66
67/// Set up a `CTRL+C` & `SIGTERM` handler for graceful shutdown and spawn a
68/// tokio task.
69fn shutdown_signal_listener(token: CancellationToken) {
70    tokio::spawn(async move {
71        #[cfg(unix)]
72        let terminate = async {
73            tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
74                .expect("Cannot listen to SIGTERM signal")
75                .recv()
76                .await;
77        };
78
79        #[cfg(not(unix))]
80        let terminate = std::future::pending::<()>();
81
82        tokio::select! {
83            _ = tokio::signal::ctrl_c() => tracing::info!("CTRL+C signal received, shutting down"),
84            _ = terminate => tracing::info!("SIGTERM signal received, shutting down")
85        };
86
87        token.cancel();
88    });
89}
iota_rest_kv/main.rs

iota_rest_kv/
main.rs
