neon/control_plane/src/storage.rs

use std::collections::HashMap;
use std::net::TcpStream;
use std::path::PathBuf;
use std::process::Command;
use std::thread;
use std::time::Duration;

use anyhow::{anyhow, bail, Result};
use nix::sys::signal::{kill, Signal};
use nix::unistd::Pid;
use postgres::{Config, NoTls};
use zenith_utils::postgres_backend::AuthType;
use zenith_utils::zid::ZTenantId;

use crate::local_env::LocalEnv;
use crate::read_pidfile;
use pageserver::branches::BranchInfo;
use zenith_utils::connstring::connection_address;

//
// Control routines for pageserver.
//
// Used in CLI and tests.
//
#[derive(Debug)]
pub struct PageServerNode {
    pub kill_on_exit: bool,
    pub connection_config: Config,
    pub env: LocalEnv,
}

impl PageServerNode {
    pub fn from_env(env: &LocalEnv) -> PageServerNode {
        let password = if matches!(env.auth_type, AuthType::ZenithJWT) {
            &env.auth_token
        } else {
            ""
        };

        PageServerNode {
            kill_on_exit: false,
            connection_config: Self::default_config(password), // default
            env: env.clone(),
        }
    }

    fn default_config(password: &str) -> Config {
        format!("postgresql://no_user:{}@localhost:64000/no_db", password)
            .parse()
            .unwrap()
    }

    pub fn init(
        &self,
        create_tenant: Option<&str>,
        enable_auth: bool,
        repository_format: Option<&str>,
    ) -> Result<()> {
        let mut cmd = Command::new(self.env.pageserver_bin()?);
        let mut args = vec![
            "--init",
            "-D",
            self.env.base_data_dir.to_str().unwrap(),
            "--postgres-distrib",
            self.env.pg_distrib_dir.to_str().unwrap(),
        ];

        if enable_auth {
            args.extend(&["--auth-validation-public-key-path", "auth_public_key.pem"]);
            args.extend(&["--auth-type", "ZenithJWT"]);
        }

        if let Some(repo_format) = repository_format {
            args.extend(&["--repository-format", repo_format]);
        }

        create_tenant.map(|tenantid| args.extend(&["--create-tenant", tenantid]));
        let status = cmd
            .args(args)
            .env_clear()
            .env("RUST_BACKTRACE", "1")
            .status()
            .expect("pageserver init failed");

        if status.success() {
            Ok(())
        } else {
            Err(anyhow!("pageserver init failed"))
        }
    }

    pub fn repo_path(&self) -> PathBuf {
        self.env.pageserver_data_dir()
    }

    pub fn pid_file(&self) -> PathBuf {
        self.repo_path().join("pageserver.pid")
    }

    pub fn start(&self) -> Result<()> {
        println!(
            "Starting pageserver at '{}' in {}",
            connection_address(&self.connection_config),
            self.repo_path().display()
        );

        let mut cmd = Command::new(self.env.pageserver_bin()?);
        cmd.args(&["-D", self.repo_path().to_str().unwrap()])
            .arg("-d")
            .env_clear()
            .env("RUST_BACKTRACE", "1");

        if !cmd.status()?.success() {
            bail!(
                "Pageserver failed to start. See '{}' for details.",
                self.repo_path().join("pageserver.log").display()
            );
        }

        // It takes a while for the page server to start up. Wait until it is
        // open for business.
        for retries in 1..15 {
            match self.page_server_psql_client() {
                Ok(_) => {
                    println!("Pageserver started");
                    return Ok(());
                }
                Err(err) => {
                    println!(
                        "Pageserver not responding yet, err {} retrying ({})...",
                        err, retries
                    );
                    thread::sleep(Duration::from_secs(1));
                }
            }
        }
        bail!("pageserver failed to start");
    }

