From 24d0395f66f20421d138e074321bb9592a2191a7 Mon Sep 17 00:00:00 2001 From: John Spray Date: Thu, 7 Dec 2023 14:24:04 +0000 Subject: [PATCH] neon_local: update various TenantId uses to TenantShardId --- control_plane/src/attachment_service.rs | 15 +++--- control_plane/src/bin/attachment_service.rs | 57 +++++---------------- control_plane/src/bin/neon_local.rs | 57 ++++++++++++++++----- control_plane/src/pageserver.rs | 4 +- control_plane/src/tenant_migration.rs | 46 +++++++++-------- libs/pageserver_api/src/models.rs | 2 +- pageserver/client/src/mgmt_api.rs | 16 ++++-- 7 files changed, 105 insertions(+), 92 deletions(-) diff --git a/control_plane/src/attachment_service.rs b/control_plane/src/attachment_service.rs index 731c05809e..aec57e75b3 100644 --- a/control_plane/src/attachment_service.rs +++ b/control_plane/src/attachment_service.rs @@ -1,9 +1,10 @@ use crate::{background_process, local_env::LocalEnv}; use anyhow::anyhow; use camino::Utf8PathBuf; +use pageserver_api::shard::TenantShardId; use serde::{Deserialize, Serialize}; use std::{path::PathBuf, process::Child}; -use utils::id::{NodeId, TenantId}; +use utils::id::NodeId; pub struct AttachmentService { env: LocalEnv, @@ -16,7 +17,7 @@ const COMMAND: &str = "attachment_service"; #[derive(Serialize, Deserialize)] pub struct AttachHookRequest { - pub tenant_id: TenantId, + pub tenant_shard_id: TenantShardId, pub node_id: Option, } @@ -27,7 +28,7 @@ pub struct AttachHookResponse { #[derive(Serialize, Deserialize)] pub struct InspectRequest { - pub tenant_id: TenantId, + pub tenant_shard_id: TenantShardId, } #[derive(Serialize, Deserialize)] @@ -87,7 +88,7 @@ impl AttachmentService { /// Call into the attach_hook API, for use before handing out attachments to pageservers pub async fn attach_hook( &self, - tenant_id: TenantId, + tenant_shard_id: TenantShardId, pageserver_id: NodeId, ) -> anyhow::Result> { use hyper::StatusCode; @@ -101,7 +102,7 @@ impl AttachmentService { .unwrap(); let request = AttachHookRequest { - tenant_id, + tenant_shard_id, node_id: Some(pageserver_id), }; @@ -114,7 +115,7 @@ impl AttachmentService { Ok(response.gen) } - pub async fn inspect(&self, tenant_id: TenantId) -> anyhow::Result> { + pub async fn inspect(&self, tenant_id: TenantShardId) -> anyhow::Result> { use hyper::StatusCode; let url = self @@ -125,7 +126,7 @@ impl AttachmentService { .join("inspect") .unwrap(); - let request = InspectRequest { tenant_id }; + let request = InspectRequest { tenant_shard_id }; let response = self.client.post(url).json(&request).send().await?; if response.status() != StatusCode::OK { diff --git a/control_plane/src/bin/attachment_service.rs b/control_plane/src/bin/attachment_service.rs index e50c8fbba0..bda02eeec4 100644 --- a/control_plane/src/bin/attachment_service.rs +++ b/control_plane/src/bin/attachment_service.rs @@ -6,13 +6,13 @@ /// use anyhow::anyhow; use clap::Parser; -use hex::FromHex; use hyper::StatusCode; use hyper::{Body, Request, Response}; use pageserver_api::shard::TenantShardId; use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; use std::path::{Path, PathBuf}; -use std::{collections::HashMap, sync::Arc}; +use std::sync::Arc; use utils::http::endpoint::request_span; use utils::logging::{self, LogFormat}; use utils::signals::{ShutdownSignals, Signal}; @@ -24,7 +24,7 @@ use utils::{ json::{json_request, json_response}, RequestExt, RouterBuilder, }, - id::{NodeId, TenantId}, + id::NodeId, tcp_listener, }; @@ -61,39 +61,10 @@ struct TenantState { generation: u32, } -fn to_hex_map(input: &HashMap, serializer: S) -> Result -where - S: serde::Serializer, - V: Clone + Serialize, -{ - let transformed = input.iter().map(|(k, v)| (hex::encode(k), v.clone())); - - transformed - .collect::>() - .serialize(serializer) -} - -fn from_hex_map<'de, D, V>(deserializer: D) -> Result, D::Error> -where - D: serde::de::Deserializer<'de>, - V: Deserialize<'de>, -{ - let hex_map = HashMap::::deserialize(deserializer)?; - hex_map - .into_iter() - .map(|(k, v)| { - TenantId::from_hex(k) - .map(|k| (k, v)) - .map_err(serde::de::Error::custom) - }) - .collect() -} - // Top level state available to all HTTP handlers #[derive(Serialize, Deserialize)] struct PersistentState { - #[serde(serialize_with = "to_hex_map", deserialize_with = "from_hex_map")] - tenants: HashMap, + tenants: BTreeMap, #[serde(skip)] path: PathBuf, @@ -127,7 +98,7 @@ impl PersistentState { { tracing::info!("Will create state file at {}", path.display()); Self { - tenants: HashMap::new(), + tenants: BTreeMap::new(), path: path.to_owned(), } } @@ -174,8 +145,7 @@ async fn handle_re_attach(mut req: Request) -> Result, ApiE if state.pageserver == Some(reattach_req.node_id) { state.generation += 1; response.tenants.push(ReAttachResponseTenant { - // TODO(sharding): make this shard-aware - id: TenantShardId::unsharded(*t), + id: *t, gen: state.generation, }); } @@ -198,8 +168,7 @@ async fn handle_validate(mut req: Request) -> Result, ApiEr }; for req_tenant in validate_req.tenants { - // TODO(sharding): make this shard-aware - if let Some(tenant_state) = locked.tenants.get(&req_tenant.id.tenant_id) { + if let Some(tenant_state) = locked.tenants.get(&req_tenant.id) { let valid = tenant_state.generation == req_tenant.gen; tracing::info!( "handle_validate: {}(gen {}): valid={valid} (latest {})", @@ -227,7 +196,7 @@ async fn handle_attach_hook(mut req: Request) -> Result, Ap let tenant_state = locked .tenants - .entry(attach_req.tenant_id) + .entry(attach_req.tenant_shard_id) .or_insert_with(|| TenantState { pageserver: attach_req.node_id, generation: 0, @@ -236,21 +205,21 @@ async fn handle_attach_hook(mut req: Request) -> Result, Ap if let Some(attaching_pageserver) = attach_req.node_id.as_ref() { tenant_state.generation += 1; tracing::info!( - tenant_id = %attach_req.tenant_id, + tenant_id = %attach_req.tenant_shard_id, ps_id = %attaching_pageserver, generation = %tenant_state.generation, "issuing", ); } else if let Some(ps_id) = tenant_state.pageserver { tracing::info!( - tenant_id = %attach_req.tenant_id, + tenant_id = %attach_req.tenant_shard_id, %ps_id, generation = %tenant_state.generation, "dropping", ); } else { tracing::info!( - tenant_id = %attach_req.tenant_id, + tenant_id = %attach_req.tenant_shard_id, "no-op: tenant already has no pageserver"); } tenant_state.pageserver = attach_req.node_id; @@ -258,7 +227,7 @@ async fn handle_attach_hook(mut req: Request) -> Result, Ap tracing::info!( "handle_attach_hook: tenant {} set generation {}, pageserver {}", - attach_req.tenant_id, + attach_req.tenant_shard_id, tenant_state.generation, attach_req.node_id.unwrap_or(utils::id::NodeId(0xfffffff)) ); @@ -278,7 +247,7 @@ async fn handle_inspect(mut req: Request) -> Result, ApiErr let state = get_state(&req).inner.clone(); let locked = state.write().await; - let tenant_state = locked.tenants.get(&inspect_req.tenant_id); + let tenant_state = locked.tenants.get(&inspect_req.tenant_shard_id); json_response( StatusCode::OK, diff --git a/control_plane/src/bin/neon_local.rs b/control_plane/src/bin/neon_local.rs index 03e69010f7..b4cea14725 100644 --- a/control_plane/src/bin/neon_local.rs +++ b/control_plane/src/bin/neon_local.rs @@ -16,6 +16,7 @@ use control_plane::safekeeper::SafekeeperNode; use control_plane::tenant_migration::migrate_tenant; use control_plane::{broker, local_env}; use pageserver_api::models::TimelineInfo; +use pageserver_api::shard::TenantShardId; use pageserver_api::{ DEFAULT_HTTP_LISTEN_PORT as DEFAULT_PAGESERVER_HTTP_PORT, DEFAULT_PG_LISTEN_PORT as DEFAULT_PAGESERVER_PG_PORT, @@ -276,10 +277,10 @@ fn print_timeline( /// Connects to the pageserver to query this information. async fn get_timeline_infos( env: &local_env::LocalEnv, - tenant_id: &TenantId, + tenant_shard_id: &TenantShardId, ) -> Result> { Ok(get_default_pageserver(env) - .timeline_list(tenant_id) + .timeline_list(tenant_shard_id) .await? .into_iter() .map(|timeline_info| (timeline_info.timeline_id, timeline_info)) @@ -297,6 +298,20 @@ fn get_tenant_id(sub_match: &ArgMatches, env: &local_env::LocalEnv) -> anyhow::R } } +// Helper function to parse --tenant_id option, for commands that accept a shard suffix +fn get_tenant_shard_id( + sub_match: &ArgMatches, + env: &local_env::LocalEnv, +) -> 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 { + Ok(TenantShardId::unsharded(default_id)) + } else { + anyhow::bail!("No tenant shard id. Use --tenant-id, or set a default tenant"); + } +} + fn parse_tenant_id(sub_match: &ArgMatches) -> anyhow::Result> { sub_match .get_one::("tenant-id") @@ -305,6 +320,14 @@ fn parse_tenant_id(sub_match: &ArgMatches) -> anyhow::Result> { .context("Failed to parse tenant id from the argument string") } +fn parse_tenant_shard_id(sub_match: &ArgMatches) -> anyhow::Result> { + sub_match + .get_one::("tenant-id") + .map(|id_str| TenantShardId::from_str(id_str)) + .transpose() + .context("Failed to parse tenant shard id from the argument string") +} + fn parse_timeline_id(sub_match: &ArgMatches) -> anyhow::Result> { sub_match .get_one::("timeline-id") @@ -399,12 +422,15 @@ async fn handle_tenant( // If tenant ID was not specified, generate one let tenant_id = parse_tenant_id(create_match)?.unwrap_or_else(TenantId::generate); + // TODO(sharding): create multiple shards if requested + let tenant_shard_id = TenantShardId::unsharded(tenant_id); + let generation = if env.control_plane_api.is_some() { // We must register the tenant with the attachment service, so // that when the pageserver restarts, it will be re-attached. let attachment_service = AttachmentService::from_env(env); attachment_service - .attach_hook(tenant_id, pageserver.conf.id) + .attach_hook(tenant_shard_id, pageserver.conf.id) .await? } else { None @@ -470,12 +496,12 @@ async fn handle_tenant( println!("tenant {tenant_id} successfully configured on the pageserver"); } Some(("migrate", matches)) => { - let tenant_id = get_tenant_id(matches, env)?; + let tenant_shard_id = get_tenant_shard_id(matches, env)?; let new_pageserver = get_pageserver(env, matches)?; let new_pageserver_id = new_pageserver.conf.id; - migrate_tenant(env, tenant_id, new_pageserver).await?; - println!("tenant {tenant_id} migrated to {}", new_pageserver_id); + migrate_tenant(env, tenant_shard_id, new_pageserver).await?; + println!("tenant {tenant_shard_id} migrated to {}", new_pageserver_id); } Some((sub_name, _)) => bail!("Unexpected tenant subcommand '{}'", sub_name), @@ -489,8 +515,10 @@ async fn handle_timeline(timeline_match: &ArgMatches, env: &mut local_env::Local match timeline_match.subcommand() { Some(("list", list_match)) => { - let tenant_id = get_tenant_id(list_match, env)?; - let timelines = pageserver.timeline_list(&tenant_id).await?; + // TODO(sharding): this command shouldn't have to specify a shard ID: we should ask the attachment service + // where shard 0 is attached, and query there. + let tenant_shard_id = get_tenant_shard_id(list_match, env)?; + let timelines = pageserver.timeline_list(&tenant_shard_id).await?; print_timelines_tree(timelines, env.timeline_name_mappings())?; } Some(("create", create_match)) => { @@ -635,8 +663,10 @@ async fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Re match sub_name { "list" => { - let tenant_id = get_tenant_id(sub_args, env)?; - let timeline_infos = get_timeline_infos(env, &tenant_id) + // TODO(sharding): this command shouldn't have to specify a shard ID: we should ask the attachment service + // where shard 0 is attached, and query there. + let tenant_shard_id = get_tenant_shard_id(sub_args, env)?; + let timeline_infos = get_timeline_infos(env, &tenant_shard_id) .await .unwrap_or_else(|e| { eprintln!("Failed to load timeline info: {}", e); @@ -661,7 +691,7 @@ async fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Re for (endpoint_id, endpoint) in cplane .endpoints .iter() - .filter(|(_, endpoint)| endpoint.tenant_id == tenant_id) + .filter(|(_, endpoint)| endpoint.tenant_id == tenant_shard_id.tenant_id) { let lsn_str = match endpoint.mode { ComputeMode::Static(lsn) => { @@ -680,7 +710,10 @@ async fn handle_endpoint(ep_match: &ArgMatches, env: &local_env::LocalEnv) -> Re }; let branch_name = timeline_name_mappings - .get(&TenantTimelineId::new(tenant_id, endpoint.timeline_id)) + .get(&TenantTimelineId::new( + tenant_shard_id.tenant_id, + endpoint.timeline_id, + )) .map(|name| name.as_str()) .unwrap_or("?"); diff --git a/control_plane/src/pageserver.rs b/control_plane/src/pageserver.rs index 7d490016bf..6dfe1934fc 100644 --- a/control_plane/src/pageserver.rs +++ b/control_plane/src/pageserver.rs @@ -471,13 +471,13 @@ impl PageServerNode { pub async fn location_config( &self, - tenant_id: TenantId, + tenant_shard_id: TenantShardId, config: LocationConfig, flush_ms: Option, ) -> anyhow::Result<()> { Ok(self .http_client - .location_config(tenant_id, config, flush_ms) + .location_config(tenant_shard_id, config, flush_ms) .await?) } diff --git a/control_plane/src/tenant_migration.rs b/control_plane/src/tenant_migration.rs index 79df108896..9aa15b1cb2 100644 --- a/control_plane/src/tenant_migration.rs +++ b/control_plane/src/tenant_migration.rs @@ -11,19 +11,17 @@ use crate::{ use pageserver_api::models::{ LocationConfig, LocationConfigMode, LocationConfigSecondary, TenantConfig, }; +use pageserver_api::shard::TenantShardId; use std::collections::HashMap; use std::time::Duration; -use utils::{ - id::{TenantId, TimelineId}, - lsn::Lsn, -}; +use utils::{id::TimelineId, lsn::Lsn}; /// Given an attached pageserver, retrieve the LSN for all timelines async fn get_lsns( - tenant_id: TenantId, + tenant_shard_id: TenantShardId, pageserver: &PageServerNode, ) -> anyhow::Result> { - let timelines = pageserver.timeline_list(&tenant_id).await?; + let timelines = pageserver.timeline_list(&tenant_shard_id).await?; Ok(timelines .into_iter() .map(|t| (t.timeline_id, t.last_record_lsn)) @@ -33,12 +31,12 @@ async fn get_lsns( /// Wait for the timeline LSNs on `pageserver` to catch up with or overtake /// `baseline`. async fn await_lsn( - tenant_id: TenantId, + tenant_shard_id: TenantShardId, pageserver: &PageServerNode, baseline: HashMap, ) -> anyhow::Result<()> { loop { - let latest = match get_lsns(tenant_id, pageserver).await { + let latest = match get_lsns(tenant_shard_id, pageserver).await { Ok(l) => l, Err(e) => { println!( @@ -86,7 +84,7 @@ async fn await_lsn( /// - reconfigure compute endpoints to point to new attached pageserver pub async fn migrate_tenant( env: &LocalEnv, - tenant_id: TenantId, + tenant_shard_id: TenantShardId, dest_ps: PageServerNode, ) -> anyhow::Result<()> { // Get a new generation @@ -108,7 +106,7 @@ pub async fn migrate_tenant( } } - let previous = attachment_service.inspect(tenant_id).await?; + let previous = attachment_service.inspect(tenant_shard_id).await?; let mut baseline_lsns = None; if let Some((generation, origin_ps_id)) = &previous { let origin_ps = PageServerNode::from_env(env, env.get_pageserver_conf(*origin_ps_id)?); @@ -116,10 +114,12 @@ pub async fn migrate_tenant( if origin_ps_id == &dest_ps.conf.id { println!("🔁 Already attached to {origin_ps_id}, freshening..."); let gen = attachment_service - .attach_hook(tenant_id, dest_ps.conf.id) + .attach_hook(tenant_shard_id, dest_ps.conf.id) .await?; let dest_conf = build_location_config(LocationConfigMode::AttachedSingle, gen, None); - dest_ps.location_config(tenant_id, dest_conf, None).await?; + dest_ps + .location_config(tenant_shard_id, dest_conf, None) + .await?; println!("✅ Migration complete"); return Ok(()); } @@ -129,28 +129,30 @@ pub async fn migrate_tenant( let stale_conf = build_location_config(LocationConfigMode::AttachedStale, Some(*generation), None); origin_ps - .location_config(tenant_id, stale_conf, Some(Duration::from_secs(10))) + .location_config(tenant_shard_id, stale_conf, Some(Duration::from_secs(10))) .await?; - baseline_lsns = Some(get_lsns(tenant_id, &origin_ps).await?); + baseline_lsns = Some(get_lsns(tenant_shard_id, &origin_ps).await?); } let gen = attachment_service - .attach_hook(tenant_id, dest_ps.conf.id) + .attach_hook(tenant_shard_id, dest_ps.conf.id) .await?; let dest_conf = build_location_config(LocationConfigMode::AttachedMulti, gen, None); println!("🔁 Attaching to pageserver {}", dest_ps.conf.id); - dest_ps.location_config(tenant_id, dest_conf, None).await?; + dest_ps + .location_config(tenant_shard_id, dest_conf, None) + .await?; if let Some(baseline) = baseline_lsns { println!("🕑 Waiting for LSN to catch up..."); - await_lsn(tenant_id, &dest_ps, baseline).await?; + await_lsn(tenant_shard_id, &dest_ps, baseline).await?; } let cplane = ComputeControlPlane::load(env.clone())?; for (endpoint_name, endpoint) in &cplane.endpoints { - if endpoint.tenant_id == tenant_id { + if endpoint.tenant_id == tenant_shard_id.tenant_id { println!( "🔁 Reconfiguring endpoint {} to use pageserver {}", endpoint_name, dest_ps.conf.id @@ -171,7 +173,7 @@ pub async fn migrate_tenant( let found = other_ps_tenants .into_iter() .map(|t| t.id) - .any(|i| i.tenant_id == tenant_id); + .any(|i| i == tenant_shard_id); if !found { continue; } @@ -188,7 +190,7 @@ pub async fn migrate_tenant( other_ps.conf.id ); other_ps - .location_config(tenant_id, secondary_conf, None) + .location_config(tenant_shard_id, secondary_conf, None) .await?; } @@ -197,7 +199,9 @@ pub async fn migrate_tenant( dest_ps.conf.id ); let dest_conf = build_location_config(LocationConfigMode::AttachedSingle, gen, None); - dest_ps.location_config(tenant_id, dest_conf, None).await?; + dest_ps + .location_config(tenant_shard_id, dest_conf, None) + .await?; println!("✅ Migration complete"); diff --git a/libs/pageserver_api/src/models.rs b/libs/pageserver_api/src/models.rs index be41b610b8..0587deb508 100644 --- a/libs/pageserver_api/src/models.rs +++ b/libs/pageserver_api/src/models.rs @@ -297,7 +297,7 @@ pub struct StatusResponse { #[derive(Serialize, Deserialize, Debug)] #[serde(deny_unknown_fields)] pub struct TenantLocationConfigRequest { - pub tenant_id: TenantId, + pub tenant_shard_id: TenantShardId, #[serde(flatten)] pub config: LocationConfig, // as we have a flattened field, we should reject all unknown fields in it } diff --git a/pageserver/client/src/mgmt_api.rs b/pageserver/client/src/mgmt_api.rs index 87e4ed8efd..d3daf68ea5 100644 --- a/pageserver/client/src/mgmt_api.rs +++ b/pageserver/client/src/mgmt_api.rs @@ -80,9 +80,12 @@ impl Client { pub async fn list_timelines( &self, - tenant_id: TenantId, + tenant_shard_id: TenantShardId, ) -> Result> { - let uri = format!("{}/v1/tenant/{tenant_id}/timeline", self.mgmt_api_endpoint); + let uri = format!( + "{}/v1/tenant/{tenant_shard_id}/timeline", + self.mgmt_api_endpoint + ); self.get(&uri) .await? .json() @@ -166,14 +169,17 @@ impl Client { pub async fn location_config( &self, - tenant_id: TenantId, + tenant_shard_id: TenantShardId, config: LocationConfig, flush_ms: Option, ) -> Result<()> { - let req_body = TenantLocationConfigRequest { tenant_id, config }; + let req_body = TenantLocationConfigRequest { + tenant_shard_id, + config, + }; let path = format!( "{}/v1/tenant/{}/location_config", - self.mgmt_api_endpoint, tenant_id + self.mgmt_api_endpoint, tenant_shard_id ); let path = if let Some(flush_ms) = flush_ms { format!("{}?flush_ms={}", path, flush_ms.as_millis())