fix(rivetkit): use startup kv preload

Author: NathanFlurry

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

@@ -96,9 +96,6 @@ pub(crate) async fn batch_preload(
 					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 {
@@ -139,14 +136,26 @@ pub(crate) async fn batch_preload(
 					let size = entry_size(&k, &v, &m);
 					if total_bytes + size <= max_total_bytes {
 						total_bytes += size;
+						requested_get_keys.push(key.clone());
 						entries.push(ep::PreloadedKvEntry {
 							key: k,
 							value: v,
 							metadata: m,
 						});
+					} else {
+						tracing::debug!(
+							?key,
+							size,
+							remaining_budget = max_total_bytes.saturating_sub(total_bytes),
+							"preload get-key skipped due to global budget exhaustion"
+						);
 					}
+				} else {
+					// Mark missing keys as scanned so the runtime can distinguish
+					// "preloaded and absent" from "not preloaded".
+					requested_get_keys.push(key.clone());
+				}
 				}
-			}
 
 			// 2. Read prefix ranges in priority order. Each prefix is bounded by
 			// its per-prefix max_bytes and the remaining global budget.

rivetkit-rust/packages/rivetkit-core/src/actor/state.rs

@@ -550,7 +550,13 @@ impl ActorContext {
 	}
 
 	pub fn set_has_initialized(&self, has_initialized: bool) {
-		self.0.persisted.write().has_initialized = has_initialized;
+		{
+			let mut persisted = self.0.persisted.write();
+			if persisted.has_initialized == has_initialized {
+				return;
+			}
+			persisted.has_initialized = has_initialized;
+		}
 		self.0
 			.metrics
 			.inc_state_mutation(StateMutationReason::HasInitialized);

rivetkit-rust/packages/rivetkit-core/src/actor/task.rs

@@ -1211,9 +1211,12 @@ impl ActorTask {
 			.await
 			.context("persist actor initialization")?;
 		let init_inspector_token_started_at = Instant::now();
-		crate::inspector::init_inspector_token(&self.ctx)
-			.await
-			.context("initialize inspector token")?;
+		crate::inspector::auth::init_inspector_token_with_preload(
+			&self.ctx,
+			self.preloaded_kv.as_ref(),
+		)
+		.await
+		.context("initialize inspector token")?;
 		tracing::debug!(
 			actor_id = %actor_id,
 			duration_ms = duration_ms_f64(init_inspector_token_started_at.elapsed()),
@@ -1250,9 +1253,10 @@ impl ActorTask {
 	async fn load_persisted_startup(&mut self) -> Result<PersistedStartup> {
 		match std::mem::take(&mut self.preload_persisted_actor) {
 			PreloadedPersistedActor::Some(preloaded) => {
+				let last_pushed_alarm = self.load_startup_last_pushed_alarm().await?;
 				return Ok(PersistedStartup {
 					actor: preloaded,
-					last_pushed_alarm: Self::load_last_pushed_alarm(self.ctx.kv().clone()).await?,
+					last_pushed_alarm,
 				});
 			}
 			PreloadedPersistedActor::BundleExistsButEmpty => {
@@ -1267,6 +1271,16 @@ impl ActorTask {
 			PreloadedPersistedActor::NoBundle => {}
 		}
 
+		if self.preloaded_kv.is_some() {
+			let actor =
+				self.decode_persisted_actor_startup(self.load_startup_key(PERSIST_DATA_KEY).await?)?;
+			let last_pushed_alarm = self.load_startup_last_pushed_alarm().await?;
+			return Ok(PersistedStartup {
+				actor,
+				last_pushed_alarm,
+			});
+		}
+
 		let mut values = self
 			.ctx
 			.kv()
@@ -1297,16 +1311,43 @@ impl ActorTask {
 		})
 	}
 
-	async fn load_last_pushed_alarm(kv: crate::kv::Kv) -> Result<Option<i64>> {
-		kv.get(LAST_PUSHED_ALARM_KEY)
+	async fn load_startup_key(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
+		if let Some(entry) = self
+			.preloaded_kv
+			.as_ref()
+			.and_then(|preloaded| preloaded.key_entry(key))
+		{
+			return Ok(entry);
+		}
+
+		self.ctx
+			.kv()
+			.get(key)
 			.await
-			.context("load persisted last pushed alarm")?
+			.context("load persisted actor startup key")
+	}
+
+	async fn load_startup_last_pushed_alarm(&self) -> Result<Option<i64>> {
+		self.load_startup_key(LAST_PUSHED_ALARM_KEY)
+			.await?
 			.map(|bytes| decode_last_pushed_alarm(&bytes))
 			.transpose()
 			.context("decode persisted last pushed alarm")
 			.map(Option::flatten)
 	}
 
+	fn decode_persisted_actor_startup(&self, encoded: Option<Vec<u8>>) -> Result<PersistedActor> {
+		match encoded {
+			Some(bytes) => {
+				decode_persisted_actor(&bytes).context("decode persisted actor startup data")
+			}
+			None => Ok(PersistedActor {
+				input: self.start_input.clone(),
+				..PersistedActor::default()
+			}),
+		}
+	}
+
 	fn ensure_actor_event_channel(&mut self) {
 		if self.actor_event_tx.is_some() && self.actor_event_rx.is_some() {
 			return;

rivetkit-rust/packages/rivetkit-core/src/inspector/auth.rs

@@ -11,6 +11,7 @@ use subtle::ConstantTimeEq;
 
 use crate::ActorContext;
 use crate::actor::keys::INSPECTOR_TOKEN_KEY;
+use crate::actor::preload::PreloadedKv;
 
 /// Test-only override. Not a public/production auth mechanism; production
 /// inspector auth goes through the per-actor KV token.
@@ -75,15 +76,26 @@ impl InspectorAuth {
 /// precedence over any KV-stored token and we do not want to pin a per-actor
 /// token that will never be consulted.
 pub async fn init_inspector_token(ctx: &ActorContext) -> Result<()> {
+	init_inspector_token_with_preload(ctx, None).await
+}
+
+pub(crate) async fn init_inspector_token_with_preload(
+	ctx: &ActorContext,
+	preloaded_kv: Option<&PreloadedKv>,
+) -> Result<()> {
 	if configured_test_token().is_some() {
 		return Ok(());
 	}
 
-	let existing = ctx
-		.kv()
-		.get(&INSPECTOR_TOKEN_KEY)
-		.await
-		.context("load inspector token")?;
+	let existing =
+		match preloaded_kv.and_then(|preloaded| preloaded.key_entry(&INSPECTOR_TOKEN_KEY)) {
+			Some(existing) => existing,
+			None => ctx
+				.kv()
+				.get(&INSPECTOR_TOKEN_KEY)
+				.await
+				.context("load inspector token")?,
+		};
 	if existing.is_some() {
 		return Ok(());
 	}

rivetkit-rust/packages/rivetkit-core/src/registry/mod.rs

@@ -550,6 +550,7 @@ impl RegistryDispatcher {
 							[1],
 							[3],
 							[5, 1, 1],
+							[6],
 						],
 						"prefixes": [
 							{

rivetkit-rust/packages/rivetkit-core/tests/task.rs

@@ -31,9 +31,12 @@ mod moved_tests {
 		ConnHandle, HibernatableConnectionMetadata, decode_persisted_connection,
 	};
 	use crate::actor::context::tests::new_with_kv;
-	use crate::actor::keys::{LAST_PUSHED_ALARM_KEY, PERSIST_DATA_KEY, make_connection_key};
+	use crate::actor::keys::{
+		CONN_PREFIX, INSPECTOR_TOKEN_KEY, LAST_PUSHED_ALARM_KEY, PERSIST_DATA_KEY,
+		QUEUE_MESSAGES_PREFIX, WORKFLOW_STORAGE_PREFIX, make_connection_key,
+	};
 	use crate::actor::messages::{ActorEvent, SerializeStateReason, StateDelta};
-	use crate::actor::preload::PreloadedPersistedActor;
+	use crate::actor::preload::{PreloadedKv, PreloadedPersistedActor};
 	use crate::actor::sleep::CanSleep;
 	use crate::actor::state::{
 		PersistedActor, PersistedScheduleEvent, RequestSaveOpts, decode_last_pushed_alarm,
@@ -47,6 +50,7 @@ mod moved_tests {
 	use crate::inspector::auth::test_inspector_env_lock;
 	use crate::kv::tests::new_in_memory;
 	use crate::{ActorConfig, ActorContext, ActorFactory};
+	use rivet_envoy_client::utils::EnvoyShutdownError;
 
 	fn test_hook_lock() -> &'static AsyncMutex<()> {
 		static LOCK: OnceLock<AsyncMutex<()>> = OnceLock::new();
@@ -2108,6 +2112,183 @@ mod moved_tests {
 		assert!(ctx.persisted_actor().has_initialized);
 	}
 
+	#[tokio::test]
+	async fn startup_uses_preloaded_last_pushed_alarm_without_live_kv() {
+		let _env_guard = test_inspector_env_lock().lock().expect("env lock poisoned");
+		unsafe {
+			std::env::remove_var("_RIVET_TEST_INSPECTOR_TOKEN");
+		}
+
+		let future_ts = 4_102_445_000_000;
+		let persisted = PersistedActor {
+			has_initialized: true,
+			scheduled_events: vec![PersistedScheduleEvent {
+				event_id: "evt-preloaded-alarm".to_owned(),
+				timestamp: future_ts,
+				action: "tick".to_owned(),
+				args: None,
+			}],
+			..PersistedActor::default()
+		};
+		let encoded_alarm =
+			encode_last_pushed_alarm(Some(future_ts)).expect("last pushed alarm should encode");
+		let preloaded_kv = PreloadedKv::new_with_requested_get_keys(
+			vec![
+				(LAST_PUSHED_ALARM_KEY.to_vec(), encoded_alarm),
+				(
+					INSPECTOR_TOKEN_KEY.to_vec(),
+					b"startup-preload-token".to_vec(),
+				),
+			],
+			vec![
+				PERSIST_DATA_KEY.to_vec(),
+				LAST_PUSHED_ALARM_KEY.to_vec(),
+				INSPECTOR_TOKEN_KEY.to_vec(),
+				vec![5, 1, 1],
+			],
+			vec![
+				WORKFLOW_STORAGE_PREFIX.to_vec(),
+				CONN_PREFIX.to_vec(),
+				QUEUE_MESSAGES_PREFIX.to_vec(),
+			],
+		);
+
+		let (handle, mut rx) = test_envoy_handle();
+		tokio::spawn(async move {
+			while let Some(message) = rx.recv().await {
+				if let ToEnvoyMessage::KvRequest { response_tx, .. } = message {
+					let _ = response_tx.send(Err(anyhow::anyhow!(EnvoyShutdownError)));
+				}
+			}
+		});
+
+		let ctx = new_with_kv(
+			"actor-preloaded-alarm",
+			"task-preloaded-alarm",
+			Vec::new(),
+			"local",
+			crate::kv::Kv::new(handle.clone(), "actor-preloaded-alarm"),
+		);
+		ctx.configure_envoy(handle, Some(31));
+		let mut task = new_task(ctx.clone())
+			.with_preloaded_persisted_actor(PreloadedPersistedActor::Some(persisted))
+			.with_preloaded_kv(Some(preloaded_kv));
+		let (start_tx, start_rx) = oneshot::channel();
+
+		task.handle_lifecycle(LifecycleCommand::Start { reply: start_tx })
+			.await;
+		start_rx
+			.await
+			.expect("start reply should send")
+			.expect("start should use preloaded alarm without live kv");
+
+		assert_eq!(ctx.last_pushed_alarm(), Some(future_ts));
+	}
+
+	#[tokio::test]
+	async fn startup_uses_preloaded_alarm_with_partial_startup_preload() {
+		let _env_guard = test_inspector_env_lock().lock().expect("env lock poisoned");
+		unsafe {
+			std::env::remove_var("_RIVET_TEST_INSPECTOR_TOKEN");
+		}
+
+		let future_ts = 4_102_445_123_000;
+		let persisted = PersistedActor {
+			has_initialized: true,
+			scheduled_events: vec![PersistedScheduleEvent {
+				event_id: "evt-partial-preload-alarm".to_owned(),
+				timestamp: future_ts,
+				action: "tick".to_owned(),
+				args: None,
+			}],
+			..PersistedActor::default()
+		};
+		let encoded_persisted =
+			encode_persisted_actor(&persisted).expect("persisted actor should encode");
+		let encoded_alarm =
+			encode_last_pushed_alarm(Some(future_ts)).expect("last pushed alarm should encode");
+		let preloaded_kv = PreloadedKv::new_with_requested_get_keys(
+			vec![
+				(LAST_PUSHED_ALARM_KEY.to_vec(), encoded_alarm),
+				(
+					INSPECTOR_TOKEN_KEY.to_vec(),
+					b"startup-partial-preload-token".to_vec(),
+				),
+			],
+			vec![
+				LAST_PUSHED_ALARM_KEY.to_vec(),
+				INSPECTOR_TOKEN_KEY.to_vec(),
+				vec![5, 1, 1],
+			],
+			vec![
+				WORKFLOW_STORAGE_PREFIX.to_vec(),
+				CONN_PREFIX.to_vec(),
+				QUEUE_MESSAGES_PREFIX.to_vec(),
+			],
+		);
+
+		let saw_persist_live_get = Arc::new(AtomicBool::new(false));
+		let saw_persist_live_get_for_task = Arc::clone(&saw_persist_live_get);
+		let (handle, mut rx) = test_envoy_handle();
+		tokio::spawn(async move {
+			while let Some(message) = rx.recv().await {
+				if let ToEnvoyMessage::KvRequest {
+					data, response_tx, ..
+				} = message
+				{
+					match data {
+						protocol::KvRequestData::KvGetRequest(request) => {
+							assert_eq!(
+								request.keys,
+								vec![PERSIST_DATA_KEY.to_vec()],
+								"partial preload should only live-fetch missing persisted actor data"
+							);
+							saw_persist_live_get_for_task.store(true, Ordering::SeqCst);
+							let _ = response_tx.send(Ok(
+								protocol::KvResponseData::KvGetResponse(protocol::KvGetResponse {
+									keys: vec![PERSIST_DATA_KEY.to_vec()],
+									values: vec![encoded_persisted.clone()],
+									metadata: vec![protocol::KvMetadata {
+										version: Vec::new(),
+										update_ts: 0,
+									}],
+								}),
+							));
+						}
+						protocol::KvRequestData::KvPutRequest(_) => {
+							let _ = response_tx.send(Ok(protocol::KvResponseData::KvPutResponse));
+						}
+						other => {
+							let _ = response_tx
+								.send(Err(anyhow::anyhow!("unexpected KV request: {other:?}")));
+						}
+					}
+				}
+			}
+		});
+
+		let ctx = new_with_kv(
+			"actor-partial-preload-alarm",
+			"task-partial-preload-alarm",
+			Vec::new(),
+			"local",
+			crate::kv::Kv::new(handle.clone(), "actor-partial-preload-alarm"),
+		);
+		ctx.configure_envoy(handle, Some(32));
+		let mut task = new_task(ctx.clone()).with_preloaded_kv(Some(preloaded_kv));
+		let (start_tx, start_rx) = oneshot::channel();
+
+		task.handle_lifecycle(LifecycleCommand::Start { reply: start_tx })
+			.await;
+		start_rx
+			.await
+			.expect("start reply should send")
+			.expect("start should use preloaded alarm in partial preload bundle");
+
+		assert!(saw_persist_live_get.load(Ordering::SeqCst));
+		assert_eq!(ctx.last_pushed_alarm(), Some(future_ts));
+	}
+
 	#[tokio::test]
 	async fn startup_skips_future_alarm_push_when_last_pushed_matches() {
 		let kv = new_in_memory();