diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs index 73eee590a6..b6059a9a63 100644 --- a/pageserver/src/tenant.rs +++ b/pageserver/src/tenant.rs @@ -194,6 +194,45 @@ struct TimelinePreload { index_part: Result, } +/// To include the HeatmapWriter in tenant shutdown, we provide a hook +/// for it to publish a barrier when upload is going on. We will take +/// this and wait on it during shutdown, ensuring that there is no +/// upload going on once shutdown() returns. +pub struct HeatmapHook { + // Mutually exclude shutdown and any in-flight uploads + // + // If this is None, we are shutting down + in_progress: Option>>, +} + +impl Default for HeatmapHook { + fn default() -> Self { + Self { + in_progress: Some(Arc::default()), + } + } +} + +impl HeatmapHook { + pub(crate) fn enter(&self) -> Option> { + self.in_progress.as_ref().map(|l| { + l.clone() + .try_lock_owned() + // expect: shutdown cannot have started yet or in_progress would have been None, + // so we expect that only one HeatmapWriter may take this lock at once. + // Depends on the invariant that HeatmapWriter is the only thing that calls + // enter(), and that it will never try and do uploads concurrently for the same + // tenant. + .expect("Tried to double-lock HeatmapHook") + }) + } + + /// Returns a lock that the caller should wait on before proceeding with shutdown + fn shutdown(&mut self) -> Arc> { + self.in_progress.take().expect("Called shutdown twice") + } +} + /// /// Tenant consists of multiple timelines. Keep them in a hash table. /// @@ -245,6 +284,8 @@ pub struct Tenant { eviction_task_tenant_state: tokio::sync::Mutex, pub(crate) delete_progress: Arc>, + + pub(crate) heatmap_hook: Mutex, } impl std::fmt::Debug for Tenant { @@ -2063,6 +2104,14 @@ impl Tenant { } } + let heatmap_hook_lock = { + let mut hook = self.heatmap_hook.lock().unwrap(); + hook.shutdown() + }; + tracing::debug!("shutdown waiting heatmap uploads..."); + // Take & drop lock to ensure any heatmap upload is complete. + drop(heatmap_hook_lock.lock().await); + tracing::debug!("shutdown waiting for tasks..."); // shutdown all tenant and timeline tasks: gc, compaction, page service // No new tasks will be started for this tenant because it's in `Stopping` state. @@ -2570,6 +2619,7 @@ impl Tenant { cached_synthetic_tenant_size: Arc::new(AtomicU64::new(0)), eviction_task_tenant_state: tokio::sync::Mutex::new(EvictionTaskTenantState::default()), delete_progress: Arc::new(tokio::sync::Mutex::new(DeleteTenantFlow::default())), + heatmap_hook: Mutex::default(), } }