@rivetkit/engine 1.0.0 → 2.2.1-pr.4600.b74ff3b

package/packages/pegboard/src/ops/envoy/drain.rs

@@ -0,0 +1,101 @@
+use anyhow::Result;
+use futures_util::TryStreamExt;
+use gas::prelude::*;
+use rivet_envoy_protocol::{self as protocol, PROTOCOL_VERSION, versioned};
+use universaldb::options::StreamingMode;
+use universaldb::utils::IsolationLevel::*;
+use universalpubsub::PublishOpts;
+use vbare::OwnedVersionedData;
+
+use crate::{keys, metrics};
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Input {
+	pub namespace_id: Id,
+	pub pool_name: String,
+	pub version: u32,
+}
+
+#[operation]
+pub async fn pegboard_envoy_drain_older_versions(ctx: &OperationCtx, input: &Input) -> Result<()> {
+	let pool_res = ctx
+		.op(crate::ops::runner_config::get::Input {
+			runners: vec![(input.namespace_id, input.pool_name.clone())],
+			bypass_cache: false,
+		})
+		.await?;
+
+	let Some(pool) = pool_res.into_iter().next() else {
+		return Ok(());
+	};
+
+	// Use config's drain_on_version_upgrade if config exists, otherwise default to false
+	if !pool.config.drain_on_version_upgrade {
+		return Ok(());
+	}
+
+	// Scan EnvoyLoadBalancerIdxKey for older versions
+	let older_envoys = ctx
+		.udb()?
+		.run(|tx| async move {
+			let tx = tx.with_subspace(keys::subspace());
+			let mut older_envoys = Vec::new();
+
+			let lb_subspace =
+				keys::subspace().subspace(&keys::ns::EnvoyLoadBalancerIdxKey::subspace(
+					input.namespace_id,
+					input.pool_name.clone(),
+				));
+
+			let mut stream = tx.get_ranges_keyvalues(
+				universaldb::RangeOption {
+					mode: StreamingMode::WantAll,
+					..(&lb_subspace).into()
+				},
+				Snapshot,
+			);
+
+			while let Some(entry) = stream.try_next().await? {
+				let (key, _) = tx.read_entry::<keys::ns::EnvoyLoadBalancerIdxKey>(&entry)?;
+
+				// Only collect envoys with older versions
+				if key.version < input.version {
+					older_envoys.push(key.envoy_key);
+				}
+			}
+
+			Ok(older_envoys)
+		})
+		.custom_instrument(tracing::info_span!("drain_older_versions_tx"))
+		.await?;
+
+	if !older_envoys.is_empty() {
+		tracing::info!(
+			namespace_id = %input.namespace_id,
+			pool_name = %input.pool_name,
+			new_version = input.version,
+			older_envoy_count = older_envoys.len(),
+			"draining older envoy versions due to drain_on_version_upgrade"
+		);
+
+		metrics::ENVOY_VERSION_UPGRADE_DRAIN_TOTAL
+			.with_label_values(&[&input.namespace_id.to_string(), &input.pool_name])
+			.inc_by(older_envoys.len() as u64);
+
+		for envoy_key in older_envoys {
+			let receiver_subject =
+				crate::pubsub_subjects::EnvoyReceiverSubject::new(input.namespace_id, envoy_key)
+					.to_string();
+
+			let message_serialized =
+				versioned::ToEnvoyConn::wrap_latest(protocol::ToEnvoyConn::ToEnvoyConnClose)
+					.serialize_with_embedded_version(PROTOCOL_VERSION)?;
+
+			ctx.ups()?
+				.publish(&receiver_subject, &message_serialized, PublishOpts::one())
+				.await?;
+		}
+	}
+
+	Ok(())
+}

package/packages/pegboard/src/ops/envoy/evict_actors.rs

