@rivetkit/engine 1.0.0 → 2.2.1-pr.4600.b74ff3b

package/packages/cache/src/driver.rs

@@ -1,21 +1,18 @@
 use std::{
 	fmt::Debug,
+	sync::OnceLock,
 	time::{Duration, Instant},
 };
 
 use moka::future::{Cache, CacheBuilder};
-use serde::{Serialize, de::DeserializeOwned};
 use tracing::Instrument;
 
-use rivet_metrics::KeyValue;
-
-use crate::{errors::Error, metrics};
+use crate::{RawCacheKey, errors::Error};
 
 /// Type alias for cache values stored as bytes
 pub type CacheValue = Vec<u8>;
 
 /// Enum wrapper for different cache driver implementations
-#[derive(Debug, Clone)]
 #[non_exhaustive]
 pub enum Driver {
 	InMemory(InMemoryDriver),
@@ -24,37 +21,37 @@ pub enum Driver {
 impl Driver {
 	/// Fetch multiple values from cache at once
 	#[tracing::instrument(skip_all, fields(driver=%self))]
-	pub async fn fetch_values<'a>(
+	pub async fn get<'a>(
 		&'a self,
 		base_key: &'a str,
-		keys: &[String],
+		keys: &[RawCacheKey],
 	) -> Result<Vec<Option<CacheValue>>, Error> {
 		match self {
-			Driver::InMemory(d) => d.fetch_values(base_key, keys).await,
+			Driver::InMemory(d) => d.get(base_key, keys).await,
 		}
 	}
 
 	/// Set multiple values in cache at once
 	#[tracing::instrument(skip_all, fields(driver=%self))]
-	pub async fn set_values<'a>(
+	pub async fn set<'a>(
 		&'a self,
 		base_key: &'a str,
-		keys_values: Vec<(String, CacheValue, i64)>,
+		keys_values: Vec<(RawCacheKey, CacheValue, i64)>,
 	) -> Result<(), Error> {
 		match self {
-			Driver::InMemory(d) => d.set_values(base_key, keys_values).await,
+			Driver::InMemory(d) => d.set(base_key, keys_values).await,
 		}
 	}
 
 	/// Delete multiple keys from cache
 	#[tracing::instrument(skip_all, fields(driver=%self))]
-	pub async fn delete_keys<'a>(
+	pub async fn delete<'a>(
 		&'a self,
 		base_key: &'a str,
-		keys: Vec<String>,
+		keys: Vec<RawCacheKey>,
 	) -> Result<(), Error> {
 		match self {
-			Driver::InMemory(d) => d.delete_keys(base_key, keys).await,
+			Driver::InMemory(d) => d.delete(base_key, keys).await,
 		}
 	}
 
