From 0aecfbb09c26b63932f3c5ed8219cd5253a9a939 Mon Sep 17 00:00:00 2001
From: Konstantin Knizhnik <knizhnik@garret.ru>
Date: Mon, 22 May 2023 16:39:16 +0300
Subject: [PATCH] Fix loading logical database size

---
 pageserver/src/tenant.rs          | 21 ++++++++++++++++++---
 pageserver/src/tenant/timeline.rs | 18 ++++++++++++------
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/pageserver/src/tenant.rs b/pageserver/src/tenant.rs
index 121b8b9505..c5415d9c2f 100644
--- a/pageserver/src/tenant.rs
+++ b/pageserver/src/tenant.rs
@@ -543,8 +543,6 @@ impl Tenant {
             }
         };
 
-        timeline.init_logical_size().await;
-
         if self.remote_storage.is_some() {
             // Reconcile local state with remote storage, downloading anything that's
             // missing locally, and scheduling uploads for anything that's missing
@@ -1043,11 +1041,29 @@ impl Tenant {
         // The loops will shut themselves down when they notice that the tenant is inactive.
         self.activate(ctx)?;
 
+        self.load_logical_sizes().await?;
+
         info!("Done");
 
         Ok(())
     }
 
+    async fn load_logical_sizes(&self) -> anyhow::Result<()> {
+        let not_broken_timelines: Vec<Arc<Timeline>>;
+        {
+            let timelines_accessor = self.timelines.lock().unwrap();
+            not_broken_timelines = timelines_accessor
+                .values()
+                .filter(|timeline| timeline.current_state() != TimelineState::Broken)
+                .cloned()
+                .collect();
+        }
+        for timeline in not_broken_timelines {
+            timeline.load_inmem_logical_size().await?;
+        }
+        Ok(())
+    }
+
     /// Subroutine of `load_tenant`, to load an individual timeline
     ///
     /// NB: The parent is assumed to be already loaded!
@@ -2516,7 +2532,6 @@ impl Tenant {
         ancestor: Option<Arc<Timeline>>,
     ) -> anyhow::Result<UninitializedTimeline> {
         let tenant_id = self.tenant_id;
-
         let remote_client = if let Some(remote_storage) = self.remote_storage.as_ref() {
             let remote_client = RemoteTimelineClient::new(
                 remote_storage.clone(),
diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs
index af0085fd49..4b27552528 100644
--- a/pageserver/src/tenant/timeline.rs
+++ b/pageserver/src/tenant/timeline.rs
@@ -713,13 +713,19 @@ impl Timeline {
         self.current_logical_size.load(AtomicOrdering::Relaxed) as u64
     }
 
-    pub async fn init_logical_size(&self) {
-        let ctx = RequestContext::todo_child(TaskKind::Startup, DownloadBehavior::Error);
-        let last_record_lsn = self.get_last_record_lsn();
-        if let Ok(size) = self.get_logical_size(last_record_lsn, &ctx).await {
-            self.current_logical_size
-                .store(size as i64, AtomicOrdering::Relaxed);
+    /// Load from KV storage value of logical timeline size and store it in inmemory atomic variable
+    pub async fn load_inmem_logical_size(&self) -> anyhow::Result<()> {
+        let lsn = self.get_disk_consistent_lsn();
+        if lsn != Lsn::INVALID {
+            let ctx = RequestContext::todo_child(TaskKind::Startup, DownloadBehavior::Error);
+            match self.get_logical_size(lsn, &ctx).await {
+                Ok(size) => self
+                    .current_logical_size
+                    .store(size as i64, AtomicOrdering::Relaxed),
+                Err(e) => info!("Failed to load logical size: {:?}", e),
+            }
         }
+        Ok(())
     }
 
     /// Check if more than 'checkpoint_distance' of WAL has been accumulated in