@@ -0,0 +1,54 @@
+use futures_util::{StreamExt, TryStreamExt};
+use gas::prelude::*;
+use universaldb::prelude::*;
+
+use crate::keys;
+
+#[derive(Debug)]
+pub struct Input {
+	pub namespace_id: Id,
+	pub envoy_key: String,
+}
+
+#[operation]
+pub async fn pegboard_envoy_evict_actors(ctx: &OperationCtx, input: &Input) -> Result<()> {
+	let actors = ctx
+		.udb()?
+		.run(|tx| async move {
+			let tx = tx.with_subspace(keys::subspace());
+
+			let actor_subspace = keys::subspace().subspace(&keys::envoy::ActorKey::subspace(
+				input.namespace_id,
+				input.envoy_key.clone(),
+			));
+
+			tx.get_ranges_keyvalues(
+				universaldb::RangeOption {
+					mode: StreamingMode::WantAll,
+					..(&actor_subspace).into()
+				},
+				Serializable,
+			)
+			.map(|res| {
+				let (key, generation) = tx.read_entry::<keys::runner::ActorKey>(&res?)?;
+
+				Ok((key.actor_id, generation))
+			})
+			.try_collect::<Vec<_>>()
+			.await
+		})
+		.custom_instrument(tracing::info_span!("envoy_list_actors_tx"))
+		.await?;
+
+	// TODO: Parallelize
+	for (actor_id, generation) in actors {
+		ctx.signal(crate::workflows::actor2::GoingAway { generation })
+			.to_workflow::<crate::workflows::actor2::Workflow>()
+			.tag("actor_id", actor_id)
+			.graceful_not_found()
+			.send()
+			.await?;
+	}
+
+	Ok(())
+}

package/packages/pegboard/src/ops/envoy/expire.rs

@@ -0,0 +1,92 @@
+use gas::prelude::*;
+use universaldb::options::ConflictRangeType;
+use universaldb::utils::IsolationLevel::*;
+
+use crate::keys;
+
+#[derive(Debug)]
+pub struct Input {
+	pub namespace_id: Id,
+	pub envoy_key: String,
+}
+
+#[operation]
+pub async fn pegboard_envoy_expire(ctx: &OperationCtx, input: &Input) -> Result<()> {
+	ctx
+		.udb()?
+		.run(|tx| {
+			async move {
+				let tx = tx.with_subspace(keys::subspace());
+
+				let pool_name_key = keys::envoy::PoolNameKey::new(input.namespace_id, input.envoy_key.clone());
+				let version_key = keys::envoy::VersionKey::new(input.namespace_id, input.envoy_key.clone());
+				let create_ts_key = keys::envoy::CreateTsKey::new(input.namespace_id, input.envoy_key.clone());
+				let last_ping_ts_key = keys::envoy::LastPingTsKey::new(input.namespace_id, input.envoy_key.clone());
+				let expired_ts_key = keys::envoy::ExpiredTsKey::new(input.namespace_id, input.envoy_key.clone());
+
+				let (
+					pool_name_entry,
+					version_entry,
+					create_ts_entry,
+					last_ping_ts_entry,
+					expired,
+				) = tokio::try_join!(
+					tx.read_opt(&pool_name_key, Serializable),
+					tx.read_opt(&version_key, Serializable),
+					tx.read_opt(&create_ts_key, Serializable),
+					tx.read_opt(&last_ping_ts_key, Serializable),
+					tx.exists(&expired_ts_key, Serializable),
+				)?;
+
+				let (
+					Some(pool_name),
+					Some(version),
+					Some(create_ts),
+					Some(old_last_ping_ts),
+				) = (
+					pool_name_entry,
+					version_entry,
+					create_ts_entry,
+					last_ping_ts_entry,
+				)
+				else {
+					tracing::debug!(namespace_id=?input.namespace_id, envoy_key=%input.envoy_key, "envoy not found");
+					return Ok(());
+				};
+
+				if !expired {
+					let idx_key = keys::ns::EnvoyLoadBalancerIdxKey::new(
+						input.namespace_id,
+						pool_name.clone(),
+						version,
+						old_last_ping_ts,
+						input.envoy_key.clone(),
+					);
+
+					tx.add_conflict_key(&idx_key, ConflictRangeType::Read)?;
+					tx.delete(&idx_key);
+
+					tx.write(&expired_ts_key, util::timestamp::now())?;
+					tx.delete(
+						&keys::ns::ActiveEnvoyKey::new(
+							input.namespace_id,
+							create_ts,
+							input.envoy_key.clone(),
+						),
+					);
+					tx.delete(
+						&keys::ns::ActiveEnvoyByNameKey::new(
+							input.namespace_id,
+							pool_name.clone(),
+							create_ts,
+							input.envoy_key.clone(),
+						),
+					);
+				}
+
+				Ok(())
+			}
+		})
+		.custom_instrument(tracing::info_span!("envoy_expire_tx"))
+		.await
+}

