neon_local: update various TenantId uses to TenantShardId

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<NodeId>,
 }
 
@@ -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<Option<u32>> {
         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<Option<(u32, NodeId)>> {
+    pub async fn inspect(&self, tenant_id: TenantShardId) -> anyhow::Result<Option<(u32, NodeId)>> {
         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 {

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<S, V>(input: &HashMap<TenantId, V>, serializer: S) -> Result<S::Ok, S::Error>
-where
-    S: serde::Serializer,
-    V: Clone + Serialize,
-{
-    let transformed = input.iter().map(|(k, v)| (hex::encode(k), v.clone()));
-
-    transformed
-        .collect::<HashMap<String, V>>()
-        .serialize(serializer)
-}
-
-fn from_hex_map<'de, D, V>(deserializer: D) -> Result<HashMap<TenantId, V>, D::Error>
-where
-    D: serde::de::Deserializer<'de>,
-    V: Deserialize<'de>,
-{
-    let hex_map = HashMap::<String, V>::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<TenantId, TenantState>,
+    tenants: BTreeMap<TenantShardId, TenantState>,
 
     #[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<Body>) -> Result<Response<Body>, 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<Body>) -> Result<Response<Body>, 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<Body>) -> Result<Response<Body>, 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<Body>) -> Result<Response<Body>, 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<Body>) -> Result<Response<Body>, 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<Body>) -> Result<Response<Body>, 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,

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<HashMap<TimelineId, TimelineInfo>> {
     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<TenantShardId> {
+    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<Option<TenantId>> {
     sub_match
         .get_one::<String>("tenant-id")
@@ -305,6 +320,14 @@ fn parse_tenant_id(sub_match: &ArgMatches) -> anyhow::Result<Option<TenantId>> {
         .context("Failed to parse tenant id from the argument string")
 }
 
+fn parse_tenant_shard_id(sub_match: &ArgMatches) -> anyhow::Result<Option<TenantShardId>> {
+    sub_match
+        .get_one::<String>("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<Option<TimelineId>> {
     sub_match
         .get_one::<String>("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("?");
 

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<Duration>,
     ) -> anyhow::Result<()> {
         Ok(self
             .http_client
-            .location_config(tenant_id, config, flush_ms)
+            .location_config(tenant_shard_id, config, flush_ms)
             .await?)
     }
 

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<HashMap<TimelineId, Lsn>> {
-    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<TimelineId, Lsn>,
 ) -> 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");
 

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
 }

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<Vec<pageserver_api::models::TimelineInfo>> {
-        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<std::time::Duration>,
     ) -> 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())