diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs index 7438215a68..c585fac5e6 100644 --- a/pageserver/src/tenant/timeline.rs +++ b/pageserver/src/tenant/timeline.rs @@ -1903,6 +1903,7 @@ impl Timeline { .set((calculated_size, metrics_guard.calculation_result_saved())) .ok() .expect("only this task sets it"); + self.current_logical_size.initialized.notify_one(); } pub fn spawn_ondemand_logical_size_calculation( @@ -3104,6 +3105,24 @@ impl Timeline { Ok(image_layers) } + + /// Wait until the background initial logical size calculation is complete, or + /// this Timeline is shut down. Calling this function will cause the initial + /// logical size calculation to skip waiting for the background jobs barrier. + pub(crate) async fn await_initial_logical_size(self: Arc) { + if let Some(await_bg_cancel) = self + .current_logical_size + .cancel_wait_for_background_loop_concurrency_limit_semaphore + .get() + { + await_bg_cancel.cancel(); + } + + tokio::select!( + _ = self.current_logical_size.initialized.notified() => {}, + _ = self.cancel.cancelled() => {} + ) + } } #[derive(Default)] diff --git a/pageserver/src/tenant/timeline/logical_size.rs b/pageserver/src/tenant/timeline/logical_size.rs index f2db8c91fc..ca7ba48e75 100644 --- a/pageserver/src/tenant/timeline/logical_size.rs +++ b/pageserver/src/tenant/timeline/logical_size.rs @@ -34,6 +34,9 @@ pub(super) struct LogicalSize { pub(crate) cancel_wait_for_background_loop_concurrency_limit_semaphore: OnceCell, + /// Once the initial logical size is initialized, this is notified. + pub(crate) initialized: tokio::sync::Notify, + /// Latest Lsn that has its size uncalculated, could be absent for freshly created timelines. pub initial_part_end: Option, @@ -125,6 +128,7 @@ impl LogicalSize { initial_part_end: None, size_added_after_initial: AtomicI64::new(0), did_return_approximate_to_walreceiver: AtomicBool::new(false), + initialized: tokio::sync::Notify::new(), } } @@ -135,6 +139,7 @@ impl LogicalSize { initial_part_end: Some(compute_to), size_added_after_initial: AtomicI64::new(0), did_return_approximate_to_walreceiver: AtomicBool::new(false), + initialized: tokio::sync::Notify::new(), } }