revert removing writes to totalSizeRounded (#25742)

GitOrigin-RevId: 0c99dba6bfe54cdbc075eebd91470591317f7613

Author: rrwang7

crates/database/src/table_summary.rs

@@ -75,18 +75,27 @@ use crate::{
     TableIterator,
 };
 
+pub const KB: u64 = 1 << 10;
+
+pub fn round_up(n: u64, k: u64) -> u64 {
+    (n + k - 1) / k * k
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct TableSummary {
     inferred_type: CountedShape<ProdConfigWithOptionalFields>,
     total_size: i64,
+    // Used for metered billing, we charge for database storage rounded up to
+    // the nearest KB per document
+    total_size_rounded: Option<i64>,
 }
 
 impl fmt::Display for TableSummary {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(
             f,
-            "TableSummary {{ inferred_type: {}, total_size: {} }}",
-            self.inferred_type, self.total_size
+            "TableSummary {{ inferred_type: {}, total_size: {} , total_size_rounded: {:?} }}",
+            self.inferred_type, self.total_size, self.total_size_rounded
         )
     }
 }
@@ -96,6 +105,7 @@ impl TableSummary {
         Self {
             inferred_type: Shape::empty(),
             total_size: 0,
+            total_size_rounded: None,
         }
     }
 
@@ -117,18 +127,36 @@ impl TableSummary {
 
     pub fn insert(&self, object: &ConvexObject) -> Self {
         let total_size = self.total_size + object.size() as i64;
+        let total_size_rounded = match self.total_size_rounded {
+            Some(total_size_rounded) => {
+                Some(total_size_rounded + round_up(object.size() as u64, KB) as i64)
+            },
+            None => None,
+        };
         Self {
             inferred_type: self.inferred_type.insert(object),
             total_size,
+            total_size_rounded,
         }
     }
 
     pub fn remove(&self, object: &ConvexObject) -> anyhow::Result<Self> {
         let size = object.size() as i64;
-        anyhow::ensure!(self.total_size >= size, "Negative size due to {object}");
+        let total_size_rounded = match self.total_size_rounded {
+            Some(total_size_rounded) => {
+                let size_rounded = round_up(object.size() as u64, KB) as i64;
+                anyhow::ensure!(
+                    total_size_rounded >= size_rounded,
+                    "Negative size due to {object}"
+                );
+                Some(total_size_rounded - size_rounded)
+            },
+            None => None,
+        };
         Ok(Self {
             inferred_type: self.inferred_type.remove(object)?,
             total_size: self.total_size - size,
+            total_size_rounded,
         })
     }
 
@@ -143,10 +171,17 @@ impl TableSummary {
 
 impl From<&TableSummary> for JsonValue {
     fn from(summary: &TableSummary) -> Self {
-        json!({
-            "totalSize": JsonInteger::encode(summary.total_size),
-            "inferredTypeWithOptionalFields": JsonValue::from(&summary.inferred_type)
-        })
+        match summary.total_size_rounded {
+            Some(total_size_rounded) => json!({
+                "totalSize": JsonInteger::encode(summary.total_size),
+                "totalSizeRounded": JsonInteger::encode(total_size_rounded),
+                "inferredTypeWithOptionalFields": JsonValue::from(&summary.inferred_type)
+            }),
+            None => json!({
+                "totalSize": JsonInteger::encode(summary.total_size),
+                "inferredTypeWithOptionalFields": JsonValue::from(&summary.inferred_type)
+            }),
+        }
     }
 }
 
@@ -161,13 +196,22 @@ impl TryFrom<JsonValue> for TableSummary {
                     _ => anyhow::bail!("Invalid totalSize"),
                 };
                 anyhow::ensure!(total_size >= 0);
+                let total_size_rounded = match v.remove("totalSizeRounded") {
+                    Some(JsonValue::String(s)) => {
+                        let total_size_rounded = JsonInteger::decode(s)?;
+                        anyhow::ensure!(total_size_rounded >= 0);
+                        Some(total_size_rounded)
+                    },
+                    _ => None,
+                };
                 let inferred_type = match v.remove("inferredTypeWithOptionalFields") {
                     Some(v) => CountedShape::<ProdConfigWithOptionalFields>::try_from(v)?,
                     None => anyhow::bail!("Missing field inferredTypeWithOptionalFields"),
                 };
                 Ok(TableSummary {
                     inferred_type,
                     total_size,
+                    total_size_rounded,
                 })
             },
             _ => anyhow::bail!("Wrong type of json value for TableSummaryJson"),