remove rounding for bandwidth and db storage (#25583)

We've decided to no longer round up usage to the nearest 1KB!

GitOrigin-RevId: dcff0ed46a69d7017f0d144060559fe476074225

Author: rrwang7

crates/application/src/export_worker.rs

@@ -877,7 +877,6 @@ mod tests {
     };
     use headers::ContentType;
     use keybroker::Identity;
-    use maplit::btreemap;
     use model::{
         exports::types::{
             Export,
@@ -1056,14 +1055,9 @@ mod tests {
         assert_eq!(zip_entries, expected_export_entries);
 
         let usage = usage.gather_user_stats();
-        assert_eq!(
-            *usage.database_egress_size,
-            btreemap! {
-               "table_0".to_string() => 1024,
-               "table_1".to_string() => 1024,
-               "table_2".to_string() => 1024,
-            }
-        );
+        assert!(usage.database_egress_size["table_0"] > 0);
+        assert!(usage.database_egress_size["table_1"] > 0);
+        assert!(usage.database_egress_size["table_2"] > 0);
 
         Ok(())
     }

crates/application/src/snapshot_import.rs

@@ -3018,10 +3018,7 @@ a
         .await?;
 
         let stats = usage.gather_user_stats();
-        assert_eq!(
-            *stats.database_ingress_size,
-            btreemap! {table_name.to_string() => 2048}
-        );
+        assert!(stats.database_ingress_size[&table_name.to_string()] > 0);
         assert_eq!(stats.storage_ingress_size, 9);
 
         Ok(())

crates/database/src/database.rs

@@ -1780,7 +1780,7 @@ impl<RT: Runtime> Database<RT> {
 
         let mut document_storage_by_table = BTreeMap::new();
         for (table_name, summary) in snapshot.iter_user_table_summaries() {
-            let table_size = summary.total_size_rounded() as usize;
+            let table_size = summary.total_size();
             document_storage_by_table.insert(table_name, (table_size, 0));
         }
 

crates/database/src/table_summary.rs

@@ -63,10 +63,6 @@ use shape_inference::{
     Shape,
     ShapeEnum,
 };
-use usage_tracking::{
-    round_up,
-    KB,
-};
 
 #[cfg(any(test, feature = "testing"))]
 use crate::IndexModel;
@@ -83,17 +79,14 @@ use crate::{
 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: i64,
 }
 
 impl fmt::Display for TableSummary {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(
             f,
-            "TableSummary {{ inferred_type: {}, total_size: {}, total_size_rounded: {} }}",
-            self.inferred_type, self.total_size, self.total_size_rounded
+            "TableSummary {{ inferred_type: {}, total_size: {} }}",
+            self.inferred_type, self.total_size
         )
     }
 }
@@ -103,22 +96,17 @@ impl TableSummary {
         Self {
             inferred_type: Shape::empty(),
             total_size: 0,
-            total_size_rounded: 0,
         }
     }
 
     pub fn is_empty(&self) -> bool {
-        self.inferred_type.is_empty() && self.total_size == 0 && self.total_size_rounded == 0
+        self.inferred_type.is_empty() && self.total_size == 0
     }
 
     pub fn total_size(&self) -> usize {
         self.total_size as usize
     }
 
-    pub fn total_size_rounded(&self) -> i64 {
-        self.total_size_rounded
-    }
-
     pub fn num_values(&self) -> usize {
         *self.inferred_type.num_values() as usize
     }
@@ -129,27 +117,18 @@ impl TableSummary {
 
     pub fn insert(&self, object: &ConvexObject) -> Self {
         let total_size = self.total_size + object.size() as i64;
-        let total_size_rounded =
-            self.total_size_rounded + round_up(object.size() as u64, KB) as i64;
         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;
-        let size_rounded = round_up(object.size() as u64, KB) as i64;
         anyhow::ensure!(self.total_size >= size, "Negative size due to {object}");
-        anyhow::ensure!(
-            self.total_size_rounded >= size_rounded,
-            "Negative size due to {object}"
-        );
         Ok(Self {
             inferred_type: self.inferred_type.remove(object)?,
             total_size: self.total_size - size,
-            total_size_rounded: self.total_size_rounded - size_rounded,
         })
     }
 
@@ -166,7 +145,6 @@ impl From<&TableSummary> for JsonValue {
     fn from(summary: &TableSummary) -> Self {
         json!({
             "totalSize": JsonInteger::encode(summary.total_size),
-            "totalSizeRounded": JsonInteger::encode(summary.total_size_rounded),
             "inferredTypeWithOptionalFields": JsonValue::from(&summary.inferred_type)
         })
     }
@@ -183,19 +161,13 @@ 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)) => JsonInteger::decode(s)?,
-                    _ => anyhow::bail!("Invalid totalSizeRounded"),
-                };
-                anyhow::ensure!(total_size_rounded >= 0);
                 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"),

crates/database/src/tests/usage_tracking.rs

@@ -25,7 +25,6 @@ use runtime::testing::TestRuntime;
 use usage_tracking::{
     CallType,
     FunctionUsageTracker,
-    KB,
 };
 use vector::VectorSearch;
 
@@ -99,8 +98,7 @@ async fn vector_insert_counts_usage_for_backfilling_indexes(rt: TestRuntime) ->
         .get(&(*index_name.table()).to_string())
         .cloned();
 
-    // round up.
-    assert_eq!(value, Some(1024));
+    assert!(value.unwrap() > 0);
     Ok(())
 }
 
@@ -133,8 +131,7 @@ async fn vector_insert_counts_usage_for_enabled_indexes(rt: TestRuntime) -> anyh
         .recent_vector_ingress_size
         .get(&(*index_name.table()).to_string())
         .cloned();
-    // We round up to the nearest KB
-    assert_eq!(value, Some(1024_u64));
+    assert!(value.unwrap() > 0);
     Ok(())
 }
 
