Line data Source code
1 : use anyhow::{anyhow, Context};
2 : use clap::Parser;
3 : use hyper::Uri;
4 : use metrics::launch_timestamp::LaunchTimestamp;
5 : use metrics::BuildInfo;
6 : use std::path::PathBuf;
7 : use std::sync::Arc;
8 : use std::time::Duration;
9 : use storage_controller::http::make_router;
10 : use storage_controller::metrics::preinitialize_metrics;
11 : use storage_controller::persistence::Persistence;
12 : use storage_controller::service::chaos_injector::ChaosInjector;
13 : use storage_controller::service::{
14 : Config, Service, MAX_OFFLINE_INTERVAL_DEFAULT, MAX_WARMING_UP_INTERVAL_DEFAULT,
15 : RECONCILER_CONCURRENCY_DEFAULT,
16 : };
17 : use tokio::signal::unix::SignalKind;
18 : use tokio_util::sync::CancellationToken;
19 : use tracing::Instrument;
20 : use utils::auth::{JwtAuth, SwappableJwtAuth};
21 : use utils::logging::{self, LogFormat};
22 :
23 : use utils::sentry_init::init_sentry;
24 : use utils::{project_build_tag, project_git_version, tcp_listener};
25 :
26 : project_git_version!(GIT_VERSION);
27 : project_build_tag!(BUILD_TAG);
28 :
29 0 : #[derive(Parser)]
30 : #[command(author, version, about, long_about = None)]
31 : #[command(arg_required_else_help(true))]
32 : struct Cli {
33 : /// Host and port to listen on, like `127.0.0.1:1234`
34 : #[arg(short, long)]
35 0 : listen: std::net::SocketAddr,
36 :
37 : /// Public key for JWT authentication of clients
38 : #[arg(long)]
39 : public_key: Option<String>,
40 :
41 : /// Token for authenticating this service with the pageservers it controls
42 : #[arg(long)]
43 : jwt_token: Option<String>,
44 :
45 : /// Token for authenticating this service with the control plane, when calling
46 : /// the compute notification endpoint
47 : #[arg(long)]
48 : control_plane_jwt_token: Option<String>,
49 :
50 : #[arg(long)]
51 : peer_jwt_token: Option<String>,
52 :
53 : /// URL to control plane compute notification endpoint
54 : #[arg(long)]
55 : compute_hook_url: Option<String>,
56 :
57 : /// URL to connect to postgres, like postgresql://localhost:1234/storage_controller
58 : #[arg(long)]
59 : database_url: Option<String>,
60 :
61 : /// Flag to enable dev mode, which permits running without auth
62 : #[arg(long, default_value = "false")]
63 0 : dev: bool,
64 :
65 : /// Grace period before marking unresponsive pageserver offline
66 : #[arg(long)]
67 : max_offline_interval: Option<humantime::Duration>,
68 :
69 : /// More tolerant grace period before marking unresponsive pagserver offline used
70 : /// around pageserver restarts
71 : #[arg(long)]
72 : max_warming_up_interval: Option<humantime::Duration>,
73 :
74 : /// Size threshold for automatically splitting shards (disabled by default)
75 : #[arg(long)]
76 : split_threshold: Option<u64>,
77 :
78 : /// Maximum number of reconcilers that may run in parallel
79 : #[arg(long)]
80 : reconciler_concurrency: Option<usize>,
81 :
82 : /// How long to wait for the initial database connection to be available.
83 : #[arg(long, default_value = "5s")]
84 0 : db_connect_timeout: humantime::Duration,
85 :
86 : #[arg(long, default_value = "false")]
87 0 : start_as_candidate: bool,
88 :
89 : // TODO: make this mandatory once the helm chart gets updated
90 : #[arg(long)]
91 : address_for_peers: Option<Uri>,
92 :
93 : /// `neon_local` sets this to the path of the neon_local repo dir.
94 : /// Only relevant for testing.
95 : // TODO: make `cfg(feature = "testing")`
96 : #[arg(long)]
97 : neon_local_repo_dir: Option<PathBuf>,
98 :
99 : /// Chaos testing
100 : #[arg(long)]
101 : chaos_interval: Option<humantime::Duration>,
102 :
103 : // Maximum acceptable lag for the secondary location while draining
104 : // a pageserver
105 : #[arg(long)]
106 : max_secondary_lag_bytes: Option<u64>,
107 : }
108 :
109 : enum StrictMode {
110 : /// In strict mode, we will require that all secrets are loaded, i.e. security features
111 : /// may not be implicitly turned off by omitting secrets in the environment.
112 : Strict,
113 : /// In dev mode, secrets are optional, and omitting a particular secret will implicitly
114 : /// disable the auth related to it (e.g. no pageserver jwt key -> send unauthenticated
115 : /// requests, no public key -> don't authenticate incoming requests).
116 : Dev,
117 : }
118 :
119 : impl Default for StrictMode {
120 0 : fn default() -> Self {
121 0 : Self::Strict
122 0 : }
123 : }
124 :
125 : /// Secrets may either be provided on the command line (for testing), or loaded from AWS SecretManager: this
126 : /// type encapsulates the logic to decide which and do the loading.
127 : struct Secrets {
128 : database_url: String,
129 : public_key: Option<JwtAuth>,
130 : jwt_token: Option<String>,
131 : control_plane_jwt_token: Option<String>,
132 : peer_jwt_token: Option<String>,
133 : }
134 :
135 : impl Secrets {
136 : const DATABASE_URL_ENV: &'static str = "DATABASE_URL";
137 : const PAGESERVER_JWT_TOKEN_ENV: &'static str = "PAGESERVER_JWT_TOKEN";
138 : const CONTROL_PLANE_JWT_TOKEN_ENV: &'static str = "CONTROL_PLANE_JWT_TOKEN";
139 : const PEER_JWT_TOKEN_ENV: &'static str = "PEER_JWT_TOKEN";
140 : const PUBLIC_KEY_ENV: &'static str = "PUBLIC_KEY";
141 :
142 : /// Load secrets from, in order of preference:
143 : /// - CLI args if database URL is provided on the CLI
144 : /// - Environment variables if DATABASE_URL is set.
145 0 : async fn load(args: &Cli) -> anyhow::Result<Self> {
146 0 : let Some(database_url) = Self::load_secret(&args.database_url, Self::DATABASE_URL_ENV)
147 : else {
148 0 : anyhow::bail!(
149 0 : "Database URL is not set (set `--database-url`, or `DATABASE_URL` environment)"
150 0 : )
151 : };
152 :
153 0 : let public_key = match Self::load_secret(&args.public_key, Self::PUBLIC_KEY_ENV) {
154 0 : Some(v) => Some(JwtAuth::from_key(v).context("Loading public key")?),
155 0 : None => None,
156 : };
157 :
158 0 : let this = Self {
159 0 : database_url,
160 0 : public_key,
161 0 : jwt_token: Self::load_secret(&args.jwt_token, Self::PAGESERVER_JWT_TOKEN_ENV),
162 0 : control_plane_jwt_token: Self::load_secret(
163 0 : &args.control_plane_jwt_token,
164 0 : Self::CONTROL_PLANE_JWT_TOKEN_ENV,
165 0 : ),
166 0 : peer_jwt_token: Self::load_secret(&args.peer_jwt_token, Self::PEER_JWT_TOKEN_ENV),
167 0 : };
168 0 :
169 0 : Ok(this)
170 0 : }
171 :
172 0 : fn load_secret(cli: &Option<String>, env_name: &str) -> Option<String> {
173 0 : if let Some(v) = cli {
174 0 : Some(v.clone())
175 0 : } else if let Ok(v) = std::env::var(env_name) {
176 0 : Some(v)
177 : } else {
178 0 : None
179 : }
180 0 : }
181 : }
182 :
183 0 : fn main() -> anyhow::Result<()> {
184 0 : logging::init(
185 0 : LogFormat::Plain,
186 0 : logging::TracingErrorLayerEnablement::Disabled,
187 0 : logging::Output::Stdout,
188 0 : )?;
189 :
190 : // log using tracing so we don't get confused output by default hook writing to stderr
191 0 : utils::logging::replace_panic_hook_with_tracing_panic_hook().forget();
192 0 :
193 0 : let _sentry_guard = init_sentry(Some(GIT_VERSION.into()), &[]);
194 0 :
195 0 : let hook = std::panic::take_hook();
196 0 : std::panic::set_hook(Box::new(move |info| {
197 0 : // let sentry send a message (and flush)
198 0 : // and trace the error
199 0 : hook(info);
200 0 :
201 0 : std::process::exit(1);
202 0 : }));
203 0 :
204 0 : tokio::runtime::Builder::new_current_thread()
205 0 : // We use spawn_blocking for database operations, so require approximately
206 0 : // as many blocking threads as we will open database connections.
207 0 : .max_blocking_threads(Persistence::MAX_CONNECTIONS as usize)
208 0 : .enable_all()
209 0 : .build()
210 0 : .unwrap()
211 0 : .block_on(async_main())
212 0 : }
213 :
214 0 : async fn async_main() -> anyhow::Result<()> {
215 0 : let launch_ts = Box::leak(Box::new(LaunchTimestamp::generate()));
216 0 :
217 0 : preinitialize_metrics();
218 0 :
219 0 : let args = Cli::parse();
220 0 : tracing::info!(
221 0 : "version: {}, launch_timestamp: {}, build_tag {}, listening on {}",
222 0 : GIT_VERSION,
223 0 : launch_ts.to_string(),
224 : BUILD_TAG,
225 : args.listen
226 : );
227 :
228 0 : let build_info = BuildInfo {
229 0 : revision: GIT_VERSION,
230 0 : build_tag: BUILD_TAG,
231 0 : };
232 :
233 0 : let strict_mode = if args.dev {
234 0 : StrictMode::Dev
235 : } else {
236 0 : StrictMode::Strict
237 : };
238 :
239 0 : let secrets = Secrets::load(&args).await?;
240 :
241 : // Validate required secrets and arguments are provided in strict mode
242 0 : match strict_mode {
243 : StrictMode::Strict
244 0 : if (secrets.public_key.is_none()
245 0 : || secrets.jwt_token.is_none()
246 0 : || secrets.control_plane_jwt_token.is_none()) =>
247 0 : {
248 0 : // Production systems should always have secrets configured: if public_key was not set
249 0 : // then we would implicitly disable auth.
250 0 : anyhow::bail!(
251 0 : "Insecure config! One or more secrets is not set. This is only permitted in `--dev` mode"
252 0 : );
253 : }
254 0 : StrictMode::Strict if args.compute_hook_url.is_none() => {
255 0 : // Production systems should always have a compute hook set, to prevent falling
256 0 : // back to trying to use neon_local.
257 0 : anyhow::bail!(
258 0 : "`--compute-hook-url` is not set: this is only permitted in `--dev` mode"
259 0 : );
260 : }
261 : StrictMode::Strict => {
262 0 : tracing::info!("Starting in strict mode: configuration is OK.")
263 : }
264 : StrictMode::Dev => {
265 0 : tracing::warn!("Starting in dev mode: this may be an insecure configuration.")
266 : }
267 : }
268 :
269 0 : let config = Config {
270 0 : jwt_token: secrets.jwt_token,
271 0 : control_plane_jwt_token: secrets.control_plane_jwt_token,
272 0 : peer_jwt_token: secrets.peer_jwt_token,
273 0 : compute_hook_url: args.compute_hook_url,
274 0 : max_offline_interval: args
275 0 : .max_offline_interval
276 0 : .map(humantime::Duration::into)
277 0 : .unwrap_or(MAX_OFFLINE_INTERVAL_DEFAULT),
278 0 : max_warming_up_interval: args
279 0 : .max_warming_up_interval
280 0 : .map(humantime::Duration::into)
281 0 : .unwrap_or(MAX_WARMING_UP_INTERVAL_DEFAULT),
282 0 : reconciler_concurrency: args
283 0 : .reconciler_concurrency
284 0 : .unwrap_or(RECONCILER_CONCURRENCY_DEFAULT),
285 0 : split_threshold: args.split_threshold,
286 0 : neon_local_repo_dir: args.neon_local_repo_dir,
287 0 : max_secondary_lag_bytes: args.max_secondary_lag_bytes,
288 0 : address_for_peers: args.address_for_peers,
289 0 : start_as_candidate: args.start_as_candidate,
290 0 : http_service_port: args.listen.port() as i32,
291 0 : };
292 0 :
293 0 : // Validate that we can connect to the database
294 0 : Persistence::await_connection(&secrets.database_url, args.db_connect_timeout.into()).await?;
295 :
296 0 : let persistence = Arc::new(Persistence::new(secrets.database_url));
297 :
298 0 : let service = Service::spawn(config, persistence.clone()).await?;
299 :
300 0 : let http_listener = tcp_listener::bind(args.listen)?;
301 :
302 0 : let auth = secrets
303 0 : .public_key
304 0 : .map(|jwt_auth| Arc::new(SwappableJwtAuth::new(jwt_auth)));
305 0 : let router = make_router(service.clone(), auth, build_info)
306 0 : .build()
307 0 : .map_err(|err| anyhow!(err))?;
308 0 : let router_service = utils::http::RouterService::new(router).unwrap();
309 0 :
310 0 : // Start HTTP server
311 0 : let server_shutdown = CancellationToken::new();
312 0 : let server = hyper::Server::from_tcp(http_listener)?
313 0 : .serve(router_service)
314 0 : .with_graceful_shutdown({
315 0 : let server_shutdown = server_shutdown.clone();
316 0 : async move {
317 0 : server_shutdown.cancelled().await;
318 0 : }
319 0 : });
320 0 : tracing::info!("Serving on {0}", args.listen);
321 0 : let server_task = tokio::task::spawn(server);
322 0 :
323 0 : let chaos_task = args.chaos_interval.map(|interval| {
324 0 : let service = service.clone();
325 0 : let cancel = CancellationToken::new();
326 0 : let cancel_bg = cancel.clone();
327 0 : (
328 0 : tokio::task::spawn(
329 0 : async move {
330 0 : let mut chaos_injector = ChaosInjector::new(service, interval.into());
331 0 : chaos_injector.run(cancel_bg).await
332 0 : }
333 0 : .instrument(tracing::info_span!("chaos_injector")),
334 : ),
335 0 : cancel,
336 0 : )
337 0 : });
338 :
339 : // Wait until we receive a signal
340 0 : let mut sigint = tokio::signal::unix::signal(SignalKind::interrupt())?;
341 0 : let mut sigquit = tokio::signal::unix::signal(SignalKind::quit())?;
342 0 : let mut sigterm = tokio::signal::unix::signal(SignalKind::terminate())?;
343 : tokio::select! {
344 : _ = sigint.recv() => {},
345 : _ = sigterm.recv() => {},
346 : _ = sigquit.recv() => {},
347 : }
348 0 : tracing::info!("Terminating on signal");
349 :
350 : // Stop HTTP server first, so that we don't have to service requests
351 : // while shutting down Service.
352 0 : server_shutdown.cancel();
353 0 : match tokio::time::timeout(Duration::from_secs(5), server_task).await {
354 : Ok(Ok(_)) => {
355 0 : tracing::info!("Joined HTTP server task");
356 : }
357 0 : Ok(Err(e)) => {
358 0 : tracing::error!("Error joining HTTP server task: {e}")
359 : }
360 : Err(_) => {
361 0 : tracing::warn!("Timed out joining HTTP server task");
362 : // We will fall through and shut down the service anyway, any request handlers
363 : // in flight will experience cancellation & their clients will see a torn connection.
364 : }
365 : }
366 :
367 : // If we were injecting chaos, stop that so that we're not calling into Service while it shuts down
368 0 : if let Some((chaos_jh, chaos_cancel)) = chaos_task {
369 0 : chaos_cancel.cancel();
370 0 : chaos_jh.await.ok();
371 0 : }
372 :
373 0 : service.shutdown().await;
374 0 : tracing::info!("Service shutdown complete");
375 :
376 0 : std::process::exit(0);
377 0 : }
|
