fix: add serverless backfill

Author: MasterPtato

Cargo.lock

@@ -1,7 +1,7 @@
 
 [workspace]
 resolver = "2"
-members = ["engine/packages/actor-kv","engine/packages/api-builder","engine/packages/api-peer","engine/packages/api-public","engine/packages/api-types","engine/packages/api-util","engine/packages/bootstrap","engine/packages/cache","engine/packages/cache-purge","engine/packages/cache-result","engine/packages/clickhouse-inserter","engine/packages/clickhouse-user-query","engine/packages/config","engine/packages/dump-openapi","engine/packages/engine","engine/packages/env","engine/packages/epoxy","engine/packages/error","engine/packages/error-macros","engine/packages/gasoline","engine/packages/gasoline-macros","engine/packages/guard","engine/packages/guard-core","engine/packages/logs","engine/packages/metrics","engine/packages/namespace","engine/packages/pegboard","engine/packages/pegboard-gateway","engine/packages/pegboard-runner","engine/packages/pools","engine/packages/runtime","engine/packages/service-manager","engine/packages/telemetry","engine/packages/test-deps","engine/packages/test-deps-docker","engine/packages/tracing-reconfigure","engine/packages/types","engine/packages/universaldb","engine/packages/universalpubsub","engine/packages/util","engine/packages/util-id","engine/packages/workflow-worker","engine/sdks/rust/api-full","engine/sdks/rust/data","engine/sdks/rust/epoxy-protocol","engine/sdks/rust/runner-protocol","engine/sdks/rust/ups-protocol"]
+members = ["engine/packages/actor-kv","engine/packages/api-builder","engine/packages/api-peer","engine/packages/api-public","engine/packages/api-types","engine/packages/api-util","engine/packages/bootstrap","engine/packages/cache","engine/packages/cache-purge","engine/packages/cache-result","engine/packages/clickhouse-inserter","engine/packages/clickhouse-user-query","engine/packages/config","engine/packages/dump-openapi","engine/packages/engine","engine/packages/env","engine/packages/epoxy","engine/packages/error","engine/packages/error-macros","engine/packages/gasoline","engine/packages/gasoline-macros","engine/packages/guard","engine/packages/guard-core","engine/packages/logs","engine/packages/metrics","engine/packages/namespace","engine/packages/pegboard","engine/packages/pegboard-gateway","engine/packages/pegboard-runner","engine/packages/pools","engine/packages/runtime","engine/packages/serverless-backfill","engine/packages/service-manager","engine/packages/telemetry","engine/packages/test-deps","engine/packages/test-deps-docker","engine/packages/tracing-reconfigure","engine/packages/tracing-utils","engine/packages/types","engine/packages/universaldb","engine/packages/universalpubsub","engine/packages/util","engine/packages/util-id","engine/packages/workflow-worker","engine/sdks/rust/api-full","engine/sdks/rust/data","engine/sdks/rust/epoxy-protocol","engine/sdks/rust/runner-protocol","engine/sdks/rust/ups-protocol"]
 
 [workspace.package]
 version = "2.0.25"
@@ -83,10 +83,13 @@ tracing = "0.1.40"
 tracing-core = "0.1"
 tracing-opentelemetry = "0.29"
 tracing-slog = "0.2"
-vergen = { version = "9.0.4", features = ["build", "cargo", "rustc"] }
 vergen-gitcl = "1.0.0"
 reqwest-eventsource = "0.6.0"
 
+[workspace.dependencies.vergen]
+version = "9.0.4"
+features = ["build","cargo","rustc"]
+
 [workspace.dependencies.sentry]
 version = "0.45.0"
 default-features = false
@@ -148,7 +151,7 @@ features = ["now"]
 
 [workspace.dependencies.clap]
 version = "4.3"
-features = ["derive", "cargo"]
+features = ["derive","cargo"]
 
 [workspace.dependencies.rivet-term]
 git = "https://github.com/rivet-dev/rivet-term"
@@ -357,6 +360,9 @@ path = "engine/packages/pools"
 [workspace.dependencies.rivet-runtime]
 path = "engine/packages/runtime"
 
+[workspace.dependencies.rivet-serverless-backfill]
+path = "engine/packages/serverless-backfill"
+
 [workspace.dependencies.rivet-service-manager]
 path = "engine/packages/service-manager"
 
@@ -425,8 +431,3 @@ debug = false
 lto = "fat"
 codegen-units = 1
 opt-level = 3
-
-# strip = true
-# panic = "abort"
-# overflow-checks = false
-# debug-assertions = false

Cargo.toml

@@ -32,6 +32,7 @@ rivet-guard.workspace = true
 rivet-logs.workspace = true
 rivet-pools.workspace = true
 rivet-runtime.workspace = true
