LCOV - code coverage report
Current view: top level - storage_controller/src - main.rs (source / functions) Coverage Total Hit
Test: 98683a8629f0f7f0031d02e04512998d589d76ea.info Lines: 0.0 % 266 0
Test Date: 2025-04-11 16:58:57 Functions: 0.0 % 39 0

            Line data    Source code
       1              : use std::num::NonZeroU32;
       2              : use std::path::PathBuf;
       3              : use std::sync::Arc;
       4              : use std::time::Duration;
       5              : 
       6              : use anyhow::{Context, anyhow};
       7              : use camino::Utf8PathBuf;
       8              : use clap::Parser;
       9              : use futures::future::OptionFuture;
      10              : use http_utils::tls_certs::ReloadingCertificateResolver;
      11              : use hyper0::Uri;
      12              : use metrics::BuildInfo;
      13              : use metrics::launch_timestamp::LaunchTimestamp;
      14              : use reqwest::Certificate;
      15              : use storage_controller::http::make_router;
      16              : use storage_controller::metrics::preinitialize_metrics;
      17              : use storage_controller::persistence::Persistence;
      18              : use storage_controller::service::chaos_injector::ChaosInjector;
      19              : use storage_controller::service::{
      20              :     Config, HEARTBEAT_INTERVAL_DEFAULT, LONG_RECONCILE_THRESHOLD_DEFAULT,
      21              :     MAX_OFFLINE_INTERVAL_DEFAULT, MAX_WARMING_UP_INTERVAL_DEFAULT,
      22              :     PRIORITY_RECONCILER_CONCURRENCY_DEFAULT, RECONCILER_CONCURRENCY_DEFAULT, Service,
      23              : };
      24              : use tokio::signal::unix::SignalKind;
      25              : use tokio_util::sync::CancellationToken;
      26              : use tracing::Instrument;
      27              : use utils::auth::{JwtAuth, SwappableJwtAuth};
      28              : use utils::logging::{self, LogFormat};
      29              : use utils::sentry_init::init_sentry;
      30              : use utils::{project_build_tag, project_git_version, tcp_listener};
      31              : 
      32              : project_git_version!(GIT_VERSION);
      33              : project_build_tag!(BUILD_TAG);
      34              : 
      35              : #[global_allocator]
      36              : static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
      37              : 
      38              : /// Configure jemalloc to profile heap allocations by sampling stack traces every 2 MB (1 << 21).
      39              : /// This adds roughly 3% overhead for allocations on average, which is acceptable considering
      40              : /// performance-sensitive code will avoid allocations as far as possible anyway.
      41              : #[allow(non_upper_case_globals)]
      42              : #[unsafe(export_name = "malloc_conf")]
      43              : pub static malloc_conf: &[u8] = b"prof:true,prof_active:true,lg_prof_sample:21\0";
      44              : 
      45              : const DEFAULT_SSL_KEY_FILE: &str = "server.key";
      46              : const DEFAULT_SSL_CERT_FILE: &str = "server.crt";
      47              : const DEFAULT_SSL_CERT_RELOAD_PERIOD: &str = "60s";
      48              : 
      49              : #[derive(Parser)]
      50              : #[command(author, version, about, long_about = None)]
      51              : #[command(arg_required_else_help(true))]
      52              : #[clap(group(
      53              :     clap::ArgGroup::new("listen-addresses")
      54              :         .required(true)
      55              :         .multiple(true)
      56              :         .args(&["listen", "listen_https"]),
      57              : ))]
      58              : struct Cli {
      59              :     /// Host and port to listen HTTP on, like `127.0.0.1:1234`.
      60              :     /// At least one of ["listen", "listen_https"] should be specified.
      61              :     // TODO: Make this option dev-only when https is out everywhere.
      62              :     #[arg(short, long)]
      63              :     listen: Option<std::net::SocketAddr>,
      64              :     /// Host and port to listen HTTPS on, like `127.0.0.1:1234`.
      65              :     /// At least one of ["listen", "listen_https"] should be specified.
      66              :     #[arg(long)]
      67              :     listen_https: Option<std::net::SocketAddr>,
      68              : 
      69              :     /// Public key for JWT authentication of clients
      70              :     #[arg(long)]
      71              :     public_key: Option<String>,
      72              : 
      73              :     /// Token for authenticating this service with the pageservers it controls
      74              :     #[arg(long)]
      75              :     jwt_token: Option<String>,
      76              : 
      77              :     /// Token for authenticating this service with the safekeepers it controls
      78              :     #[arg(long)]
      79              :     safekeeper_jwt_token: Option<String>,
      80              : 
      81              :     /// Token for authenticating this service with the control plane, when calling
      82              :     /// the compute notification endpoint
      83              :     #[arg(long)]
      84              :     control_plane_jwt_token: Option<String>,
      85              : 
      86              :     #[arg(long)]
      87              :     peer_jwt_token: Option<String>,
      88              : 
      89              :     /// URL to control plane compute notification endpoint
      90              :     #[arg(long)]
      91              :     compute_hook_url: Option<String>,
      92              : 
      93              :     /// URL to control plane storage API prefix
      94              :     #[arg(long)]
      95              :     control_plane_url: Option<String>,
      96              : 
      97              :     /// URL to connect to postgres, like postgresql://localhost:1234/storage_controller
      98              :     #[arg(long)]
      99              :     database_url: Option<String>,
     100              : 
     101              :     /// Flag to enable dev mode, which permits running without auth
     102              :     #[arg(long, default_value = "false")]
     103            0 :     dev: bool,
     104              : 
     105              :     /// Grace period before marking unresponsive pageserver offline
     106              :     #[arg(long)]
     107              :     max_offline_interval: Option<humantime::Duration>,
     108              : 
     109              :     /// More tolerant grace period before marking unresponsive pagserver offline used
     110              :     /// around pageserver restarts
     111              :     #[arg(long)]
     112              :     max_warming_up_interval: Option<humantime::Duration>,
     113              : 
     114              :     /// Size threshold for automatically splitting shards (disabled by default)
     115              :     #[arg(long)]
     116              :     split_threshold: Option<u64>,
     117              : 
     118              :     /// Maximum number of shards during autosplits. 0 disables autosplits. Defaults
     119              :     /// to 16 as a safety to avoid too many shards by accident.
     120              :     #[arg(long, default_value = "16")]
     121            0 :     max_split_shards: u8,
     122              : 
     123              :     /// Size threshold for initial shard splits of unsharded tenants. 0 disables initial splits.
     124              :     #[arg(long)]
     125              :     initial_split_threshold: Option<u64>,
     126              : 
     127              :     /// Number of target shards for initial splits. 0 or 1 disables initial splits. Defaults to 2.
     128              :     #[arg(long, default_value = "2")]
     129            0 :     initial_split_shards: u8,
     130              : 
     131              :     /// Maximum number of normal-priority reconcilers that may run in parallel
     132              :     #[arg(long)]
     133              :     reconciler_concurrency: Option<usize>,
     134              : 
     135              :     /// Maximum number of high-priority reconcilers that may run in parallel
     136              :     #[arg(long)]
     137              :     priority_reconciler_concurrency: Option<usize>,
     138              : 
     139              :     /// Tenant API rate limit, as requests per second per tenant.
     140              :     #[arg(long, default_value = "10")]
     141            0 :     tenant_rate_limit: NonZeroU32,
     142              : 
     143              :     /// How long to wait for the initial database connection to be available.
     144              :     #[arg(long, default_value = "5s")]
     145            0 :     db_connect_timeout: humantime::Duration,
     146              : 
     147              :     #[arg(long, default_value = "false")]
     148            0 :     start_as_candidate: bool,
     149              : 
     150              :     // TODO: make this mandatory once the helm chart gets updated
     151              :     #[arg(long)]
     152              :     address_for_peers: Option<Uri>,
     153              : 
     154              :     /// `neon_local` sets this to the path of the neon_local repo dir.
     155              :     /// Only relevant for testing.
     156              :     // TODO: make `cfg(feature = "testing")`
     157              :     #[arg(long)]
     158              :     neon_local_repo_dir: Option<PathBuf>,
     159              : 
     160              :     /// Chaos testing: exercise tenant migrations
     161              :     #[arg(long)]
     162              :     chaos_interval: Option<humantime::Duration>,
     163              : 
     164              :     /// Chaos testing: exercise an immediate exit
     165              :     #[arg(long)]
     166              :     chaos_exit_crontab: Option<cron::Schedule>,
     167              : 
     168              :     /// Maximum acceptable lag for the secondary location while draining
     169              :     /// a pageserver
     170              :     #[arg(long)]
     171              :     max_secondary_lag_bytes: Option<u64>,
     172              : 
     173              :     /// Period with which to send heartbeats to registered nodes
     174              :     #[arg(long)]
     175              :     heartbeat_interval: Option<humantime::Duration>,
     176              : 
     177              :     #[arg(long)]
     178              :     long_reconcile_threshold: Option<humantime::Duration>,
     179              : 
     180              :     /// Flag to use https for requests to pageserver API.
     181              :     #[arg(long, default_value = "false")]
     182            0 :     use_https_pageserver_api: bool,
     183              : 
     184              :     // Whether to put timelines onto safekeepers
     185              :     #[arg(long, default_value = "false")]
     186            0 :     timelines_onto_safekeepers: bool,
     187              : 
     188              :     /// Flag to use https for requests to safekeeper API.
     189              :     #[arg(long, default_value = "false")]
     190            0 :     use_https_safekeeper_api: bool,
     191              : 
     192              :     /// Path to a file with certificate's private key for https API.
     193              :     #[arg(long, default_value = DEFAULT_SSL_KEY_FILE)]
     194            0 :     ssl_key_file: Utf8PathBuf,
     195              :     /// Path to a file with a X509 certificate for https API.
     196              :     #[arg(long, default_value = DEFAULT_SSL_CERT_FILE)]
     197            0 :     ssl_cert_file: Utf8PathBuf,
     198              :     /// Period to reload certificate and private key from files.
     199              :     #[arg(long, default_value = DEFAULT_SSL_CERT_RELOAD_PERIOD)]
     200            0 :     ssl_cert_reload_period: humantime::Duration,
     201              :     /// Trusted root CA certificates to use in https APIs.
     202              :     #[arg(long)]
     203              :     ssl_ca_file: Option<PathBuf>,
     204              : 
     205              :     /// Neon local specific flag. When set, ignore [`Cli::control_plane_url`] and deliver
     206              :     /// the compute notification directly (instead of via control plane).
     207              :     #[arg(long, default_value = "false")]
     208            0 :     use_local_compute_notifications: bool,
     209              : }
     210              : 
     211              : enum StrictMode {
     212              :     /// In strict mode, we will require that all secrets are loaded, i.e. security features
     213              :     /// may not be implicitly turned off by omitting secrets in the environment.
     214              :     Strict,
     215              :     /// In dev mode, secrets are optional, and omitting a particular secret will implicitly
     216              :     /// disable the auth related to it (e.g. no pageserver jwt key -> send unauthenticated
     217              :     /// requests, no public key -> don't authenticate incoming requests).
     218              :     Dev,
     219              : }
     220              : 
     221              : impl Default for StrictMode {
     222            0 :     fn default() -> Self {
     223            0 :         Self::Strict
     224            0 :     }
     225              : }
     226              : 
     227              : /// Secrets may either be provided on the command line (for testing), or loaded from AWS SecretManager: this
     228              : /// type encapsulates the logic to decide which and do the loading.
     229              : struct Secrets {
     230              :     database_url: String,
     231              :     public_key: Option<JwtAuth>,
     232              :     pageserver_jwt_token: Option<String>,
     233              :     safekeeper_jwt_token: Option<String>,
     234              :     control_plane_jwt_token: Option<String>,
     235              :     peer_jwt_token: Option<String>,
     236              : }
     237              : 
     238              : impl Secrets {
     239              :     const DATABASE_URL_ENV: &'static str = "DATABASE_URL";
     240              :     const PAGESERVER_JWT_TOKEN_ENV: &'static str = "PAGESERVER_JWT_TOKEN";
     241              :     const SAFEKEEPER_JWT_TOKEN_ENV: &'static str = "SAFEKEEPER_JWT_TOKEN";
     242              :     const CONTROL_PLANE_JWT_TOKEN_ENV: &'static str = "CONTROL_PLANE_JWT_TOKEN";
     243              :     const PEER_JWT_TOKEN_ENV: &'static str = "PEER_JWT_TOKEN";
     244              :     const PUBLIC_KEY_ENV: &'static str = "PUBLIC_KEY";
     245              : 
     246              :     /// Load secrets from, in order of preference:
     247              :     /// - CLI args if database URL is provided on the CLI
     248              :     /// - Environment variables if DATABASE_URL is set.
     249            0 :     async fn load(args: &Cli) -> anyhow::Result<Self> {
     250            0 :         let Some(database_url) = Self::load_secret(&args.database_url, Self::DATABASE_URL_ENV)
     251              :         else {
     252            0 :             anyhow::bail!(
     253            0 :                 "Database URL is not set (set `--database-url`, or `DATABASE_URL` environment)"
     254            0 :             )
     255              :         };
     256              : 
     257            0 :         let public_key = match Self::load_secret(&args.public_key, Self::PUBLIC_KEY_ENV) {
     258            0 :             Some(v) => Some(JwtAuth::from_key(v).context("Loading public key")?),
     259            0 :             None => None,
     260              :         };
     261              : 
     262            0 :         let this = Self {
     263            0 :             database_url,
     264            0 :             public_key,
     265            0 :             pageserver_jwt_token: Self::load_secret(
     266            0 :                 &args.jwt_token,
     267            0 :                 Self::PAGESERVER_JWT_TOKEN_ENV,
     268            0 :             ),
     269            0 :             safekeeper_jwt_token: Self::load_secret(
     270            0 :                 &args.safekeeper_jwt_token,
     271            0 :                 Self::SAFEKEEPER_JWT_TOKEN_ENV,
     272            0 :             ),
     273            0 :             control_plane_jwt_token: Self::load_secret(
     274            0 :                 &args.control_plane_jwt_token,
     275            0 :                 Self::CONTROL_PLANE_JWT_TOKEN_ENV,
     276            0 :             ),
     277            0 :             peer_jwt_token: Self::load_secret(&args.peer_jwt_token, Self::PEER_JWT_TOKEN_ENV),
     278            0 :         };
     279            0 : 
     280            0 :         Ok(this)
     281            0 :     }
     282              : 
     283            0 :     fn load_secret(cli: &Option<String>, env_name: &str) -> Option<String> {
     284            0 :         if let Some(v) = cli {
     285            0 :             Some(v.clone())
     286              :         } else {
     287            0 :             std::env::var(env_name).ok()
     288              :         }
     289            0 :     }
     290              : }
     291              : 
     292            0 : fn main() -> anyhow::Result<()> {
     293            0 :     logging::init(
     294            0 :         LogFormat::Plain,
     295            0 :         logging::TracingErrorLayerEnablement::Disabled,
     296            0 :         logging::Output::Stdout,
     297            0 :     )?;
     298              : 
     299              :     // log using tracing so we don't get confused output by default hook writing to stderr
     300            0 :     utils::logging::replace_panic_hook_with_tracing_panic_hook().forget();
     301            0 : 
     302            0 :     let _sentry_guard = init_sentry(Some(GIT_VERSION.into()), &[]);
     303            0 : 
     304            0 :     let hook = std::panic::take_hook();
     305            0 :     std::panic::set_hook(Box::new(move |info| {
     306            0 :         // let sentry send a message (and flush)
     307            0 :         // and trace the error
     308            0 :         hook(info);
     309            0 : 
     310            0 :         std::process::exit(1);
     311            0 :     }));
     312            0 : 
     313            0 :     tokio::runtime::Builder::new_current_thread()
     314            0 :         // We use spawn_blocking for database operations, so require approximately
     315            0 :         // as many blocking threads as we will open database connections.
     316            0 :         .max_blocking_threads(Persistence::MAX_CONNECTIONS as usize)
     317            0 :         .enable_all()
     318            0 :         .build()
     319            0 :         .unwrap()
     320            0 :         .block_on(async_main())
     321            0 : }
     322              : 
     323            0 : async fn async_main() -> anyhow::Result<()> {
     324            0 :     let launch_ts = Box::leak(Box::new(LaunchTimestamp::generate()));
     325            0 : 
     326            0 :     preinitialize_metrics();
     327            0 : 
     328            0 :     let args = Cli::parse();
     329            0 :     tracing::info!(
     330            0 :         "version: {}, launch_timestamp: {}, build_tag {}",
     331            0 :         GIT_VERSION,
     332            0 :         launch_ts.to_string(),
     333              :         BUILD_TAG,
     334              :     );
     335              : 
     336            0 :     let build_info = BuildInfo {
     337            0 :         revision: GIT_VERSION,
     338            0 :         build_tag: BUILD_TAG,
     339            0 :     };
     340              : 
     341            0 :     let strict_mode = if args.dev {
     342            0 :         StrictMode::Dev
     343              :     } else {
     344            0 :         StrictMode::Strict
     345              :     };
     346              : 
     347            0 :     let secrets = Secrets::load(&args).await?;
     348              : 
     349              :     // Validate required secrets and arguments are provided in strict mode
     350            0 :     match strict_mode {
     351              :         StrictMode::Strict
     352            0 :             if (secrets.public_key.is_none()
     353            0 :                 || secrets.pageserver_jwt_token.is_none()
     354            0 :                 || secrets.control_plane_jwt_token.is_none()
     355            0 :                 || secrets.safekeeper_jwt_token.is_none()) =>
     356            0 :         {
     357            0 :             // Production systems should always have secrets configured: if public_key was not set
     358            0 :             // then we would implicitly disable auth.
     359            0 :             anyhow::bail!(
     360            0 :                 "Insecure config!  One or more secrets is not set.  This is only permitted in `--dev` mode"
     361            0 :             );
     362              :         }
     363              :         StrictMode::Strict
     364            0 :             if args.compute_hook_url.is_none() && args.control_plane_url.is_none() =>
     365            0 :         {
     366            0 :             // Production systems should always have a control plane URL set, to prevent falling
     367            0 :             // back to trying to use neon_local.
     368            0 :             anyhow::bail!(
     369            0 :                 "neither `--compute-hook-url` nor `--control-plane-url` are set: this is only permitted in `--dev` mode"
     370            0 :             );
     371              :         }
     372            0 :         StrictMode::Strict if args.use_local_compute_notifications => {
     373            0 :             anyhow::bail!("`--use-local-compute-notifications` is only permitted in `--dev` mode");
     374              :         }
     375              :         StrictMode::Strict => {
     376            0 :             tracing::info!("Starting in strict mode: configuration is OK.")
     377              :         }
     378              :         StrictMode::Dev => {
     379            0 :             tracing::warn!("Starting in dev mode: this may be an insecure configuration.")
     380              :         }
     381              :     }
     382              : 
     383            0 :     let ssl_ca_certs = match args.ssl_ca_file.as_ref() {
     384            0 :         Some(ssl_ca_file) => {
     385            0 :             tracing::info!("Using ssl root CA file: {ssl_ca_file:?}");
     386            0 :             let buf = tokio::fs::read(ssl_ca_file).await?;
     387            0 :             Certificate::from_pem_bundle(&buf)?
     388              :         }
     389            0 :         None => Vec::new(),
     390              :     };
     391              : 
     392            0 :     let config = Config {
     393            0 :         pageserver_jwt_token: secrets.pageserver_jwt_token,
     394            0 :         safekeeper_jwt_token: secrets.safekeeper_jwt_token,
     395            0 :         control_plane_jwt_token: secrets.control_plane_jwt_token,
     396            0 :         peer_jwt_token: secrets.peer_jwt_token,
     397            0 :         compute_hook_url: args.compute_hook_url,
     398            0 :         control_plane_url: args.control_plane_url,
     399            0 :         max_offline_interval: args
     400            0 :             .max_offline_interval
     401            0 :             .map(humantime::Duration::into)
     402            0 :             .unwrap_or(MAX_OFFLINE_INTERVAL_DEFAULT),
     403            0 :         max_warming_up_interval: args
     404            0 :             .max_warming_up_interval
     405            0 :             .map(humantime::Duration::into)
     406            0 :             .unwrap_or(MAX_WARMING_UP_INTERVAL_DEFAULT),
     407            0 :         reconciler_concurrency: args
     408            0 :             .reconciler_concurrency
     409            0 :             .unwrap_or(RECONCILER_CONCURRENCY_DEFAULT),
     410            0 :         priority_reconciler_concurrency: args
     411            0 :             .priority_reconciler_concurrency
     412            0 :             .unwrap_or(PRIORITY_RECONCILER_CONCURRENCY_DEFAULT),
     413            0 :         tenant_rate_limit: args.tenant_rate_limit,
     414            0 :         split_threshold: args.split_threshold,
     415            0 :         max_split_shards: args.max_split_shards,
     416            0 :         initial_split_threshold: args.initial_split_threshold,
     417            0 :         initial_split_shards: args.initial_split_shards,
     418            0 :         neon_local_repo_dir: args.neon_local_repo_dir,
     419            0 :         max_secondary_lag_bytes: args.max_secondary_lag_bytes,
     420            0 :         heartbeat_interval: args
     421            0 :             .heartbeat_interval
     422            0 :             .map(humantime::Duration::into)
     423            0 :             .unwrap_or(HEARTBEAT_INTERVAL_DEFAULT),
     424            0 :         long_reconcile_threshold: args
     425            0 :             .long_reconcile_threshold
     426            0 :             .map(humantime::Duration::into)
     427            0 :             .unwrap_or(LONG_RECONCILE_THRESHOLD_DEFAULT),
     428            0 :         address_for_peers: args.address_for_peers,
     429            0 :         start_as_candidate: args.start_as_candidate,
     430            0 :         use_https_pageserver_api: args.use_https_pageserver_api,
     431            0 :         use_https_safekeeper_api: args.use_https_safekeeper_api,
     432            0 :         ssl_ca_certs,
     433            0 :         timelines_onto_safekeepers: args.timelines_onto_safekeepers,
     434            0 :         use_local_compute_notifications: args.use_local_compute_notifications,
     435            0 :     };
     436            0 : 
     437            0 :     // Validate that we can connect to the database
     438            0 :     Persistence::await_connection(&secrets.database_url, args.db_connect_timeout.into()).await?;
     439              : 
     440            0 :     let persistence = Arc::new(Persistence::new(secrets.database_url).await);
     441              : 
     442            0 :     let service = Service::spawn(config, persistence.clone()).await?;
     443              : 
     444            0 :     let auth = secrets
     445            0 :         .public_key
     446            0 :         .map(|jwt_auth| Arc::new(SwappableJwtAuth::new(jwt_auth)));
     447            0 :     let router = make_router(service.clone(), auth, build_info)
     448            0 :         .build()
     449            0 :         .map_err(|err| anyhow!(err))?;
     450            0 :     let http_service =
     451            0 :         Arc::new(http_utils::RequestServiceBuilder::new(router).map_err(|err| anyhow!(err))?);
     452              : 
     453            0 :     let api_shutdown = CancellationToken::new();
     454              : 
     455              :     // Start HTTP server
     456            0 :     let http_server_task: OptionFuture<_> = match args.listen {
     457            0 :         Some(http_addr) => {
     458            0 :             let http_listener = tcp_listener::bind(http_addr)?;
     459            0 :             let http_server =
     460            0 :                 http_utils::server::Server::new(Arc::clone(&http_service), http_listener, None)?;
     461              : 
     462            0 :             tracing::info!("Serving HTTP on {}", http_addr);
     463            0 :             Some(tokio::task::spawn(http_server.serve(api_shutdown.clone())))
     464              :         }
     465            0 :         None => None,
     466              :     }
     467            0 :     .into();
     468              : 
     469              :     // Start HTTPS server
     470            0 :     let https_server_task: OptionFuture<_> = match args.listen_https {
     471            0 :         Some(https_addr) => {
     472            0 :             let https_listener = tcp_listener::bind(https_addr)?;
     473              : 
     474            0 :             let resolver = ReloadingCertificateResolver::new(
     475            0 :                 "main",
     476            0 :                 &args.ssl_key_file,
     477            0 :                 &args.ssl_cert_file,
     478            0 :                 *args.ssl_cert_reload_period,
     479            0 :             )
     480            0 :             .await?;
     481              : 
     482            0 :             let server_config = rustls::ServerConfig::builder()
     483            0 :                 .with_no_client_auth()
     484            0 :                 .with_cert_resolver(resolver);
     485            0 : 
     486            0 :             let tls_acceptor = tokio_rustls::TlsAcceptor::from(Arc::new(server_config));
     487            0 :             let https_server =
     488            0 :                 http_utils::server::Server::new(http_service, https_listener, Some(tls_acceptor))?;
     489              : 
     490            0 :             tracing::info!("Serving HTTPS on {}", https_addr);
     491            0 :             Some(tokio::task::spawn(https_server.serve(api_shutdown.clone())))
     492              :         }
     493            0 :         None => None,
     494              :     }
     495            0 :     .into();
     496            0 : 
     497            0 :     let chaos_task = args.chaos_interval.map(|interval| {
     498            0 :         let service = service.clone();
     499            0 :         let cancel = CancellationToken::new();
     500            0 :         let cancel_bg = cancel.clone();
     501            0 :         let chaos_exit_crontab = args.chaos_exit_crontab;
     502            0 :         (
     503            0 :             tokio::task::spawn(
     504            0 :                 async move {
     505            0 :                     let mut chaos_injector =
     506            0 :                         ChaosInjector::new(service, interval.into(), chaos_exit_crontab);
     507            0 :                     chaos_injector.run(cancel_bg).await
     508            0 :                 }
     509            0 :                 .instrument(tracing::info_span!("chaos_injector")),
     510              :             ),
     511            0 :             cancel,
     512            0 :         )
     513            0 :     });
     514              : 
     515              :     // Wait until we receive a signal
     516            0 :     let mut sigint = tokio::signal::unix::signal(SignalKind::interrupt())?;
     517            0 :     let mut sigquit = tokio::signal::unix::signal(SignalKind::quit())?;
     518            0 :     let mut sigterm = tokio::signal::unix::signal(SignalKind::terminate())?;
     519            0 :     tokio::pin!(http_server_task, https_server_task);
     520            0 :     tokio::select! {
     521            0 :         _ = sigint.recv() => {},
     522            0 :         _ = sigterm.recv() => {},
     523            0 :         _ = sigquit.recv() => {},
     524            0 :         Some(err) = &mut http_server_task => {
     525            0 :             panic!("HTTP server task failed: {err:#?}");
     526              :         }
     527            0 :         Some(err) = &mut https_server_task => {
     528            0 :             panic!("HTTPS server task failed: {err:#?}");
     529              :         }
     530              :     }
     531            0 :     tracing::info!("Terminating on signal");
     532              : 
     533              :     // Stop HTTP and HTTPS servers first, so that we don't have to service requests
     534              :     // while shutting down Service.
     535            0 :     api_shutdown.cancel();
     536            0 : 
     537            0 :     // If the deadline is exceeded, we will fall through and shut down the service anyway,
     538            0 :     // any request handlers in flight will experience cancellation & their clients will
     539            0 :     // see a torn connection.
     540            0 :     let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
     541            0 : 
     542            0 :     match tokio::time::timeout_at(deadline, http_server_task).await {
     543            0 :         Ok(Some(Ok(_))) => tracing::info!("Joined HTTP server task"),
     544            0 :         Ok(Some(Err(e))) => tracing::error!("Error joining HTTP server task: {e}"),
     545            0 :         Ok(None) => {} // HTTP is disabled.
     546            0 :         Err(_) => tracing::warn!("Timed out joining HTTP server task"),
     547              :     }
     548              : 
     549            0 :     match tokio::time::timeout_at(deadline, https_server_task).await {
     550            0 :         Ok(Some(Ok(_))) => tracing::info!("Joined HTTPS server task"),
     551            0 :         Ok(Some(Err(e))) => tracing::error!("Error joining HTTPS server task: {e}"),
     552            0 :         Ok(None) => {} // HTTPS is disabled.
     553            0 :         Err(_) => tracing::warn!("Timed out joining HTTPS server task"),
     554              :     }
     555              : 
     556              :     // If we were injecting chaos, stop that so that we're not calling into Service while it shuts down
     557            0 :     if let Some((chaos_jh, chaos_cancel)) = chaos_task {
     558            0 :         chaos_cancel.cancel();
     559            0 :         chaos_jh.await.ok();
     560            0 :     }
     561              : 
     562            0 :     service.shutdown().await;
     563            0 :     tracing::info!("Service shutdown complete");
     564              : 
     565            0 :     std::process::exit(0);
     566            0 : }
        

Generated by: LCOV version 2.1-beta