Skip to content

Commit e4476fd

Browse files
NathanFlurryNathanFlurry
committed
feat(engine): optionally automatically drain old runners on new runner version connected
1 parent 99737e0 commit e4476fd

File tree

24 files changed

+945
-154
lines changed

24 files changed

+945
-154
lines changed

‎engine/artifacts/openapi.json‎

Lines changed: 9 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

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

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ pub struct RunnerConfig {
99
pub kind: RunnerConfigKind,
1010
#[serde(default, skip_serializing_if = "Option::is_none")]
1111
pub metadata: Option<serde_json::Value>,
12+
#[serde(default = "default_drain_on_version_upgrade")]
13+
pub drain_on_version_upgrade: bool,
1214
}
1315

1416
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
@@ -27,9 +29,17 @@ pub enum RunnerConfigKind {
2729
},
2830
}
2931

32+
fn default_drain_on_version_upgrade() -> bool {
33+
false
34+
}
35+
3036
impl Into<rivet_types::runner_configs::RunnerConfig> for RunnerConfig {
3137
fn into(self) -> rivet_types::runner_configs::RunnerConfig {
32-
let RunnerConfig { kind, metadata } = self;
38+
let RunnerConfig {
39+
kind,
40+
metadata,
41+
drain_on_version_upgrade,
42+
} = self;
3343
let kind = match kind {
3444
RunnerConfigKind::Normal {} => rivet_types::runner_configs::RunnerConfigKind::Normal {},
3545
RunnerConfigKind::Serverless {
@@ -50,7 +60,10 @@ impl Into<rivet_types::runner_configs::RunnerConfig> for RunnerConfig {
5060
runners_margin: runners_margin.unwrap_or_default(),
5161
},
5262
};
53-
54-
rivet_types::runner_configs::RunnerConfig { kind, metadata }
63+
rivet_types::runner_configs::RunnerConfig {
64+
kind,
65+
metadata,
66+
drain_on_version_upgrade,
67+
}
5568
}
5669
}

‎engine/packages/engine/tests/actors_kv_drop.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use anyhow::*;
22
use async_trait::async_trait;
33
use common::test_runner::*;
4-
use rivet_runner_protocol as rp;
4+
use rivet_runner_protocol::mk2 as rp;
55
use std::sync::{Arc, Mutex};
66

77
mod common;

‎engine/packages/engine/tests/actors_kv_list.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use anyhow::*;
22
use async_trait::async_trait;
33
use common::test_runner::*;
4-
use rivet_runner_protocol as rp;
4+
use rivet_runner_protocol::mk2 as rp;
55
use std::sync::{Arc, Mutex};
66

77
mod common;

‎engine/packages/engine/tests/actors_kv_misc.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
use anyhow::*;
22
use async_trait::async_trait;
33
use common::test_runner::*;
4-
use rivet_runner_protocol as rp;
4+
use rivet_runner_protocol::mk2 as rp;
55
use std::sync::{Arc, Mutex};
66

77
mod common;

‎engine/packages/engine/tests/api_runner_configs_list.rs‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ fn list_runner_configs_single_runner() {
4141
rivet_api_types::namespaces::runner_configs::RunnerConfig {
4242
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
4343
metadata: None,
44+
drain_on_version_upgrade: true,
4445
},
4546
);
4647

@@ -95,6 +96,7 @@ fn list_runner_configs_multiple_runners() {
9596
rivet_api_types::namespaces::runner_configs::RunnerConfig {
9697
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
9798
metadata: None,
99+
drain_on_version_upgrade: true,
98100
},
99101
);
100102

@@ -147,13 +149,15 @@ fn list_runner_configs_multiple_dcs() {
147149
rivet_api_types::namespaces::runner_configs::RunnerConfig {
148150
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
149151
metadata: None,
152+
drain_on_version_upgrade: true,
150153
},
151154
);
152155
datacenters.insert(
153156
"dc-2".to_string(),
154157
rivet_api_types::namespaces::runner_configs::RunnerConfig {
155158
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
156159
metadata: None,
160+
drain_on_version_upgrade: true,
157161
},
158162
);
159163

@@ -210,6 +214,7 @@ fn list_runner_configs_filter_by_name() {
210214
rivet_api_types::namespaces::runner_configs::RunnerConfig {
211215
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
212216
metadata: None,
217+
drain_on_version_upgrade: true,
213218
},
214219
);
215220

@@ -260,6 +265,7 @@ fn list_runner_configs_filter_by_variant_normal() {
260265
rivet_api_types::namespaces::runner_configs::RunnerConfig {
261266
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
262267
metadata: None,
268+
drain_on_version_upgrade: true,
263269
},
264270
);
265271

@@ -320,6 +326,7 @@ fn list_runner_configs_filter_by_variant_serverless() {
320326
runners_margin: Some(2),
321327
},
322328
metadata: None,
329+
drain_on_version_upgrade: true,
323330
},
324331
);
325332

