From 7565939dcec4d820bd745f4380625fc177a8129f Mon Sep 17 00:00:00 2001 From: Heikki Linnakangas Date: Tue, 10 Sep 2024 01:29:51 +0300 Subject: [PATCH] Store branch name mappings in separate branches.toml file --- control_plane/src/bin/neon_local.rs | 41 +++++-- control_plane/src/local_env.rs | 164 +++++++++++++++++-------- storage_controller/src/compute_hook.rs | 2 +- test_runner/fixtures/neon_fixtures.py | 11 +- 4 files changed, 151 insertions(+), 67 deletions(-) diff --git a/control_plane/src/bin/neon_local.rs b/control_plane/src/bin/neon_local.rs index f53f433b7c..5a8064cf25 100644 --- a/control_plane/src/bin/neon_local.rs +++ b/control_plane/src/bin/neon_local.rs @@ -90,8 +90,11 @@ fn main() -> Result<()> { handle_init(sub_args).map(Some) } else { // all other commands need an existing config - let mut env = - LocalEnv::load_config(&local_env::base_path()).context("Error loading config")?; + let base_path = local_env::base_path(); + let branch_mappings_path = local_env::branch_mappings_path(); + + let mut env = LocalEnv::load_config(&base_path, Some(&branch_mappings_path)) + .context("Error loading config")?; let original_env = env.clone(); let rt = tokio::runtime::Builder::new_current_thread() @@ -264,7 +267,7 @@ async fn get_timeline_infos( fn get_tenant_id(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> anyhow::Result { if let Some(tenant_id_from_arguments) = parse_tenant_id(sub_match).transpose() { tenant_id_from_arguments - } else if let Some(default_id) = env.default_tenant_id { + } else if let Some(default_id) = env.branch_mappings.default_tenant_id { Ok(default_id) } else { anyhow::bail!("No tenant id. Use --tenant-id, or set a default tenant"); @@ -278,7 +281,7 @@ fn get_tenant_shard_id( ) -> anyhow::Result { if let Some(tenant_id_from_arguments) = parse_tenant_shard_id(sub_match).transpose() { tenant_id_from_arguments - } else if let Some(default_id) = env.default_tenant_id { + } else if let Some(default_id) = env.branch_mappings.default_tenant_id { Ok(TenantShardId::unsharded(default_id)) } else { anyhow::bail!("No tenant shard id. Use --tenant-id, or set a default tenant"); @@ -360,7 +363,6 @@ fn handle_init(init_match: &ArgMatches) -> anyhow::Result { .collect(), pg_distrib_dir: None, neon_distrib_dir: None, - default_tenant_id: TenantId::from_array(std::array::from_fn(|_| 0)), storage_controller: None, control_plane_compute_hook_api: None, } @@ -368,8 +370,12 @@ fn handle_init(init_match: &ArgMatches) -> anyhow::Result { LocalEnv::init(init_conf, force) .context("materialize initial neon_local environment on disk")?; - Ok(LocalEnv::load_config(&local_env::base_path()) - .expect("freshly written config should be loadable")) + let base_path = local_env::base_path(); + let branch_mappings_path = local_env::branch_mappings_path(); + let env = LocalEnv::load_config(&base_path, Some(&branch_mappings_path)) + .expect("freshly written config should be loadable"); + + Ok(env) } /// The default pageserver is the one where CLI tenant/timeline operations are sent by default. @@ -525,14 +531,14 @@ async fn handle_tenant( if create_match.get_flag("set-default") { println!("Setting tenant {tenant_id} as a default one"); - env.default_tenant_id = Some(tenant_id); + env.branch_mappings.default_tenant_id = Some(tenant_id); } } Some(("set-default", set_default_match)) => { let tenant_id = parse_tenant_id(set_default_match)?.context("No tenant id specified")?; println!("Setting tenant {tenant_id} as a default one"); - env.default_tenant_id = Some(tenant_id); + env.branch_mappings.default_tenant_id = Some(tenant_id); } Some(("config", create_match)) => { let tenant_id = get_tenant_id(create_match, env)?; @@ -693,6 +699,7 @@ async fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Re Some(ep_subcommand_data) => ep_subcommand_data, None => bail!("no endpoint subcommand provided"), }; + let mut cplane = ComputeControlPlane::load(env.clone())?; match sub_name { @@ -1360,6 +1367,7 @@ async fn handle_stop_all(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> R async fn try_stop_all(env: &local_env::LocalEnv, immediate: bool) { // Stop all endpoints + // NOTE: This only knows about endpoints in the default endpoints dir match ComputeControlPlane::load(env.clone()) { Ok(cplane) => { for (_k, node) in cplane.endpoints { @@ -1420,6 +1428,18 @@ fn cli() -> Command { .default_value("10s") .required(false); + let branch_mappings_arg = Arg::new("branch-mappings") + .long("branch-mappings") + .help("File holding all branch names. Default is /branches.toml") + .value_parser(value_parser!(PathBuf)) + .required(false); + + let endpoints_dir_arg = Arg::new("endpoints-dir") + .long("endpoints-dir") + .help("Path to directory holding all endpoints. Default is /endpoints") + .value_parser(value_parser!(PathBuf)) + .required(false); + let branch_name_arg = Arg::new("branch-name") .long("branch-name") .help("Name of the branch to be created or used as an alias for other services") @@ -1560,6 +1580,8 @@ fn cli() -> Command { Command::new("Neon CLI") .arg_required_else_help(true) .version(GIT_VERSION) + .arg(branch_mappings_arg) + .arg(endpoints_dir_arg) .subcommand( Command::new("init") .about("Initialize a new Neon repository, preparing configs for services to start with") @@ -1571,7 +1593,6 @@ fn cli() -> Command { .value_parser(value_parser!(PathBuf)) .value_name("config") ) - .arg(pg_version_arg.clone()) .arg(force_arg) ) .subcommand( diff --git a/control_plane/src/local_env.rs b/control_plane/src/local_env.rs index 5dbc3bcbbc..4508fe3f7d 100644 --- a/control_plane/src/local_env.rs +++ b/control_plane/src/local_env.rs @@ -46,6 +46,11 @@ pub struct LocalEnv { // must be an absolute path. If the env var is not set, $PWD/.neon is used. pub base_data_dir: PathBuf, + // Similarly, path to branch mappings file. Not stored in the config file but + // read from the NEON_BRANCH_MAPPINGS env variable. "None" means no mappings + // are loaded, and they cannot be saved either. + pub branch_name_mappings_path: Option, + // Path to postgres distribution. It's expected that "bin", "include", // "lib", "share" from postgres distribution are there. If at some point // in time we will be able to run against vanilla postgres we may split that @@ -55,10 +60,6 @@ pub struct LocalEnv { // Path to pageserver binary. pub neon_distrib_dir: PathBuf, - // Default tenant ID to use with the 'neon_local' command line utility, when - // --tenant_id is not explicitly specified. - pub default_tenant_id: Option, - // used to issue tokens during e.g pg start pub private_key_path: PathBuf, @@ -82,11 +83,7 @@ pub struct LocalEnv { // storage controller's configuration. pub control_plane_compute_hook_api: Option, - /// Keep human-readable aliases in memory (and persist them to config), to hide ZId hex strings from the user. - // A `HashMap>` would be more appropriate here, - // but deserialization into a generic toml object as `toml::Value::try_from` fails with an error. - // https://toml.io/en/v1.0.0 does not contain a concept of "a table inside another table". - pub branch_name_mappings: HashMap>, + pub branch_mappings: BranchMappings, } /// On-disk state stored in `.neon/config`. @@ -95,7 +92,6 @@ pub struct LocalEnv { pub struct OnDiskConfig { pub pg_distrib_dir: PathBuf, pub neon_distrib_dir: PathBuf, - pub default_tenant_id: Option, pub private_key_path: PathBuf, pub broker: NeonBroker, pub storage_controller: NeonStorageControllerConf, @@ -107,7 +103,20 @@ pub struct OnDiskConfig { pub safekeepers: Vec, pub control_plane_api: Option, pub control_plane_compute_hook_api: Option, - branch_name_mappings: HashMap>, +} + +#[derive(PartialEq, Eq, Clone, Debug, Default, Serialize, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct BranchMappings { + // Default tenant ID to use with the 'neon_local' command line utility, when + // --tenant_id is not explicitly specified. This comes from the branches. + pub default_tenant_id: Option, + + /// Keep human-readable aliases in memory (and persist them to config XXX), to hide ZId hex strings from the user. + // A `HashMap>` would be more appropriate here, + // but deserialization into a generic toml object as `toml::Value::try_from` fails with an error. + // https://toml.io/en/v1.0.0 does not contain a concept of "a table inside another table". + pub mappings: HashMap>, } fn fail_if_pageservers_field_specified<'de, D>(_: D) -> Result, D::Error> @@ -128,7 +137,6 @@ pub struct NeonLocalInitConf { pub pg_distrib_dir: Option, // TODO: do we need this? Seems unused pub neon_distrib_dir: Option, - pub default_tenant_id: TenantId, pub broker: NeonBroker, pub storage_controller: Option, pub pageservers: Vec, @@ -443,7 +451,8 @@ impl LocalEnv { timeline_id: TimelineId, ) -> anyhow::Result<()> { let existing_values = self - .branch_name_mappings + .branch_mappings + .mappings .entry(branch_name.clone()) .or_default(); @@ -468,7 +477,8 @@ impl LocalEnv { branch_name: &str, tenant_id: TenantId, ) -> Option { - self.branch_name_mappings + self.branch_mappings + .mappings .get(branch_name)? .iter() .find(|(mapped_tenant_id, _)| mapped_tenant_id == &tenant_id) @@ -477,7 +487,8 @@ impl LocalEnv { } pub fn timeline_name_mappings(&self) -> HashMap { - self.branch_name_mappings + self.branch_mappings + .mappings .iter() .flat_map(|(name, tenant_timelines)| { tenant_timelines.iter().map(|&(tenant_id, timeline_id)| { @@ -488,7 +499,10 @@ impl LocalEnv { } /// Construct `Self` from on-disk state. - pub fn load_config(repopath: &Path) -> anyhow::Result { + pub fn load_config( + repopath: &Path, + branch_name_mappings_path: Option<&Path>, + ) -> anyhow::Result { if !repopath.exists() { bail!( "Neon config is not found in {}. You need to run 'neon_local init' first", @@ -498,37 +512,43 @@ impl LocalEnv { // TODO: check that it looks like a neon repository - // load and parse file + // load and parse config file let config_file_contents = fs::read_to_string(repopath.join("config"))?; let on_disk_config: OnDiskConfig = toml::from_str(config_file_contents.as_str())?; - let mut env = { - let OnDiskConfig { - pg_distrib_dir, - neon_distrib_dir, - default_tenant_id, - private_key_path, - broker, - storage_controller, - pageservers, - safekeepers, - control_plane_api, - control_plane_compute_hook_api, - branch_name_mappings, - } = on_disk_config; - LocalEnv { - base_data_dir: repopath.to_owned(), - pg_distrib_dir, - neon_distrib_dir, - default_tenant_id, - private_key_path, - broker, - storage_controller, - pageservers, - safekeepers, - control_plane_api, - control_plane_compute_hook_api, - branch_name_mappings, - } + let OnDiskConfig { + pg_distrib_dir, + neon_distrib_dir, + private_key_path, + broker, + storage_controller, + pageservers, + safekeepers, + control_plane_api, + control_plane_compute_hook_api, + } = on_disk_config; + + // load and parse "branches.toml" file + let branch_mappings = if let Some(path) = branch_name_mappings_path { + let contents = fs::read_to_string(path) + .context(format!("load branch mappings file {}", path.display()))?; + toml::from_str::(contents.as_str())? + } else { + BranchMappings::default() + }; + + let mut env = LocalEnv { + base_data_dir: repopath.to_owned(), + branch_name_mappings_path: branch_name_mappings_path.map(|p| p.to_owned()), + pg_distrib_dir, + neon_distrib_dir, + private_key_path, + broker, + storage_controller, + pageservers, + safekeepers, + control_plane_api, + control_plane_compute_hook_api, + branch_mappings, }; // The source of truth for pageserver configuration is the pageserver.toml. @@ -618,7 +638,6 @@ impl LocalEnv { &OnDiskConfig { pg_distrib_dir: self.pg_distrib_dir.clone(), neon_distrib_dir: self.neon_distrib_dir.clone(), - default_tenant_id: self.default_tenant_id, private_key_path: self.private_key_path.clone(), broker: self.broker.clone(), storage_controller: self.storage_controller.clone(), @@ -626,9 +645,20 @@ impl LocalEnv { safekeepers: self.safekeepers.clone(), control_plane_api: self.control_plane_api.clone(), control_plane_compute_hook_api: self.control_plane_compute_hook_api.clone(), - branch_name_mappings: self.branch_name_mappings.clone(), }, - ) + )?; + + if let Some(path) = &self.branch_name_mappings_path { + Self::persist_branches_impl(path, &self.branch_mappings)?; + } else { + if !self.branch_mappings.mappings.is_empty() { + tracing::warn!("command created a branch mapping, but it was not saved because no mappings file was configured") + } else if self.branch_mappings.default_tenant_id.is_some() { + tracing::warn!("command created a tenant default, but it was not saved because no mappings file was configured") + } + } + + Ok(()) } pub fn persist_config_impl(base_path: &Path, config: &OnDiskConfig) -> anyhow::Result<()> { @@ -642,6 +672,19 @@ impl LocalEnv { }) } + pub fn persist_branches_impl( + branch_name_mappings_path: &Path, + branch_mappings: &BranchMappings, + ) -> anyhow::Result<()> { + let content = &toml::to_string_pretty(branch_mappings)?; + fs::write(branch_name_mappings_path, content).with_context(|| { + format!( + "Failed to write branch information into path '{}'", + branch_name_mappings_path.display() + ) + }) + } + // this function is used only for testing purposes in CLI e g generate tokens during init pub fn generate_auth_token(&self, claims: &Claims) -> anyhow::Result { let private_key_path = self.get_private_key_path(); @@ -702,7 +745,6 @@ impl LocalEnv { let NeonLocalInitConf { pg_distrib_dir, neon_distrib_dir, - default_tenant_id, broker, storage_controller, pageservers, @@ -746,9 +788,9 @@ impl LocalEnv { // TODO: refactor to avoid this, LocalEnv should only be constructed from on-disk state let env = LocalEnv { base_data_dir: base_path.clone(), + branch_name_mappings_path: Some(base_path.join("branches.toml")), pg_distrib_dir, neon_distrib_dir, - default_tenant_id: Some(default_tenant_id), private_key_path, broker, storage_controller: storage_controller.unwrap_or_default(), @@ -756,10 +798,10 @@ impl LocalEnv { safekeepers, control_plane_api: control_plane_api.unwrap_or_default(), control_plane_compute_hook_api: control_plane_compute_hook_api.unwrap_or_default(), - branch_name_mappings: Default::default(), + branch_mappings: Default::default(), }; - // create endpoints dir + // create the default endpoints dir fs::create_dir_all(env.endpoints_path())?; // create safekeeper dirs @@ -806,6 +848,24 @@ pub fn base_path() -> PathBuf { path } +pub fn branch_mappings_path() -> PathBuf { + let path = match std::env::var_os("NEON_BRANCH_MAPPINGS") { + Some(val) => { + let path = PathBuf::from(val); + + // a relative path is relative to repo dir + if !path.is_absolute() { + base_path().join(path) + } else { + path + } + } + None => base_path().join("branches.toml"), + }; + assert!(path.is_absolute()); + path +} + /// Generate a public/private key pair for JWT authentication fn generate_auth_keys(private_key_path: &Path, public_key_path: &Path) -> anyhow::Result<()> { // Generate the key pair diff --git a/storage_controller/src/compute_hook.rs b/storage_controller/src/compute_hook.rs index c46539485c..91ce48929c 100644 --- a/storage_controller/src/compute_hook.rs +++ b/storage_controller/src/compute_hook.rs @@ -292,7 +292,7 @@ impl ComputeHook { ); return Ok(()); }; - let env = match LocalEnv::load_config(repo_dir) { + let env = match LocalEnv::load_config(repo_dir, None) { Ok(e) => e, Err(e) => { tracing::warn!("Couldn't load neon_local config, skipping compute update ({e})"); diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py index 7386e25740..a5fcb70c79 100644 --- a/test_runner/fixtures/neon_fixtures.py +++ b/test_runner/fixtures/neon_fixtures.py @@ -654,6 +654,10 @@ class NeonEnvBuilder: with snapshot_config_toml.open("r") as f: snapshot_config = toml.load(f) + snapshot_branches_toml = repo_dir / "branches.toml" + with snapshot_branches_toml.open("r") as f: + snapshot_branch_mappings = toml.load(f) + self.initial_tenant = TenantId(snapshot_config["default_tenant_id"]) self.initial_timeline = TimelineId( dict(snapshot_config["branch_name_mappings"][DEFAULT_BRANCH_NAME])[ @@ -729,9 +733,6 @@ class NeonEnvBuilder: with (self.repo_dir / "config").open("r") as f: config = toml.load(f) - config["default_tenant_id"] = snapshot_config["default_tenant_id"] - config["branch_name_mappings"] = snapshot_config["branch_name_mappings"] - # Update the config with new neon + postgres path in case of compat test config["pg_distrib_dir"] = str(self.pg_distrib_dir) config["neon_distrib_dir"] = str(self.neon_binpath) @@ -739,6 +740,9 @@ class NeonEnvBuilder: with (self.repo_dir / "config").open("w") as f: toml.dump(config, f) + with (self.repo_dir / "branches.toml").open("w") as f: + toml.dump(snapshot_branch_mappings, f) + return self.env def overlay_mount(self, ident: str, srcdir: Path, dstdir: Path): @@ -1118,7 +1122,6 @@ class NeonEnv: # Create the neon_local's `NeonLocalInitConf` cfg: Dict[str, Any] = { - "default_tenant_id": str(self.initial_tenant), "broker": { "listen_addr": self.broker.listen_addr(), },