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 : }
|
