LCOV - code coverage report
Current view: top level - control_plane/src - storage_controller.rs (source / functions) Coverage Total Hit
Test: 1e20c4f2b28aa592527961bb32170ebbd2c9172f.info Lines: 0.0 % 570 0
Test Date: 2025-07-16 12:29:03 Functions: 0.0 % 92 0

            Line data    Source code
       1              : use std::ffi::OsStr;
       2              : use std::fs;
       3              : use std::path::PathBuf;
       4              : use std::process::ExitStatus;
       5              : use std::str::FromStr;
       6              : use std::sync::OnceLock;
       7              : use std::time::{Duration, Instant};
       8              : 
       9              : use crate::background_process;
      10              : use crate::local_env::{LocalEnv, NeonStorageControllerConf};
      11              : use camino::{Utf8Path, Utf8PathBuf};
      12              : use hyper0::Uri;
      13              : use nix::unistd::Pid;
      14              : use pageserver_api::controller_api::{
      15              :     NodeConfigureRequest, NodeDescribeResponse, NodeRegisterRequest,
      16              :     SafekeeperSchedulingPolicyRequest, SkSchedulingPolicy, TenantCreateRequest,
      17              :     TenantCreateResponse, TenantLocateResponse,
      18              : };
      19              : use pageserver_api::models::{
      20              :     TenantConfig, TenantConfigRequest, TimelineCreateRequest, TimelineInfo,
      21              : };
      22              : use pageserver_api::shard::TenantShardId;
      23              : use pageserver_client::mgmt_api::ResponseErrorMessageExt;
      24              : use pem::Pem;
      25              : use postgres_backend::AuthType;
      26              : use reqwest::{Method, Response};
      27              : use safekeeper_api::PgMajorVersion;
      28              : use serde::de::DeserializeOwned;
      29              : use serde::{Deserialize, Serialize};
      30              : use tokio::process::Command;
      31              : use tracing::instrument;
      32              : use url::Url;
      33              : use utils::auth::{Claims, Scope, encode_from_key_file};
      34              : use utils::id::{NodeId, TenantId};
      35              : use whoami::username;
      36              : 
      37              : pub struct StorageController {
      38              :     env: LocalEnv,
      39              :     private_key: Option<Pem>,
      40              :     public_key: Option<Pem>,
      41              :     client: reqwest::Client,
      42              :     config: NeonStorageControllerConf,
      43              : 
      44              :     // The listen port is learned when starting the storage controller,
      45              :     // hence the use of OnceLock to init it at the right time.
      46              :     listen_port: OnceLock<u16>,
      47              : }
      48              : 
      49              : const COMMAND: &str = "storage_controller";
      50              : 
      51              : const STORAGE_CONTROLLER_POSTGRES_VERSION: PgMajorVersion = PgMajorVersion::PG16;
      52              : 
      53              : const DB_NAME: &str = "storage_controller";
      54              : 
      55              : pub struct NeonStorageControllerStartArgs {
      56              :     pub instance_id: u8,
      57              :     pub base_port: Option<u16>,
      58              :     pub start_timeout: humantime::Duration,
      59              : }
      60              : 
      61              : impl NeonStorageControllerStartArgs {
      62            0 :     pub fn with_default_instance_id(start_timeout: humantime::Duration) -> Self {
      63            0 :         Self {
      64            0 :             instance_id: 1,
      65            0 :             base_port: None,
      66            0 :             start_timeout,
      67            0 :         }
      68            0 :     }
      69              : }
      70              : 
      71              : pub struct NeonStorageControllerStopArgs {
      72              :     pub instance_id: u8,
      73              :     pub immediate: bool,
      74              : }
      75              : 
      76              : impl NeonStorageControllerStopArgs {
      77            0 :     pub fn with_default_instance_id(immediate: bool) -> Self {
      78            0 :         Self {
      79            0 :             instance_id: 1,
      80            0 :             immediate,
      81            0 :         }
      82            0 :     }
      83              : }
      84              : 
      85            0 : #[derive(Serialize, Deserialize)]
      86              : pub struct AttachHookRequest {
      87              :     pub tenant_shard_id: TenantShardId,
      88              :     pub node_id: Option<NodeId>,
      89              :     pub generation_override: Option<i32>, // only new tenants
      90              :     pub config: Option<TenantConfig>,     // only new tenants
      91              : }
      92              : 
      93            0 : #[derive(Serialize, Deserialize)]
      94              : pub struct AttachHookResponse {
      95              :     #[serde(rename = "gen")]
      96              :     pub generation: Option<u32>,
      97              : }
      98              : 
      99            0 : #[derive(Serialize, Deserialize)]
     100              : pub struct InspectRequest {
     101              :     pub tenant_shard_id: TenantShardId,
     102              : }
     103              : 
     104            0 : #[derive(Serialize, Deserialize)]
     105              : pub struct InspectResponse {
     106              :     pub attachment: Option<(u32, NodeId)>,
     107              : }
     108              : 
     109              : impl StorageController {
     110            0 :     pub fn from_env(env: &LocalEnv) -> Self {
     111              :         // Assume all pageservers have symmetric auth configuration: this service
     112              :         // expects to use one JWT token to talk to all of them.
     113            0 :         let ps_conf = env
     114            0 :             .pageservers
     115            0 :             .first()
     116            0 :             .expect("Config is validated to contain at least one pageserver");
     117            0 :         let (private_key, public_key) = match ps_conf.http_auth_type {
     118            0 :             AuthType::Trust => (None, None),
     119              :             AuthType::NeonJWT => {
     120            0 :                 let private_key_path = env.get_private_key_path();
     121            0 :                 let private_key =
     122            0 :                     pem::parse(fs::read(private_key_path).expect("failed to read private key"))
     123            0 :                         .expect("failed to parse PEM file");
     124              : 
     125              :                 // If pageserver auth is enabled, this implicitly enables auth for this service,
     126              :                 // using the same credentials.
     127            0 :                 let public_key_path =
     128            0 :                     camino::Utf8PathBuf::try_from(env.base_data_dir.join("auth_public_key.pem"))
     129            0 :                         .unwrap();
     130              : 
     131              :                 // This service takes keys as a string rather than as a path to a file/dir: read the key into memory.
     132            0 :                 let public_key = if std::fs::metadata(&public_key_path)
     133            0 :                     .expect("Can't stat public key")
     134            0 :                     .is_dir()
     135              :                 {
     136              :                     // Our config may specify a directory: this is for the pageserver's ability to handle multiple
     137              :                     // keys.  We only use one key at a time, so, arbitrarily load the first one in the directory.
     138            0 :                     let mut dir =
     139            0 :                         std::fs::read_dir(&public_key_path).expect("Can't readdir public key path");
     140            0 :                     let dent = dir
     141            0 :                         .next()
     142            0 :                         .expect("Empty key dir")
     143            0 :                         .expect("Error reading key dir");
     144              : 
     145            0 :                     pem::parse(std::fs::read_to_string(dent.path()).expect("Can't read public key"))
     146            0 :                         .expect("Failed to parse PEM file")
     147              :                 } else {
     148            0 :                     pem::parse(
     149            0 :                         std::fs::read_to_string(&public_key_path).expect("Can't read public key"),
     150              :                     )
     151            0 :                     .expect("Failed to parse PEM file")
     152              :                 };
     153            0 :                 (Some(private_key), Some(public_key))
     154              :             }
     155              :         };
     156              : 
     157            0 :         Self {
     158            0 :             env: env.clone(),
     159            0 :             private_key,
     160            0 :             public_key,
     161            0 :             client: env.create_http_client(),
     162            0 :             config: env.storage_controller.clone(),
     163            0 :             listen_port: OnceLock::default(),
     164            0 :         }
     165            0 :     }
     166              : 
     167            0 :     fn storage_controller_instance_dir(&self, instance_id: u8) -> PathBuf {
     168            0 :         self.env
     169            0 :             .base_data_dir
     170            0 :             .join(format!("storage_controller_{instance_id}"))
     171            0 :     }
     172              : 
     173            0 :     fn pid_file(&self, instance_id: u8) -> Utf8PathBuf {
     174            0 :         Utf8PathBuf::from_path_buf(
     175            0 :             self.storage_controller_instance_dir(instance_id)
     176            0 :                 .join("storage_controller.pid"),
     177              :         )
     178            0 :         .expect("non-Unicode path")
     179            0 :     }
     180              : 
     181              :     /// Find the directory containing postgres subdirectories, such `bin` and `lib`
     182              :     ///
     183              :     /// This usually uses STORAGE_CONTROLLER_POSTGRES_VERSION of postgres, but will fall back
     184              :     /// to other versions if that one isn't found.  Some automated tests create circumstances
     185              :     /// where only one version is available in pg_distrib_dir, such as `test_remote_extensions`.
     186            0 :     async fn get_pg_dir(&self, dir_name: &str) -> anyhow::Result<Utf8PathBuf> {
     187              :         const PREFER_VERSIONS: [PgMajorVersion; 5] = [
     188              :             STORAGE_CONTROLLER_POSTGRES_VERSION,
     189              :             PgMajorVersion::PG16,
     190              :             PgMajorVersion::PG15,
     191              :             PgMajorVersion::PG14,
     192              :             PgMajorVersion::PG17,
     193              :         ];
     194              : 
     195            0 :         for v in PREFER_VERSIONS {
     196            0 :             let path = Utf8PathBuf::from_path_buf(self.env.pg_dir(v, dir_name)?).unwrap();
     197            0 :             if tokio::fs::try_exists(&path).await? {
     198            0 :                 return Ok(path);
     199            0 :             }
     200              :         }
     201              : 
     202              :         // Fall through
     203            0 :         anyhow::bail!(
     204            0 :             "Postgres directory '{}' not found in {}",
     205              :             dir_name,
     206            0 :             self.env.pg_distrib_dir.display(),
     207              :         );
     208            0 :     }
     209              : 
     210            0 :     pub async fn get_pg_bin_dir(&self) -> anyhow::Result<Utf8PathBuf> {
     211            0 :         self.get_pg_dir("bin").await
     212            0 :     }
     213              : 
     214            0 :     pub async fn get_pg_lib_dir(&self) -> anyhow::Result<Utf8PathBuf> {
     215            0 :         self.get_pg_dir("lib").await
     216            0 :     }
     217              : 
     218              :     /// Readiness check for our postgres process
     219            0 :     async fn pg_isready(&self, pg_bin_dir: &Utf8Path, postgres_port: u16) -> anyhow::Result<bool> {
     220            0 :         let bin_path = pg_bin_dir.join("pg_isready");
     221            0 :         let args = [
     222            0 :             "-h",
     223            0 :             "localhost",
     224            0 :             "-U",
     225            0 :             &username(),
     226            0 :             "-d",
     227            0 :             DB_NAME,
     228            0 :             "-p",
     229            0 :             &format!("{postgres_port}"),
     230            0 :         ];
     231            0 :         let pg_lib_dir = self.get_pg_lib_dir().await.unwrap();
     232            0 :         let envs = [
     233            0 :             ("LD_LIBRARY_PATH".to_owned(), pg_lib_dir.to_string()),
     234            0 :             ("DYLD_LIBRARY_PATH".to_owned(), pg_lib_dir.to_string()),
     235            0 :         ];
     236            0 :         let exitcode = Command::new(bin_path)
     237            0 :             .args(args)
     238            0 :             .envs(envs)
     239            0 :             .spawn()?
     240            0 :             .wait()
     241            0 :             .await?;
     242              : 
     243            0 :         Ok(exitcode.success())
     244            0 :     }
     245              : 
     246              :     /// Create our database if it doesn't exist
     247              :     ///
     248              :     /// This function is equivalent to the `diesel setup` command in the diesel CLI.  We implement
     249              :     /// the same steps by hand to avoid imposing a dependency on installing diesel-cli for developers
     250              :     /// who just want to run `cargo neon_local` without knowing about diesel.
     251              :     ///
     252              :     /// Returns the database url
     253            0 :     pub async fn setup_database(&self, postgres_port: u16) -> anyhow::Result<String> {
     254            0 :         let database_url = format!(
     255            0 :             "postgresql://{}@localhost:{}/{DB_NAME}",
     256            0 :             &username(),
     257              :             postgres_port
     258              :         );
     259              : 
     260            0 :         let pg_bin_dir = self.get_pg_bin_dir().await?;
     261            0 :         let createdb_path = pg_bin_dir.join("createdb");
     262            0 :         let pg_lib_dir = self.get_pg_lib_dir().await.unwrap();
     263            0 :         let envs = [
     264            0 :             ("LD_LIBRARY_PATH".to_owned(), pg_lib_dir.to_string()),
     265            0 :             ("DYLD_LIBRARY_PATH".to_owned(), pg_lib_dir.to_string()),
     266            0 :         ];
     267            0 :         let output = Command::new(&createdb_path)
     268            0 :             .args([
     269            0 :                 "-h",
     270            0 :                 "localhost",
     271            0 :                 "-p",
     272            0 :                 &format!("{postgres_port}"),
     273            0 :                 "-U",
     274            0 :                 &username(),
     275            0 :                 "-O",
     276            0 :                 &username(),
     277            0 :                 DB_NAME,
     278            0 :             ])
     279            0 :             .envs(envs)
     280            0 :             .output()
     281            0 :             .await
     282            0 :             .expect("Failed to spawn createdb");
     283              : 
     284            0 :         if !output.status.success() {
     285            0 :             let stderr = String::from_utf8(output.stderr).expect("Non-UTF8 output from createdb");
     286            0 :             if stderr.contains("already exists") {
     287            0 :                 tracing::info!("Database {DB_NAME} already exists");
     288              :             } else {
     289            0 :                 anyhow::bail!("createdb failed with status {}: {stderr}", output.status);
     290              :             }
     291            0 :         }
     292              : 
     293            0 :         Ok(database_url)
     294            0 :     }
     295              : 
     296            0 :     pub async fn connect_to_database(
     297            0 :         &self,
     298            0 :         postgres_port: u16,
     299            0 :     ) -> anyhow::Result<(
     300            0 :         tokio_postgres::Client,
     301            0 :         tokio_postgres::Connection<tokio_postgres::Socket, tokio_postgres::tls::NoTlsStream>,
     302            0 :     )> {
     303            0 :         tokio_postgres::Config::new()
     304            0 :             .host("localhost")
     305            0 :             .port(postgres_port)
     306            0 :             // The user is the ambient operating system user name.
     307            0 :             // That is an impurity which we want to fix in => TODO https://github.com/neondatabase/neon/issues/8400
     308            0 :             //
     309            0 :             // Until we get there, use the ambient operating system user name.
     310            0 :             // Recent tokio-postgres versions default to this if the user isn't specified.
     311            0 :             // But tokio-postgres fork doesn't have this upstream commit:
     312            0 :             // https://github.com/sfackler/rust-postgres/commit/cb609be758f3fb5af537f04b584a2ee0cebd5e79
     313            0 :             // => we should rebase our fork => TODO https://github.com/neondatabase/neon/issues/8399
     314            0 :             .user(&username())
     315            0 :             .dbname(DB_NAME)
     316            0 :             .connect(tokio_postgres::NoTls)
     317            0 :             .await
     318            0 :             .map_err(anyhow::Error::new)
     319            0 :     }
     320              : 
     321              :     /// Wrapper for the pg_ctl binary, which we spawn as a short-lived subprocess when starting and stopping postgres
     322            0 :     async fn pg_ctl<I, S>(&self, args: I) -> ExitStatus
     323            0 :     where
     324            0 :         I: IntoIterator<Item = S>,
     325            0 :         S: AsRef<OsStr>,
     326            0 :     {
     327            0 :         let pg_bin_dir = self.get_pg_bin_dir().await.unwrap();
     328            0 :         let bin_path = pg_bin_dir.join("pg_ctl");
     329              : 
     330            0 :         let pg_lib_dir = self.get_pg_lib_dir().await.unwrap();
     331            0 :         let envs = [
     332            0 :             ("LD_LIBRARY_PATH".to_owned(), pg_lib_dir.to_string()),
     333            0 :             ("DYLD_LIBRARY_PATH".to_owned(), pg_lib_dir.to_string()),
     334            0 :         ];
     335              : 
     336            0 :         Command::new(bin_path)
     337            0 :             .args(args)
     338            0 :             .envs(envs)
     339            0 :             .spawn()
     340            0 :             .expect("Failed to spawn pg_ctl, binary_missing?")
     341            0 :             .wait()
     342            0 :             .await
     343            0 :             .expect("Failed to wait for pg_ctl termination")
     344            0 :     }
     345              : 
     346            0 :     pub async fn start(&self, start_args: NeonStorageControllerStartArgs) -> anyhow::Result<()> {
     347            0 :         let instance_dir = self.storage_controller_instance_dir(start_args.instance_id);
     348            0 :         if let Err(err) = tokio::fs::create_dir(&instance_dir).await {
     349            0 :             if err.kind() != std::io::ErrorKind::AlreadyExists {
     350            0 :                 panic!("Failed to create instance dir {instance_dir:?}");
     351            0 :             }
     352            0 :         }
     353              : 
     354            0 :         if self.env.generate_local_ssl_certs {
     355            0 :             self.env.generate_ssl_cert(
     356            0 :                 &instance_dir.join("server.crt"),
     357            0 :                 &instance_dir.join("server.key"),
     358            0 :             )?;
     359            0 :         }
     360              : 
     361            0 :         let listen_url = &self.env.control_plane_api;
     362              : 
     363            0 :         let scheme = listen_url.scheme();
     364            0 :         let host = listen_url.host_str().unwrap();
     365              : 
     366            0 :         let (listen_port, postgres_port) = if let Some(base_port) = start_args.base_port {
     367            0 :             (
     368            0 :                 base_port,
     369            0 :                 self.config
     370            0 :                     .database_url
     371            0 :                     .expect("--base-port requires NeonStorageControllerConf::database_url")
     372            0 :                     .port(),
     373            0 :             )
     374              :         } else {
     375            0 :             let port = listen_url.port().unwrap();
     376            0 :             (port, port + 1)
     377              :         };
     378              : 
     379            0 :         self.listen_port
     380            0 :             .set(listen_port)
     381            0 :             .expect("StorageController::listen_port is only set here");
     382              : 
     383              :         // Do we remove the pid file on stop?
     384            0 :         let pg_started = self.is_postgres_running().await?;
     385            0 :         let pg_lib_dir = self.get_pg_lib_dir().await?;
     386              : 
     387            0 :         if !pg_started {
     388              :             // Start a vanilla Postgres process used by the storage controller for persistence.
     389            0 :             let pg_data_path = Utf8PathBuf::from_path_buf(self.env.base_data_dir.clone())
     390            0 :                 .unwrap()
     391            0 :                 .join("storage_controller_db");
     392            0 :             let pg_bin_dir = self.get_pg_bin_dir().await?;
     393            0 :             let pg_log_path = pg_data_path.join("postgres.log");
     394              : 
     395            0 :             if !tokio::fs::try_exists(&pg_data_path).await? {
     396            0 :                 let initdb_args = [
     397            0 :                     "--pgdata",
     398            0 :                     pg_data_path.as_ref(),
     399            0 :                     "--username",
     400            0 :                     &username(),
     401            0 :                     "--no-sync",
     402            0 :                     "--no-instructions",
     403            0 :                 ];
     404            0 :                 tracing::info!(
     405            0 :                     "Initializing storage controller database with args: {:?}",
     406              :                     initdb_args
     407              :                 );
     408              : 
     409              :                 // Initialize empty database
     410            0 :                 let initdb_path = pg_bin_dir.join("initdb");
     411            0 :                 let mut child = Command::new(&initdb_path)
     412            0 :                     .envs(vec![
     413            0 :                         ("LD_LIBRARY_PATH".to_owned(), pg_lib_dir.to_string()),
     414            0 :                         ("DYLD_LIBRARY_PATH".to_owned(), pg_lib_dir.to_string()),
     415            0 :                     ])
     416            0 :                     .args(initdb_args)
     417            0 :                     .spawn()
     418            0 :                     .expect("Failed to spawn initdb");
     419            0 :                 let status = child.wait().await?;
     420            0 :                 if !status.success() {
     421            0 :                     anyhow::bail!("initdb failed with status {status}");
     422            0 :                 }
     423            0 :             };
     424              : 
     425              :             // Write a minimal config file:
     426              :             // - Specify the port, since this is chosen dynamically
     427              :             // - Switch off fsync, since we're running on lightweight test environments and when e.g. scale testing
     428              :             //   the storage controller we don't want a slow local disk to interfere with that.
     429              :             //
     430              :             // NB: it's important that we rewrite this file on each start command so we propagate changes
     431              :             // from `LocalEnv`'s config file (`.neon/config`).
     432            0 :             tokio::fs::write(
     433            0 :                 &pg_data_path.join("postgresql.conf"),
     434            0 :                 format!("port = {postgres_port}\nfsync=off\n"),
     435            0 :             )
     436            0 :             .await?;
     437              : 
     438            0 :             println!("Starting storage controller database...");
     439            0 :             let db_start_args = [
     440            0 :                 "-w",
     441            0 :                 "-D",
     442            0 :                 pg_data_path.as_ref(),
     443            0 :                 "-l",
     444            0 :                 pg_log_path.as_ref(),
     445            0 :                 "-U",
     446            0 :                 &username(),
     447            0 :                 "start",
     448            0 :             ];
     449            0 :             tracing::info!(
     450            0 :                 "Starting storage controller database with args: {:?}",
     451              :                 db_start_args
     452              :             );
     453              : 
     454            0 :             let db_start_status = self.pg_ctl(db_start_args).await;
     455            0 :             let start_timeout: Duration = start_args.start_timeout.into();
     456            0 :             let db_start_deadline = Instant::now() + start_timeout;
     457            0 :             if !db_start_status.success() {
     458            0 :                 return Err(anyhow::anyhow!(
     459            0 :                     "Failed to start postgres {}",
     460            0 :                     db_start_status.code().unwrap()
     461            0 :                 ));
     462            0 :             }
     463              : 
     464              :             loop {
     465            0 :                 if Instant::now() > db_start_deadline {
     466            0 :                     return Err(anyhow::anyhow!("Timed out waiting for postgres to start"));
     467            0 :                 }
     468              : 
     469            0 :                 match self.pg_isready(&pg_bin_dir, postgres_port).await {
     470              :                     Ok(true) => {
     471            0 :                         tracing::info!("storage controller postgres is now ready");
     472            0 :                         break;
     473              :                     }
     474              :                     Ok(false) => {
     475            0 :                         tokio::time::sleep(Duration::from_millis(100)).await;
     476              :                     }
     477            0 :                     Err(e) => {
     478            0 :                         tracing::warn!("Failed to check postgres status: {e}")
     479              :                     }
     480              :                 }
     481              :             }
     482              : 
     483            0 :             self.setup_database(postgres_port).await?;
     484            0 :         }
     485              : 
     486            0 :         let database_url = format!("postgresql://localhost:{postgres_port}/{DB_NAME}");
     487              : 
     488              :         // We support running a startup SQL script to fiddle with the database before we launch storcon.
     489              :         // This is used by the test suite.
     490            0 :         let startup_script_path = self
     491            0 :             .env
     492            0 :             .base_data_dir
     493            0 :             .join("storage_controller_db.startup.sql");
     494            0 :         let startup_script = match tokio::fs::read_to_string(&startup_script_path).await {
     495            0 :             Ok(script) => {
     496            0 :                 tokio::fs::remove_file(startup_script_path).await?;
     497            0 :                 script
     498              :             }
     499            0 :             Err(e) => {
     500            0 :                 if e.kind() == std::io::ErrorKind::NotFound {
     501              :                     // always run some startup script so that this code path doesn't bit rot
     502            0 :                     "BEGIN; COMMIT;".to_string()
     503              :                 } else {
     504            0 :                     anyhow::bail!("Failed to read startup script: {e}")
     505              :                 }
     506              :             }
     507              :         };
     508            0 :         let (mut client, conn) = self.connect_to_database(postgres_port).await?;
     509            0 :         let conn = tokio::spawn(conn);
     510            0 :         let tx = client.build_transaction();
     511            0 :         let tx = tx.start().await?;
     512            0 :         tx.batch_execute(&startup_script).await?;
     513            0 :         tx.commit().await?;
     514            0 :         drop(client);
     515            0 :         conn.await??;
     516              : 
     517            0 :         let addr = format!("{host}:{listen_port}");
     518            0 :         let address_for_peers = Uri::builder()
     519            0 :             .scheme(scheme)
     520            0 :             .authority(addr.clone())
     521            0 :             .path_and_query("")
     522            0 :             .build()
     523            0 :             .unwrap();
     524              : 
     525            0 :         let mut args = vec![
     526              :             "--dev",
     527            0 :             "--database-url",
     528            0 :             &database_url,
     529            0 :             "--max-offline-interval",
     530            0 :             &humantime::Duration::from(self.config.max_offline).to_string(),
     531            0 :             "--max-warming-up-interval",
     532            0 :             &humantime::Duration::from(self.config.max_warming_up).to_string(),
     533            0 :             "--heartbeat-interval",
     534            0 :             &humantime::Duration::from(self.config.heartbeat_interval).to_string(),
     535            0 :             "--address-for-peers",
     536            0 :             &address_for_peers.to_string(),
     537              :         ]
     538            0 :         .into_iter()
     539            0 :         .map(|s| s.to_string())
     540            0 :         .collect::<Vec<_>>();
     541              : 
     542            0 :         match scheme {
     543            0 :             "http" => args.extend(["--listen".to_string(), addr]),
     544            0 :             "https" => args.extend(["--listen-https".to_string(), addr]),
     545              :             _ => {
     546            0 :                 panic!("Unexpected url scheme in control_plane_api: {scheme}");
     547              :             }
     548              :         }
     549              : 
     550            0 :         if self.config.start_as_candidate {
     551            0 :             args.push("--start-as-candidate".to_string());
     552            0 :         }
     553              : 
     554            0 :         if self.config.use_https_pageserver_api {
     555            0 :             args.push("--use-https-pageserver-api".to_string());
     556            0 :         }
     557              : 
     558            0 :         if self.config.use_https_safekeeper_api {
     559            0 :             args.push("--use-https-safekeeper-api".to_string());
     560            0 :         }
     561              : 
     562            0 :         if self.config.use_local_compute_notifications {
     563            0 :             args.push("--use-local-compute-notifications".to_string());
     564            0 :         }
     565              : 
     566            0 :         if let Some(value) = self.config.kick_secondary_downloads {
     567            0 :             args.push(format!("--kick-secondary-downloads={value}"));
     568            0 :         }
     569              : 
     570            0 :         if let Some(ssl_ca_file) = self.env.ssl_ca_cert_path() {
     571            0 :             args.push(format!("--ssl-ca-file={}", ssl_ca_file.to_str().unwrap()));
     572            0 :         }
     573              : 
     574            0 :         if let Some(private_key) = &self.private_key {
     575            0 :             let claims = Claims::new(None, Scope::PageServerApi);
     576            0 :             let jwt_token =
     577            0 :                 encode_from_key_file(&claims, private_key).expect("failed to generate jwt token");
     578            0 :             args.push(format!("--jwt-token={jwt_token}"));
     579            0 : 
     580            0 :             let peer_claims = Claims::new(None, Scope::Admin);
     581            0 :             let peer_jwt_token = encode_from_key_file(&peer_claims, private_key)
     582            0 :                 .expect("failed to generate jwt token");
     583            0 :             args.push(format!("--peer-jwt-token={peer_jwt_token}"));
     584            0 : 
     585            0 :             let claims = Claims::new(None, Scope::SafekeeperData);
     586            0 :             let jwt_token =
     587            0 :                 encode_from_key_file(&claims, private_key).expect("failed to generate jwt token");
     588            0 :             args.push(format!("--safekeeper-jwt-token={jwt_token}"));
     589            0 :         }
     590              : 
     591            0 :         if let Some(public_key) = &self.public_key {
     592            0 :             args.push(format!("--public-key=\"{public_key}\""));
     593            0 :         }
     594              : 
     595            0 :         if let Some(control_plane_hooks_api) = &self.env.control_plane_hooks_api {
     596            0 :             args.push(format!("--control-plane-url={control_plane_hooks_api}"));
     597            0 :         }
     598              : 
     599            0 :         if let Some(split_threshold) = self.config.split_threshold.as_ref() {
     600            0 :             args.push(format!("--split-threshold={split_threshold}"))
     601            0 :         }
     602              : 
     603            0 :         if let Some(max_split_shards) = self.config.max_split_shards.as_ref() {
     604            0 :             args.push(format!("--max-split-shards={max_split_shards}"))
     605            0 :         }
     606              : 
     607            0 :         if let Some(initial_split_threshold) = self.config.initial_split_threshold.as_ref() {
     608            0 :             args.push(format!(
     609            0 :                 "--initial-split-threshold={initial_split_threshold}"
     610              :             ))
     611            0 :         }
     612              : 
     613            0 :         if let Some(initial_split_shards) = self.config.initial_split_shards.as_ref() {
     614            0 :             args.push(format!("--initial-split-shards={initial_split_shards}"))
     615            0 :         }
     616              : 
     617            0 :         if let Some(lag) = self.config.max_secondary_lag_bytes.as_ref() {
     618            0 :             args.push(format!("--max-secondary-lag-bytes={lag}"))
     619            0 :         }
     620              : 
     621            0 :         if let Some(threshold) = self.config.long_reconcile_threshold {
     622            0 :             args.push(format!(
     623            0 :                 "--long-reconcile-threshold={}",
     624            0 :                 humantime::Duration::from(threshold)
     625              :             ))
     626            0 :         }
     627              : 
     628            0 :         args.push(format!(
     629            0 :             "--neon-local-repo-dir={}",
     630            0 :             self.env.base_data_dir.display()
     631              :         ));
     632              : 
     633            0 :         if self.env.safekeepers.iter().any(|sk| sk.auth_enabled) && self.private_key.is_none() {
     634            0 :             anyhow::bail!("Safekeeper set up for auth but no private key specified");
     635            0 :         }
     636              : 
     637            0 :         if self.config.timelines_onto_safekeepers {
     638            0 :             args.push("--timelines-onto-safekeepers".to_string());
     639            0 :         }
     640              : 
     641              :         // neon_local is used in test environments where we often have less than 3 safekeepers.
     642            0 :         if self.config.timeline_safekeeper_count.is_some() || self.env.safekeepers.len() < 3 {
     643            0 :             let sk_cnt = self
     644            0 :                 .config
     645            0 :                 .timeline_safekeeper_count
     646            0 :                 .unwrap_or(self.env.safekeepers.len());
     647            0 : 
     648            0 :             args.push(format!("--timeline-safekeeper-count={sk_cnt}"));
     649            0 :         }
     650              : 
     651            0 :         if let Some(duration) = self.config.shard_split_request_timeout {
     652            0 :             args.push(format!(
     653            0 :                 "--shard-split-request-timeout={}",
     654            0 :                 humantime::Duration::from(duration)
     655            0 :             ));
     656            0 :         }
     657              : 
     658            0 :         let mut envs = vec![
     659            0 :             ("LD_LIBRARY_PATH".to_owned(), pg_lib_dir.to_string()),
     660            0 :             ("DYLD_LIBRARY_PATH".to_owned(), pg_lib_dir.to_string()),
     661              :         ];
     662              : 
     663            0 :         if let Some(posthog_config) = &self.config.posthog_config {
     664            0 :             envs.push((
     665            0 :                 "POSTHOG_CONFIG".to_string(),
     666            0 :                 serde_json::to_string(posthog_config)?,
     667              :             ));
     668            0 :         }
     669              : 
     670            0 :         println!("Starting storage controller at {scheme}://{host}:{listen_port}");
     671              : 
     672            0 :         background_process::start_process(
     673            0 :             COMMAND,
     674            0 :             &instance_dir,
     675            0 :             &self.env.storage_controller_bin(),
     676            0 :             args,
     677            0 :             envs,
     678            0 :             background_process::InitialPidFile::Create(self.pid_file(start_args.instance_id)),
     679            0 :             &start_args.start_timeout,
     680            0 :             || async {
     681            0 :                 match self.ready().await {
     682            0 :                     Ok(_) => Ok(true),
     683            0 :                     Err(_) => Ok(false),
     684              :                 }
     685            0 :             },
     686              :         )
     687            0 :         .await?;
     688              : 
     689            0 :         if self.config.timelines_onto_safekeepers {
     690            0 :             self.register_safekeepers().await?;
     691            0 :         }
     692              : 
     693            0 :         Ok(())
     694            0 :     }
     695              : 
     696            0 :     pub async fn stop(&self, stop_args: NeonStorageControllerStopArgs) -> anyhow::Result<()> {
     697            0 :         background_process::stop_process(
     698            0 :             stop_args.immediate,
     699            0 :             COMMAND,
     700            0 :             &self.pid_file(stop_args.instance_id),
     701            0 :         )?;
     702              : 
     703            0 :         let storcon_instances = self.env.storage_controller_instances().await?;
     704            0 :         for (instance_id, instanced_dir_path) in storcon_instances {
     705            0 :             if instance_id == stop_args.instance_id {
     706            0 :                 continue;
     707            0 :             }
     708              : 
     709            0 :             let pid_file = instanced_dir_path.join("storage_controller.pid");
     710            0 :             let pid = tokio::fs::read_to_string(&pid_file)
     711            0 :                 .await
     712            0 :                 .map_err(|err| {
     713            0 :                     anyhow::anyhow!("Failed to read storcon pid file at {pid_file:?}: {err}")
     714            0 :                 })?
     715            0 :                 .parse::<i32>()
     716            0 :                 .expect("pid is valid i32");
     717              : 
     718            0 :             let other_proc_alive = !background_process::process_has_stopped(Pid::from_raw(pid))?;
     719            0 :             if other_proc_alive {
     720              :                 // There is another storage controller instance running, so we return
     721              :                 // and leave the database running.
     722            0 :                 return Ok(());
     723            0 :             }
     724              :         }
     725              : 
     726            0 :         let pg_data_path = self.env.base_data_dir.join("storage_controller_db");
     727              : 
     728            0 :         println!("Stopping storage controller database...");
     729            0 :         let pg_stop_args = ["-D", &pg_data_path.to_string_lossy(), "stop"];
     730            0 :         let stop_status = self.pg_ctl(pg_stop_args).await;
     731            0 :         if !stop_status.success() {
     732            0 :             match self.is_postgres_running().await {
     733              :                 Ok(false) => {
     734            0 :                     println!("Storage controller database is already stopped");
     735            0 :                     return Ok(());
     736              :                 }
     737              :                 Ok(true) => {
     738            0 :                     anyhow::bail!("Failed to stop storage controller database");
     739              :                 }
     740            0 :                 Err(err) => {
     741            0 :                     anyhow::bail!("Failed to stop storage controller database: {err}");
     742              :                 }
     743              :             }
     744            0 :         }
     745              : 
     746            0 :         Ok(())
     747            0 :     }
     748              : 
     749            0 :     async fn is_postgres_running(&self) -> anyhow::Result<bool> {
     750            0 :         let pg_data_path = self.env.base_data_dir.join("storage_controller_db");
     751              : 
     752            0 :         let pg_status_args = ["-D", &pg_data_path.to_string_lossy(), "status"];
     753            0 :         let status_exitcode = self.pg_ctl(pg_status_args).await;
     754              : 
     755              :         // pg_ctl status returns this exit code if postgres is not running: in this case it is
     756              :         // fine that stop failed.  Otherwise it is an error that stop failed.
     757              :         const PG_STATUS_NOT_RUNNING: i32 = 3;
     758              :         const PG_NO_DATA_DIR: i32 = 4;
     759              :         const PG_STATUS_RUNNING: i32 = 0;
     760            0 :         match status_exitcode.code() {
     761            0 :             Some(PG_STATUS_NOT_RUNNING) => Ok(false),
     762            0 :             Some(PG_NO_DATA_DIR) => Ok(false),
     763            0 :             Some(PG_STATUS_RUNNING) => Ok(true),
     764            0 :             Some(code) => Err(anyhow::anyhow!(
     765            0 :                 "pg_ctl status returned unexpected status code: {:?}",
     766            0 :                 code
     767            0 :             )),
     768            0 :             None => Err(anyhow::anyhow!("pg_ctl status returned no status code")),
     769              :         }
     770            0 :     }
     771              : 
     772            0 :     fn get_claims_for_path(path: &str) -> anyhow::Result<Option<Claims>> {
     773            0 :         let category = match path.find('/') {
     774            0 :             Some(idx) => &path[..idx],
     775            0 :             None => path,
     776              :         };
     777              : 
     778            0 :         match category {
     779            0 :             "status" | "ready" => Ok(None),
     780            0 :             "control" | "debug" => Ok(Some(Claims::new(None, Scope::Admin))),
     781            0 :             "v1" => Ok(Some(Claims::new(None, Scope::PageServerApi))),
     782            0 :             _ => Err(anyhow::anyhow!("Failed to determine claims for {}", path)),
     783              :         }
     784            0 :     }
     785              : 
     786              :     /// Simple HTTP request wrapper for calling into storage controller
     787            0 :     async fn dispatch<RQ, RS>(
     788            0 :         &self,
     789            0 :         method: reqwest::Method,
     790            0 :         path: String,
     791            0 :         body: Option<RQ>,
     792            0 :     ) -> anyhow::Result<RS>
     793            0 :     where
     794            0 :         RQ: Serialize + Sized,
     795            0 :         RS: DeserializeOwned + Sized,
     796            0 :     {
     797            0 :         let response = self.dispatch_inner(method, path, body).await?;
     798            0 :         Ok(response
     799            0 :             .json()
     800            0 :             .await
     801            0 :             .map_err(pageserver_client::mgmt_api::Error::ReceiveBody)?)
     802            0 :     }
     803              : 
     804              :     /// Simple HTTP request wrapper for calling into storage controller
     805            0 :     async fn dispatch_inner<RQ>(
     806            0 :         &self,
     807            0 :         method: reqwest::Method,
     808            0 :         path: String,
     809            0 :         body: Option<RQ>,
     810            0 :     ) -> anyhow::Result<Response>
     811            0 :     where
     812            0 :         RQ: Serialize + Sized,
     813            0 :     {
     814              :         // In the special case of the `storage_controller start` subcommand, we wish
     815              :         // to use the API endpoint of the newly started storage controller in order
     816              :         // to pass the readiness check. In this scenario [`Self::listen_port`] will
     817              :         // be set (see [`Self::start`]).
     818              :         //
     819              :         // Otherwise, we infer the storage controller api endpoint from the configured
     820              :         // control plane API.
     821            0 :         let port = if let Some(port) = self.listen_port.get() {
     822            0 :             *port
     823              :         } else {
     824            0 :             self.env.control_plane_api.port().unwrap()
     825              :         };
     826              : 
     827              :         // The configured URL has the /upcall path prefix for pageservers to use: we will strip that out
     828              :         // for general purpose API access.
     829            0 :         let url = Url::from_str(&format!(
     830            0 :             "{}://{}:{port}/{path}",
     831            0 :             self.env.control_plane_api.scheme(),
     832            0 :             self.env.control_plane_api.host_str().unwrap(),
     833            0 :         ))
     834            0 :         .unwrap();
     835              : 
     836            0 :         let mut builder = self.client.request(method, url);
     837            0 :         if let Some(body) = body {
     838            0 :             builder = builder.json(&body)
     839            0 :         }
     840            0 :         if let Some(private_key) = &self.private_key {
     841            0 :             println!("Getting claims for path {path}");
     842            0 :             if let Some(required_claims) = Self::get_claims_for_path(&path)? {
     843            0 :                 println!("Got claims {required_claims:?} for path {path}");
     844            0 :                 let jwt_token = encode_from_key_file(&required_claims, private_key)?;
     845            0 :                 builder = builder.header(
     846            0 :                     reqwest::header::AUTHORIZATION,
     847            0 :                     format!("Bearer {jwt_token}"),
     848              :                 );
     849            0 :             }
     850            0 :         }
     851              : 
     852            0 :         let response = builder.send().await?;
     853            0 :         let response = response.error_from_body().await?;
     854              : 
     855            0 :         Ok(response)
     856            0 :     }
     857              : 
     858              :     /// Register the safekeepers in the storage controller
     859              :     #[instrument(skip(self))]
     860              :     async fn register_safekeepers(&self) -> anyhow::Result<()> {
     861              :         for sk in self.env.safekeepers.iter() {
     862              :             let sk_id = sk.id;
     863              :             let body = serde_json::json!({
     864              :                 "id": sk_id,
     865              :                 "created_at": "2023-10-25T09:11:25Z",
     866              :                 "updated_at": "2024-08-28T11:32:43Z",
     867              :                 "region_id": "aws-us-east-2",
     868              :                 "host": "127.0.0.1",
     869              :                 "port": sk.pg_port,
     870              :                 "http_port": sk.http_port,
     871              :                 "https_port": sk.https_port,
     872              :                 "version": 5957,
     873              :                 "availability_zone_id": format!("us-east-2b-{sk_id}"),
     874              :             });
     875              :             self.upsert_safekeeper(sk_id, body).await?;
     876              :             self.safekeeper_scheduling_policy(sk_id, SkSchedulingPolicy::Active)
     877              :                 .await?;
     878              :         }
     879              :         Ok(())
     880              :     }
     881              : 
     882              :     /// Call into the attach_hook API, for use before handing out attachments to pageservers
     883              :     #[instrument(skip(self))]
     884              :     pub async fn attach_hook(
     885              :         &self,
     886              :         tenant_shard_id: TenantShardId,
     887              :         pageserver_id: NodeId,
     888              :     ) -> anyhow::Result<Option<u32>> {
     889              :         let request = AttachHookRequest {
     890              :             tenant_shard_id,
     891              :             node_id: Some(pageserver_id),
     892              :             generation_override: None,
     893              :             config: None,
     894              :         };
     895              : 
     896              :         let response = self
     897              :             .dispatch::<_, AttachHookResponse>(
     898              :                 Method::POST,
     899              :                 "debug/v1/attach-hook".to_string(),
     900              :                 Some(request),
     901              :             )
     902              :             .await?;
     903              : 
     904              :         Ok(response.generation)
     905              :     }
     906              : 
     907              :     #[instrument(skip(self))]
     908              :     pub async fn upsert_safekeeper(
     909              :         &self,
     910              :         node_id: NodeId,
     911              :         request: serde_json::Value,
     912              :     ) -> anyhow::Result<()> {
     913              :         let resp = self
     914              :             .dispatch_inner::<serde_json::Value>(
     915              :                 Method::POST,
     916              :                 format!("control/v1/safekeeper/{node_id}"),
     917              :                 Some(request),
     918              :             )
     919              :             .await?;
     920              :         if !resp.status().is_success() {
     921              :             anyhow::bail!(
     922              :                 "setting scheduling policy unsuccessful for safekeeper {node_id}: {}",
     923              :                 resp.status()
     924              :             );
     925              :         }
     926              :         Ok(())
     927              :     }
     928              : 
     929              :     #[instrument(skip(self))]
     930              :     pub async fn safekeeper_scheduling_policy(
     931              :         &self,
     932              :         node_id: NodeId,
     933              :         scheduling_policy: SkSchedulingPolicy,
     934              :     ) -> anyhow::Result<()> {
     935              :         self.dispatch::<SafekeeperSchedulingPolicyRequest, ()>(
     936              :             Method::POST,
     937              :             format!("control/v1/safekeeper/{node_id}/scheduling_policy"),
     938              :             Some(SafekeeperSchedulingPolicyRequest { scheduling_policy }),
     939              :         )
     940              :         .await
     941              :     }
     942              : 
     943              :     #[instrument(skip(self))]
     944              :     pub async fn inspect(
     945              :         &self,
     946              :         tenant_shard_id: TenantShardId,
     947              :     ) -> anyhow::Result<Option<(u32, NodeId)>> {
     948              :         let request = InspectRequest { tenant_shard_id };
     949              : 
     950              :         let response = self
     951              :             .dispatch::<_, InspectResponse>(
     952              :                 Method::POST,
     953              :                 "debug/v1/inspect".to_string(),
     954              :                 Some(request),
     955              :             )
     956              :             .await?;
     957              : 
     958              :         Ok(response.attachment)
     959              :     }
     960              : 
     961              :     #[instrument(skip(self))]
     962              :     pub async fn tenant_create(
     963              :         &self,
     964              :         req: TenantCreateRequest,
     965              :     ) -> anyhow::Result<TenantCreateResponse> {
     966              :         self.dispatch(Method::POST, "v1/tenant".to_string(), Some(req))
     967              :             .await
     968              :     }
     969              : 
     970              :     #[instrument(skip(self))]
     971              :     pub async fn tenant_import(&self, tenant_id: TenantId) -> anyhow::Result<TenantCreateResponse> {
     972              :         self.dispatch::<(), TenantCreateResponse>(
     973              :             Method::POST,
     974              :             format!("debug/v1/tenant/{tenant_id}/import"),
     975              :             None,
     976              :         )
     977              :         .await
     978              :     }
     979              : 
     980              :     #[instrument(skip(self))]
     981              :     pub async fn tenant_locate(&self, tenant_id: TenantId) -> anyhow::Result<TenantLocateResponse> {
     982              :         self.dispatch::<(), _>(
     983              :             Method::GET,
     984              :             format!("debug/v1/tenant/{tenant_id}/locate"),
     985              :             None,
     986              :         )
     987              :         .await
     988              :     }
     989              : 
     990              :     #[instrument(skip_all, fields(node_id=%req.node_id))]
     991              :     pub async fn node_register(&self, req: NodeRegisterRequest) -> anyhow::Result<()> {
     992              :         self.dispatch::<_, ()>(Method::POST, "control/v1/node".to_string(), Some(req))
     993              :             .await
     994              :     }
     995              : 
     996              :     #[instrument(skip_all, fields(node_id=%req.node_id))]
     997              :     pub async fn node_configure(&self, req: NodeConfigureRequest) -> anyhow::Result<()> {
     998              :         self.dispatch::<_, ()>(
     999              :             Method::PUT,
    1000              :             format!("control/v1/node/{}/config", req.node_id),
    1001              :             Some(req),
    1002              :         )
    1003              :         .await
    1004              :     }
    1005              : 
    1006            0 :     pub async fn node_list(&self) -> anyhow::Result<Vec<NodeDescribeResponse>> {
    1007            0 :         self.dispatch::<(), Vec<NodeDescribeResponse>>(
    1008            0 :             Method::GET,
    1009            0 :             "control/v1/node".to_string(),
    1010            0 :             None,
    1011            0 :         )
    1012            0 :         .await
    1013            0 :     }
    1014              : 
    1015              :     #[instrument(skip(self))]
    1016              :     pub async fn ready(&self) -> anyhow::Result<()> {
    1017              :         self.dispatch::<(), ()>(Method::GET, "ready".to_string(), None)
    1018              :             .await
    1019              :     }
    1020              : 
    1021              :     #[instrument(skip_all, fields(%tenant_id, timeline_id=%req.new_timeline_id))]
    1022              :     pub async fn tenant_timeline_create(
    1023              :         &self,
    1024              :         tenant_id: TenantId,
    1025              :         req: TimelineCreateRequest,
    1026              :     ) -> anyhow::Result<TimelineInfo> {
    1027              :         self.dispatch(
    1028              :             Method::POST,
    1029              :             format!("v1/tenant/{tenant_id}/timeline"),
    1030              :             Some(req),
    1031              :         )
    1032              :         .await
    1033              :     }
    1034              : 
    1035            0 :     pub async fn set_tenant_config(&self, req: &TenantConfigRequest) -> anyhow::Result<()> {
    1036            0 :         self.dispatch(Method::PUT, "v1/tenant/config".to_string(), Some(req))
    1037            0 :             .await
    1038            0 :     }
    1039              : }
        

Generated by: LCOV version 2.1-beta