Skip to content

Commit a37cb03

Browse files
rrwang7Convex, Inc.
authored and
Convex, Inc.
committed
revert removing writes to totalSizeRounded (#25742)
GitOrigin-RevId: 0c99dba6bfe54cdbc075eebd91470591317f7613
1 parent b28ceff commit a37cb03

File tree

1 file changed

+51
-7
lines changed

1 file changed

+51
-7
lines changed

crates/database/src/table_summary.rs

Lines changed: 51 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -75,18 +75,27 @@ use crate::{
7575
TableIterator,
7676
};
7777

78+
pub const KB: u64 = 1 << 10;
79+
80+
pub fn round_up(n: u64, k: u64) -> u64 {
81+
(n + k - 1) / k * k
82+
}
83+
7884
#[derive(Debug, Clone, PartialEq, Eq)]
7985
pub struct TableSummary {
8086
inferred_type: CountedShape<ProdConfigWithOptionalFields>,
8187
total_size: i64,
88+
// Used for metered billing, we charge for database storage rounded up to
89+
// the nearest KB per document
90+
total_size_rounded: Option<i64>,
8291
}
8392

8493
impl fmt::Display for TableSummary {
8594
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
8695
write!(
8796
f,
88-
"TableSummary {{ inferred_type: {}, total_size: {} }}",
89-
self.inferred_type, self.total_size
97+
"TableSummary {{ inferred_type: {}, total_size: {} , total_size_rounded: {:?} }}",
98+
self.inferred_type, self.total_size, self.total_size_rounded
9099
)
91100
}
92101
}
@@ -96,6 +105,7 @@ impl TableSummary {
96105
Self {
97106
inferred_type: Shape::empty(),
98107
total_size: 0,
108+
total_size_rounded: None,
99109
}
100110
}
101111

@@ -117,18 +127,36 @@ impl TableSummary {
117127

118128
pub fn insert(&self, object: &ConvexObject) -> Self {
119129
let total_size = self.total_size + object.size() as i64;
130+
let total_size_rounded = match self.total_size_rounded {
131+
Some(total_size_rounded) => {
132+
Some(total_size_rounded + round_up(object.size() as u64, KB) as i64)
133+
},
134+
None => None,
135+
};
120136
Self {
121137
inferred_type: self.inferred_type.insert(object),
122138
total_size,
139+
total_size_rounded,
123140
}
124141
}
125142

126143
pub fn remove(&self, object: &ConvexObject) -> anyhow::Result<Self> {
127144
let size = object.size() as i64;
128-
anyhow::ensure!(self.total_size >= size, "Negative size due to {object}");
145+
let total_size_rounded = match self.total_size_rounded {
146+
Some(total_size_rounded) => {
147+
let size_rounded = round_up(object.size() as u64, KB) as i64;
148+
anyhow::ensure!(
149+
total_size_rounded >= size_rounded,
150+
"Negative size due to {object}"
151+
);
152+
Some(total_size_rounded - size_rounded)
153+
},
154+
None => None,
155+
};
129156
Ok(Self {
130157
inferred_type: self.inferred_type.remove(object)?,
131158
total_size: self.total_size - size,
159+
total_size_rounded,
132160
})
133161
}
134162

@@ -143,10 +171,17 @@ impl TableSummary {
143171

144172
impl From<&TableSummary> for JsonValue {
145173
fn from(summary: &TableSummary) -> Self {
146-
json!({
147-
"totalSize": JsonInteger::encode(summary.total_size),
148-
"inferredTypeWithOptionalFields": JsonValue::from(&summary.inferred_type)
149-
})
174+
match summary.total_size_rounded {
175+
Some(total_size_rounded) => json!({
176+
"totalSize": JsonInteger::encode(summary.total_size),
177+
"totalSizeRounded": JsonInteger::encode(total_size_rounded),
178+
"inferredTypeWithOptionalFields": JsonValue::from(&summary.inferred_type)
179+
}),
180+
None => json!({
181+
"totalSize": JsonInteger::encode(summary.total_size),
182+
"inferredTypeWithOptionalFields": JsonValue::from(&summary.inferred_type)
183+
}),
184+
}
150185
}
151186
}
152187

@@ -161,13 +196,22 @@ impl TryFrom<JsonValue> for TableSummary {
161196
_ => anyhow::bail!("Invalid totalSize"),
162197
};
163198
anyhow::ensure!(total_size >= 0);
199+
let total_size_rounded = match v.remove("totalSizeRounded") {
200+
Some(JsonValue::String(s)) => {
201+
let total_size_rounded = JsonInteger::decode(s)?;
202+
anyhow::ensure!(total_size_rounded >= 0);
203+
Some(total_size_rounded)
204+
},
205+
_ => None,
206+
};
164207
let inferred_type = match v.remove("inferredTypeWithOptionalFields") {
165208
Some(v) => CountedShape::<ProdConfigWithOptionalFields>::try_from(v)?,
166209
None => anyhow::bail!("Missing field inferredTypeWithOptionalFields"),
167210
};
168211
Ok(TableSummary {
169212
inferred_type,
170213
total_size,
214+
total_size_rounded,
171215
})
172216
},
173217
_ => anyhow::bail!("Wrong type of json value for TableSummaryJson"),

0 commit comments

Comments
 (0)