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