iota_rest_kv/
main.rs

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

iota_rest_kv/
main.rs
