Fix update_stats and add all tests

cmu-db · Nov 8, 2024 · 52d290a · 52d290a
1 parent 43e66c3
commit 52d290a
Show file tree

Hide file tree

Showing 4 changed files with 221 additions and 65 deletions.
diff --git a/optd-persistent/init.db b/optd-persistent/init.db
diff --git a/optd-persistent/src/cost_model/interface.rs b/optd-persistent/src/cost_model/interface.rs
@@ -42,6 +42,14 @@ pub enum StatType {
     Max,
 }
 
+#[derive(PartialEq)]
+pub enum EpochOption {
+    // TODO(lanlou): Could I make i32 -> EpochId?
+    Existed(i32),
+    New(String, String),
+}
+
+#[derive(Clone)]
 pub struct Stat {
     pub stat_type: i32,
     pub stat_value: Json,
@@ -72,7 +80,11 @@ pub trait CostModelStorageLayer {
         epoch_id: Self::EpochId,
     ) -> StorageResult<()>;
 
-    async fn update_stats(&self, stat: Stat, epoch_id: Self::EpochId) -> StorageResult<()>;
+    async fn update_stats(
+        &mut self,
+        stat: Stat,
+        epoch_option: EpochOption,
+    ) -> StorageResult<Option<Self::EpochId>>;
 
     async fn store_cost(
         &self,