Sharded in-memory store

crates/store/src/dispatch/lookup.rs

@@ -20,9 +20,9 @@ use crate::{
 };
 
 pub struct KeyValue<T> {
-    key: Vec<u8>,
-    value: T,
-    expires: Option<u64>,
+    pub key: Vec<u8>,
+    pub value: T,
+    pub expires: Option<u64>,
 }
 
 impl InMemoryStore {
@@ -44,6 +44,8 @@ impl InMemoryStore {
             }
             #[cfg(feature = "redis")]
             InMemoryStore::Redis(store) => store.key_set(&kv.key, &kv.value, kv.expires).await,
+            #[cfg(feature = "enterprise")]
+            InMemoryStore::Sharded(store) => store.key_set(kv).await,
             InMemoryStore::Static(_) | InMemoryStore::Http(_) => {
                 Err(trc::StoreEvent::NotSupported.into_err())
             }
@@ -81,6 +83,8 @@ impl InMemoryStore {
             }
             #[cfg(feature = "redis")]
             InMemoryStore::Redis(store) => store.key_incr(&kv.key, kv.value, kv.expires).await,
+            #[cfg(feature = "enterprise")]
+            InMemoryStore::Sharded(store) => store.counter_incr(kv).await,
             InMemoryStore::Static(_) | InMemoryStore::Http(_) => {
                 Err(trc::StoreEvent::NotSupported.into_err())
             }
@@ -100,6 +104,8 @@ impl InMemoryStore {
             }
             #[cfg(feature = "redis")]
             InMemoryStore::Redis(store) => store.key_delete(key.into().as_bytes()).await,
+            #[cfg(feature = "enterprise")]
+            InMemoryStore::Sharded(store) => store.key_delete(key).await,
             InMemoryStore::Static(_) | InMemoryStore::Http(_) => {
                 Err(trc::StoreEvent::NotSupported.into_err())
             }
@@ -119,6 +125,8 @@ impl InMemoryStore {
             }
             #[cfg(feature = "redis")]
             InMemoryStore::Redis(store) => store.key_delete(key.into().as_bytes()).await,
+            #[cfg(feature = "enterprise")]
+            InMemoryStore::Sharded(store) => store.counter_delete(key).await,
             InMemoryStore::Static(_) | InMemoryStore::Http(_) => {
                 Err(trc::StoreEvent::NotSupported.into_err())
             }
@@ -156,6 +164,8 @@ impl InMemoryStore {
             }
             #[cfg(feature = "redis")]
             InMemoryStore::Redis(store) => store.key_delete_prefix(prefix).await,
+            #[cfg(feature = "enterprise")]
+            InMemoryStore::Sharded(store) => store.key_delete_prefix(prefix).await,
             InMemoryStore::Static(_) | InMemoryStore::Http(_) => {
                 Err(trc::StoreEvent::NotSupported.into_err())
             }
@@ -176,6 +186,8 @@ impl InMemoryStore {
                 .map(|value| value.and_then(|v| v.into())),
             #[cfg(feature = "redis")]
             InMemoryStore::Redis(store) => store.key_get(key.into().as_bytes()).await,
+            #[cfg(feature = "enterprise")]
+            InMemoryStore::Sharded(store) => store.key_get(key).await,
             InMemoryStore::Static(store) => Ok(store
                 .get(key.into().as_str())
                 .map(|value| T::from(value.clone()))),
@@ -197,6 +209,8 @@ impl InMemoryStore {
             }
             #[cfg(feature = "redis")]
             InMemoryStore::Redis(store) => store.counter_get(key.into().as_bytes()).await,
+            #[cfg(feature = "enterprise")]
+            InMemoryStore::Sharded(store) => store.counter_get(key).await,
             InMemoryStore::Static(_) | InMemoryStore::Http(_) => {
                 Err(trc::StoreEvent::NotSupported.into_err())
             }
@@ -214,6 +228,8 @@ impl InMemoryStore {
                 .map(|value| matches!(value, Some(LookupValue::Value(())))),
             #[cfg(feature = "redis")]
             InMemoryStore::Redis(store) => store.key_exists(key.into().as_bytes()).await,
+            #[cfg(feature = "enterprise")]
+            InMemoryStore::Sharded(store) => store.key_exists(key).await,
             InMemoryStore::Static(store) => Ok(store.get(key.into().as_str()).is_some()),
             InMemoryStore::Http(store) => Ok(store.contains(key.into().as_str())),
         }
@@ -345,6 +361,8 @@ impl InMemoryStore {
             }
             #[cfg(feature = "redis")]
             InMemoryStore::Redis(_) => {}
+            #[cfg(feature = "enterprise")]
+            InMemoryStore::Sharded(_) => {}
             InMemoryStore::Static(_) | InMemoryStore::Http(_) => {}
         }
 
@@ -357,6 +375,15 @@ impl InMemoryStore {
             _ => false,
         }
     }
+
+    pub fn is_redis(&self) -> bool {
+        match self {
+            #[cfg(feature = "redis")]
+            InMemoryStore::Redis(_) => true,
+            InMemoryStore::Static(_) => false,
+            _ => false,
+        }
+    }
 }
 
 pub enum LookupKey<'x> {