feat(engine): optionally automatically drain old runners on new runner version connected

engine/packages/api-types/src/namespaces/runner_configs.rs

@@ -9,6 +9,8 @@ pub struct RunnerConfig {
 	pub kind: RunnerConfigKind,
 	#[serde(default, skip_serializing_if = "Option::is_none")]
 	pub metadata: Option<serde_json::Value>,
+	#[serde(default = "default_drain_on_version_upgrade")]
+	pub drain_on_version_upgrade: bool,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
@@ -27,9 +29,17 @@ pub enum RunnerConfigKind {
 	},
 }
 
+fn default_drain_on_version_upgrade() -> bool {
+	false
+}
+
 impl Into<rivet_types::runner_configs::RunnerConfig> for RunnerConfig {
 	fn into(self) -> rivet_types::runner_configs::RunnerConfig {
-		let RunnerConfig { kind, metadata } = self;
+		let RunnerConfig {
+			kind,
+			metadata,
+			drain_on_version_upgrade,
+		} = self;
 		let kind = match kind {
 			RunnerConfigKind::Normal {} => rivet_types::runner_configs::RunnerConfigKind::Normal {},
 			RunnerConfigKind::Serverless {
@@ -50,7 +60,10 @@ impl Into<rivet_types::runner_configs::RunnerConfig> for RunnerConfig {
 				runners_margin: runners_margin.unwrap_or_default(),
 			},
 		};
-
-		rivet_types::runner_configs::RunnerConfig { kind, metadata }
+		rivet_types::runner_configs::RunnerConfig {
+			kind,
+			metadata,
+			drain_on_version_upgrade,
+		}
 	}
 }

engine/packages/engine/tests/actors_kv_drop.rs

@@ -1,7 +1,7 @@
 use anyhow::*;
 use async_trait::async_trait;
 use common::test_runner::*;
-use rivet_runner_protocol as rp;
+use rivet_runner_protocol::mk2 as rp;
 use std::sync::{Arc, Mutex};
 
 mod common;

engine/packages/engine/tests/actors_kv_list.rs

@@ -1,7 +1,7 @@
 use anyhow::*;
 use async_trait::async_trait;
 use common::test_runner::*;
-use rivet_runner_protocol as rp;
+use rivet_runner_protocol::mk2 as rp;
 use std::sync::{Arc, Mutex};
 
 mod common;

engine/packages/engine/tests/actors_kv_misc.rs

@@ -1,7 +1,7 @@
 use anyhow::*;
 use async_trait::async_trait;
 use common::test_runner::*;
-use rivet_runner_protocol as rp;
+use rivet_runner_protocol::mk2 as rp;
 use std::sync::{Arc, Mutex};
 
 mod common;

engine/packages/engine/tests/api_runner_configs_list.rs

@@ -41,6 +41,7 @@ fn list_runner_configs_single_runner() {
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -95,6 +96,7 @@ fn list_runner_configs_multiple_runners() {
 				rivet_api_types::namespaces::runner_configs::RunnerConfig {
 					kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 					metadata: None,
+					drain_on_version_upgrade: true,
 				},
 			);
 
@@ -147,13 +149,15 @@ fn list_runner_configs_multiple_dcs() {
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 		datacenters.insert(
 			"dc-2".to_string(),
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -210,6 +214,7 @@ fn list_runner_configs_filter_by_name() {
 				rivet_api_types::namespaces::runner_configs::RunnerConfig {
 					kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 					metadata: None,
+					drain_on_version_upgrade: true,
 				},
 			);
 
@@ -260,6 +265,7 @@ fn list_runner_configs_filter_by_variant_normal() {
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -320,6 +326,7 @@ fn list_runner_configs_filter_by_variant_serverless() {
 					runners_margin: Some(2),
 				},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -453,6 +460,7 @@ fn list_runner_configs_validates_returned_data() {
 					runners_margin: Some(3),
 				},
 				metadata: Some(serde_json::json!({"key": "value"})),
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -499,6 +507,7 @@ fn list_runner_configs_validates_returned_data() {
 			min_runners,
 			max_runners,
 			runners_margin,
+			..
 		} = &dc_config.config.kind
 		{
 			assert_eq!(url, "http://localhost:9000");
@@ -530,6 +539,7 @@ fn list_runner_configs_mixed_variants() {
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -565,6 +575,7 @@ fn list_runner_configs_mixed_variants() {
 					runners_margin: Some(2),
 				},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 

engine/packages/engine/tests/api_runner_configs_upsert.rs

@@ -16,6 +16,7 @@ fn upsert_runner_config_normal_single_dc() {
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -50,13 +51,15 @@ fn upsert_runner_config_normal_multiple_dcs() {
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 		datacenters.insert(
 			"dc-2".to_string(),
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -97,6 +100,7 @@ fn upsert_runner_config_serverless() {
 					runners_margin: Some(2),
 				},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -129,6 +133,7 @@ fn upsert_runner_config_update_existing() {
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -156,6 +161,7 @@ fn upsert_runner_config_update_existing() {
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: Some(serde_json::json!({"test": "value"})),
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -189,6 +195,7 @@ fn upsert_runner_config_returns_endpoint_changed() {
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -232,6 +239,7 @@ fn upsert_runner_config_with_metadata() {
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: Some(metadata_value),
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -268,13 +276,15 @@ fn upsert_runner_config_removes_missing_dcs() {
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 		datacenters.insert(
 			"dc-2".to_string(),
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -300,6 +310,7 @@ fn upsert_runner_config_removes_missing_dcs() {
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -365,6 +376,7 @@ fn upsert_runner_config_empty_map_deletes_all() {
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -438,6 +450,7 @@ fn upsert_runner_config_non_existent_namespace() {
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -473,6 +486,7 @@ fn upsert_runner_config_overwrites_different_variant() {
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -506,6 +520,7 @@ fn upsert_runner_config_overwrites_different_variant() {
 					runners_margin: Some(2),
 				},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -540,6 +555,7 @@ fn upsert_runner_config_idempotent() {
 			rivet_api_types::namespaces::runner_configs::RunnerConfig {
 				kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 
@@ -603,6 +619,7 @@ fn upsert_runner_config_serverless_slots_per_runner_zero() {
 					runners_margin: Some(2),
 				},
 				metadata: None,
+				drain_on_version_upgrade: true,
 			},
 		);
 

engine/packages/engine/tests/common/test_runner/actor.rs

@@ -1,6 +1,6 @@
 use anyhow::*;
 use async_trait::async_trait;
-use rivet_runner_protocol as rp;
+use rivet_runner_protocol::mk2 as rp;
 use std::time::Duration;
 use tokio::sync::{mpsc, oneshot};
 
@@ -47,33 +47,25 @@ impl ActorConfig {
 impl ActorConfig {
 	/// Send a sleep intent
 	pub fn send_sleep_intent(&self) {
-		let event = protocol::make_actor_intent(
-			&self.actor_id,
-			self.generation,
-			rp::ActorIntent::ActorIntentSleep,
-		);
+		let event = protocol::make_actor_intent(rp::ActorIntent::ActorIntentSleep);
 		self.send_event(event);
 	}
 
 	/// Send a stop intent
 	pub fn send_stop_intent(&self) {
-		let event = protocol::make_actor_intent(
-			&self.actor_id,
-			self.generation,
-			rp::ActorIntent::ActorIntentStop,
-		);
+		let event = protocol::make_actor_intent(rp::ActorIntent::ActorIntentStop);
 		self.send_event(event);
 	}
 
 	/// Set an alarm to wake at specified timestamp (milliseconds)
 	pub fn send_set_alarm(&self, alarm_ts: i64) {
-		let event = protocol::make_set_alarm(&self.actor_id, self.generation, Some(alarm_ts));
+		let event = protocol::make_set_alarm(Some(alarm_ts));
 		self.send_event(event);
 	}
 
 	/// Clear the alarm
 	pub fn send_clear_alarm(&self) {
-		let event = protocol::make_set_alarm(&self.actor_id, self.generation, None);
+		let event = protocol::make_set_alarm(None);
 		self.send_event(event);
 	}