+rivet-serverless-backfill.workspace = true
 rivet-service-manager.workspace = true
 rivet-telemetry.workspace = true
 rivet-term.workspace = true

engine/packages/engine/Cargo.toml

@@ -40,6 +40,12 @@ pub fn config(_rivet_config: rivet_config::Config) -> Result<RunConfigData> {
 			|config, pools| Box::pin(rivet_cache_purge::start(config, pools)),
 			false,
 		),
+		Service::new(
+			"serverless_backfill",
+			ServiceKind::Oneshot,
+			|config, pools| Box::pin(rivet_serverless_backfill::start(config, pools)),
+			false,
+		),
 	];
 
 	Ok(RunConfigData { services })

engine/packages/engine/src/run_config.rs

@@ -0,0 +1,13 @@
+[package]
+name = "rivet-serverless-backfill"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+anyhow.workspace = true
+gas.workspace = true
+pegboard.workspace = true
+rivet-config.workspace = true
+rivet-types.workspace = true
+tracing.workspace = true
+universaldb.workspace = true

engine/packages/serverless-backfill/Cargo.toml

@@ -0,0 +1,81 @@
+use anyhow::Result;
+use futures_util::{StreamExt, TryStreamExt};
+use gas::prelude::*;
+use universaldb::options::StreamingMode;
+use universaldb::utils::IsolationLevel::*;
+
+#[tracing::instrument(skip_all)]
+pub async fn start(config: rivet_config::Config, pools: rivet_pools::Pools) -> Result<()> {
+	let cache = rivet_cache::CacheInner::from_env(&config, pools.clone())?;
+	let ctx = StandaloneCtx::new(
+		db::DatabaseKv::from_pools(pools.clone()).await?,
+		config.clone(),
+		pools,
+		cache,
+		"serverless_backfill",
+		Id::new_v1(config.dc_label()),
+		Id::new_v1(config.dc_label()),
+	)?;
+
+	let serverless_data = ctx
+		.udb()?
+		.run(|tx| async move {
+			let tx = tx.with_subspace(pegboard::keys::subspace());
+
+			let serverless_desired_subspace = pegboard::keys::subspace().subspace(
+				&rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey::entire_subspace(),
+			);
+
+			tx.get_ranges_keyvalues(
+				universaldb::RangeOption {
+					mode: StreamingMode::WantAll,
+					..(&serverless_desired_subspace).into()
+				},
+				// NOTE: This is a snapshot to prevent conflict with updates to this subspace
+				Snapshot,
+			)
+			.map(|res| {
+				tx.unpack::<rivet_types::keys::pegboard::ns::ServerlessDesiredSlotsKey>(res?.key())
+			})
+			.try_collect::<Vec<_>>()
+			.await
+		})
+		.custom_instrument(tracing::info_span!("read_serverless_tx"))
+		.await?;
+
+	let runner_configs = ctx
+		.op(pegboard::ops::runner_config::get::Input {
+			runners: serverless_data
+				.iter()
+				.map(|key| (key.namespace_id, key.runner_name.clone()))
+				.collect(),
+			bypass_cache: true,
+		})
+		.await?;
+
+	for key in &serverless_data {
+		if !runner_configs
+			.iter()
+			.any(|rc| rc.namespace_id == key.namespace_id)
+		{
+			tracing::debug!(
+				namespace_id=?key.namespace_id,
+				runner_name=?key.runner_name,
+				"runner config not found, likely deleted"
+			);
+			continue;
+		};
+
+		ctx.workflow(pegboard::workflows::runner_pool::Input {
+			namespace_id: key.namespace_id,
+			runner_name: key.runner_name.clone(),
+		})
+		.tag("namespace_id", key.namespace_id)
+		.tag("runner_name", key.runner_name.clone())
+		.unique()
+		.dispatch()
+		.await?;
+	}
+
+	Ok(())
+}

engine/packages/serverless-backfill/src/lib.rs

@@ -54,16 +54,11 @@ kubectl apply -f 02-engine-configmap.yaml
 kubectl apply -f 03-rivet-engine-deployment.yaml
 kubectl apply -f 04-rivet-engine-service.yaml
 kubectl apply -f 05-rivet-engine-hpa.yaml
-kubectl apply -f 06-rivet-engine-singleton-deployment.yaml
 
 # Wait for engine to be ready
 echo "Waiting for engine to be ready..."
 kubectl -n "${NAMESPACE}" wait --for=condition=ready pod -l app=rivet-engine --timeout=300s
 
-# Wait for singleton to be ready
-echo "Waiting for singleton to be ready..."
-kubectl -n "${NAMESPACE}" wait --for=condition=ready pod -l app=rivet-engine-singleton --timeout=300s
-
 echo ""
 echo "Deployment complete."
 echo ""