@@ -62,46 +59,11 @@ impl Driver {
 	///
 	/// Different implementations use different key formats:
 	/// - In-memory uses simpler keys
-	pub fn process_key(&self, base_key: &str, key: &impl crate::CacheKey) -> String {
+	pub fn process_key(&self, base_key: &str, key: &impl crate::CacheKey) -> RawCacheKey {
 		match self {
 			Driver::InMemory(d) => d.process_key(base_key, key),
 		}
 	}
-
-	/// Process a rate limit key into a driver-specific format
-	pub fn process_rate_limit_key(
-		&self,
-		key: &impl crate::CacheKey,
-		remote_address: impl AsRef<str>,
-		bucket: i64,
-		bucket_duration_ms: i64,
-	) -> String {
-		match self {
-			Driver::InMemory(d) => {
-				d.process_rate_limit_key(key, remote_address, bucket, bucket_duration_ms)
-			}
-		}
-	}
-
-	/// Increment a rate limit counter and return the new count
-	#[tracing::instrument(skip_all, fields(driver=%self))]
-	pub async fn rate_limit_increment<'a>(
-		&'a self,
-		key: &'a str,
-		ttl_ms: i64,
-	) -> Result<i64, Error> {
-		match self {
-			Driver::InMemory(d) => d.rate_limit_increment(key, ttl_ms).await,
-		}
-	}
-
-	pub fn encode_value<T: Serialize>(&self, value: &T) -> Result<CacheValue, Error> {
-		serde_json::to_vec(value).map_err(Error::SerdeEncode)
-	}
-
-	pub fn decode_value<T: DeserializeOwned>(&self, bytes: &[u8]) -> Result<T, Error> {
-		serde_json::from_slice(bytes).map_err(Error::SerdeDecode)
-	}
 }
 
 impl std::fmt::Display for Driver {
@@ -122,7 +84,7 @@ struct ExpiringValue {
 }
 
 /// Cache expiry implementation for Moka
-#[derive(Clone, Debug)]
+#[derive(Debug)]
 struct ValueExpiry;
 
 impl moka::Expiry<String, ExpiringValue> for ValueExpiry {
@@ -171,56 +133,47 @@ impl moka::Expiry<String, ExpiringValue> for ValueExpiry {
 	}
 }
 
+static CACHE: OnceLock<Cache<String, ExpiringValue>> = OnceLock::new();
+
 /// In-memory cache driver implementation using the moka crate
-#[derive(Clone)]
-pub struct InMemoryDriver {
-	service_name: String,
-	cache: Cache<String, ExpiringValue>,
-	/// In-memory rate limiting store - maps keys to hit counts with expiration
-	rate_limits: Cache<String, ExpiringValue>,
-}
+pub struct InMemoryDriver {}
 
 impl Debug for InMemoryDriver {
 	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-		f.debug_struct("InMemoryDriver")
-			.field("service_name", &self.service_name)
-			.finish()
+		f.debug_struct("InMemoryDriver").finish()
 	}
 }
 
 impl InMemoryDriver {
-	pub fn new(service_name: String, max_capacity: u64) -> Self {
+	pub fn new(max_capacity: u64) -> Self {
 		// Create a cache with ValueExpiry implementation for custom expiration times
-		let cache = CacheBuilder::new(max_capacity)
-			.expire_after(ValueExpiry)
-			.build();
+		CACHE.get_or_init(|| {
+			CacheBuilder::new(max_capacity)
+				.expire_after(ValueExpiry)
+				.eviction_listener(|key, _value, cause| {
+					tracing::debug!(?key, ?cause, "cache eviction");
+				})
+				.build()
+		});
 
-		// Create a separate cache for rate limiting with the same expiry mechanism
-		let rate_limits = CacheBuilder::new(max_capacity)
-			.expire_after(ValueExpiry)
-			.build();
+		Self {}
+	}
 
-		Self {
-			service_name,
-			cache,
-			rate_limits,
-		}
+	fn cache(&self) -> &Cache<String, ExpiringValue> {
+		CACHE.get().expect("should be initialized")
 	}
 
-	pub async fn fetch_values<'a>(
+	pub async fn get<'a>(
 		&'a self,
 		_base_key: &'a str,
-		keys: &[String],
+		keys: &[RawCacheKey],
 	) -> Result<Vec<Option<CacheValue>>, Error> {
-		let keys = keys.to_vec();
-		let cache = self.cache.clone();
-
 		let mut result = Vec::with_capacity(keys.len());
 
 		// Async block for metrics
 		async {
 			for key in keys {
-				result.push(cache.get(&key).await.map(|x| x.value.clone()));
+				result.push(self.cache().get(&**key).await.map(|x| x.value.clone()));
 			}
 		}
 		.instrument(tracing::info_span!("get"))
@@ -235,13 +188,11 @@ impl InMemoryDriver {
 		Ok(result)
 	}
 
-	pub async fn set_values<'a>(
+	pub async fn set<'a>(
 		&'a self,
 		_base_key: &'a str,
-		keys_values: Vec<(String, CacheValue, i64)>,
+		keys_values: Vec<(RawCacheKey, CacheValue, i64)>,
 	) -> Result<(), Error> {
-		let cache = self.cache.clone();
-
 		// Async block for metrics
 		async {
 			for (key, value, expire_at) in keys_values {
@@ -252,7 +203,7 @@ impl InMemoryDriver {
 				};
 
 				// Store in cache - expiry will be handled by ValueExpiry
-				cache.insert(key, entry).await;
+				self.cache().insert(key.into(), entry).await;
 			}
 		}
 		.instrument(tracing::info_span!("set"))
@@ -262,85 +213,26 @@ impl InMemoryDriver {
 		Ok(())
 	}
 
-	pub async fn delete_keys<'a>(
+	pub async fn delete<'a>(
 		&'a self,
-		base_key: &'a str,
-		keys: Vec<String>,
+		_base_key: &'a str,
+		keys: Vec<RawCacheKey>,
 	) -> Result<(), Error> {
-		let cache = self.cache.clone();
-		let base_key = base_key.to_string();
-
-		metrics::CACHE_PURGE_REQUEST_TOTAL.add(1, &[KeyValue::new("key", base_key.clone())]);
-		metrics::CACHE_PURGE_VALUE_TOTAL
-			.add(keys.len() as u64, &[KeyValue::new("key", base_key.clone())]);
-
 		// Async block for metrics
 		async {
 			for key in keys {
 				// Use remove instead of invalidate to ensure it's actually removed
-				cache.remove(&key).await;
+				self.cache().remove(&*key).await;
 			}
 		}
-		.instrument(tracing::info_span!("remove"))
+		.instrument(tracing::info_span!("delete"))
 		.await;
 
 		tracing::trace!("successfully deleted keys from in-memory cache");
 		Ok(())
 	}
 
-	pub fn process_key(&self, base_key: &str, key: &impl crate::CacheKey) -> String {
-		format!("{}:{}", base_key, key.cache_key())
-	}
-
-	pub fn process_rate_limit_key(
-		&self,
-		key: &impl crate::CacheKey,
-		remote_address: impl AsRef<str>,
-		bucket: i64,
-		bucket_duration_ms: i64,
-	) -> String {
-		format!(
-			"rate_limit:{}:{}:{}:{}",
-			key.cache_key(),
-			remote_address.as_ref(),
-			bucket_duration_ms,
-			bucket,
-		)
-	}
-
-	/// Increment a rate limit counter for in-memory storage
-	pub async fn rate_limit_increment<'a>(
-		&'a self,
-		key: &'a str,
-		ttl_ms: i64,
-	) -> Result<i64, Error> {
-		let rate_limits = self.rate_limits.clone();
-
-		// Get current value or default to 0
-		let current_value = match rate_limits.get(key).await {
-			Some(value) => {
-				// Try to decode the value as an integer
-				match serde_json::from_slice::<i64>(&value.value).map_err(Error::SerdeDecode) {
-					Ok(count) => count,
-					Err(_) => 0, // If we can't decode, reset to 0
-				}
-			}
-			None => 0,
-		};
-
-		// Increment the value
-		let new_value = current_value + 1;
-		let encoded = serde_json::to_vec(&new_value).map_err(Error::SerdeEncode)?;
-
-		// Store with expiration
-		let entry = ExpiringValue {
-			value: encoded,
-			expiry_time: rivet_util::timestamp::now() + ttl_ms,
-		};
-
-		// Update the rate limit cache
-		rate_limits.insert(key.to_string(), entry).await;
-
-		Ok(new_value)
+	pub fn process_key(&self, base_key: &str, key: &impl crate::CacheKey) -> RawCacheKey {
+		RawCacheKey::from(format!("{}:{}", base_key, key.cache_key()))
 	}
 }

