From e51f2be3d0d2172421016ece22332c079328e851 Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki@neon.tech>
Date: Fri, 25 Nov 2022 19:06:50 +0200
Subject: [PATCH] Fix physical size tracking when files are downloaded.

Includes a test.
---
 pageserver/src/storage_sync.rs                   |  6 ++++--
 pageserver/src/tenant/timeline.rs                | 11 +++++------
 .../regress/test_tenants_with_remote_storage.py  | 16 ++++++++++++++--
 3 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/pageserver/src/storage_sync.rs b/pageserver/src/storage_sync.rs
index 7d90f11ac4..c7a451c3a3 100644
--- a/pageserver/src/storage_sync.rs
+++ b/pageserver/src/storage_sync.rs
@@ -544,11 +544,13 @@ impl RemoteTimelineClient {
     /// Download a (layer) file from `path`, into local filesystem.
     ///
     /// 'layer_metadata' is the metadata from the remote index file.
+    ///
+    /// On success, returns the size of the downloaded file.
     pub async fn download_layer_file(
         &self,
         path: &RelativePath,
         layer_metadata: &LayerFileMetadata,
-    ) -> anyhow::Result<()> {
+    ) -> anyhow::Result<u64> {
         let downloaded_size = download::download_layer_file(
             self.conf,
             &self.storage_impl,
@@ -576,7 +578,7 @@ impl RemoteTimelineClient {
                 );
             }
         }
-        Ok(())
+        Ok(downloaded_size)
     }
 
     //
diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs
index 75394145bd..ebc49d2dc2 100644
--- a/pageserver/src/tenant/timeline.rs
+++ b/pageserver/src/tenant/timeline.rs
@@ -1024,6 +1024,7 @@ impl Timeline {
                                 error!("could not rename file \"{}\": {:?}",
                                        local_path.display(), err);
                             }
+                            self.metrics.current_physical_size_gauge.sub(local_size);
                             false
                         } else {
                             true
@@ -1067,7 +1068,7 @@ impl Timeline {
                 }
 
                 trace!("downloading image file: {}", file = path.display());
-                remote_client
+                let sz = remote_client
                     .download_layer_file(&RelativePath::from_filename(path), &layer_metadata)
                     .await
                     .context("download image layer")?;
@@ -1076,11 +1077,11 @@ impl Timeline {
                 let image_layer =
                     ImageLayer::new(self.conf, self.timeline_id, self.tenant_id, &imgfilename);
 
-                // FIXME: when to update physical size?
                 self.layers
                     .write()
                     .unwrap()
                     .insert_historic(Arc::new(image_layer));
+                self.metrics.current_physical_size_gauge.add(sz);
             } else if let Some(deltafilename) = DeltaFileName::parse_str(fname) {
                 // Create a DeltaLayer struct for each delta file.
                 // The end-LSN is exclusive, while disk_consistent_lsn is
@@ -1097,7 +1098,7 @@ impl Timeline {
                 }
 
                 trace!("downloading image file: {}", file = path.display());
-                remote_client
+                let sz = remote_client
                     .download_layer_file(&RelativePath::from_filename(path), &layer_metadata)
                     .await
                     .context("download delta layer")?;
@@ -1106,11 +1107,11 @@ impl Timeline {
                 let delta_layer =
                     DeltaLayer::new(self.conf, self.timeline_id, self.tenant_id, &deltafilename);
 
-                // FIXME: when to update physical size?
                 self.layers
                     .write()
                     .unwrap()
                     .insert_historic(Arc::new(delta_layer));
+                self.metrics.current_physical_size_gauge.add(sz);
             } else {
                 bail!("unexpected layer filename in remote storage: {}", fname);
             }
@@ -1184,8 +1185,6 @@ impl Timeline {
             }
         };
 
-        // TODO what to do with physical size?
-
         // Are there local files that don't exist remotely? Schedule uploads for them
         let timeline_path = self.conf.timeline_path(&self.timeline_id, &self.tenant_id);
         for fname in &local_only_filenames {
diff --git a/test_runner/regress/test_tenants_with_remote_storage.py b/test_runner/regress/test_tenants_with_remote_storage.py
index 5024b9b782..1a1f9da41e 100644
--- a/test_runner/regress/test_tenants_with_remote_storage.py
+++ b/test_runner/regress/test_tenants_with_remote_storage.py
@@ -160,6 +160,12 @@ def test_tenants_attached_after_download(
 
     ##### Stop the pageserver, erase its layer file to force it being downloaded from S3
     env.postgres.stop_all()
+
+    detail_before = client.timeline_detail(
+        tenant_id, timeline_id, include_non_incremental_physical_size=True
+    )
+    assert detail_before["current_physical_size_non_incremental"] == detail_before["current_physical_size"]
+
     env.pageserver.stop()
 
     timeline_dir = Path(env.repo_dir) / "tenants" / str(tenant_id) / "timelines" / str(timeline_id)
@@ -186,11 +192,17 @@ def test_tenants_attached_after_download(
     assert (
         len(restored_timelines) == 1
     ), f"Tenant {tenant_id} should have its timeline reattached after its layer is downloaded from the remote storage"
-    retored_timeline = restored_timelines[0]
-    assert retored_timeline["timeline_id"] == str(
+    restored_timeline = restored_timelines[0]
+    assert restored_timeline["timeline_id"] == str(
         timeline_id
     ), f"Tenant {tenant_id} should have its old timeline {timeline_id} restored from the remote storage"
 
+    # Check that the physical size matches after re-downloading
+    detail_after = client.timeline_detail(
+        tenant_id, timeline_id, include_non_incremental_physical_size=True
+    )
+    assert detail_before["current_physical_size"] == detail_after["current_physical_size"]
+
 
 @pytest.mark.parametrize("remote_storage_kind", [RemoteStorageKind.LOCAL_FS])
 def test_tenant_upgrades_index_json_from_v0(