fix(rivetkit-napi): key actor runtime state by generation

Author: abcxff

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

@@ -953,6 +953,12 @@ impl ActorContext {
 		&self.0.actor_id
 	}
 
+	/// Returns the envoy generation this context was started for, or `None` before an envoy is
+	/// configured.
+	pub fn generation(&self) -> Option<u32> {
+		self.sleep_generation()
+	}
+
 	pub fn name(&self) -> &str {
 		&self.0.name
 	}

rivetkit-rust/packages/rivetkit-core/src/testing.rs

@@ -71,9 +71,23 @@ impl ActorContextHarness {
 		key: ActorKey,
 		region: impl Into<String>,
 		config: ActorConfig,
+	) -> ActorContext {
+		self.context_with_generation(actor_id, name, key, region, config, 1)
+	}
+
+	/// Constructs a context for a specific actor generation, for tests where generations of the
+	/// same actor coexist.
+	pub fn context_with_generation(
+		&self,
+		actor_id: impl Into<String>,
+		name: impl Into<String>,
+		key: ActorKey,
+		region: impl Into<String>,
+		config: ActorConfig,
+		generation: u32,
 	) -> ActorContext {
 		let actor_id = actor_id.into();
-		let generation = Some(1);
+		let generation = Some(generation);
 		let sql = SqliteDb::new_with_remote_sqlite(
 			self.handle.clone(),
 			actor_id.clone(),

rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs

@@ -43,8 +43,13 @@ type DisconnectPredicateTsfn =
 type RunRestartHook = Arc<dyn Fn(Option<(i64, u64)>) -> anyhow::Result<()> + Send + Sync + 'static>;
 pub(crate) type RegisteredTask = Pin<Box<dyn Future<Output = ()> + Send + 'static>>;
 
-static ACTOR_CONTEXT_SHARED: LazyLock<SccHashMap<String, Weak<ActorContextShared>>> =
-	LazyLock::new(SccHashMap::new);
+// Keyed by generation so a newer generation of the same actor never shares, resets, or clears the
+// JS runtime state of an older generation that is still shutting down on this host.
+static ACTOR_CONTEXT_SHARED: LazyLock<
+	SccHashMap<ActorContextSharedKey, Weak<ActorContextShared>>,
+> = LazyLock::new(SccHashMap::new);
+
+type ActorContextSharedKey = (String, Option<u32>);
 
 /// N-API wrapper around `rivetkit-core::ActorContext`.
 #[derive(Clone)]
@@ -193,10 +198,12 @@ struct DisconnectPredicatePayload {
 impl ActorContext {
 	pub(crate) fn new(inner: CoreActorContext) -> Self {
 		let actor_id = inner.actor_id().to_owned();
-		let shared = actor_context_shared(&actor_id);
+		let generation = inner.generation();
+		let shared = actor_context_shared(&actor_id, generation);
 		tracing::debug!(
 			class = "ActorContext",
 			%actor_id,
+			?generation,
 			shared_strong_count = Arc::strong_count(&shared),
 			"constructed napi class"
 		);
@@ -994,14 +1001,15 @@ impl Drop for ActorContextShared {
 	}
 }
 
-fn actor_context_shared(actor_id: &str) -> Arc<ActorContextShared> {
+fn actor_context_shared(actor_id: &str, generation: Option<u32>) -> Arc<ActorContextShared> {
 	ACTOR_CONTEXT_SHARED.retain_sync(|_, shared| shared.strong_count() > 0);
 
-	match ACTOR_CONTEXT_SHARED.entry_sync(actor_id.to_owned()) {
+	match ACTOR_CONTEXT_SHARED.entry_sync((actor_id.to_owned(), generation)) {
 		scc::hash_map::Entry::Occupied(mut entry) => {
 			if let Some(shared) = entry.get().upgrade() {
 				tracing::debug!(
 					%actor_id,
+					?generation,
 					outcome = "hit",
 					strong_count = Arc::strong_count(&shared),
 					"actor context shared-state cache lookup"
@@ -1013,6 +1021,7 @@ fn actor_context_shared(actor_id: &str) -> Arc<ActorContextShared> {
 			*entry.get_mut() = Arc::downgrade(&shared);
 			tracing::debug!(
 				%actor_id,
+				?generation,
 				outcome = "stale",
 				"actor context shared-state cache lookup"
 			);
@@ -1023,6 +1032,7 @@ fn actor_context_shared(actor_id: &str) -> Arc<ActorContextShared> {
 			entry.insert_entry(Arc::downgrade(&shared));
 			tracing::debug!(
 				%actor_id,
+				?generation,
 				outcome = "miss",
 				"actor context shared-state cache lookup"
 			);

rivetkit-typescript/packages/rivetkit-napi/tests/actor_context.rs

@@ -18,4 +18,32 @@ mod moved_tests {
 		assert!(!shared.has_end_reason());
 		assert!(ctx.started());
 	}
+
+	#[test]
+	fn generations_of_the_same_actor_do_not_share_runtime_state() {
+		let harness = rivetkit_core::testing::ActorContextHarness::new();
+		let context_for = |generation| {
+			ActorContext::new(harness.context_with_generation(
+				"actor-generations",
+				"actor",
+				Vec::new(),
+				"local",
+				rivetkit_core::ActorConfig::default(),
+				generation,
+			))
+		};
+
+		let old = context_for(1);
+		let old_again = context_for(1);
+		let new = context_for(2);
+
+		assert!(Arc::ptr_eq(&old.shared, &old_again.shared));
+		assert!(!Arc::ptr_eq(&old.shared, &new.shared));
+
+		// An older generation that finishes shutting down after a newer generation started must
+		// not reset the newer generation's runtime state.
+		new.shared.set_end_reason(EndReason::Sleep);
+		old.shared.reset_runtime_state();
+		assert!(new.shared.has_end_reason());
+	}
 }