@@ -214,16 +211,12 @@ async fn vector_query_counts_bandwidth(rt: TestRuntime) -> anyhow::Result<()> {
     let stats = fixtures.db.usage_counter().collect();
     let vector_egress = stats.recent_vector_egress_size;
     let bandwidth_egress = stats.recent_database_egress_size;
-    // Rounded up.
-    assert_eq!(
-        *vector_egress.get(&index_name.table().to_string()).unwrap(),
-        KB
-    );
-    assert_eq!(
+    assert!(*vector_egress.get(&index_name.table().to_string()).unwrap() > 0);
+    assert!(
         *bandwidth_egress
             .get(&index_name.table().to_string())
-            .unwrap(),
-        KB
+            .unwrap()
+            > 0
     );
     Ok(())
 }
@@ -254,7 +247,7 @@ async fn test_usage_tracking_basic_insert_and_get(rt: TestRuntime) -> anyhow::Re
     let database_ingress = stats.recent_database_ingress_size;
     assert_eq!(database_ingress.len(), 1);
     assert!(database_ingress.contains_key("my_table"));
-    assert_eq!(*database_ingress.get("my_table").unwrap(), KB);
+    assert!(*database_ingress.get("my_table").unwrap() > 0);
     let database_egress = stats.recent_database_egress_size;
     assert!(database_egress.is_empty());
 
@@ -280,7 +273,7 @@ async fn test_usage_tracking_basic_insert_and_get(rt: TestRuntime) -> anyhow::Re
     let database_egress = stats.recent_database_egress_size;
     assert_eq!(database_egress.len(), 1);
     assert!(database_egress.contains_key("my_table"));
-    assert_eq!(*database_egress.get("my_table").unwrap(), KB);
+    assert!(*database_egress.get("my_table").unwrap() > 0);
 
     Ok(())
 }
@@ -336,11 +329,10 @@ async fn test_usage_tracking_insert_with_index(rt: TestRuntime) -> anyhow::Resul
     );
 
     let stats = db.usage_counter().collect();
