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

Generated by: LCOV version 2.1-beta