rust/neon
1
0
Fork 0
mirror of https://github.com/neondatabase/neon.git synced 2026-07-07 14:10:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

neon_local: implement Locate API for attachment service

Browse Source
This commit is contained in:
John Spray
2023-12-08 09:20:08 +00:00
parent 78e673fbb3
commit 58f64339f3
2 changed files with 85 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
control_plane/src/attachment_service.rs
View File

@@ -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<TenantLocateResponseShard>,
}
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<TenantLocateResponse> {
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,

50
control_plane/src/bin/attachment_service.rs
View File

@@ -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<Body>) -> Result<Respons
json_response(StatusCode::OK, ())
}
async fn handle_tenant_locate(req: Request<Body>) -> Result<Response<Body>, 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<Body>) -> Result<Response<Body>, ApiError> {
let register_req = json_request::<NodeRegisterRequest>(&mut req).await?;
let state = get_state(&req).inner.clone();
@@ -658,6 +701,8 @@ async fn handle_node_register(mut req: Request<Body>) -> Result<Response<Body>,
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<hyper::Body,
.post("/tenant/:tenant_id/timeline", |r| {
request_span(r, handle_tenant_timeline_create)
})
.get("/tenant/:tenant_id/locate", |r| {
request_span(r, handle_tenant_locate)
})
}
#[tokio::main]