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

package/packages/pegboard/src/actor_kv/preload.rs

@@ -0,0 +1,363 @@
+use anyhow::Result;
+use futures_util::TryStreamExt;
+use gas::prelude::*;
+use rivet_config::config::pegboard::Pegboard;
+use rivet_envoy_protocol as ep;
+use serde::Deserialize;
+use universaldb::prelude::*;
+use universaldb::tuple::Subspace;
+
+use super::entry::EntryBuilder;
+use crate::keys;
+
+/// Request to preload a prefix range from the actor's KV store.
+pub struct PreloadPrefixRequest {
+	/// The raw key prefix bytes (e.g., [2] for connections, [8] for SQLite).
+	pub prefix: ep::KvKey,
+	/// Maximum bytes to preload for this prefix.
+	pub max_bytes: u64,
+	/// If true, return whatever fits even if truncated (for per-key lookup subsystems
+	/// like SQLite VFS). If false, return nothing if the total data exceeds max_bytes
+	/// (for list-based subsystems like connections and workflows).
+	pub partial: bool,
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct PreloadConfig {
+	#[serde(default)]
+	keys: Vec<Vec<u8>>,
+	#[serde(default)]
+	prefixes: Vec<PreloadPrefixConfig>,
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct PreloadPrefixConfig {
+	prefix: Vec<u8>,
+	max_bytes: u64,
+	#[serde(default)]
+	partial: bool,
+}
+
+/// Fetches all preload data for an actor in a single FDB snapshot transaction.
+///
+/// Reads exact get-keys and prefix ranges, reassembles chunked FDB values using
+/// EntryBuilder, strips tuple encoding via KeyWrapper::unpack, and returns raw
+/// byte key-value pairs ready for TypeScript consumption.
+///
+/// Prefix requests should be passed in descending priority order (highest priority
+/// first). When the global byte cap is reached, lower-priority prefixes are
+/// truncated first.
+#[tracing::instrument(skip_all)]
+pub(crate) async fn batch_preload(
+	db: &universaldb::Database,
+	actor_id: Id,
+	get_keys: Vec<ep::KvKey>,
+	prefix_requests: Vec<PreloadPrefixRequest>,
+	max_total_bytes: u64,
+) -> Result<ep::PreloadedKv> {
+	let subspace = keys::actor_kv::subspace(actor_id);
+
+	// Break prefix_requests into separate vectors so they can be cloned for the
+	// FDB transaction closure (which may retry on conflicts).
+	let prefix_keys: Vec<ep::KvKey> = prefix_requests.iter().map(|r| r.prefix.clone()).collect();
+	let prefix_params: Vec<(u64, bool)> = prefix_requests
+		.iter()
+		.map(|r| (r.max_bytes, r.partial))
+		.collect();
+
+	db.run(|tx| {
+		let subspace = subspace.clone();
+		let get_keys = get_keys.clone();
+		let prefix_keys = prefix_keys.clone();
+		let prefix_params = prefix_params.clone();
+
+		async move {
+			let tx = tx.with_subspace(subspace.clone());
+			let mut entries = Vec::new();
+			let mut total_bytes: u64 = 0;
+
+			// Build requested lists dynamically so they only contain keys/prefixes
+			// that were actually scanned. Keys or prefixes skipped due to budget
+			// exhaustion or disabled config must not appear, otherwise the actor
+			// would mistake "not scanned" for "scanned and not found".
+			let mut requested_get_keys: Vec<ep::KvKey> = Vec::new();
+			let mut requested_prefixes: Vec<ep::KvKey> = Vec::new();
+
+			// 1. Read exact get-keys. Each key maps to a single logical entry
+			// (or nothing if the key doesn't exist in FDB).
+			for key in &get_keys {
+				if total_bytes >= max_total_bytes {
+					tracing::debug!(
+						skipped_keys = get_keys.len() - requested_get_keys.len(),
+						"preload get-keys skipped due to global budget exhaustion"
+					);
+					break;
+				}
+
+				// Mark this key as scanned regardless of whether it exists in FDB.
+				requested_get_keys.push(key.clone());
+
+				let key_subspace = subspace.subspace(&keys::actor_kv::KeyWrapper(key.clone()));
+				let mut stream = tx.get_ranges_keyvalues(
+					universaldb::RangeOption {
+						mode: universaldb::options::StreamingMode::WantAll,
+						..key_subspace.range().into()
+					},
+					Snapshot,
+				);
+
+				let mut builder: Option<EntryBuilder> = None;
+
+				while let Some(fdb_kv) = stream.try_next().await? {
+					if builder.is_none() {
+						let parsed_key = tx
+							.unpack::<keys::actor_kv::EntryBaseKey>(&fdb_kv.key())?
+							.key;
+						builder = Some(EntryBuilder::new(parsed_key));
+					}
+
+					let b = builder.as_mut().unwrap();
+
+					if let Ok(chunk_key) =
+						tx.unpack::<keys::actor_kv::EntryValueChunkKey>(&fdb_kv.key())
+					{
+						b.append_chunk(chunk_key.chunk, fdb_kv.value());
+					} else if let Ok(metadata_key) =
+						tx.unpack::<keys::actor_kv::EntryMetadataKey>(&fdb_kv.key())
+					{
+						let metadata = metadata_key.deserialize(fdb_kv.value())?;
+						b.append_metadata(metadata);
+					} else {
+						bail!("unexpected sub key in preload get");
+					}
+				}
+
+				if let Some(b) = builder {
+					let (k, v, m) = b.build()?;
+					let size = entry_size(&k, &v, &m);
+					if total_bytes + size <= max_total_bytes {
+						total_bytes += size;
+						entries.push(ep::PreloadedKvEntry {
+							key: k,
+							value: v,
+							metadata: m,
+						});
+					}
+				}
+			}
+
+			// 2. Read prefix ranges in priority order. Each prefix is bounded by
+			// its per-prefix max_bytes and the remaining global budget.
+			for (i, prefix) in prefix_keys.iter().enumerate() {
+				let (max_bytes, partial) = prefix_params[i];
+
+				// Skip prefixes disabled by config (max_bytes == 0) or when
+				// global budget is exhausted. Do not include in requested_prefixes
+				// so the actor falls back to kvListPrefix.
+				let remaining_budget = max_total_bytes.saturating_sub(total_bytes);
+				let effective_limit = max_bytes.min(remaining_budget);
+
+				if effective_limit == 0 {
+					tracing::debug!(
+						?prefix,
+						max_bytes,
+						remaining_budget,
+						"preload prefix skipped, effective limit is 0"
+					);
+					continue;
+				}
+
+				let range = prefix_range(prefix, &subspace);
+				let mut stream = tx.get_ranges_keyvalues(
+					universaldb::RangeOption {
+						mode: universaldb::options::StreamingMode::Iterator,
+						..range.into()
+					},
+					Snapshot,
+				);
+
+				let mut prefix_entries: Vec<ep::PreloadedKvEntry> = Vec::new();
+				let mut prefix_bytes: u64 = 0;
+				let mut current_entry: Option<EntryBuilder> = None;
+				let mut exceeded = false;
+
+				while let Some(fdb_kv) = stream.try_next().await? {
+					let key = tx
+						.unpack::<keys::actor_kv::EntryBaseKey>(&fdb_kv.key())?
+						.key;
+
+					let curr = if let Some(inner) = &mut current_entry {
+						if inner.key != key {
+							// Finalize the previous entry.
+							let prev = std::mem::replace(inner, EntryBuilder::new(key));
+							let (k, v, m) = prev.build()?;
+							let size = entry_size(&k, &v, &m);
+
+							if prefix_bytes + size > effective_limit {
+								exceeded = true;
+								break;
+							}
+
+							prefix_bytes += size;
+							prefix_entries.push(ep::PreloadedKvEntry {
+								key: k,
+								value: v,
+								metadata: m,
+							});
+						}
+
+						inner
+					} else {
+						current_entry = Some(EntryBuilder::new(key));
+						current_entry.as_mut().expect("just set")
+					};
+
+					if let Ok(chunk_key) =
+						tx.unpack::<keys::actor_kv::EntryValueChunkKey>(&fdb_kv.key())
+					{
+						curr.append_chunk(chunk_key.chunk, fdb_kv.value());
+					} else if let Ok(metadata_key) =
+						tx.unpack::<keys::actor_kv::EntryMetadataKey>(&fdb_kv.key())
+					{
+						let metadata = metadata_key.deserialize(fdb_kv.value())?;
+						curr.append_metadata(metadata);
+					} else {
+						bail!("unexpected sub key in preload prefix scan");
+					}
+				}
+
+				// Finalize the last entry if the stream ended without exceeding.
+				if !exceeded {
+					if let Some(b) = current_entry {
+						let (k, v, m) = b.build()?;
+						let size = entry_size(&k, &v, &m);
+						if prefix_bytes + size > effective_limit {
+							exceeded = true;
+						} else {
+							prefix_bytes += size;
+							prefix_entries.push(ep::PreloadedKvEntry {
+								key: k,
+								value: v,
+								metadata: m,
+							});
+						}
+					}
+				}
+
+				// For non-partial prefixes, discard all entries if the data exceeded
+				// the limit. The subsystem will fall back to a full kvListPrefix.
+				// Do not include in requested_prefixes so the actor knows to fall back.
+				if exceeded && !partial {
+					tracing::debug!(
+						?prefix,
+						prefix_bytes,
+						effective_limit,
+						"preload prefix truncated (partial: false), discarding entries"
+					);
+					continue;
+				}
+
+				if exceeded {
+					tracing::debug!(
+						?prefix,
+						prefix_bytes,
+						effective_limit,
+						"preload prefix truncated (partial: true), keeping partial data"
+					);
+				}
+
+				requested_prefixes.push(prefix.clone());
+				total_bytes += prefix_bytes;
+				entries.extend(prefix_entries);
+			}
+
+			Ok(ep::PreloadedKv {
+				entries,
+				requested_get_keys,
+				requested_prefixes,
+			})
+		}
+	})
+	.custom_instrument(tracing::info_span!("kv_batch_preload_tx"))
+	.await
+	.map_err(Into::<anyhow::Error>::into)
+}
+
+/// Fetches preloaded KV data for an actor using engine config and actor name
+/// metadata. Returns `None` if preloading is disabled or missing from actor
+/// metadata. Fails if the FDB transaction fails (no silent fallback).
+#[tracing::instrument(skip_all)]
+pub async fn fetch_preloaded_kv(
+	db: &universaldb::Database,
+	config: &Pegboard,
+	actor_id: Id,
+	namespace_id: Id,
+	actor_name: &str,
+) -> Result<Option<ep::PreloadedKv>> {
+	// Read actor name metadata from FDB.
+	let metadata = db
+		.run(|tx| {
+			let tx = tx.with_subspace(keys::subspace());
+			let name_key = keys::ns::ActorNameKey::new(namespace_id, actor_name.to_string());
+			async move { tx.read_opt(&name_key, Snapshot).await }
+		})
+		.instrument(tracing::info_span!("read_actor_metadata_tx"))
+		.await?;
+
+	let metadata_map = metadata
+		.map(|d: rivet_data::converted::ActorNameKeyData| d.metadata)
+		.unwrap_or_default();
+
+	let Some(preload_config) = metadata_map
+		.get("preload")
+		.and_then(|value| serde_json::from_value::<PreloadConfig>(value.clone()).ok())
+	else {
+		return Ok(None);
+	};
+
+	let preload_max_total_bytes = config.preload_max_total_bytes();
+	if preload_max_total_bytes == 0 {
+		return Ok(None);
+	};
+
+	let prefix_requests = preload_config
+		.prefixes
+		.into_iter()
+		.map(|prefix| PreloadPrefixRequest {
+			prefix: prefix.prefix,
+			max_bytes: prefix.max_bytes,
+			partial: prefix.partial,
+		})
+		.collect();
+
+	let preloaded = batch_preload(
+		db,
+		actor_id,
+		preload_config.keys,
+		prefix_requests,
+		preload_max_total_bytes,
+	)
+	.await?;
+
+	Ok(Some(preloaded))
+}
+
+/// Computes the serialized size of a preloaded KV entry, including key, value,
+/// and metadata (version bytes + i64 timestamp).
+fn entry_size(key: &ep::KvKey, value: &ep::KvValue, metadata: &ep::KvMetadata) -> u64 {
+	(key.len() + value.len() + metadata.version.len() + std::mem::size_of::<i64>()) as u64
+}
+
+/// Computes the FDB key range for a prefix scan within the actor KV subspace.
+fn prefix_range(prefix: &ep::KvKey, subspace: &Subspace) -> (Vec<u8>, Vec<u8>) {
+	let mut start = subspace.pack(&keys::actor_kv::ListKeyWrapper(prefix.clone()));
+	// Remove the trailing 0 byte that tuple encoding adds to bytes.
+	if let Some(&0) = start.last() {
+		start.pop();
+	}
+	let mut end = start.clone();
+	end.push(0xFF);
+	(start, end)
+}

package/packages/{actor-kv/src → pegboard/src/actor_kv}/utils.rs

@@ -1,35 +1,19 @@
-use std::result::Result::Ok;
+use anyhow::{Result, ensure};
+use rivet_envoy_protocol as ep;
 
-use anyhow::*;
-use rivet_runner_protocol as rp;
-
-use crate::{
-	MAX_KEY_SIZE, MAX_KEYS, MAX_PUT_PAYLOAD_SIZE, MAX_STORAGE_SIZE, MAX_VALUE_SIZE, key::KeyWrapper,
+use super::{
+	MAX_KEY_SIZE, MAX_KEYS, MAX_PUT_PAYLOAD_SIZE, MAX_STORAGE_SIZE, MAX_VALUE_SIZE,
+	keys::actor_kv::KeyWrapper,
 };
+use crate::errors;
 
-pub fn now() -> i64 {
-	std::time::SystemTime::now()
-		.duration_since(std::time::UNIX_EPOCH)
-		.unwrap_or_else(|err| unreachable!("time is broken: {}", err))
-		.as_millis()
-		.try_into()
-		.expect("now doesn't fit in i64")
-}
-
-pub fn validate_list_query(query: &rp::KvListQuery) -> Result<()> {
+pub fn validate_list_query(query: &ep::KvListQuery) -> Result<()> {
 	match query {
-		rp::KvListQuery::KvListAllQuery => {}
-		rp::KvListQuery::KvListRangeQuery(range) => {
-			ensure!(
-				KeyWrapper::tuple_len(&range.start) <= MAX_KEY_SIZE,
-				"start key is too long (max 2048 bytes)"
-			);
-			ensure!(
-				KeyWrapper::tuple_len(&range.end) <= MAX_KEY_SIZE,
-				"end key is too long (max 2048 bytes)"
-			);
+		ep::KvListQuery::KvListAllQuery => {}
+		ep::KvListQuery::KvListRangeQuery(range) => {
+			validate_range(&range.start, &range.end)?;
 		}
-		rp::KvListQuery::KvListPrefixQuery(prefix) => {
+		ep::KvListQuery::KvListPrefixQuery(prefix) => {
 			ensure!(
 				KeyWrapper::tuple_len(&prefix.key) <= MAX_KEY_SIZE,
 				"prefix key is too long (max 2048 bytes)"
@@ -40,7 +24,20 @@ pub fn validate_list_query(query: &rp::KvListQuery) -> Result<()> {
 	Ok(())
 }
 
-pub fn validate_keys(keys: &[rp::KvKey]) -> Result<()> {
+pub fn validate_range(start: &ep::KvKey, end: &ep::KvKey) -> Result<()> {
+	ensure!(
+		KeyWrapper::tuple_len(start) <= MAX_KEY_SIZE,
+		"start key is too long (max 2048 bytes)"
+	);
+	ensure!(
+		KeyWrapper::tuple_len(end) <= MAX_KEY_SIZE,
+		"end key is too long (max 2048 bytes)"
+	);
+
+	Ok(())
+}
+
+pub fn validate_keys(keys: &[ep::KvKey]) -> Result<()> {
 	ensure!(keys.len() <= MAX_KEYS, "a maximum of 128 keys is allowed");
 
 	for key in keys {
@@ -54,8 +51,8 @@ pub fn validate_keys(keys: &[rp::KvKey]) -> Result<()> {
 }
 
 pub fn validate_entries(
-	keys: &[rp::KvKey],
-	values: &[rp::KvValue],
+	keys: &[ep::KvKey],
+	values: &[ep::KvValue],
 	total_size: usize,
 ) -> Result<()> {
 	ensure!(
@@ -78,10 +75,13 @@ pub fn validate_entries(
 	);
 
 	let storage_remaining = MAX_STORAGE_SIZE.saturating_sub(total_size);
-	ensure!(
-		payload_size <= storage_remaining,
-		"not enough space left in storage ({storage_remaining} bytes remaining, current payload is {payload_size} bytes)"
-	);
+	if payload_size > storage_remaining {
+		return Err(errors::Actor::KvStorageQuotaExceeded {
+			remaining: storage_remaining,
+			payload_size,
+		}
+		.build());
+	}
 
 	for key in keys {
 		ensure!(
@@ -93,7 +93,8 @@ pub fn validate_entries(
 	for value in values {
 		ensure!(
 			value.len() <= MAX_VALUE_SIZE,
-			"value is too large (max 128 KiB)"
+			"value is too large (max {} KiB)",
+			MAX_VALUE_SIZE / 1024
 		);
 	}
 

package/packages/pegboard/src/errors.rs

@@ -2,6 +2,8 @@ use gas::prelude::*;
 use rivet_error::*;
 use serde::{Deserialize, Serialize};
 
+use crate::ops::serverless_metadata::fetch::ServerlessMetadataError;
+
 #[derive(RivetError, Debug, Clone, Deserialize, Serialize)]
 #[error("actor")]
 pub enum Actor {
@@ -55,17 +57,27 @@ pub enum Actor {
 	KeyReservedInDifferentDatacenter { datacenter_label: u16 },
 
 	#[error(
-		"no_runners_available",
-		"No runners are available in any datacenter. Validate the runner is listed in the Connect tab and that the runner's name matches the requested runner name.",
-		"No runners with name '{runner_name}' are available in any datacenter for the namespace '{namespace}'. Validate the runner is listed in the Connect tab and that the runner's name matches the requested runner name."
+		"no_runner_config_configured",
+		"No runner config configured in any datacenter. Validate a provider is listed that matches requested pool name.",
+		"No runner config with name '{pool_name}' are available in any datacenter for the namespace '{namespace}'. Validate a provider is listed that matches the requested pool name."
 	)]
-	NoRunnersAvailable {
+	NoRunnerConfigConfigured {
 		namespace: String,
-		runner_name: String,
+		pool_name: String,
 	},
 
 	#[error("kv_key_not_found", "The KV key does not exist for this actor.")]
 	KvKeyNotFound,
+
+	#[error(
+		"kv_storage_quota_exceeded",
+		"Not enough space left in storage.",
+		"Not enough space left in storage ({remaining} bytes remaining, current payload is {payload_size} bytes)."
+	)]
+	KvStorageQuotaExceeded {
+		remaining: usize,
+		payload_size: usize,
+	},
 }
 
 #[derive(RivetError, Debug, Clone, Deserialize, Serialize)]
@@ -74,3 +86,25 @@ pub enum Runner {
 	#[error("not_found", "The runner does not exist.")]
 	NotFound,
 }
+
+#[derive(RivetError, Debug, Deserialize, Serialize)]
+#[error("runner_config")]
+pub enum RunnerConfig {
+	#[error("invalid", "Invalid runner config.", "Invalid runner config: {reason}")]
+	Invalid { reason: String },
+
+	#[error("not_found", "No config for this runner exists.")]
+	NotFound,
+}
+
+#[derive(RivetError, Debug, Deserialize, Serialize)]
+#[error("serverless_runner_pool")]
+pub enum ServerlessRunnerPool {
+	#[error("not_found", "No serverless pool for this runner exists.")]
+	NotFound,
+	#[error(
+		"failed_to_fetch_metadata",
+		"Failed to fetch serverless metadata: {reason:?}"
+	)]
+	FailedToFetchMetadata { reason: ServerlessMetadataError },
+}