LCOV - code coverage report
Current view: top level - control_plane/src - storage_controller.rs (source / functions) Coverage Total Hit
Test: 07bee600374ccd486c69370d0972d9035964fe68.info Lines: 0.0 % 534 0
Test Date: 2025-02-20 13:11:02 Functions: 0.0 % 76 0

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

Generated by: LCOV version 2.1-beta