@@ -453,6 +460,7 @@ fn list_runner_configs_validates_returned_data() {
453460
runners_margin: Some(3),
454461
},
455462
metadata: Some(serde_json::json!({"key": "value"})),
463+
drain_on_version_upgrade: true,
456464
},
457465
);
458466

@@ -499,6 +507,7 @@ fn list_runner_configs_validates_returned_data() {
499507
min_runners,
500508
max_runners,
501509
runners_margin,
510+
..
502511
} = &dc_config.config.kind
503512
{
504513
assert_eq!(url, "http://localhost:9000");
@@ -530,6 +539,7 @@ fn list_runner_configs_mixed_variants() {
530539
rivet_api_types::namespaces::runner_configs::RunnerConfig {
531540
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
532541
metadata: None,
542+
drain_on_version_upgrade: true,
533543
},
534544
);
535545

@@ -565,6 +575,7 @@ fn list_runner_configs_mixed_variants() {
565575
runners_margin: Some(2),
566576
},
567577
metadata: None,
578+
drain_on_version_upgrade: true,
568579
},
569580
);
570581

‎engine/packages/engine/tests/api_runner_configs_upsert.rs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ fn upsert_runner_config_normal_single_dc() {
1616
rivet_api_types::namespaces::runner_configs::RunnerConfig {
1717
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
1818
metadata: None,
19+
drain_on_version_upgrade: true,
1920
},
2021
);
2122

@@ -50,13 +51,15 @@ fn upsert_runner_config_normal_multiple_dcs() {
5051
rivet_api_types::namespaces::runner_configs::RunnerConfig {
5152
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
5253
metadata: None,
54+
drain_on_version_upgrade: true,
5355
},
5456
);
5557
datacenters.insert(
5658
"dc-2".to_string(),
5759
rivet_api_types::namespaces::runner_configs::RunnerConfig {
5860
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
5961
metadata: None,
62+
drain_on_version_upgrade: true,
6063
},
6164
);
6265

@@ -97,6 +100,7 @@ fn upsert_runner_config_serverless() {
97100
runners_margin: Some(2),
98101
},
99102
metadata: None,
103+
drain_on_version_upgrade: true,
100104
},
101105
);
102106

@@ -129,6 +133,7 @@ fn upsert_runner_config_update_existing() {
129133
rivet_api_types::namespaces::runner_configs::RunnerConfig {
130134
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
131135
metadata: None,
136+
drain_on_version_upgrade: true,
132137
},
133138
);
134139

@@ -156,6 +161,7 @@ fn upsert_runner_config_update_existing() {
156161
rivet_api_types::namespaces::runner_configs::RunnerConfig {
157162
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
158163
metadata: Some(serde_json::json!({"test": "value"})),
164+
drain_on_version_upgrade: true,
159165
},
160166
);
161167

@@ -189,6 +195,7 @@ fn upsert_runner_config_returns_endpoint_changed() {
189195
rivet_api_types::namespaces::runner_configs::RunnerConfig {
190196
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
191197
metadata: None,
198+
drain_on_version_upgrade: true,
192199
},
193200
);
194201

@@ -232,6 +239,7 @@ fn upsert_runner_config_with_metadata() {
232239
rivet_api_types::namespaces::runner_configs::RunnerConfig {
233240
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
234241
metadata: Some(metadata_value),
242+
drain_on_version_upgrade: true,
235243
},
236244
);
237245

