diff --git a/libs/pageserver_api/src/controller_api.rs b/libs/pageserver_api/src/controller_api.rs index 91f9c03ba4..c5b49edba0 100644 --- a/libs/pageserver_api/src/controller_api.rs +++ b/libs/pageserver_api/src/controller_api.rs @@ -169,6 +169,8 @@ pub struct TenantDescribeResponseShard { pub is_pending_compute_notification: bool, /// A shard split is currently underway pub is_splitting: bool, + /// A timeline is being imported into this tenant + pub is_importing: bool, pub scheduling_policy: ShardSchedulingPolicy, diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs index 997fc24052..698579e8fb 100644 --- a/pageserver/src/tenant.rs +++ b/pageserver/src/tenant.rs @@ -3816,6 +3816,24 @@ impl TenantShard { MaybeDeletedIndexPart::IndexPart(p) => p, }; + // A shard split may not take place while a timeline import is on-going + // for the tenant. Timeline imports run as part of each tenant shard + // and rely on the sharding scheme to split the work among pageservers. + // If we were to split in the middle of this process, we would have to + // either ensure that it's driven to completion on the old shard set + // or transfer it to the new shard set. It's technically possible, but complex. + match index_part.import_pgdata { + Some(ref import) if !import.is_done() => { + anyhow::bail!( + "Cannot split due to import with idempotency key: {:?}", + import.idempotency_key() + ); + } + Some(_) | None => { + // fallthrough + } + } + for child_shard in child_shards { tracing::info!(%timeline_id, "Uploading index_part for child {}", child_shard.to_index()); upload_index_part( diff --git a/storage_controller/src/persistence.rs b/storage_controller/src/persistence.rs index 64a8846a9d..9ffcf9b9e6 100644 --- a/storage_controller/src/persistence.rs +++ b/storage_controller/src/persistence.rs @@ -134,6 +134,7 @@ pub(crate) enum DatabaseOperation { UpdateTimelineImport, DeleteTimelineImport, ListTimelineImports, + IsTenantImportingTimeline, } #[must_use] @@ -1641,9 +1642,7 @@ impl Persistence { .await } - pub(crate) async fn list_complete_timeline_imports( - &self, - ) -> DatabaseResult> { + pub(crate) async fn list_timeline_imports(&self) -> DatabaseResult> { use crate::schema::timeline_imports::dsl; let persistent = self .with_measured_conn(DatabaseOperation::ListTimelineImports, move |conn| { @@ -1660,10 +1659,7 @@ impl Persistence { .map(TimelineImport::from_persistent) .collect(); match imports { - Ok(ok) => Ok(ok - .into_iter() - .filter(|import| import.is_complete()) - .collect()), + Ok(ok) => Ok(ok.into_iter().collect()), Err(err) => Err(DatabaseError::Logical(format!( "failed to deserialize import: {err}" ))), @@ -1773,6 +1769,25 @@ impl Persistence { }) .await } + + pub(crate) async fn is_tenant_importing_timeline( + &self, + tenant_id: TenantId, + ) -> DatabaseResult { + use crate::schema::timeline_imports::dsl; + self.with_measured_conn(DatabaseOperation::IsTenantImportingTimeline, move |conn| { + Box::pin(async move { + let imports: i64 = dsl::timeline_imports + .filter(dsl::tenant_id.eq(tenant_id.to_string())) + .count() + .get_result(conn) + .await?; + + Ok(imports > 0) + }) + }) + .await + } } pub(crate) fn load_certs() -> anyhow::Result> { diff --git a/storage_controller/src/service.rs b/storage_controller/src/service.rs index 7e5e3fd8f4..ca9b911c4d 100644 --- a/storage_controller/src/service.rs +++ b/storage_controller/src/service.rs @@ -97,7 +97,9 @@ use crate::tenant_shard::{ ReconcileNeeded, ReconcileResult, ReconcileWaitError, ReconcilerStatus, ReconcilerWaiter, ScheduleOptimization, ScheduleOptimizationAction, TenantShard, }; -use crate::timeline_import::{ShardImportStatuses, TimelineImport, UpcallClient}; +use crate::timeline_import::{ + ShardImportStatuses, TimelineImport, TimelineImportState, UpcallClient, +}; const WAITER_FILL_DRAIN_POLL_TIMEOUT: Duration = Duration::from_millis(500); @@ -878,15 +880,33 @@ impl Service { }); } - // Fetch the list of completed imports and attempt to finalize them in the background. - // This handles the case where the previous storage controller instance shut down - // whilst finalizing imports. - let complete_imports = self.persistence.list_complete_timeline_imports().await; - match complete_imports { - Ok(ok) => { + // Reconcile the timeline imports: + // 1. Mark each tenant shard of tenants with an importing timeline as importing. + // 2. Finalize the completed imports in the background. This handles the case where + // the previous storage controller instance shut down whilst finalizing imports. + let imports = self.persistence.list_timeline_imports().await; + match imports { + Ok(mut imports) => { + { + let mut locked = self.inner.write().unwrap(); + for import in &imports { + locked + .tenants + .range_mut(TenantShardId::tenant_range(import.tenant_id)) + .for_each(|(_id, shard)| { + shard.importing = TimelineImportState::Importing + }); + } + } + + imports.retain(|import| import.is_complete()); tokio::task::spawn({ let finalize_imports_self = self.clone(); - async move { finalize_imports_self.finalize_timeline_imports(ok).await } + async move { + finalize_imports_self + .finalize_timeline_imports(imports) + .await + } }); } Err(err) => { @@ -3772,6 +3792,22 @@ impl Service { failpoint_support::sleep_millis_async!("tenant-create-timeline-shared-lock"); let is_import = create_req.is_import(); + if is_import { + // Ensure that there is no split on-going. + // [`Self::tenant_shard_split`] holds the exclusive tenant lock + // for the duration of the split, but here we handle the case + // where we restarted and the split is being aborted. + let locked = self.inner.read().unwrap(); + let splitting = locked + .tenants + .range(TenantShardId::tenant_range(tenant_id)) + .any(|(_id, shard)| shard.splitting != SplitState::Idle); + + if splitting { + return Err(ApiError::Conflict("Tenant is splitting shard".to_string())); + } + } + let timeline_info = self .tenant_timeline_create_pageservers(tenant_id, create_req) .await?; @@ -3809,6 +3845,14 @@ impl Service { .context("timeline import insert") .map_err(ApiError::InternalServerError)?; + // Set the importing flag on the tenant shards + self.inner + .write() + .unwrap() + .tenants + .range_mut(TenantShardId::tenant_range(tenant_id)) + .for_each(|(_id, shard)| shard.importing = TimelineImportState::Importing); + match inserted { true => { tracing::info!(%tenant_id, %timeline_id, "Inserted timeline import"); @@ -3931,6 +3975,13 @@ impl Service { tracing::warn!("Failed to delete timeline import entry from database: {err}"); } + self.inner + .write() + .unwrap() + .tenants + .range_mut(TenantShardId::tenant_range(import.tenant_id)) + .for_each(|(_id, shard)| shard.importing = TimelineImportState::Idle); + // TODO(vlad): Timeline creations in import mode do not return a correct initdb lsn, // so we can't create the timeline on the safekeepers. Fix by moving creation here. // https://github.com/neondatabase/neon/issues/11569 @@ -4914,6 +4965,7 @@ impl Service { is_reconciling: shard.reconciler.is_some(), is_pending_compute_notification: shard.pending_compute_notification, is_splitting: matches!(shard.splitting, SplitState::Splitting), + is_importing: shard.importing == TimelineImportState::Importing, scheduling_policy: shard.get_scheduling_policy(), preferred_az_id: shard.preferred_az().map(ToString::to_string), }) @@ -5404,6 +5456,27 @@ impl Service { .enter() .map_err(|_| ApiError::ShuttingDown)?; + // Timeline imports on the pageserver side can't handle shard-splits. + // If the tenant is importing a timeline, dont't shard split it. + match self + .persistence + .is_tenant_importing_timeline(tenant_id) + .await + { + Ok(importing) => { + if importing { + return Err(ApiError::Conflict( + "Cannot shard split during timeline import".to_string(), + )); + } + } + Err(err) => { + return Err(ApiError::InternalServerError(anyhow::anyhow!( + "Failed to check for running imports: {err}" + ))); + } + } + let new_shard_count = ShardCount::new(split_req.new_shard_count); let new_stripe_size = split_req.new_stripe_size; @@ -8076,12 +8149,25 @@ impl Service { candidates.extend(size_candidates); } - // Filter out tenants in a prohibiting scheduling mode. + // Filter out tenants in a prohibiting scheduling modes + // and tenants with an ongoing import. + // + // Note that the import check here is oportunistic. An import might start + // after the check before we actually update [`TenantShard::splitting`]. + // [`Self::tenant_shard_split`] checks the database whilst holding the exclusive + // tenant lock. Imports might take a long time, so the check here allows us + // to split something else instead of trying the same shard over and over. { let state = self.inner.read().unwrap(); candidates.retain(|i| { - let policy = state.tenants.get(&i.id).map(|s| s.get_scheduling_policy()); - policy == Some(ShardSchedulingPolicy::Active) + let shard = state.tenants.get(&i.id); + match shard { + Some(t) => { + t.get_scheduling_policy() == ShardSchedulingPolicy::Active + && t.importing == TimelineImportState::Idle + } + None => false, + } }); } diff --git a/storage_controller/src/tenant_shard.rs b/storage_controller/src/tenant_shard.rs index 3a75e96cb2..c7b2628ec4 100644 --- a/storage_controller/src/tenant_shard.rs +++ b/storage_controller/src/tenant_shard.rs @@ -33,6 +33,7 @@ use crate::scheduler::{ RefCountUpdate, ScheduleContext, ScheduleError, Scheduler, SecondaryShardTag, ShardTag, }; use crate::service::ReconcileResultRequest; +use crate::timeline_import::TimelineImportState; use crate::{Sequence, service}; /// Serialization helper @@ -100,6 +101,10 @@ pub(crate) struct TenantShard { /// reconciliation, and timeline creation. pub(crate) splitting: SplitState, + /// Flag indicating whether the tenant has an in-progress timeline import. + /// Used to disallow shard splits while an import is in progress. + pub(crate) importing: TimelineImportState, + /// If a tenant was enqueued for later reconcile due to hitting concurrency limit, this flag /// is set. This flag is cleared when the tenant is popped off the delay queue. pub(crate) delayed_reconcile: bool, @@ -583,6 +588,7 @@ impl TenantShard { config: TenantConfig::default(), reconciler: None, splitting: SplitState::Idle, + importing: TimelineImportState::Idle, sequence: Sequence(1), delayed_reconcile: false, waiter: Arc::new(SeqWait::new(Sequence(0))), @@ -1844,6 +1850,8 @@ impl TenantShard { config: serde_json::from_str(&tsp.config).unwrap(), reconciler: None, splitting: tsp.splitting, + // Filled in during [`Service::startup_reconcile`] + importing: TimelineImportState::Idle, waiter: Arc::new(SeqWait::new(Sequence::initial())), error_waiter: Arc::new(SeqWait::new(Sequence::initial())), last_error: Arc::default(), diff --git a/storage_controller/src/timeline_import.rs b/storage_controller/src/timeline_import.rs index b6dd4b252e..6dcc538c4b 100644 --- a/storage_controller/src/timeline_import.rs +++ b/storage_controller/src/timeline_import.rs @@ -14,6 +14,12 @@ use utils::{ use crate::{persistence::TimelineImportPersistence, service::Config}; +#[derive(Deserialize, Serialize, PartialEq, Eq)] +pub(crate) enum TimelineImportState { + Importing, + Idle, +} + #[derive(Serialize, Deserialize, Clone, Debug)] pub(crate) struct ShardImportStatuses(pub(crate) HashMap);