patches/neon-controller-placement.patch
Rama main · réplica de lectura del repositorio original.
diff --git a/storage_controller/src/persistence.rs b/storage_controller/src/persistence.rs
index 0693dcd..31405c2 100644
--- a/storage_controller/src/persistence.rs
+++ b/storage_controller/src/persistence.rs
@@ -80,6 +80,31 @@ pub struct Persistence {
}
impl Persistence {
+ // Read the generation issued to the destination before a compute callback
+ // returns. In-memory shard state is updated only after reconciliation.
+ pub(crate) async fn owned_tenant_placement(
+ &self,
+ filter_tenant_id: TenantId,
+ ) -> DatabaseResult<Option<JsonValue>> {
+ self.with_conn(move |conn| {
+ Box::pin(async move {
+ let rows = diesel::sql_query(
+ "SELECT s.tenant_id, s.shard_number, s.shard_count, s.generation, \
+ s.generation_pageserver, s.placement_policy, o.ownership_token, \
+ o.state::text AS ownership_state FROM tenant_shards s \
+ JOIN hakopod_resource_ownership o ON o.resource_kind = 'tenant' \
+ AND o.external_key = s.tenant_id WHERE s.tenant_id = $1 \
+ ORDER BY s.shard_number LIMIT 2",
+ )
+ .bind::<diesel::sql_types::Text, _>(filter_tenant_id.to_string())
+ .load::<OwnedTenantPlacementRow>(conn)
+ .await?;
+ owned_tenant_placement_response(filter_tenant_id, rows)
+ })
+ })
+ .await
+ }
+
pub(crate) async fn reserve_ownership(
&self,
token: Uuid,
@@ -612,6 +637,115 @@ struct OwnershipStateRow {
state: String,
}
+#[derive(QueryableByName, Clone)]
+struct OwnedTenantPlacementRow {
+ #[diesel(sql_type = diesel::sql_types::Text)]
+ tenant_id: String,
+ #[diesel(sql_type = diesel::sql_types::Integer)]
+ shard_number: i32,
+ #[diesel(sql_type = diesel::sql_types::Integer)]
+ shard_count: i32,
+ #[diesel(sql_type = diesel::sql_types::Nullable<diesel::sql_types::Integer>)]
+ generation: Option<i32>,
+ #[diesel(sql_type = diesel::sql_types::Nullable<diesel::sql_types::BigInt>)]
+ generation_pageserver: Option<i64>,
+ #[diesel(sql_type = diesel::sql_types::Text)]
+ placement_policy: String,
+ #[diesel(sql_type = diesel::sql_types::Uuid)]
+ ownership_token: Uuid,
+ #[diesel(sql_type = diesel::sql_types::Text)]
+ ownership_state: String,
+}
+
+fn owned_tenant_placement_response(
+ tenant_id: TenantId,
+ mut rows: Vec<OwnedTenantPlacementRow>,
+) -> DatabaseResult<Option<JsonValue>> {
+ if rows.is_empty() {
+ return Ok(None);
+ }
+ if rows.len() != 1 {
+ return Err(DatabaseError::Cas(
+ "owned placement requires one unsharded tenant".into(),
+ ));
+ }
+ let row = rows.pop().unwrap();
+ let policy = serde_json::from_str::<PlacementPolicy>(&row.placement_policy);
+ if row.tenant_id != tenant_id.to_string()
+ || row.shard_number != 0
+ || row.shard_count != 0
+ || !matches!(policy, Ok(PlacementPolicy::Attached(_)))
+ || row.ownership_state != "completed"
+ || row.ownership_token.is_nil()
+ || row.generation.is_none_or(|value| value < 1)
+ || row.generation_pageserver.is_none_or(|value| value < 1)
+ {
+ return Err(DatabaseError::Cas(
+ "owned tenant placement is not attached and complete".into(),
+ ));
+ }
+ Ok(Some(serde_json::json!({
+ "schema_version": 1,
+ "tenant_id": row.tenant_id,
+ "node_id": row.generation_pageserver.unwrap(),
+ "generation": row.generation.unwrap(),
+ "ownership_token": row.ownership_token.simple().to_string(),
+ "ownership_state": row.ownership_state,
+ })))
+}
+
+#[test]
+fn test_owned_tenant_placement_response() {
+ let tenant = TenantId::from_str("11111111111111111111111111111111").unwrap();
+ let row = OwnedTenantPlacementRow {
+ tenant_id: tenant.to_string(),
+ shard_number: 0,
+ shard_count: 0,
+ generation: Some(8),
+ generation_pageserver: Some(2),
+ placement_policy: serde_json::to_string(&PlacementPolicy::Attached(1)).unwrap(),
+ ownership_token: Uuid::from_u128(1),
+ ownership_state: "completed".into(),
+ };
+ let observed = owned_tenant_placement_response(tenant, vec![row.clone()])
+ .unwrap()
+ .unwrap();
+ assert_eq!(observed["generation"], 8);
+ assert_eq!(observed["node_id"], 2);
+ assert_eq!(
+ observed["ownership_token"],
+ row.ownership_token.simple().to_string()
+ );
+ assert!(
+ owned_tenant_placement_response(tenant, vec![])
+ .unwrap()
+ .is_none()
+ );
+ assert!(owned_tenant_placement_response(tenant, vec![row.clone(), row.clone()]).is_err());
+ for state in ["reserved", "applying", "deleting", "tombstoned"] {
+ let mut invalid = row.clone();
+ invalid.ownership_state = state.into();
+ assert!(owned_tenant_placement_response(tenant, vec![invalid]).is_err());
+ }
+ for case in 0..8 {
+ let mut invalid = row.clone();
+ match case {
+ 0 => invalid.shard_count = 1,
+ 1 => invalid.shard_number = 1,
+ 2 => invalid.generation = None,
+ 3 => invalid.generation = Some(0),
+ 4 => invalid.generation_pageserver = None,
+ 5 => invalid.generation_pageserver = Some(0),
+ 6 => {
+ invalid.placement_policy =
+ serde_json::to_string(&PlacementPolicy::Detached).unwrap()
+ }
+ _ => invalid.tenant_id = "22222222222222222222222222222222".into(),
+ }
+ assert!(owned_tenant_placement_response(tenant, vec![invalid]).is_err());
+ }
+}
+
fn ownership_resume_decision(row: OwnershipRow) -> DatabaseResult<OwnershipDecision> {
match row.state.as_str() {
"reserved" | "applying" => Ok(OwnershipDecision::Resume),
diff --git a/storage_controller/src/service.rs b/storage_controller/src/service.rs
index 925906a..a7a2664 100644
--- a/storage_controller/src/service.rs
+++ b/storage_controller/src/service.rs
@@ -842,6 +842,16 @@ impl Service {
.map_err(Into::into)
}
+ pub(crate) async fn owned_tenant_placement(
+ &self,
+ tenant_id: TenantId,
+ ) -> Result<Option<serde_json::Value>, ApiError> {
+ self.persistence
+ .owned_tenant_placement(tenant_id)
+ .await
+ .map_err(Into::into)
+ }
+
pub(crate) async fn preflight_owned_timelines_for_tenant(
&self,
tenant_id: TenantId,
diff --git a/storage_controller/src/http.rs b/storage_controller/src/http.rs
index 80c7be3..e833167 100644
--- a/storage_controller/src/http.rs
+++ b/storage_controller/src/http.rs
@@ -1121,6 +1121,33 @@ async fn handle_tenant_locate(
json_response(StatusCode::OK, service.tenant_locate(tenant_id)?)
}
+async fn handle_owned_tenant_placement(req: Request<Body>) -> Result<Response<Body>, ApiError> {
+ check_permissions(&req, Scope::Admin)?;
+ let tenant_id: TenantId = parse_request_param(&req, "tenant_id")?;
+ let service = get_state(&req).service.clone();
+ if !service.get_config().hakopod_ownership_v1 {
+ return Err(ApiError::NotFound(
+ anyhow::anyhow!("owned placement is disabled").into(),
+ ));
+ }
+ match maybe_forward(req).await {
+ ForwardOutcome::Forwarded(res) => return res,
+ ForwardOutcome::NotForwarded(_req) => {}
+ };
+ // This callback read must not wait for startup or a tenant operation lock:
+ // startup and live migration may themselves be waiting for the callback.
+ let placement = tokio::time::timeout(
+ Duration::from_secs(2),
+ service.owned_tenant_placement(tenant_id),
+ )
+ .await
+ .map_err(|_| ApiError::Timeout("owned placement read timed out".into()))??
+ .ok_or_else(|| {
+ ApiError::NotFound(anyhow::anyhow!("owned tenant placement is unavailable").into())
+ })?;
+ json_response(StatusCode::OK, placement)
+}
+
async fn handle_tenant_describe(
service: Arc<Service>,
req: Request<Body>,
@@ -2910,6 +2937,13 @@ pub fn make_router(
RequestName("control_v1_tenant_shard_split"),
)
})
+ .get("/control/v1/tenant/:tenant_id/placement", |r| {
+ named_request_span(
+ r,
+ handle_owned_tenant_placement,
+ RequestName("control_v1_owned_tenant_placement"),
+ )
+ })
.get("/control/v1/tenant/:tenant_id", |r| {
tenant_service_handler(
r,