fix(epoxy): fix sending msgs to self

Author: MasterPtato

Cargo.lock

@@ -49,6 +49,26 @@ impl ApiCtx {
 		})
 	}
 
+	pub fn new_from_activity(ctx: &ActivityCtx) -> Result<Self> {
+		let req_id = Id::new_v1(ctx.config().dc_label());
+
+		Ok(Self {
+			ray_id: ctx.ray_id(),
+			req_id,
+			standalone_ctx: StandaloneCtx::new_from_activity(ctx, req_id)?,
+		})
+	}
+
+	pub fn new_from_operation(ctx: &OperationCtx) -> Result<Self> {
+		let req_id = Id::new_v1(ctx.config().dc_label());
+
+		Ok(Self {
+			ray_id: ctx.ray_id(),
+			req_id,
+			standalone_ctx: StandaloneCtx::new_from_operation(ctx, req_id)?,
+		})
+	}
+
 	pub fn ray_id(&self) -> Id {
 		self.ray_id
 	}

out/openapi.json

@@ -144,6 +144,10 @@ impl ActivityCtx {
 }
 
 impl ActivityCtx {
+	pub(crate) fn db(&self) -> &DatabaseHandle {
+		&self.db
+	}
+
 	pub fn name(&self) -> &str {
 		self.name
 	}

packages/common/api-builder/src/context.rs

@@ -152,6 +152,10 @@ impl OperationCtx {
 }
 
 impl OperationCtx {
+	pub(crate) fn db(&self) -> &DatabaseHandle {
+		&self.db
+	}
+
 	pub fn name(&self) -> &str {
 		self.name
 	}

packages/common/gasoline/core/src/ctx/activity.rs

@@ -7,7 +7,7 @@ use tracing::Instrument;
 
 use crate::{
 	builder::{WorkflowRepr, common as builder},
-	ctx::{MessageCtx, common, message::SubscriptionHandle},
+	ctx::{ActivityCtx, MessageCtx, OperationCtx, common, message::SubscriptionHandle},
 	db::{DatabaseHandle, WorkflowData},
 	error::WorkflowResult,
 	message::Message,
@@ -21,7 +21,7 @@ use crate::{
 pub struct StandaloneCtx {
 	ray_id: Id,
 	req_id: Id,
-	name: &'static str,
+	name: String,
 	ts: i64,
 
 	db: DatabaseHandle,
@@ -39,7 +39,7 @@ impl StandaloneCtx {
 		config: rivet_config::Config,
 		pools: rivet_pools::Pools,
 		cache: rivet_cache::Cache,
-		name: &'static str,
+		name: &str,
 		ray_id: Id,
 		req_id: Id,
 	) -> WorkflowResult<Self> {
@@ -54,7 +54,7 @@ impl StandaloneCtx {
 		Ok(StandaloneCtx {
 			ray_id,
 			req_id,
-			name,
+			name: name.to_string(),
 			ts,
 			db,
 			config,
@@ -63,6 +63,32 @@ impl StandaloneCtx {
 			msg_ctx,
 		})
 	}
+
+	#[tracing::instrument(skip_all)]
+	pub fn new_from_activity(ctx: &ActivityCtx, req_id: Id) -> WorkflowResult<Self> {
+		StandaloneCtx::new(
+			ctx.db().clone(),
+			ctx.config().clone(),
+			ctx.pools().clone(),
+			ctx.cache().clone(),
+			ctx.name(),
+			ctx.ray_id(),
+			req_id,
+		)
+	}
+
+	#[tracing::instrument(skip_all)]
+	pub fn new_from_operation(ctx: &OperationCtx, req_id: Id) -> WorkflowResult<Self> {
+		StandaloneCtx::new(
+			ctx.db().clone(),
+			ctx.config().clone(),
+			ctx.pools().clone(),
+			ctx.cache().clone(),
+			ctx.name(),
+			ctx.ray_id(),
+			req_id,
+		)
+	}
 }
 
 impl StandaloneCtx {
@@ -154,13 +180,9 @@ impl StandaloneCtx {
 
 impl StandaloneCtx {
 	pub fn name(&self) -> &str {
-		self.name
+		&self.name
 	}
 
-	// pub fn timeout(&self) -> Duration {
-	// 	self.timeout
-	// }
-
 	pub fn ray_id(&self) -> Id {
 		self.ray_id
 	}

packages/common/gasoline/core/src/ctx/operation.rs

@@ -1,3 +1,13 @@
+use std::result::Result::Ok as ResultOk;
+use std::{
+	collections::HashMap,
+	sync::{
+		Arc,
+		atomic::{AtomicU64, Ordering},
+	},
+	time::Duration,
+};
+
 use anyhow::*;
 use async_trait::async_trait;
 use bytes::Bytes;
@@ -108,9 +118,9 @@ impl PegboardGateway {
 		let actor_id = req
 			.headers()
 			.get("x-rivet-actor")
-			.ok_or_else(|| anyhow!("missing x-rivet-actor header"))?
+			.context("missing x-rivet-actor header")?
 			.to_str()
-			.map_err(|_| anyhow!("invalid x-rivet-actor header"))?
+			.context("invalid x-rivet-actor header")?
 			.to_string();
 
 		// Extract request parts
@@ -132,7 +142,7 @@ impl PegboardGateway {
 			.into_body()
 			.collect()
 			.await
-			.map_err(|e| anyhow!("failed to read body: {}", e))?
+			.context("failed to read body")?
 			.to_bytes();
 
 		// Build subject to publish to
@@ -212,11 +222,9 @@ impl PegboardGateway {
 		// Extract actor ID for the message
 		let actor_id = match headers
 			.get("x-rivet-actor")
-			.ok_or_else(|| anyhow!("missing x-rivet-actor header"))
-			.and_then(|v| {
-				v.to_str()
-					.map_err(|_| anyhow!("invalid x-rivet-actor header"))
-			}) {
+			.context("missing x-rivet-actor header")
+			.and_then(|v| v.to_str().context("invalid x-rivet-actor header"))
+		{
 			Result::Ok(v) => v.to_string(),
 			Err(err) => return Err((client_ws, err)),
 		};

packages/common/gasoline/core/src/ctx/standalone.rs

@@ -5,6 +5,7 @@ use epoxy_protocol::{
 	versioned,
 };
 use futures_util::{StreamExt, stream::FuturesUnordered};
+use rivet_api_builder::ApiCtx;
 use std::future::Future;
 use versioned_data_util::OwnedVersionedData;
 
@@ -56,7 +57,7 @@ where
 	.await;
 	tracing::info!(?quorum_size, len = ?responses.len(), ?quorum_type, "fanout quorum size");
 
-	// Choow how many successful responses we need before considering a success
+	// Choose how many successful responses we need before considering a success
 	let target_responses = match quorum_type {
 		// Only require 1 response
 		utils::QuorumType::Any => 1,
@@ -93,19 +94,32 @@ where
 }
 
 pub async fn send_message(
+	ctx: &ApiCtx,
 	config: &protocol::ClusterConfig,
-	to_replica_id: ReplicaId,
 	request: protocol::Request,
 ) -> Result<protocol::Response> {
-	let replica_url = find_replica_address(config, to_replica_id)?;
-	send_message_to_address(replica_url, to_replica_id, request).await
+	let replica_url = find_replica_address(config, request.to_replica_id)?;
+	send_message_to_address(ctx, replica_url, request).await
 }
 
 pub async fn send_message_to_address(
+	ctx: &ApiCtx,
 	replica_url: String,
-	to_replica_id: ReplicaId,
 	request: protocol::Request,
 ) -> Result<protocol::Response> {
+	let from_replica_id = request.from_replica_id;
+	let to_replica_id = request.to_replica_id;
+
+	if from_replica_id == to_replica_id {
+		tracing::info!(
+			to_replica = to_replica_id,
+			"sending message to replica directly"
+		);
+
+		return crate::replica::message_request::message_request(&ctx, from_replica_id, request)
+			.await;
+	}
+
 	let mut replica_url = url::Url::parse(&replica_url)?;
 	replica_url.set_path(&format!("/v{PROTOCOL_VERSION}/epoxy/message"));
 

packages/core/pegboard-gateway/src/lib.rs

@@ -1,6 +1,7 @@
 use anyhow::*;
 use epoxy_protocol::protocol::{self, ReplicaId};
 use gas::prelude::*;
+use rivet_api_builder::ApiCtx;
 
 use crate::{http_client, replica, types, utils};
 
@@ -48,8 +49,15 @@ pub async fn epoxy_explicit_prepare(
 	let quorum_members = utils::get_quorum_members(&config);
 
 	// EPaxos Step 26: Send Prepare to all replicas and wait for quorum
-	let prepare_responses =
-		send_prepares(&config, replica_id, &quorum_members, &new_ballot, instance).await?;
+	let prepare_responses = send_prepares(
+		ctx,
+		&config,
+		replica_id,
+		&quorum_members,
+		&new_ballot,
+		instance,
+	)
+	.await?;
 
 	// Check if we got enough responses for a quorum
 	let required_quorum = utils::calculate_quorum(quorum_members.len(), utils::QuorumType::Slow);
@@ -241,6 +249,7 @@ fn compare_ballots(a: &protocol::Ballot, b: &protocol::Ballot) -> std::cmp::Orde
 }
 
 async fn send_prepares(
+	ctx: &OperationCtx,
 	config: &protocol::ClusterConfig,
 	from_replica_id: ReplicaId,
 	replica_ids: &[ReplicaId],
@@ -257,8 +266,8 @@ async fn send_prepares(
 			let instance = instance.clone();
 			async move {
 				let response = http_client::send_message(
+					&ApiCtx::new_from_operation(&ctx)?,
 					&config,
-					to_replica_id,
 					protocol::Request {
 						from_replica_id,
 						to_replica_id,

packages/services/epoxy/src/http_client.rs

@@ -1,6 +1,7 @@
 use anyhow::*;
 use epoxy_protocol::protocol::{self, ReplicaId};
 use gas::prelude::*;
+use rivet_api_builder::ApiCtx;
 use universaldb::utils::{FormalKey, IsolationLevel::*};
 
 use crate::{http_client, keys, utils};
@@ -35,7 +36,7 @@ pub struct Output {
 ///
 /// We cannot use quorum reads for the fanout read because of the constraints of Epaxos.
 #[operation]
-pub async fn epoxy_kv_get_optimistic(ctx: &OperationCtx, input: &Input) -> Result<Output> {
+pub async fn epoxy_get_optimistic(ctx: &OperationCtx, input: &Input) -> Result<Output> {
 	// Try to read locally
 	let kv_key = keys::keys::KvValueKey::new(input.key.clone());
 	let cache_key = keys::keys::KvOptimisticCacheKey::new(input.key.clone());
@@ -113,7 +114,9 @@ pub async fn epoxy_kv_get_optimistic(ctx: &OperationCtx, input: &Input) -> Resul
 				};
 
 				// Send the message and extract the KV response
-				let response = http_client::send_message(&config, replica_id, request).await?;
+				let response =
+					http_client::send_message(&ApiCtx::new_from_operation(&ctx)?, &config, request)
+						.await?;
 
 				match response.kind {
 					protocol::ResponseKind::KvGetResponse(kv_response) => Ok(kv_response.value),