@@ -268,13 +276,15 @@ fn upsert_runner_config_removes_missing_dcs() {
268276
rivet_api_types::namespaces::runner_configs::RunnerConfig {
269277
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
270278
metadata: None,
279+
drain_on_version_upgrade: true,
271280
},
272281
);
273282
datacenters.insert(
274283
"dc-2".to_string(),
275284
rivet_api_types::namespaces::runner_configs::RunnerConfig {
276285
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
277286
metadata: None,
287+
drain_on_version_upgrade: true,
278288
},
279289
);
280290

@@ -300,6 +310,7 @@ fn upsert_runner_config_removes_missing_dcs() {
300310
rivet_api_types::namespaces::runner_configs::RunnerConfig {
301311
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
302312
metadata: None,
313+
drain_on_version_upgrade: true,
303314
},
304315
);
305316

@@ -365,6 +376,7 @@ fn upsert_runner_config_empty_map_deletes_all() {
365376
rivet_api_types::namespaces::runner_configs::RunnerConfig {
366377
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
367378
metadata: None,
379+
drain_on_version_upgrade: true,
368380
},
369381
);
370382

@@ -438,6 +450,7 @@ fn upsert_runner_config_non_existent_namespace() {
438450
rivet_api_types::namespaces::runner_configs::RunnerConfig {
439451
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
440452
metadata: None,
453+
drain_on_version_upgrade: true,
441454
},
442455
);
443456

@@ -473,6 +486,7 @@ fn upsert_runner_config_overwrites_different_variant() {
473486
rivet_api_types::namespaces::runner_configs::RunnerConfig {
474487
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
475488
metadata: None,
489+
drain_on_version_upgrade: true,
476490
},
477491
);
478492

@@ -506,6 +520,7 @@ fn upsert_runner_config_overwrites_different_variant() {
506520
runners_margin: Some(2),
507521
},
508522
metadata: None,
523+
drain_on_version_upgrade: true,
509524
},
510525
);
511526

@@ -540,6 +555,7 @@ fn upsert_runner_config_idempotent() {
540555
rivet_api_types::namespaces::runner_configs::RunnerConfig {
541556
kind: rivet_api_types::namespaces::runner_configs::RunnerConfigKind::Normal {},
542557
metadata: None,
558+
drain_on_version_upgrade: true,
543559
},
544560
);
545561

@@ -603,6 +619,7 @@ fn upsert_runner_config_serverless_slots_per_runner_zero() {
603619
runners_margin: Some(2),
604620
},
605621
metadata: None,
622+
drain_on_version_upgrade: true,
606623
},
607624
);
608625

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

Lines changed: 5 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use anyhow::*;
22
use async_trait::async_trait;
3-
use rivet_runner_protocol as rp;
3+
use rivet_runner_protocol::mk2 as rp;
44
use std::time::Duration;
55
use tokio::sync::{mpsc, oneshot};
66

@@ -47,33 +47,25 @@ impl ActorConfig {
4747
impl ActorConfig {
4848
/// Send a sleep intent
4949
pub fn send_sleep_intent(&self) {
50-
let event = protocol::make_actor_intent(
51-
&self.actor_id,
52-
self.generation,
53-
rp::ActorIntent::ActorIntentSleep,
54-
);
50+
let event = protocol::make_actor_intent(rp::ActorIntent::ActorIntentSleep);
5551
self.send_event(event);
5652
}
5753

5854
/// Send a stop intent
5955
pub fn send_stop_intent(&self) {
60-
let event = protocol::make_actor_intent(
61-
&self.actor_id,
62-
self.generation,
63-
rp::ActorIntent::ActorIntentStop,
64-
);
56+
let event = protocol::make_actor_intent(rp::ActorIntent::ActorIntentStop);
6557
self.send_event(event);
6658
}
6759

6860
/// Set an alarm to wake at specified timestamp (milliseconds)
6961
pub fn send_set_alarm(&self, alarm_ts: i64) {
70-
let event = protocol::make_set_alarm(&self.actor_id, self.generation, Some(alarm_ts));
62+
let event = protocol::make_set_alarm(Some(alarm_ts));
7163
self.send_event(event);
7264
}
7365

7466
/// Clear the alarm
7567
pub fn send_clear_alarm(&self) {
76-
let event = protocol::make_set_alarm(&self.actor_id, self.generation, None);
68+
let event = protocol::make_set_alarm(None);
7769
self.send_event(event);
7870
}
7971

0 commit comments

Comments
 (0)