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

package/packages/cache/src/metrics.rs

@@ -1,34 +1,46 @@
-use rivet_metrics::otel::{global::*, metrics::*};
+use rivet_metrics::{REGISTRY, prometheus::*};
 
 lazy_static::lazy_static! {
-	static ref METER: Meter = meter("rivet-cache");
-
-	/// Expected attributes: "key"
-	pub static ref CACHE_REQUEST_TOTAL: Counter<u64> = METER.u64_counter("rivet_cache_request_total")
-		.with_description("Total number of cache requests.")
-		.build();
-	/// Expected attributes: "key"
-	pub static ref CACHE_REQUEST_ERRORS: Counter<u64> = METER.u64_counter("rivet_cache_request_errors")
-		.with_description("Total number of cache request errors.")
-		.build();
-	/// Expected attributes: "key"
-	pub static ref CACHE_PURGE_REQUEST_TOTAL: Counter<u64> = METER.u64_counter("rivet_cache_purge_request_total")
-		.with_description("Total number of cache purge requests.")
-		.build();
-	/// Expected attributes: "key"
-	pub static ref CACHE_VALUE_TOTAL: Counter<u64> = METER.u64_counter("rivet_cache_value_total")
-		.with_description("Total number of cache values requested.")
-		.build();
-	/// Expected attributes: "key"
-	pub static ref CACHE_VALUE_MISS_TOTAL: Counter<u64> = METER.u64_counter("rivet_cache_value_miss_total")
-		.with_description("Total number of cache value misses.")
-		.build();
-	/// Expected attributes: "key"
-	pub static ref CACHE_VALUE_EMPTY_TOTAL: Counter<u64> = METER.u64_counter("rivet_cache_value_empty_total")
-		.with_description("Total number of cache values that were requested but not resolved by cache nor getter.")
-		.build();
-	/// Expected attributes: "key"
-	pub static ref CACHE_PURGE_VALUE_TOTAL: Counter<u64> = METER.u64_counter("rivet_cache_purge_value_total")
-		.with_description("Total number of cache values purged.")
-		.build();
+	pub static ref CACHE_REQUEST_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"cache_request_total",
+		"Total number of cache requests.",
+		&["key"],
+		*REGISTRY
+	).unwrap();
+	pub static ref CACHE_REQUEST_ERRORS: IntCounterVec = register_int_counter_vec_with_registry!(
+		"cache_request_errors",
+		"Total number of cache request errors.",
+		&["key"],
+		*REGISTRY
+	).unwrap();
+	pub static ref CACHE_PURGE_REQUEST_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"cache_purge_request_total",
+		"Total number of cache purge requests.",
+		&["key"],
+		*REGISTRY
+	).unwrap();
+	pub static ref CACHE_VALUE_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"cache_value_total",
+		"Total number of cache values requested.",
+		&["key"],
+		*REGISTRY
+	).unwrap();
+	pub static ref CACHE_VALUE_MISS_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"cache_value_miss_total",
+		"Total number of cache value misses.",
+		&["key"],
+		*REGISTRY
+	).unwrap();
+	pub static ref CACHE_VALUE_EMPTY_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"cache_value_empty_total",
+		"Total number of cache values that were requested but not resolved by cache nor getter.",
+		&["key"],
+		*REGISTRY
+	).unwrap();
+	pub static ref CACHE_PURGE_VALUE_TOTAL: IntCounterVec = register_int_counter_vec_with_registry!(
+		"cache_purge_value_total",
+		"Total number of cache values purged.",
+		&["key"],
+		*REGISTRY
+	).unwrap();
 }

package/packages/cache/src/req_config.rs

@@ -1,12 +1,16 @@
-use std::{fmt::Debug, future::Future, result::Result::Ok};
+use std::{fmt::Debug, future::Future, time::Duration};
 
-use anyhow::*;
+use anyhow::Result;
+use futures_util::StreamExt;
+use itertools::{Either, Itertools};
 use serde::{Serialize, de::DeserializeOwned};
-use tracing::Instrument;
+use tokio::sync::broadcast;
 
 use super::*;
 use crate::{errors::Error, metrics};
