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

pageserver: implement shard splitting

Browse Source
This commit is contained in:
John Spray
2023-12-08 12:56:44 +00:00
parent 9cd72caabf
commit fc2f9fa3fe
10 changed files with 366 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
control_plane/src/attachment_service.rs
View File

@@ -3,7 +3,7 @@ use anyhow::anyhow;
use camino::Utf8PathBuf;
use hyper::{Method, StatusCode};
use pageserver_api::{
models::{TenantCreateRequest, TimelineCreateRequest},
models::{TenantCreateRequest, TenantShardSplitRequest, TimelineCreateRequest},
shard::TenantShardId,
};
use postgres_connection::parse_host_port;
@@ -247,6 +247,15 @@ impl AttachmentService {
self.dispatch::<(), _>(Method::GET, format!("tenant/{tenant_id}/locate"), None)
}
#[instrument(skip(self), fields(%tenant_id, %new_shard_count))]
pub fn tenant_split(&self, tenant_id: TenantId, new_shard_count: u8) -> anyhow::Result<()> {
self.dispatch::<_, ()>(
Method::POST,
format!("tenant/{tenant_id}/shard_split"),
Some(TenantShardSplitRequest { new_shard_count }),
)
}
#[instrument(skip_all, fields(node_id=%req.node_id))]
pub fn node_register(&self, req: NodeRegisterRequest) -> anyhow::Result<()> {
self.dispatch::<_, ()>(Method::POST, "node".to_string(), Some(req))

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

@@ -10,7 +10,8 @@ use hyper::{Body, Request, Response};
use hyper::{Method, StatusCode};
use pageserver_api::models::{
LocationConfig, LocationConfigMode, TenantConfig, TenantCreateRequest,
TenantLocationConfigRequest, TimelineCreateRequest,
TenantLocationConfigRequest, TenantShardSplitRequest, TenantShardSplitResponse,
TimelineCreateRequest,
};
use pageserver_api::shard::{ShardCount, ShardIdentity, ShardNumber, TenantShardId};
use reqwest::Client;
@@ -722,6 +723,121 @@ async fn handle_node_register(mut req: Request<Body>) -> Result<Response<Body>,
json_response(StatusCode::OK, ())
}
async fn handle_tenant_shard_split(mut req: Request<Body>) -> Result<Response<Body>, ApiError> {
let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?;
let split_req = json_request::<TenantShardSplitRequest>(&mut req).await?;
let state = get_state(&req).inner.clone();
let mut locked = state.write().await;
let pageservers = locked.pageservers.clone();
let mut replacements = HashMap::new();
for (tenant_shard_id, shard) in locked
.tenants
.range_mut(TenantShardId::tenant_range(tenant_id))
{
if tenant_shard_id.shard_count == ShardCount(split_req.new_shard_count) {
tracing::warn!(
"Tenant shard {} already has shard count {}",
tenant_shard_id,
split_req.new_shard_count
);
continue;
}
let node_id = shard
.pageserver
.ok_or(ApiError::BadRequest(anyhow::anyhow!(
"Cannot split a tenant that is not attached"
)))?;
let node = pageservers
.get(&node_id)
.expect("Pageservers may not be deleted while referenced");
let client = Client::new();
let response = client
.request(
Method::POST,
format!("{}/tenant/{}/shard_split", node.base_url(), tenant_shard_id),
)
.json(&TenantShardSplitRequest {
new_shard_count: split_req.new_shard_count,
})
.send()
.await
.map_err(|e| {
ApiError::Conflict(format!("Failed to split {}: {}", tenant_shard_id, e))
})?;
// response.error_for_status().map_err(|e| {
// ApiError::Conflict(format!("Failed to split {}: {}", tenant_shard_id, e))
// })?;
let response: TenantShardSplitResponse = response.json().await.map_err(|e| {
ApiError::InternalServerError(anyhow::anyhow!(
"Malformed response from pageserver: {}",
e
))
})?;
replacements.insert(*tenant_shard_id, response.new_shards);
}
// Replace all the shards we just split with their children
for (replaced, children) in replacements.into_iter() {
let (pageserver, generation, shard_ident, config) = {
let old_state = locked
.tenants
.remove(&replaced)
.expect("It was present, we just split it");
(
old_state.pageserver.unwrap(),
old_state.generation,
old_state.shard,
old_state.config.clone(),
)
};
for child in children {
let mut child_shard = shard_ident;
child_shard.number = child.shard_number;
child_shard.count = child.shard_count;
let mut child_observed: HashMap<NodeId, ObservedStateLocation> = HashMap::new();
child_observed.insert(
pageserver,
ObservedStateLocation {
conf: Some(LocationConfig {
mode: LocationConfigMode::AttachedSingle,
generation: Some(generation),
secondary_conf: None,
shard_number: child.shard_number.0,
shard_count: child.shard_count.0,
shard_stripe_size: shard_ident.stripe_size.0,
tenant_conf: config.clone(),
}),
},
);
locked.tenants.insert(
child,
TenantState {
tenant_shard_id: child,
shard: child_shard,
pageserver: Some(pageserver),
generation: generation,
observed: ObservedState {
locations: child_observed,
},
config: config.clone(),
},
);
}
}
json_response(StatusCode::OK, ())
}
/// Status endpoint is just used for checking that our HTTP listener is up
async fn handle_status(_req: Request<Body>) -> Result<Response<Body>, ApiError> {
json_response(StatusCode::OK, ())
@@ -743,6 +859,9 @@ fn make_router(persistent_state: PersistentState) -> RouterBuilder<hyper::Body,
.get("/tenant/:tenant_id/locate", |r| {
request_span(r, handle_tenant_locate)
})
.put("/tenant/:tenant_id/shard_split", |r| {
request_span(r, handle_tenant_shard_split)
})
}
#[tokio::main]

12
control_plane/src/bin/neon_local.rs
View File

@@ -515,6 +515,18 @@ async fn handle_tenant(
migrate_tenant(env, tenant_shard_id, new_pageserver).await?;
println!("tenant {tenant_shard_id} migrated to {}", new_pageserver_id);
}
Some(("split", matches)) => {
let tenant_id = get_tenant_id(matches, env)?;
let attachment_service = AttachmentService::from_env(env);
let old_shards = attachment_service.tenant_locate(tenant_id)?.shards;
let new_shard_count = old_shards.len() * 2;
if old_shards.len() > 127 {
bail!("Cannot split further");
}
attachment_service.tenant_split(tenant_id, new_shard_count as u8)?;
println!("Split {}->{}", old_shards.len(), new_shard_count);
}
Some(("status", matches)) => {
let tenant_id = get_tenant_id(matches, env)?;