package/packages/cache/src/getter_ctx.rs

@@ -1,13 +1,10 @@
-use std::fmt::Debug;
+use std::{collections::HashMap, fmt::Debug};
 
 use super::*;
 
 /// Entry for a single value that is going to be read/written to the cache.
 #[derive(Debug)]
-pub(super) struct GetterCtxKey<K, V> {
-	/// `CacheKey` that will be used to build Redis keys.
-	pub(super) key: K,
-
+pub(super) struct GetterCtxEntry<V> {
 	/// The value that was read from the cache or getter.
 	value: Option<V>,
 
@@ -23,81 +20,68 @@ pub struct GetterCtx<K, V>
 where
 	K: CacheKey,
 {
-	/// The name of the service-specific key to write this cached value to. For
-	/// example, a team get service would use the "team_profile" key to store
-	/// the profile a "team_members" to store a cache of members.
-	///
-	/// This is local to the service & source hash that caches this value.
-	#[allow(unused)]
-	base_key: String,
-
-	/// The keys to get/populate from the cache.
-	keys: Vec<GetterCtxKey<K, V>>,
+	/// The entries to get/populate from the cache.
+	entries: HashMap<K, GetterCtxEntry<V>>,
 }
 
 impl<K, V> GetterCtx<K, V>
 where
 	K: CacheKey,
 {
-	pub(super) fn new(base_key: String, keys: Vec<K>) -> Self {
+	pub(super) fn new(keys: Vec<K>) -> Self {
 		GetterCtx {
-			base_key,
-			keys: {
-				// Create deduplicated ctx keys
-				let mut ctx_keys = Vec::<GetterCtxKey<K, V>>::new();
-				for key in keys {
-					if !ctx_keys.iter().any(|x| x.key == key) {
-						ctx_keys.push(GetterCtxKey {
-							key,
+			entries: keys
+				.into_iter()
+				.map(|k| {
+					(
+						k,
+						GetterCtxEntry {
 							value: None,
 							from_cache: false,
-						});
-					}
-				}
-				ctx_keys
-			},
+						},
+					)
+				})
+				.collect(),
 		}
 	}
 
+	pub(super) fn merge(&mut self, other: GetterCtx<K, V>) {
+		self.entries.extend(other.entries);
+	}
+
 	pub(super) fn into_values(self) -> Vec<(K, V)> {
-		self.keys
+		self.entries
 			.into_iter()
-			.filter_map(|k| {
-				if let Some(v) = k.value {
-					Some((k.key, v))
-				} else {
-					None
-				}
-			})
+			.filter_map(|(k, x)| x.value.map(|v| (k, v)))
 			.collect()
 	}
 
-	/// All keys.
-	pub(super) fn keys(&self) -> &[GetterCtxKey<K, V>] {
-		&self.keys[..]
+	/// All entries.
+	pub(super) fn entries(&self) -> impl Iterator<Item = (&K, &GetterCtxEntry<V>)> {
+		self.entries.iter()
 	}
 
-	/// If all keys have an associated value.
-	pub(super) fn all_keys_have_value(&self) -> bool {
-		self.keys.iter().all(|x| x.value.is_some())
+	/// If all entries have an associated value.
+	pub(super) fn all_entries_have_value(&self) -> bool {
+		self.entries.iter().all(|(_, x)| x.value.is_some())
 	}
 
 	/// Keys that do not have a value yet.
 	pub(super) fn unresolved_keys(&self) -> Vec<K> {
-		self.keys
+		self.entries
 			.iter()
-			.filter(|x| x.value.is_none())
-			.map(|x| x.key.clone())
+			.filter(|(_, x)| x.value.is_none())
+			.map(|(k, _)| k.clone())
 			.collect()
 	}
 
-	/// Keys that have been resolved in a getter and need to be written to the
+	/// Entries that have been resolved in a getter and need to be written to the
 	/// cache.
-	pub(super) fn values_needing_cache_write(&self) -> Vec<(&GetterCtxKey<K, V>, &V)> {
-		self.keys
+	pub(super) fn entries_needing_cache_write(&self) -> Vec<(&K, &V)> {
+		self.entries
 			.iter()
-			.filter(|x| !x.from_cache)
-			.filter_map(|k| k.value.as_ref().map(|v| (k, v)))
+			.filter(|(_, x)| !x.from_cache)
+			.filter_map(|(k, x)| x.value.as_ref().map(|v| (k, v)))
 			.collect()
 	}
 }
@@ -107,32 +91,26 @@ where
 	K: CacheKey,
 	V: Debug,
 {
-	/// Sets a value with the value provided from the cache.
-	pub(super) fn resolve_from_cache(&mut self, idx: usize, value: V) {
-		if let Some(key) = self.keys.get_mut(idx) {
-			key.value = Some(value);
-			key.from_cache = true;
+	/// Sets an entry with the value provided from the cache.
+	pub(super) fn resolve_from_cache(&mut self, key: &K, value: V) {
+		if let Some(entry) = self.entries.get_mut(key) {
+			entry.value = Some(value);
+			entry.from_cache = true;
 		} else {
-			tracing::warn!(?idx, ?value, "resolving cache key index out of range");
+			tracing::warn!(?key, ?value, "resolving nonexistent cache entry");
 		}
 	}
 
-	/// Calls the callback with a mutable reference to a given key. Validates
-	/// that the key does not already have a value.
-	fn get_key_for_resolve(&mut self, key: &K, cb: impl FnOnce(&mut GetterCtxKey<K, V>)) {
-		if let Some(key) = self.keys.iter_mut().find(|x| x.key == *key) {
-			if key.value.is_some() {
-				tracing::warn!(?key, "cache key already has value");
+	/// Sets a value with the value provided from the getter function.
+	pub fn resolve(&mut self, key: &K, value: V) {
+		if let Some(entry) = self.entries.get_mut(key) {
+			if entry.value.is_some() {
+				tracing::warn!(?entry, "cache entry already has value");
 			} else {
-				cb(key);
+				entry.value = Some(value);
 			}
 		} else {
-			tracing::warn!(?key, "resolved value for nonexistent cache key");
+			tracing::warn!(?key, "resolved value for nonexistent cache entry");
 		}
 	}
-
-	/// Sets a value with the value provided from the getter function.
-	pub fn resolve(&mut self, key: &K, value: V) {
-		self.get_key_for_resolve(key, |key| key.value = Some(value));
-	}
 }

package/packages/cache/src/inner.rs

@@ -1,22 +1,26 @@
-use std::{fmt::Debug, sync::Arc};
+use std::{
+	fmt::Debug,
+	sync::{Arc, OnceLock},
+};
+
+use tokio::sync::broadcast;
 
 use super::*;
 use crate::driver::{Driver, InMemoryDriver};
 
+static IN_FLIGHT: OnceLock<scc::HashMap<RawCacheKey, broadcast::Sender<()>>> = OnceLock::new();
+
 pub type Cache = Arc<CacheInner>;
 
 /// Utility type used to hold information relating to caching.
 pub struct CacheInner {
-	service_name: String,
-	pub(crate) driver: Driver,
+	pub(crate) driver: Option<Driver>,
 	pub(crate) ups: Option<universalpubsub::PubSub>,
 }
 
 impl Debug for CacheInner {
 	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-		f.debug_struct("CacheInner")
-			.field("service_name", &self.service_name)
-			.finish()
+		f.debug_struct("CacheInner").finish()
 	}
 }
 
@@ -26,30 +30,36 @@ impl CacheInner {
 		config: &rivet_config::Config,
 		pools: rivet_pools::Pools,
 	) -> Result<Cache, Error> {
-		let service_name = rivet_env::service_name();
 		let ups = pools.ups().ok();
 
 		match &config.cache().driver {
-			rivet_config::config::CacheDriver::Redis => todo!(),
-			rivet_config::config::CacheDriver::InMemory => {
-				Ok(Self::new_in_memory(service_name.to_string(), 1000, ups))
+			Some(rivet_config::config::CacheDriver::InMemory) => {
+				Ok(Self::new_in_memory(10000, ups))
 			}
+			None => Ok(Self::new_disabled()),
 		}
 	}
 
 	#[tracing::instrument(skip(ups))]
-	pub fn new_in_memory(
-		service_name: String,
-		max_capacity: u64,
-		ups: Option<universalpubsub::PubSub>,
-	) -> Cache {
-		let driver = Driver::InMemory(InMemoryDriver::new(service_name.clone(), max_capacity));
+	pub fn new_in_memory(max_capacity: u64, ups: Option<universalpubsub::PubSub>) -> Cache {
+		let driver = Driver::InMemory(InMemoryDriver::new(max_capacity));
+
 		Arc::new(CacheInner {
-			service_name,
-			driver,
+			driver: Some(driver),
 			ups,
 		})
 	}
+
+	pub fn new_disabled() -> Cache {
+		Arc::new(CacheInner {
+			driver: None,
+			ups: None,
+		})
+	}
+
+	pub(crate) fn in_flight(&self) -> &scc::HashMap<RawCacheKey, broadcast::Sender<()>> {
+		IN_FLIGHT.get_or_init(scc::HashMap::new)
+	}
 }
 
 impl CacheInner {

package/packages/cache/src/key.rs

@@ -1,8 +1,8 @@
 use serde::{Deserialize, Serialize};
-use std::fmt::Debug;
+use std::{fmt::Debug, hash::Hash, ops::Deref};
 
 /// A type that can be serialized in to a key that can be used in the cache.
-pub trait CacheKey: Clone + Debug + PartialEq {
+pub trait CacheKey: Clone + Debug + Eq + PartialEq + Hash {
 	fn cache_key(&self) -> String;
 }
 
@@ -85,7 +85,7 @@ impl_to_string!(isize);
 /// Unlike other types that implement `CacheKey` (which escape special characters like `:` and `\`),
 /// `RawCacheKey` uses the provided string as-is. This is useful when you already have a properly
 /// formatted cache key string or need to preserve the exact format without transformations.
-#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
+#[derive(Clone, Debug, Serialize, Deserialize, Hash, Eq, PartialEq)]
 pub struct RawCacheKey(String);
 
 impl CacheKey for RawCacheKey {
@@ -99,3 +99,17 @@ impl From<String> for RawCacheKey {
 		RawCacheKey(value)
 	}
 }
+
+impl From<RawCacheKey> for String {
+	fn from(value: RawCacheKey) -> Self {
+		value.0
+	}
+}
+
+impl Deref for RawCacheKey {
+	type Target = String;
+
+	fn deref(&self) -> &Self::Target {
+		&self.0
+	}
+}

package/packages/cache/src/lib.rs

@@ -5,7 +5,6 @@ mod inner;
 mod key;
 mod metrics;
 mod purge;
-mod rate_limit;
 mod req_config;
 
 // Re-export public API
@@ -15,5 +14,4 @@ pub use getter_ctx::*;
 pub use inner::*;
 pub use key::*;
 pub use purge::*;
-pub use rate_limit::*;
 pub use req_config::*;