package/packages/pegboard/src/ops/envoy/get.rs

@@ -0,0 +1,135 @@
+use anyhow::Result;
+use futures_util::TryStreamExt;
+use gas::prelude::*;
+use rivet_types::envoys::Envoy;
+use universaldb::options::StreamingMode;
+use universaldb::utils::{FormalChunkedKey, IsolationLevel::*};
+
+use crate::keys;
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Input {
+	pub namespace_id: Id,
+	pub envoy_keys: Vec<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Output {
+	pub envoys: Vec<Envoy>,
+}
+
+#[operation]
+pub async fn pegboard_envoy_get(ctx: &OperationCtx, input: &Input) -> Result<Output> {
+	let dc_name = ctx.config().dc_name()?;
+
+	let envoys = ctx
+		.udb()?
+		.run(|tx| {
+			let dc_name = dc_name.to_string();
+			async move {
+				let mut envoys = Vec::new();
+
+				for envoy_key in &input.envoy_keys {
+					if let Some(envoy) =
+						get_inner(&dc_name, &tx, input.namespace_id, envoy_key).await?
+					{
+						envoys.push(envoy);
+					}
+				}
+
+				Ok(envoys)
+			}
+		})
+		.custom_instrument(tracing::info_span!("envoy_get_tx"))
+		.await?;
+
+	Ok(Output { envoys })
+}
+
+pub(crate) async fn get_inner(
+	dc_name: &str,
+	tx: &universaldb::Transaction,
+	namespace_id: Id,
+	envoy_key: &str,
+) -> Result<Option<Envoy>> {
+	let tx = tx.with_subspace(keys::subspace());
+
+	// TODO: Make this part of the below try join to reduce round trip count
+	// Check if envoy exists by looking for workflow ID
+	if !tx
+		.exists(
+			&keys::envoy::PoolNameKey::new(namespace_id, envoy_key.to_string()),
+			Serializable,
+		)
+		.await?
+	{
+		return Ok(None);
+	}
+
+	let pool_name_key = keys::envoy::PoolNameKey::new(namespace_id, envoy_key.to_string());
+	let version_key = keys::envoy::VersionKey::new(namespace_id, envoy_key.to_string());
+	let slots_key = keys::envoy::SlotsKey::new(namespace_id, envoy_key.to_string());
+	let create_ts_key = keys::envoy::CreateTsKey::new(namespace_id, envoy_key.to_string());
+	let connected_ts_key = keys::envoy::ConnectedTsKey::new(namespace_id, envoy_key.to_string());
+	let stop_ts_key = keys::envoy::StopTsKey::new(namespace_id, envoy_key.to_string());
+	let last_ping_ts_key = keys::envoy::LastPingTsKey::new(namespace_id, envoy_key.to_string());
+	let last_rtt_key = keys::envoy::LastRttKey::new(namespace_id, envoy_key.to_string());
+	let metadata_key = keys::envoy::MetadataKey::new(namespace_id, envoy_key.to_string());
+	let metadata_subspace = keys::subspace().subspace(&metadata_key);
+
+	let (
+		pool_name,
+		version,
+		slots,
+		create_ts,
+		connected_ts,
+		stop_ts,
+		last_ping_ts,
+		last_rtt,
+		metadata_chunks,
+	) = tokio::try_join!(
+		// NOTE: These are not Serializable because this op is meant for basic information (i.e. data for the
+		// API)
+		tx.read(&pool_name_key, Snapshot),
+		tx.read(&version_key, Snapshot),
+		tx.read(&slots_key, Snapshot),
+		tx.read(&create_ts_key, Snapshot),
+		tx.read_opt(&connected_ts_key, Snapshot),
+		tx.read_opt(&stop_ts_key, Snapshot),
+		tx.read_opt(&last_ping_ts_key, Snapshot),
+		tx.read_opt(&last_rtt_key, Snapshot),
+		async {
+			tx.get_ranges_keyvalues(
+				universaldb::RangeOption {
+					mode: StreamingMode::WantAll,
+					..(&metadata_subspace).into()
+				},
+				Snapshot,
+			)
+			.try_collect::<Vec<_>>()
+			.await
+			.map_err(Into::into)
+		},
+	)?;
+
+	let metadata = if metadata_chunks.is_empty() {
+		None
+	} else {
+		Some(metadata_key.combine(metadata_chunks)?.metadata)
+	};
+
+	Ok(Some(Envoy {
+		envoy_key: envoy_key.to_string(),
+		namespace_id: namespace_id,
+		datacenter: dc_name.to_string(),
+		pool_name,
+		version,
+		slots: slots.try_into()?,
+		create_ts,
+		last_connected_ts: connected_ts,
+		stop_ts,
+		last_ping_ts: last_ping_ts.unwrap_or_default(),
+		last_rtt: last_rtt.unwrap_or_default(),
+		metadata,
+	}))
+}

package/packages/pegboard/src/ops/envoy/list.rs

@@ -0,0 +1,131 @@
+use anyhow::Result;
+use futures_util::{StreamExt, TryStreamExt};
+use gas::prelude::*;
+use rivet_types::envoys::Envoy;
+use universaldb::options::StreamingMode;
+use universaldb::utils::IsolationLevel::*;
+
+use crate::keys;
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Input {
+	pub namespace_id: Id,
+	pub pool_name: Option<String>,
+	pub created_before: Option<i64>,
+	pub limit: usize,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Output {
+	pub envoys: Vec<Envoy>,
+}
+
+#[operation]
+pub async fn pegboard_envoy_list(ctx: &OperationCtx, input: &Input) -> Result<Output> {
+	let dc_name = ctx.config().dc_name()?;
+
+	let envoys = ctx
+		.udb()?
+		.run(|tx| {
+			let dc_name = dc_name.to_string();
+			async move {
+				let tx = tx.with_subspace(keys::subspace());
+				let mut results = Vec::new();
+
+				if let Some(pool_name) = &input.pool_name {
+					let envoy_subspace =
+						keys::subspace().subspace(&keys::ns::ActiveEnvoyByNameKey::subspace(
+							input.namespace_id,
+							pool_name.clone(),
+						));
+					let (start, end) = envoy_subspace.range();
+
+					let end = if let Some(created_before) = input.created_before {
+						universaldb::utils::end_of_key_range(&tx.pack(
+							&keys::ns::ActiveEnvoyByNameKey::subspace_with_create_ts(
+								input.namespace_id,
+								pool_name.clone(),
+								created_before,
+							),
+						))
+					} else {
+						end
+					};
+
+					let mut stream = tx.get_ranges_keyvalues(
+						universaldb::RangeOption {
+							mode: StreamingMode::Iterator,
+							reverse: true,
+							..(start, end).into()
+						},
+						// NOTE: Does not have to be serializable because we are listing, stale data does not matter
+						Snapshot,
+					);
+
+					while let Some(entry) = stream.try_next().await? {
+						let idx_key = tx.unpack::<keys::ns::ActiveEnvoyByNameKey>(entry.key())?;
+
+						results.push(idx_key.envoy_key);
+
+						if results.len() >= input.limit {
+							break;
+						}
+					}
+				} else {
+					let envoy_subspace = keys::subspace()
+						.subspace(&keys::ns::ActiveEnvoyKey::subspace(input.namespace_id));
+					let (start, end) = envoy_subspace.range();
+
+					let end = if let Some(created_before) = input.created_before {
+						universaldb::utils::end_of_key_range(&tx.pack(
+							&keys::ns::ActiveEnvoyKey::subspace_with_create_ts(
+								input.namespace_id,
+								created_before,
+							),
+						))
+					} else {
+						end
+					};
+
+					let mut stream = tx.get_ranges_keyvalues(
+						universaldb::RangeOption {
+							mode: StreamingMode::Iterator,
+							reverse: true,
+							..(start, end).into()
+						},
+						// NOTE: Does not have to be serializable because we are listing, stale data does not matter
+						Snapshot,
+					);
+
+					while let Some(entry) = stream.try_next().await? {
+						let idx_key = tx.unpack::<keys::ns::ActiveEnvoyKey>(entry.key())?;
+
+						results.push(idx_key.envoy_key);
+
+						if results.len() >= input.limit {
+							break;
+						}
+					}
+				}
+
+				futures_util::stream::iter(results)
+					.map(|envoy_key| {
+						let tx = tx.clone();
+						let dc_name = dc_name.clone();
+
+						async move {
+							super::get::get_inner(&dc_name, &tx, input.namespace_id, &envoy_key)
+								.await
+						}
+					})
+					.buffered(512)
+					.try_filter_map(|result| async move { Ok(result) })
+					.try_collect::<Vec<_>>()
+					.await
+			}
+		})
+		.custom_instrument(tracing::info_span!("envoy_list_tx"))
+		.await?;
+
+	Ok(Output { envoys })
+}

package/packages/pegboard/src/ops/envoy/mod.rs

@@ -0,0 +1,6 @@
+pub mod drain;
+pub mod evict_actors;
+pub mod expire;
+pub mod get;
+pub mod list;
+pub mod update_ping;

package/packages/pegboard/src/ops/envoy/update_ping.rs

@@ -0,0 +1,92 @@
+use gas::prelude::*;
+use universaldb::options::ConflictRangeType;
+use universaldb::utils::IsolationLevel::*;
+
+use crate::keys;
+
+#[derive(Debug)]
+pub struct Input {
+	pub namespace_id: Id,
+	pub envoy_key: String,
+	pub update_lb: bool,
+	pub rtt: u32,
+}
+
+#[operation]
+pub async fn pegboard_envoy_update_ping(ctx: &OperationCtx, input: &Input) -> Result<()> {
+	ctx
+		.udb()?
+		.run(|tx| {
+			async move {
+				let tx = tx.with_subspace(keys::subspace());
+
+				let pool_name_key = keys::envoy::PoolNameKey::new(input.namespace_id, input.envoy_key.clone());
+				let version_key = keys::envoy::VersionKey::new(input.namespace_id, input.envoy_key.clone());
+				let last_ping_ts_key = keys::envoy::LastPingTsKey::new(input.namespace_id, input.envoy_key.clone());
+				let expired_ts_key = keys::envoy::ExpiredTsKey::new(input.namespace_id, input.envoy_key.clone());
+
+				let (
+					pool_name_entry,
+					version_entry,
+					last_ping_ts_entry,
+					expired,
+				) = tokio::try_join!(
+					tx.read_opt(&pool_name_key, Serializable),
+					tx.read_opt(&version_key, Serializable),
+					tx.read_opt(&last_ping_ts_key, Serializable),
+					tx.exists(&expired_ts_key, Serializable),
+				)?;
+
+				let (
+					Some(pool_name),
+					Some(version),
+					Some(old_last_ping_ts),
+				) = (
+					pool_name_entry,
+					version_entry,
+					last_ping_ts_entry,
+				)
+				else {
+					tracing::debug!(namespace_id=?input.namespace_id, envoy_key=%input.envoy_key, "envoy has not initiated yet");
+					return Ok(());
+				};
+
+				let last_ping_ts = util::timestamp::now();
+
+				// Write new ping
+				tx.write(&last_ping_ts_key, last_ping_ts)?;
+
+				let last_rtt_key = keys::envoy::LastRttKey::new(input.namespace_id, input.envoy_key.clone());
+				tx.write(&last_rtt_key, input.rtt)?;
+
+				if input.update_lb && !expired {
+					let old_lb_key = keys::ns::EnvoyLoadBalancerIdxKey::new(
+						input.namespace_id,
+						pool_name.clone(),
+						version,
+						old_last_ping_ts,
+						input.envoy_key.clone(),
+					);
+
+					// Clear old key
+					tx.add_conflict_key(&old_lb_key, ConflictRangeType::Read)?;
+					tx.delete(&old_lb_key);
+
+					tx.write(
+						&keys::ns::EnvoyLoadBalancerIdxKey::new(
+							input.namespace_id,
+							pool_name.clone(),
+							version,
+							last_ping_ts,
+							input.envoy_key.clone(),
+						),
+						(),
+					)?;
+				}
+
+				Ok(())
+			}
+		})
+		.custom_instrument(tracing::info_span!("envoy_update_ping_tx"))
+		.await
+}

package/packages/pegboard/src/ops/mod.rs

@@ -1,2 +1,5 @@
 pub mod actor;
+pub mod envoy;
 pub mod runner;
+pub mod runner_config;
+pub mod serverless_metadata;