From 58f64339f321cf85214b432d30fe252365f3873f Mon Sep 17 00:00:00 2001 From: John Spray Date: Fri, 8 Dec 2023 09:20:08 +0000 Subject: [PATCH] neon_local: implement Locate API for attachment service --- control_plane/src/attachment_service.rs | 36 +++++++++++++++ control_plane/src/bin/attachment_service.rs | 50 ++++++++++++++++++++- 2 files changed, 85 insertions(+), 1 deletion(-) diff --git a/control_plane/src/attachment_service.rs b/control_plane/src/attachment_service.rs index 4295659739..5c18cd6e5c 100644 --- a/control_plane/src/attachment_service.rs +++ b/control_plane/src/attachment_service.rs @@ -55,10 +55,29 @@ pub struct TenantCreateResponse { pub struct NodeRegisterRequest { pub node_id: NodeId, + pub listen_pg_addr: String, + pub listen_pg_port: u16, + pub listen_http_addr: String, pub listen_http_port: u16, } +#[derive(Serialize, Deserialize)] +pub struct TenantLocateResponseShard { + pub node_id: NodeId, + + pub listen_pg_addr: String, + pub listen_pg_port: u16, + + pub listen_http_addr: String, + pub listen_http_port: u16, +} + +#[derive(Serialize, Deserialize)] +pub struct TenantLocateResponse { + pub shards: Vec, +} + impl AttachmentService { pub fn from_env(env: &LocalEnv) -> Self { let path = env.base_data_dir.join("attachments.json"); @@ -178,6 +197,23 @@ impl AttachmentService { Ok(()) } + pub fn tenant_locate(&self, tenant_id: TenantId) -> anyhow::Result { + let url = self + .env + .control_plane_api + .clone() + .unwrap() + .join(&format!("tenant/{tenant_id}")) + .unwrap(); + + let response = self.client.get(url).send()?; + if response.status() != StatusCode::OK { + return Err(anyhow!("Unexpected status {}", response.status())); + } + + Ok(response.json()?) + } + pub fn tenant_timeline_create( &self, tenant_id: TenantId, diff --git a/control_plane/src/bin/attachment_service.rs b/control_plane/src/bin/attachment_service.rs index 30377d970b..a9e96f7efd 100644 --- a/control_plane/src/bin/attachment_service.rs +++ b/control_plane/src/bin/attachment_service.rs @@ -42,7 +42,8 @@ use pageserver_api::control_api::{ use control_plane::attachment_service::{ AttachHookRequest, AttachHookResponse, InspectRequest, InspectResponse, NodeRegisterRequest, - TenantCreateResponse, TenantCreateResponseShard, + TenantCreateResponse, TenantCreateResponseShard, TenantLocateResponse, + TenantLocateResponseShard, }; #[derive(Parser)] @@ -103,6 +104,9 @@ struct NodeState { listen_http_addr: String, listen_http_port: u16, + + listen_pg_addr: String, + listen_pg_port: u16, } impl NodeState { @@ -647,6 +651,45 @@ async fn handle_tenant_timeline_create(mut req: Request) -> Result) -> Result, ApiError> { + let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?; + + let state = get_state(&req).inner.clone(); + let mut locked = state.write().await; + + tracing::info!("Locating shards for tenant {tenant_id}"); + + // Take a snapshot of pageservers + let pageservers = locked.pageservers.clone(); + + let mut result = Vec::new(); + + for (_tenant_shard_id, shard) in locked + .tenants + .range_mut(TenantShardId::tenant_range(tenant_id)) + { + let node_id = shard + .pageserver + .ok_or(ApiError::BadRequest(anyhow::anyhow!( + "Cannot locate a tenant that is not attached" + )))?; + + let node = pageservers + .get(&node_id) + .expect("Pageservers may not be deleted while referenced"); + + result.push(TenantLocateResponseShard { + node_id, + listen_http_addr: node.listen_http_addr.clone(), + listen_http_port: node.listen_http_port, + listen_pg_addr: node.listen_pg_addr.clone(), + listen_pg_port: node.listen_pg_port, + }); + } + + json_response(StatusCode::OK, TenantLocateResponse { shards: result }) +} + async fn handle_node_register(mut req: Request) -> Result, ApiError> { let register_req = json_request::(&mut req).await?; let state = get_state(&req).inner.clone(); @@ -658,6 +701,8 @@ async fn handle_node_register(mut req: Request) -> Result, id: register_req.node_id, listen_http_addr: register_req.listen_http_addr, listen_http_port: register_req.listen_http_port, + listen_pg_addr: register_req.listen_pg_addr, + listen_pg_port: register_req.listen_pg_port, }, ); @@ -682,6 +727,9 @@ fn make_router(persistent_state: PersistentState) -> RouterBuilder