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