fix: make memdisk safer

cmu-db · May 1, 2024 · 88a913a · 88a913a
1 parent b85f491
commit 88a913a
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/storage-node/src/cache/data_store_cache/memdisk/mod.rs b/storage-node/src/cache/data_store_cache/memdisk/mod.rs
@@ -154,6 +154,12 @@ impl<R: DataStoreReplacer<MemDiskStoreReplacerKey, MemDiskStoreReplacerValue>>
         let notify;
         {
             let mut status_of_keys = self.status_of_keys.write().await;
+            if status_of_keys.get(key).is_none() {
+                warn!("Notify waiters mem for key {}: key not found", key);
+                drop(status_of_keys);
+                self.notify_waiters_error(key).await;
+                return;
+            }
             let ((status, size), notify_ref) = status_of_keys.get_mut(key).unwrap();
             *status = Status::MemCompleted;
             debug!(
@@ -186,6 +192,12 @@ impl<R: DataStoreReplacer<MemDiskStoreReplacerKey, MemDiskStoreReplacerValue>>
         let notify;
         {
             let mut status_of_keys = self.status_of_keys.write().await;
+            if status_of_keys.get(key).is_none() {
+                warn!("Notify waiters mem for key {}: key not found", key);
+                drop(status_of_keys);
+                self.notify_waiters_error(key).await;
+                return;
+            }
             let ((status, size), notify_ref) = status_of_keys.get_mut(key).unwrap();
             *status = Status::DiskCompleted;
             debug!(
@@ -561,6 +573,7 @@ impl<R: DataStoreReplacer<MemDiskStoreReplacerKey, MemDiskStoreReplacerValue>> D
                                 "MemDiskStore put_data_to_cache: 3 remove status {:?} for key {}",
                                 status, evicted_key
                             );
+                            assert!(status != Status::Incompleted);
                         }
                         self.disk_store
                             .clean_data(&self.disk_store.data_store_key(&evicted_key))