-use rivet_metrics::KeyValue;
+
+/// How long to wait for an in flight cache req before proceeding to execute the same req anyway.
+const IN_FLIGHT_TIMEOUT: Duration = Duration::from_secs(5);
 
 /// Config specifying how cached values will behave.
 #[derive(Clone)]
@@ -43,66 +47,8 @@ impl RequestConfig {
 
 // MARK: Fetch
 impl RequestConfig {
-	/// Attempts to fetch a given key from the database and falls back to the
-	/// getter if not found. If both the cache and the getter return `None`,
-	/// then this returns `None`.
-	///
-	/// See `fetch_all` for more details.
-	#[tracing::instrument(err, skip(key, getter))]
-	pub async fn fetch_one<K, V, Getter, Fut>(
-		self,
-		base_key: impl ToString + Debug,
-		key: K,
-		getter: Getter,
-	) -> Result<Option<V>>
-	where
-		K: CacheKey + Send + Sync,
-		V: Serialize + DeserializeOwned + Clone + Debug + Send + Sync,
-		Getter: Fn(GetterCtx<K, V>, K) -> Fut + Clone,
-		Fut: Future<Output = Result<GetterCtx<K, V>>>,
-	{
-		let values = self
-			.fetch_all(base_key, [key], move |cache, keys| {
-				let getter = getter.clone();
-				async move {
-					debug_assert_eq!(1, keys.len());
-					if let Some(key) = keys.into_iter().next() {
-						getter(cache, key).await
-					} else {
-						tracing::error!("no keys provided to fetch one");
-						Ok(cache)
-					}
-				}
-			})
-			.await?;
-		Ok(values.into_iter().next().map(|(_, v)| v))
-	}
-
-	#[tracing::instrument(err, skip(keys, getter))]
-	pub async fn fetch_all<Key, Value, Getter, Fut>(
-		self,
-		base_key: impl ToString + Debug,
-		keys: impl IntoIterator<Item = Key>,
-		getter: Getter,
-	) -> Result<Vec<(Key, Value)>>
-	where
-		Key: CacheKey + Send + Sync,
-		Value: Serialize + DeserializeOwned + Clone + Debug + Send + Sync,
-		Getter: Fn(GetterCtx<Key, Value>, Vec<Key>) -> Fut + Clone,
-		Fut: Future<Output = Result<GetterCtx<Key, Value>>>,
-	{
-		self.fetch_all_convert(
-			base_key,
-			keys,
-			getter,
-			|x: &Value| Ok(x.clone()),
-			|x: &Value| Ok(x.clone()),
-		)
-		.await
-	}
-
-	#[tracing::instrument(err, skip(keys, getter, encoder, decoder))]
-	async fn fetch_all_convert<Key, Value, ValueSerde, Getter, Fut, Encoder, Decoder>(
+	#[tracing::instrument(err, skip_all, fields(?base_key))]
+	async fn fetch_all_convert<Key, Value, Getter, Fut, Encoder, Decoder>(
 		self,
 		base_key: impl ToString + Debug,
 		keys: impl IntoIterator<Item = Key>,
@@ -113,11 +59,10 @@ impl RequestConfig {
 	where
 		Key: CacheKey + Send + Sync,
 		Value: Debug + Send + Sync,
-		ValueSerde: Serialize + DeserializeOwned + Debug + Send + Sync,
 		Getter: Fn(GetterCtx<Key, Value>, Vec<Key>) -> Fut + Clone,
 		Fut: Future<Output = Result<GetterCtx<Key, Value>>>,
-		Encoder: Fn(&Value) -> Result<ValueSerde> + Clone,
-		Decoder: Fn(&ValueSerde) -> Result<Value> + Clone,
+		Encoder: Fn(&Value) -> Result<Vec<u8>> + Clone,
+		Decoder: Fn(&[u8]) -> Result<Value> + Clone,
 	{
 		let base_key = base_key.to_string();
 		let keys = keys.into_iter().collect::<Vec<Key>>();
@@ -127,45 +72,51 @@ impl RequestConfig {
 			return Ok(Vec::new());
 		}
 
-		metrics::CACHE_REQUEST_TOTAL.add(1, &[KeyValue::new("key", base_key.clone())]);
+		metrics::CACHE_REQUEST_TOTAL
+			.with_label_values(&[base_key.as_str()])
+			.inc();
 		metrics::CACHE_VALUE_TOTAL
-			.add(keys.len() as u64, &[KeyValue::new("key", base_key.clone())]);
+			.with_label_values(&[base_key.as_str()])
+			.inc_by(keys.len() as u64);
 
-		// Build context.
-		//
-		// Drop `keys` bc this is not the same as the keys list in `ctx`, so it should not be used
-		// again.
-		let mut ctx = GetterCtx::new(base_key.clone(), keys);
+		let mut ctx = GetterCtx::new(keys);
+
+		// No driver (cache disabled)
+		let Some(driver) = &self.cache.driver else {
+			let keys = ctx.unresolved_keys();
+			let ctx = getter(ctx, keys).await.map_err(Error::Getter)?;
+
+			metrics::CACHE_VALUE_EMPTY_TOTAL
+				.with_label_values(&[&base_key])
+				.inc_by(ctx.unresolved_keys().len() as u64);
+
+			return Ok(ctx.into_values());
+		};
 
 		// Build driver-specific cache keys
-		let cache_keys = ctx
-			.keys()
-			.iter()
-			.map(|key| self.cache.driver.process_key(&base_key, &key.key))
-			.collect::<Vec<_>>();
+		let (keys, cache_keys): (Vec<_>, Vec<_>) = ctx
+			.entries()
+			.map(|(key, _)| (key.clone(), driver.process_key(&base_key, key)))
+			.unzip();
+		let cache_keys_len = cache_keys.len();
 
 		// Attempt to fetch value from cache, fall back to getter
-		match self.cache.driver.fetch_values(&base_key, &cache_keys).await {
+		match driver.get(&base_key, &cache_keys).await {
 			Ok(cached_values) => {
 				debug_assert_eq!(
-					cache_keys.len(),
+					cache_keys_len,
 					cached_values.len(),
 					"cache returned wrong number of values"
 				);
 
-				// Create the getter ctx and resolve the cached values
-				for (i, value) in cached_values.into_iter().enumerate() {
+				// Resolve the cached values
+				for (key, value) in keys.iter().zip(cached_values.into_iter()) {
 					if let Some(value_bytes) = value {
 						// Try to decode the value using the driver
-						match self.cache.driver.decode_value(&value_bytes) {
-							Ok(value_serde) => match decoder(&value_serde) {
-								Ok(value) => {
-									ctx.resolve_from_cache(i, value);
-								}
-								Err(err) => {
-									tracing::error!(?err, "Failed to decode value");
-								}
-							},
+						match decoder(&value_bytes) {
+							Ok(value) => {
+								ctx.resolve_from_cache(key, value);
+							}
 							Err(err) => {
 								tracing::error!(?err, "Failed to decode value");
 							}
@@ -174,48 +125,150 @@ impl RequestConfig {
 				}
 
 				// Fetch remaining values and add to the cached list
-				if !ctx.all_keys_have_value() {
+				if !ctx.all_entries_have_value() {
 					// Call the getter
 					let remaining_keys = ctx.unresolved_keys();
 					let unresolved_len = remaining_keys.len();
 
-					metrics::CACHE_VALUE_MISS_TOTAL.add(
-						unresolved_len as u64,
-						&[KeyValue::new("key", base_key.clone())],
-					);
+					metrics::CACHE_VALUE_MISS_TOTAL
+						.with_label_values(&[base_key.as_str()])
+						.inc_by(unresolved_len as u64);
+
+					let mut waiting_keys = Vec::new();
+					let mut leased_keys = Vec::new();
+					let (broadcast_tx, _) = broadcast::channel::<()>(16);
+
+					// Determine which keys are currently being fetched and not
+					for key in remaining_keys {
+						let cache_key = driver.process_key(&base_key, &key);
+						match self.cache.in_flight().entry_async(cache_key).await {
+							scc::hash_map::Entry::Occupied(broadcast) => {
+								waiting_keys.push((key, broadcast.subscribe()));
+							}
+							scc::hash_map::Entry::Vacant(entry) => {
+								entry.insert_entry(broadcast_tx.clone());
+								leased_keys.push(key);
+							}
+						}
+					}
 
-					ctx = getter(ctx, remaining_keys).await.map_err(Error::Getter)?;
+					let getter2 = getter.clone();
+					let ctx2 = GetterCtx::new(leased_keys.clone());
+					let base_key2 = base_key.clone();
+					let leased_keys2 = leased_keys.clone();
+					let (ctx2, ctx3) = tokio::try_join!(
+						async move {
+							if leased_keys2.is_empty() {
+								Ok(ctx2)
+							} else {
+								getter2(ctx2, leased_keys2).await.map_err(Error::Getter)
+							}
+						},
+						async move {
+							let ctx3 = GetterCtx::new(
+								waiting_keys.iter().map(|(key, _)| key.clone()).collect(),
+							);
+
+							// Wait on keys that are being fetched by another cache req
+							let (succeeded_keys, failed_keys): (Vec<_>, Vec<_>) =
+								futures_util::stream::iter(waiting_keys)
+									.map(|(key, mut rx)| async move {
+										(
+											key,
+											tokio::time::timeout(IN_FLIGHT_TIMEOUT, rx.recv())
+												.await
+												.ok()
+												.map(|x| x.ok())
+												.flatten()
+												.is_some(),
+										)
+									})
+									.buffer_unordered(1024)
+									.collect::<Vec<_>>()
+									.await
+									.into_iter()
+									.partition_map(|(key, succeeded)| {
+										if succeeded {
+											let cache_key = driver.process_key(&base_key2, &key);
+											Either::Left((key, cache_key))
+										} else {
+											Either::Right(key)
+										}
+									});
+							let (succeeded_keys, succeeded_cache_keys): (Vec<_>, Vec<_>) =
+								succeeded_keys.into_iter().unzip();
+
+							let (cached_values_res, ctx3_res) = tokio::join!(
+								async {
+									if succeeded_cache_keys.is_empty() {
+										Ok(Vec::new())
+									} else {
+										driver.get(&base_key2, &succeeded_cache_keys).await
+									}
+								},
+								async {
+									if failed_keys.is_empty() {
+										Ok(ctx3)
+									} else {
+										getter(ctx3, failed_keys).await.map_err(Error::Getter)
+									}
+								},
+							);
+							let mut ctx3 = ctx3_res?;
+
+							match cached_values_res {
+								Ok(cached_values) => {
+									for (key, value) in
+										succeeded_keys.iter().zip(cached_values.into_iter())
+									{
+										if let Some(value_bytes) = value {
+											// Try to decode the value using the driver
+											match decoder(&value_bytes) {
+												Ok(value) => {
+													ctx3.resolve_from_cache(key, value);
+												}
+												Err(err) => {
+													tracing::error!(?err, "Failed to decode value");
+												}
+											}
+										}
+									}
+								}
+								Err(err) => {
+									tracing::error!(?err, "failed to read batch keys from cache");
+
+									metrics::CACHE_REQUEST_ERRORS
+										.with_label_values(&[&base_key2])
+										.inc();
+								}
+							}
+
+							Ok(ctx3)
+						}
+					)?;
+
+					ctx.merge(ctx2);
+					ctx.merge(ctx3);
 
 					// Write the values to cache
 					let expire_at = rivet_util::timestamp::now() + self.ttl;
-					let values_needing_cache_write = ctx.values_needing_cache_write();
+					let entries_needing_cache_write = ctx.entries_needing_cache_write();
 
 					tracing::trace!(
 						unresolved_len,
-						fetched_len = values_needing_cache_write.len(),
+						fetched_len = entries_needing_cache_write.len(),
 						"writing new values to cache"
 					);
 
 					// Convert values to cache bytes
-					let keys_values = values_needing_cache_write
+					let entries_values = entries_needing_cache_write
 						.into_iter()
 						.filter_map(|(key, value)| {
 							// Process the key with the appropriate driver
-							let driver_key = self.cache.driver.process_key(&base_key, &key.key);
+							let cache_key = driver.process_key(&base_key, key);
 							// Try to decode the value using the driver
-							match encoder(&value) {
-								Ok(value_bytes) => {
-									match self.cache.driver.encode_value(&value_bytes) {
-										Ok(value_serde) => {
-											Some((driver_key, value_serde, expire_at))
-										}
-										Err(err) => {
-											tracing::error!(?err, "Failed to encode value");
-
-											None
-										}
-									}
-								}
+							match encoder(value) {
+								Ok(value_bytes) => Some((cache_key, value_bytes, expire_at)),
 								Err(err) => {
 									tracing::error!(?err, "Failed to encode value");
 
@@ -225,30 +278,26 @@ impl RequestConfig {
 						})
 						.collect::<Vec<_>>();
 
-					if !keys_values.is_empty() {
-						let driver = self.cache.driver.clone();
+					if !entries_values.is_empty() {
 						let base_key_clone = base_key.clone();
 
-						let spawn_res = tokio::task::Builder::new().name("cache::write").spawn(
-							async move {
-								if let Err(err) =
-									driver.set_values(&base_key_clone, keys_values).await
-								{
-									tracing::error!(?err, "failed to write to cache");
-								}
-							}
-							.in_current_span(),
-						);
-						if let Err(err) = spawn_res {
-							tracing::error!(?err, "failed to spawn cache::write task");
+						if let Err(err) = driver.set(&base_key_clone, entries_values).await {
+							tracing::error!(?err, "failed to write to cache");
 						}
+
+						let _ = broadcast_tx.send(());
+					}
+
+					// Release leases
+					for key in leased_keys {
+						let cache_key = driver.process_key(&base_key, &key);
+						self.cache.in_flight().remove_async(&cache_key).await;
 					}
 				}
 
-				metrics::CACHE_VALUE_EMPTY_TOTAL.add(
-					ctx.unresolved_keys().len() as u64,
-					&[KeyValue::new("key", base_key.clone())],
-				);
+				metrics::CACHE_VALUE_EMPTY_TOTAL
+					.with_label_values(&[&base_key])
+					.inc_by(ctx.unresolved_keys().len() as u64);
 
 				Ok(ctx.into_values())
 			}
@@ -258,19 +307,29 @@ impl RequestConfig {
 					"failed to read batch keys from cache, falling back to getter"
 				);
 
-				metrics::CACHE_REQUEST_ERRORS.add(1, &[KeyValue::new("key", base_key.clone())]);
+				let keys = ctx.unresolved_keys();
+
+				metrics::CACHE_REQUEST_ERRORS
+					.with_label_values(&[&base_key])
+					.inc();
+				metrics::CACHE_VALUE_MISS_TOTAL
+					.with_label_values(&[base_key.as_str()])
+					.inc_by(keys.len() as u64);
 
 				// Fall back to the getter since we can't fetch the value from
 				// the cache
-				let keys = ctx.unresolved_keys();
 				let ctx = getter(ctx, keys).await.map_err(Error::Getter)?;
 
+				metrics::CACHE_VALUE_EMPTY_TOTAL
+					.with_label_values(&[&base_key])
+					.inc_by(ctx.unresolved_keys().len() as u64);
+
 				Ok(ctx.into_values())
 			}
 		}
 	}
 
-	#[tracing::instrument(err, skip(keys))]
+	#[tracing::instrument(err, skip_all, fields(?base_key))]
 	pub async fn purge<Key>(
 		self,
 		base_key: impl AsRef<str> + Debug,
@@ -279,11 +338,16 @@ impl RequestConfig {
 	where
 		Key: CacheKey + Send + Sync,
 	{
+		// Cache disabled
+		let Some(driver) = &self.cache.driver else {
+			return Ok(());
+		};
+
 		// Build keys
 		let base_key = base_key.as_ref();
 		let cache_keys = keys
 			.into_iter()
-			.map(|key| self.cache.driver.process_key(base_key, &key))
+			.map(|key| driver.process_key(base_key, &key))
 			.collect::<Vec<_>>();
 
 		if cache_keys.is_empty() {
@@ -294,10 +358,7 @@ impl RequestConfig {
 		if let Some(ups) = &self.cache.ups {
 			let message = CachePurgeMessage {
 				base_key: base_key.to_string(),
-				keys: cache_keys
-					.iter()
-					.map(|k| RawCacheKey::from(k.clone()))
-					.collect(),
+				keys: cache_keys.clone(),
 			};
 
 			let payload = serde_json::to_vec(&message)?;
@@ -321,32 +382,37 @@ impl RequestConfig {
 		}
 
 		// Delete keys locally
-		let raw_keys = cache_keys
-			.into_iter()
-			.map(RawCacheKey::from)
-			.collect::<Vec<_>>();
-		self.purge_local(base_key, raw_keys).await
+		self.purge_local(base_key, cache_keys).await
 	}
 
 	/// Purges keys from the local cache without publishing to NATS.
 	/// This is used by the cache-purge service to avoid recursive publishing.
-	#[tracing::instrument(err, skip(keys))]
+	#[tracing::instrument(err, skip_all, fields(?base_key))]
 	pub async fn purge_local(
 		self,
 		base_key: impl AsRef<str> + Debug,
 		keys: Vec<RawCacheKey>,
 	) -> Result<()> {
+		// Cache disabled
+		let Some(driver) = &self.cache.driver else {
+			return Ok(());
+		};
+
 		let base_key = base_key.as_ref();
 
 		if keys.is_empty() {
 			return Ok(());
 		}
 
-		// Convert RawCacheKey to String for driver
-		let cache_keys = keys.into_iter().map(|k| k.cache_key()).collect::<Vec<_>>();
+		metrics::CACHE_PURGE_REQUEST_TOTAL
+			.with_label_values(&[base_key])
+			.inc();
+		metrics::CACHE_PURGE_VALUE_TOTAL
+			.with_label_values(&[base_key])
+			.inc_by(keys.len() as u64);
 
 		// Delete keys locally
-		match self.cache.driver.delete_keys(base_key, cache_keys).await {
+		match driver.delete(base_key, keys).await {
 			Ok(_) => {
 				tracing::trace!("successfully deleted keys");
 			}
@@ -361,7 +427,6 @@ impl RequestConfig {
 
 // MARK: JSON fetch
 impl RequestConfig {
-	#[tracing::instrument(err, skip(key, getter))]
 	pub async fn fetch_one_json<Key, Value, Getter, Fut>(
 		self,
 		base_key: impl ToString + Debug,
@@ -391,7 +456,6 @@ impl RequestConfig {
 		Ok(values.into_iter().next().map(|(_, v)| v))
 	}
 
-	#[tracing::instrument(err, skip(keys, getter))]
 	pub async fn fetch_all_json<Key, Value, Getter, Fut>(
 		self,
 		base_key: impl ToString + Debug,
@@ -410,7 +474,6 @@ impl RequestConfig {
 			.map(|x| x.into_iter().map(|(_, v)| v).collect::<Vec<_>>())
 	}
 
-	#[tracing::instrument(err, skip(keys, getter))]
 	pub async fn fetch_all_json_with_keys<Key, Value, Getter, Fut>(
 		self,
 		base_key: impl ToString + Debug,
@@ -428,12 +491,12 @@ impl RequestConfig {
 			keys,
 			getter,
 			|value: &Value| -> Result<Vec<u8>> {
-				serde_json::to_vec(value)
+				serde_json::to_vec(&value)
 					.map_err(Error::SerdeEncode)
 					.map_err(Into::into)
 			},
-			|value: &Vec<u8>| -> Result<Value> {
-				serde_json::from_slice(value.as_slice())
+			|value: &[u8]| -> Result<Value> {
+				serde_json::from_slice(value)
 					.map_err(Error::SerdeDecode)
 					.map_err(Into::into)
 			},