-    // Database bandwidth ingress is 6 KB from 3 document writes and 3 index writes
     let database_ingress = stats.recent_database_ingress_size;
     assert_eq!(database_ingress.len(), 1);
     assert!(database_ingress.contains_key("my_table"));
-    assert_eq!(*database_ingress.get("my_table").unwrap(), 6 * KB);
+    assert!(*database_ingress.get("my_table").unwrap() > 0);
     let database_egress = stats.recent_database_egress_size;
     assert!(database_egress.is_empty());
 
@@ -364,13 +356,12 @@ async fn test_usage_tracking_insert_with_index(rt: TestRuntime) -> anyhow::Resul
     );
 
     let stats = db.usage_counter().collect();
-    // Database bandwidth egress is 4 KB from 2 document writes and 2 index writes
     let database_ingress = stats.recent_database_ingress_size;
     assert!(database_ingress.is_empty());
     let database_egress = stats.recent_database_egress_size;
     assert_eq!(database_egress.len(), 1);
     assert!(database_egress.contains_key("my_table"));
-    assert_eq!(*database_egress.get("my_table").unwrap(), 4 * KB);
+    assert!(*database_egress.get("my_table").unwrap() > 0);
 
     Ok(())
 }

crates/usage_tracking/src/lib.rs

@@ -111,12 +111,6 @@ impl HeapSize for UsageCounter {
     }
 }
 
-pub const KB: u64 = 1 << 10;
-
-pub fn round_up(n: u64, k: u64) -> u64 {
-    (n + k - 1) / k * k
-}
-
 impl UsageCounter {
     pub fn new(usage_logger: Arc<dyn UsageEventLogger>) -> Self {
         let state = Arc::new(Mutex::new(UsageCounterState::default()));
@@ -570,9 +564,7 @@ impl FunctionUsageTracker {
         let mut state = self.state.lock();
         state
             .database_ingress_size
-            .mutate_entry_or_default(table_name.clone(), |count| {
-                *count += round_up(ingress_size, KB)
-            });
+            .mutate_entry_or_default(table_name.clone(), |count| *count += ingress_size);
     }
 
     pub fn track_database_egress_size(
@@ -588,9 +580,7 @@ impl FunctionUsageTracker {
         let mut state = self.state.lock();
         state
             .database_egress_size
-            .mutate_entry_or_default(table_name.clone(), |count| {
-                *count += round_up(egress_size, KB)
-            });
+            .mutate_entry_or_default(table_name.clone(), |count| *count += egress_size);
     }
 
     // Tracks the vector ingress surcharge and database usage for documents
@@ -618,16 +608,15 @@ impl FunctionUsageTracker {
         // Note that vector search counts as both database and vector bandwidth
         // per the comment above.
         let mut state = self.state.lock();
-        let rounded_size = round_up(ingress_size, KB);
         state
             .database_ingress_size
             .mutate_entry_or_default(table_name.clone(), |count| {
-                *count += rounded_size;
+                *count += ingress_size;
             });
         state
             .vector_ingress_size
             .mutate_entry_or_default(table_name.clone(), |count| {
-                *count += rounded_size;
+                *count += ingress_size;
             });
     }
 
@@ -657,13 +646,12 @@ impl FunctionUsageTracker {
         // Note that vector search counts as both database and vector bandwidth
         // per the comment above.
         let mut state = self.state.lock();
-        let rounded_size = round_up(egress_size, KB);
         state
             .database_egress_size
-            .mutate_entry_or_default(table_name.clone(), |count| *count += rounded_size);
+            .mutate_entry_or_default(table_name.clone(), |count| *count += egress_size);
         state
             .vector_egress_size
-            .mutate_entry_or_default(table_name.clone(), |count| *count += rounded_size);
+            .mutate_entry_or_default(table_name.clone(), |count| *count += egress_size);
     }
 }