    pub fn stop(&self) -> Result<()> {
        let pid = read_pidfile(&self.pid_file())?;
        let pid = Pid::from_raw(pid);
        if kill(pid, Signal::SIGTERM).is_err() {
            bail!("Failed to kill pageserver with pid {}", pid);
        }

        // wait for pageserver stop
        let address = connection_address(&self.connection_config);
        for _ in 0..5 {
            let stream = TcpStream::connect(&address);
            thread::sleep(Duration::from_secs(1));
            if let Err(_e) = stream {
                println!("Pageserver stopped");
                return Ok(());
            }
            println!("Stopping pageserver on {}", address);
        }

        bail!("Failed to stop pageserver with pid {}", pid);
    }

    pub fn page_server_psql(&self, sql: &str) -> Vec<postgres::SimpleQueryMessage> {
        let mut client = self.connection_config.connect(NoTls).unwrap();

        println!("Pageserver query: '{}'", sql);
        client.simple_query(sql).unwrap()
    }

    pub fn page_server_psql_client(&self) -> Result<postgres::Client, postgres::Error> {
        self.connection_config.connect(NoTls)
    }

    pub fn tenants_list(&self) -> Result<Vec<String>> {
        let mut client = self.page_server_psql_client()?;
        let query_result = client.simple_query("tenant_list")?;
        let tenants_json = query_result
            .first()
            .map(|msg| match msg {
                postgres::SimpleQueryMessage::Row(row) => row.get(0),
                _ => None,
            })
            .flatten()
            .ok_or_else(|| anyhow!("missing tenants"))?;

        Ok(serde_json::from_str(tenants_json)?)
    }

    pub fn tenant_create(&self, tenantid: &ZTenantId) -> Result<()> {
        let mut client = self.page_server_psql_client()?;
        client.simple_query(format!("tenant_create {}", tenantid).as_str())?;
        Ok(())
    }

    pub fn branches_list(&self, tenantid: &ZTenantId) -> Result<Vec<BranchInfo>> {
        let mut client = self.page_server_psql_client()?;
        let query_result = client.simple_query(&format!("branch_list {}", tenantid))?;
        let branches_json = query_result
            .first()
            .map(|msg| match msg {
                postgres::SimpleQueryMessage::Row(row) => row.get(0),
                _ => None,
            })
            .flatten()
            .ok_or_else(|| anyhow!("missing branches"))?;

        Ok(serde_json::from_str(branches_json)?)
    }

    pub fn branch_create(
        &self,
        branch_name: &str,
        startpoint: &str,
        tenantid: &ZTenantId,
    ) -> Result<BranchInfo> {
        let mut client = self.page_server_psql_client()?;
        let query_result = client.simple_query(
            format!("branch_create {} {} {}", tenantid, branch_name, startpoint).as_str(),
        )?;

        let branch_json = query_result
            .first()
            .map(|msg| match msg {
                postgres::SimpleQueryMessage::Row(row) => row.get(0),
                _ => None,
            })
            .flatten()
            .ok_or_else(|| anyhow!("missing branch"))?;

        let res: BranchInfo = serde_json::from_str(branch_json).map_err(|e| {
            anyhow!(
                "failed to parse branch_create response: {}: {}",
                branch_json,
                e
            )
        })?;

        Ok(res)
    }

    // TODO: make this a separate request type and avoid loading all the branches
    pub fn branch_get_by_name(
        &self,
        tenantid: &ZTenantId,
        branch_name: &str,
    ) -> Result<BranchInfo> {
        let branch_infos = self.branches_list(tenantid)?;
        let branche_by_name: Result<HashMap<String, BranchInfo>> = branch_infos
            .into_iter()
            .map(|branch_info| Ok((branch_info.name.clone(), branch_info)))
            .collect();
        let branche_by_name = branche_by_name?;

        let branch = branche_by_name
            .get(branch_name)
            .ok_or_else(|| anyhow!("Branch {} not found", branch_name))?;

        Ok(branch.clone())
    }
}

impl Drop for PageServerNode {
    fn drop(&mut self) {
        if self.kill_on_exit {
            let _ = self.stop();
        }
    }
}