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

Generated by: LCOV version 2.1-beta