From 1ff96154ee444365471802b69aa1da48cc82552d Mon Sep 17 00:00:00 2001
From: Yueh-Hsuan Chiang <yueh-hsuan@solana.com>
Date: Fri, 22 Mar 2024 14:55:39 -0700
Subject: [PATCH 1/2] [TieredStorage] Add capacity() API and limit hot-storage
 file size to 16GB.

---
 accounts-db/src/tiered_storage.rs          | 7 +++++++
 accounts-db/src/tiered_storage/hot.rs      | 9 ++++++++-
 accounts-db/src/tiered_storage/readable.rs | 6 ++++++
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/accounts-db/src/tiered_storage.rs b/accounts-db/src/tiered_storage.rs
index a0d8eea4010b94..cbca5c93d0041e 100644
--- a/accounts-db/src/tiered_storage.rs
+++ b/accounts-db/src/tiered_storage.rs
@@ -38,6 +38,8 @@ use {
 
 pub type TieredStorageResult<T> = Result<T, TieredStorageError>;
 
+const MAX_TIERED_STORAGE_FILE_SIZE: u64 = 16 * 1024 * 1024 * 1024; // 16 GiB;
+
 /// The struct that defines the formats of all building blocks of a
 /// TieredStorage.
 #[derive(Clone, Debug, PartialEq)]
@@ -163,6 +165,11 @@ impl TieredStorage {
     pub fn is_empty(&self) -> bool {
         self.len() == 0
     }
+
+    pub fn capacity(&self) -> u64 {
+        self.reader()
+            .map_or(MAX_TIERED_STORAGE_FILE_SIZE, |reader| reader.capacity())
+    }
 }
 
 #[cfg(test)]
diff --git a/accounts-db/src/tiered_storage/hot.rs b/accounts-db/src/tiered_storage/hot.rs
index 414d74b2eb81b7..6c8763f200ab13 100644
--- a/accounts-db/src/tiered_storage/hot.rs
+++ b/accounts-db/src/tiered_storage/hot.rs
@@ -16,7 +16,7 @@ use {
             mmap_utils::{get_pod, get_slice},
             owners::{OwnerOffset, OwnersBlockFormat, OwnersTable, OWNER_NO_OWNER},
             StorableAccounts, StorableAccountsWithHashesAndWriteVersions, TieredStorageError,
-            TieredStorageFormat, TieredStorageResult,
+            TieredStorageFormat, TieredStorageResult, MAX_TIERED_STORAGE_FILE_SIZE,
         },
     },
     bytemuck::{Pod, Zeroable},
@@ -369,6 +369,13 @@ impl HotStorageReader {
         self.len() == 0
     }
 
+    pub fn capacity(&self) -> u64 {
+        if self.is_empty() {
+            return MAX_TIERED_STORAGE_FILE_SIZE;
+        }
+        self.len() as u64
+    }
+
     /// Returns the footer of the underlying tiered-storage accounts file.
     pub fn footer(&self) -> &TieredStorageFooter {
         &self.footer
diff --git a/accounts-db/src/tiered_storage/readable.rs b/accounts-db/src/tiered_storage/readable.rs
index 008e805689df57..0191dad4903578 100644
--- a/accounts-db/src/tiered_storage/readable.rs
+++ b/accounts-db/src/tiered_storage/readable.rs
@@ -44,6 +44,12 @@ impl TieredStorageReader {
         }
     }
 
+    pub fn capacity(&self) -> u64 {
+        match self {
+            Self::Hot(hot) => hot.capacity(),
+        }
+    }
+
     /// Returns the footer of the associated HotAccountsFile.
     pub fn footer(&self) -> &TieredStorageFooter {
         match self {

From a960163c947be90207c346024dac9f395e635f04 Mon Sep 17 00:00:00 2001
From: Yueh-Hsuan Chiang <yueh-hsuan@solana.com>
Date: Fri, 22 Mar 2024 17:37:35 -0700
Subject: [PATCH 2/2] Have HotStorageReader::capacity() simply return len()

---
 accounts-db/src/tiered_storage/hot.rs | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/accounts-db/src/tiered_storage/hot.rs b/accounts-db/src/tiered_storage/hot.rs
index 6c8763f200ab13..260548897f66e2 100644
--- a/accounts-db/src/tiered_storage/hot.rs
+++ b/accounts-db/src/tiered_storage/hot.rs
@@ -16,7 +16,7 @@ use {
             mmap_utils::{get_pod, get_slice},
             owners::{OwnerOffset, OwnersBlockFormat, OwnersTable, OWNER_NO_OWNER},
             StorableAccounts, StorableAccountsWithHashesAndWriteVersions, TieredStorageError,
-            TieredStorageFormat, TieredStorageResult, MAX_TIERED_STORAGE_FILE_SIZE,
+            TieredStorageFormat, TieredStorageResult,
         },
     },
     bytemuck::{Pod, Zeroable},
@@ -370,9 +370,6 @@ impl HotStorageReader {
     }
 
     pub fn capacity(&self) -> u64 {
-        if self.is_empty() {
-            return MAX_TIERED_STORAGE_FILE_SIZE;
-        }
         self.len() as u64
     }