feat(epoxy): add ability to get & set coordinator state (#3294)

* chore(guard): add metadata to rate limit error

* chore(tracing-utils): add tracing utils for custom instrument

* feat(epoxy): add ability to get & set coordinator state

Author: NathanFlurry

engine/packages/api-peer/src/internal.rs

@@ -1,4 +1,4 @@
-use anyhow::Result;
+use anyhow::*;
 use gas::prelude::*;
 use rivet_api_builder::ApiCtx;
 use serde::{Deserialize, Serialize};
@@ -91,6 +91,12 @@ pub struct ReplicaReconfigureRequest {}
 #[serde(deny_unknown_fields)]
 pub struct ReplicaReconfigureResponse {}
 
+/// Triggers the epoxy coordinator to reconfigure all replicas.
+///
+/// Useful when a replica's configuration is outdated for any reason and needs to be re-notified of
+/// changes.
+///
+/// This should never need to be called manually if everything is operating correctly.
 pub async fn epoxy_replica_reconfigure(
 	ctx: ApiCtx,
 	_path: (),
@@ -107,3 +113,74 @@ pub async fn epoxy_replica_reconfigure(
 
 	Ok(ReplicaReconfigureResponse {})
 }
+
+#[derive(Serialize)]
+#[serde(deny_unknown_fields)]
+pub struct GetEpoxyStateResponse {
+	pub config: epoxy::types::ClusterConfig,
+}
+
+/// Returns the current epoxy coordinator cluster configuration.
+///
+/// Useful for inspecting the current state of the epoxy cluster, including all replicas and their statuses.
+pub async fn get_epoxy_state(ctx: ApiCtx, _path: (), _query: ()) -> Result<GetEpoxyStateResponse> {
+	let workflow_id = ctx
+		.find_workflow::<epoxy::workflows::coordinator::Workflow>((
+			"replica",
+			ctx.config().epoxy_replica_id(),
+		))
+		.await?
+		.ok_or_else(|| anyhow!("epoxy coordinator workflow not found"))?;
+
+	let wfs = ctx.get_workflows(vec![workflow_id]).await?;
+	let wf = wfs.first().ok_or_else(|| anyhow!("workflow not found"))?;
+
+	let state: epoxy::workflows::coordinator::State =
+		wf.parse_state().context("failed to parse workflow state")?;
+
+	Ok(GetEpoxyStateResponse {
+		config: state.config,
+	})
+}
+
+#[derive(Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct SetEpoxyStateRequest {
+	pub config: epoxy::types::ClusterConfig,
+}
+
+#[derive(Serialize)]
+#[serde(deny_unknown_fields)]
+pub struct SetEpoxyStateResponse {}
+
+/// Overrides the epoxy coordinator cluster configuration and triggers reconfiguration.
+///
+/// Useful for manually adjusting the cluster state in case the replica status drifts from the
+/// state in the coordinator. This will automatically trigger a reconfigure when called.
+///
+/// This should never need to be called manually if everything is operating correctly.
+pub async fn set_epoxy_state(
+	ctx: ApiCtx,
+	_path: (),
+	_query: (),
+	body: SetEpoxyStateRequest,
+) -> Result<SetEpoxyStateResponse> {
+	ensure!(
+		body.config.coordinator_replica_id == ctx.config().epoxy_replica_id(),
+		"config coordinator_replica_id ({}) does not match current replica id ({})",
+		body.config.coordinator_replica_id,
+		ctx.config().epoxy_replica_id()
+	);
+
+	if ctx.config().is_leader() {
+		ctx.signal(epoxy::workflows::coordinator::OverrideState {
+			config: body.config,
+		})
+		.to_workflow::<epoxy::workflows::coordinator::Workflow>()
+		.tag("replica", ctx.config().epoxy_replica_id())
+		.send()
+		.await?;
+	}
+
+	Ok(SetEpoxyStateResponse {})
+}

engine/packages/api-peer/src/router.rs

@@ -36,9 +36,11 @@ pub async fn router(
 				post(internal::bump_serverless_autoscaler),
 			)
 			.route(
-				"/epoxy/replica-reconfigure",
+				"/epoxy/coordinator/replica-reconfigure",
 				post(internal::epoxy_replica_reconfigure),
 			)
+			.route("/epoxy/coordinator/state", get(internal::get_epoxy_state))
+			.route("/epoxy/coordinator/state", post(internal::set_epoxy_state))
 			.route("/debug/tracing/config", put(internal::set_tracing_config))
 	})
 	.await

engine/packages/epoxy/src/workflows/coordinator/mod.rs

@@ -40,6 +40,12 @@ pub async fn epoxy_coordinator(ctx: &mut WorkflowCtx, _input: &Input) -> Result<
 				Main::ReplicaReconfigure(_) => {
 					replica_status_change::replica_reconfigure(ctx).await?;
 				}
+				Main::OverrideState(sig) => {
+					ctx.activity(OverrideStateActivityInput { config: sig.config })
+						.await?;
+
+					reconfigure::reconfigure(ctx).await?;
+				}
 			}
 
 			Ok(Loop::<()>::Continue)
@@ -67,6 +73,21 @@ pub async fn check_config_changes(ctx: &ActivityCtx, _input: &InitInput) -> Resu
 	Ok(())
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, Hash)]
+pub struct OverrideStateActivityInput {
+	pub config: types::ClusterConfig,
+}
+
+#[activity(OverrideStateActivity)]
+pub async fn override_state_activity(
+	ctx: &ActivityCtx,
+	input: &OverrideStateActivityInput,
+) -> Result<()> {
+	let mut state = ctx.state::<State>()?;
+	state.config = input.config.clone();
+	Ok(())
+}
+
 #[message("epoxy_coordinator_config_update")]
 pub struct ConfigChangeMessage {
 	pub config: types::ClusterConfig,
@@ -87,8 +108,14 @@ pub struct ReplicaStatusChange {
 #[signal("epoxy_coordinator_replica_reconfigure")]
 pub struct ReplicaReconfigure {}
 
+#[signal("epoxy_coordinator_override_state")]
+pub struct OverrideState {
+	pub config: types::ClusterConfig,
+}
+
 join_signal!(Main {
 	Reconfigure,
 	ReplicaStatusChange,
 	ReplicaReconfigure,
+	OverrideState,
 });