Line data Source code
1 : use anyhow::{anyhow, Context};
2 : use camino::Utf8PathBuf;
3 : use clap::Parser;
4 : use diesel::Connection;
5 : use metrics::launch_timestamp::LaunchTimestamp;
6 : use metrics::BuildInfo;
7 : use std::sync::Arc;
8 : use storage_controller::http::make_router;
9 : use storage_controller::metrics::preinitialize_metrics;
10 : use storage_controller::persistence::Persistence;
11 : use storage_controller::service::{Config, Service, MAX_UNAVAILABLE_INTERVAL_DEFAULT};
12 : use tokio::signal::unix::SignalKind;
13 : use tokio_util::sync::CancellationToken;
14 : use utils::auth::{JwtAuth, SwappableJwtAuth};
15 : use utils::logging::{self, LogFormat};
16 :
17 : use utils::sentry_init::init_sentry;
18 : use utils::{project_build_tag, project_git_version, tcp_listener};
19 :
20 : project_git_version!(GIT_VERSION);
21 : project_build_tag!(BUILD_TAG);
22 :
23 : use diesel_migrations::{embed_migrations, EmbeddedMigrations};
24 : pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("./migrations");
25 :
26 0 : #[derive(Parser)]
27 : #[command(author, version, about, long_about = None)]
28 : #[command(arg_required_else_help(true))]
29 : struct Cli {
30 : /// Host and port to listen on, like `127.0.0.1:1234`
31 : #[arg(short, long)]
32 0 : listen: std::net::SocketAddr,
33 :
34 : /// Public key for JWT authentication of clients
35 : #[arg(long)]
36 : public_key: Option<String>,
37 :
38 : /// Token for authenticating this service with the pageservers it controls
39 : #[arg(long)]
40 : jwt_token: Option<String>,
41 :
42 : /// Token for authenticating this service with the control plane, when calling
43 : /// the compute notification endpoint
44 : #[arg(long)]
45 : control_plane_jwt_token: Option<String>,
46 :
47 : /// URL to control plane compute notification endpoint
48 : #[arg(long)]
49 : compute_hook_url: Option<String>,
50 :
51 : /// Path to the .json file to store state (will be created if it doesn't exist)
52 : #[arg(short, long)]
53 : path: Option<Utf8PathBuf>,
54 :
55 : /// URL to connect to postgres, like postgresql://localhost:1234/storage_controller
56 : #[arg(long)]
57 : database_url: Option<String>,
58 :
59 : /// Flag to enable dev mode, which permits running without auth
60 : #[arg(long, default_value = "false")]
61 0 : dev: bool,
62 :
63 : /// Grace period before marking unresponsive pageserver offline
64 : #[arg(long)]
65 : max_unavailable_interval: Option<humantime::Duration>,
66 : }
67 :
68 : enum StrictMode {
69 : /// In strict mode, we will require that all secrets are loaded, i.e. security features
70 : /// may not be implicitly turned off by omitting secrets in the environment.
71 : Strict,
72 : /// In dev mode, secrets are optional, and omitting a particular secret will implicitly
73 : /// disable the auth related to it (e.g. no pageserver jwt key -> send unauthenticated
74 : /// requests, no public key -> don't authenticate incoming requests).
75 : Dev,
76 : }
77 :
78 : impl Default for StrictMode {
79 0 : fn default() -> Self {
80 0 : Self::Strict
81 0 : }
82 : }
83 :
84 : /// Secrets may either be provided on the command line (for testing), or loaded from AWS SecretManager: this
85 : /// type encapsulates the logic to decide which and do the loading.
86 : struct Secrets {
87 : database_url: String,
88 : public_key: Option<JwtAuth>,
89 : jwt_token: Option<String>,
90 : control_plane_jwt_token: Option<String>,
91 : }
92 :
93 : impl Secrets {
94 : const DATABASE_URL_ENV: &'static str = "DATABASE_URL";
95 : const PAGESERVER_JWT_TOKEN_ENV: &'static str = "PAGESERVER_JWT_TOKEN";
96 : const CONTROL_PLANE_JWT_TOKEN_ENV: &'static str = "CONTROL_PLANE_JWT_TOKEN";
97 : const PUBLIC_KEY_ENV: &'static str = "PUBLIC_KEY";
98 :
99 : /// Load secrets from, in order of preference:
100 : /// - CLI args if database URL is provided on the CLI
101 : /// - Environment variables if DATABASE_URL is set.
102 : /// - AWS Secrets Manager secrets
103 0 : async fn load(args: &Cli) -> anyhow::Result<Self> {
104 0 : let Some(database_url) =
105 0 : Self::load_secret(&args.database_url, Self::DATABASE_URL_ENV).await
106 : else {
107 0 : anyhow::bail!(
108 0 : "Database URL is not set (set `--database-url`, or `DATABASE_URL` environment)"
109 0 : )
110 : };
111 :
112 0 : let public_key = match Self::load_secret(&args.public_key, Self::PUBLIC_KEY_ENV).await {
113 0 : Some(v) => Some(JwtAuth::from_key(v).context("Loading public key")?),
114 0 : None => None,
115 : };
116 :
117 0 : let this = Self {
118 0 : database_url,
119 0 : public_key,
120 0 : jwt_token: Self::load_secret(&args.jwt_token, Self::PAGESERVER_JWT_TOKEN_ENV).await,
121 0 : control_plane_jwt_token: Self::load_secret(
122 0 : &args.control_plane_jwt_token,
123 0 : Self::CONTROL_PLANE_JWT_TOKEN_ENV,
124 0 : )
125 0 : .await,
126 : };
127 :
128 0 : Ok(this)
129 0 : }
130 :
131 0 : async fn load_secret(cli: &Option<String>, env_name: &str) -> Option<String> {
132 0 : if let Some(v) = cli {
133 0 : Some(v.clone())
134 0 : } else if let Ok(v) = std::env::var(env_name) {
135 0 : Some(v)
136 : } else {
137 0 : None
138 : }
139 0 : }
140 : }
141 :
142 : /// Execute the diesel migrations that are built into this binary
143 0 : async fn migration_run(database_url: &str) -> anyhow::Result<()> {
144 : use diesel::PgConnection;
145 : use diesel_migrations::{HarnessWithOutput, MigrationHarness};
146 0 : let mut conn = PgConnection::establish(database_url)?;
147 :
148 0 : HarnessWithOutput::write_to_stdout(&mut conn)
149 0 : .run_pending_migrations(MIGRATIONS)
150 0 : .map(|_| ())
151 0 : .map_err(|e| anyhow::anyhow!(e))?;
152 :
153 0 : Ok(())
154 0 : }
155 :
156 0 : fn main() -> anyhow::Result<()> {
157 0 : let default_panic = std::panic::take_hook();
158 0 : std::panic::set_hook(Box::new(move |info| {
159 0 : default_panic(info);
160 0 : std::process::exit(1);
161 0 : }));
162 0 :
163 0 : let _sentry_guard = init_sentry(Some(GIT_VERSION.into()), &[]);
164 0 :
165 0 : tokio::runtime::Builder::new_current_thread()
166 0 : // We use spawn_blocking for database operations, so require approximately
167 0 : // as many blocking threads as we will open database connections.
168 0 : .max_blocking_threads(Persistence::MAX_CONNECTIONS as usize)
169 0 : .enable_all()
170 0 : .build()
171 0 : .unwrap()
172 0 : .block_on(async_main())
173 0 : }
174 :
175 0 : async fn async_main() -> anyhow::Result<()> {
176 0 : let launch_ts = Box::leak(Box::new(LaunchTimestamp::generate()));
177 0 :
178 0 : logging::init(
179 0 : LogFormat::Plain,
180 0 : logging::TracingErrorLayerEnablement::Disabled,
181 0 : logging::Output::Stdout,
182 0 : )?;
183 :
184 0 : preinitialize_metrics();
185 0 :
186 0 : let args = Cli::parse();
187 0 : tracing::info!(
188 0 : "version: {}, launch_timestamp: {}, build_tag {}, state at {}, listening on {}",
189 0 : GIT_VERSION,
190 0 : launch_ts.to_string(),
191 0 : BUILD_TAG,
192 0 : args.path.as_ref().unwrap_or(&Utf8PathBuf::from("<none>")),
193 0 : args.listen
194 0 : );
195 :
196 0 : let build_info = BuildInfo {
197 0 : revision: GIT_VERSION,
198 0 : build_tag: BUILD_TAG,
199 0 : };
200 :
201 0 : let strict_mode = if args.dev {
202 0 : StrictMode::Dev
203 : } else {
204 0 : StrictMode::Strict
205 : };
206 :
207 0 : let secrets = Secrets::load(&args).await?;
208 :
209 : // Validate required secrets and arguments are provided in strict mode
210 0 : match strict_mode {
211 : StrictMode::Strict
212 0 : if (secrets.public_key.is_none()
213 0 : || secrets.jwt_token.is_none()
214 0 : || secrets.control_plane_jwt_token.is_none()) =>
215 0 : {
216 0 : // Production systems should always have secrets configured: if public_key was not set
217 0 : // then we would implicitly disable auth.
218 0 : anyhow::bail!(
219 0 : "Insecure config! One or more secrets is not set. This is only permitted in `--dev` mode"
220 0 : );
221 : }
222 0 : StrictMode::Strict if args.compute_hook_url.is_none() => {
223 0 : // Production systems should always have a compute hook set, to prevent falling
224 0 : // back to trying to use neon_local.
225 0 : anyhow::bail!(
226 0 : "`--compute-hook-url` is not set: this is only permitted in `--dev` mode"
227 0 : );
228 : }
229 : StrictMode::Strict => {
230 0 : tracing::info!("Starting in strict mode: configuration is OK.")
231 : }
232 : StrictMode::Dev => {
233 0 : tracing::warn!("Starting in dev mode: this may be an insecure configuration.")
234 : }
235 : }
236 :
237 0 : let config = Config {
238 0 : jwt_token: secrets.jwt_token,
239 0 : control_plane_jwt_token: secrets.control_plane_jwt_token,
240 0 : compute_hook_url: args.compute_hook_url,
241 0 : max_unavailable_interval: args
242 0 : .max_unavailable_interval
243 0 : .map(humantime::Duration::into)
244 0 : .unwrap_or(MAX_UNAVAILABLE_INTERVAL_DEFAULT),
245 0 : };
246 0 :
247 0 : // After loading secrets & config, but before starting anything else, apply database migrations
248 0 : migration_run(&secrets.database_url)
249 0 : .await
250 0 : .context("Running database migrations")?;
251 :
252 0 : let json_path = args.path;
253 0 : let persistence = Arc::new(Persistence::new(secrets.database_url, json_path.clone()));
254 :
255 0 : let service = Service::spawn(config, persistence.clone()).await?;
256 :
257 0 : let http_listener = tcp_listener::bind(args.listen)?;
258 :
259 0 : let auth = secrets
260 0 : .public_key
261 0 : .map(|jwt_auth| Arc::new(SwappableJwtAuth::new(jwt_auth)));
262 0 : let router = make_router(service.clone(), auth, build_info)
263 0 : .build()
264 0 : .map_err(|err| anyhow!(err))?;
265 0 : let router_service = utils::http::RouterService::new(router).unwrap();
266 0 :
267 0 : // Start HTTP server
268 0 : let server_shutdown = CancellationToken::new();
269 0 : let server = hyper::Server::from_tcp(http_listener)?
270 0 : .serve(router_service)
271 0 : .with_graceful_shutdown({
272 0 : let server_shutdown = server_shutdown.clone();
273 0 : async move {
274 0 : server_shutdown.cancelled().await;
275 0 : }
276 0 : });
277 0 : tracing::info!("Serving on {0}", args.listen);
278 0 : let server_task = tokio::task::spawn(server);
279 :
280 : // Wait until we receive a signal
281 0 : let mut sigint = tokio::signal::unix::signal(SignalKind::interrupt())?;
282 0 : let mut sigquit = tokio::signal::unix::signal(SignalKind::quit())?;
283 0 : let mut sigterm = tokio::signal::unix::signal(SignalKind::terminate())?;
284 0 : tokio::select! {
285 0 : _ = sigint.recv() => {},
286 0 : _ = sigterm.recv() => {},
287 0 : _ = sigquit.recv() => {},
288 0 : }
289 0 : tracing::info!("Terminating on signal");
290 :
291 0 : if json_path.is_some() {
292 : // Write out a JSON dump on shutdown: this is used in compat tests to avoid passing
293 : // full postgres dumps around.
294 0 : if let Err(e) = persistence.write_tenants_json().await {
295 0 : tracing::error!("Failed to write JSON on shutdown: {e}")
296 0 : }
297 0 : }
298 :
299 : // Stop HTTP server first, so that we don't have to service requests
300 : // while shutting down Service
301 0 : server_shutdown.cancel();
302 0 : if let Err(e) = server_task.await {
303 0 : tracing::error!("Error joining HTTP server task: {e}")
304 0 : }
305 0 : tracing::info!("Joined HTTP server task");
306 :
307 0 : service.shutdown().await;
308 0 : tracing::info!("Service shutdown complete");
309 :
310 0 : std::process::exit(0);
311 0 : }
|
