use super::{
service::{ActorTarget, ControlPlaneService},
socket_bus::SocketBus,
socket_inventory::SocketInventory,
socket_ticket::SocketTicket,
};
use crate::request_tracking::HostState;
use crate::{
actor::{ActorKey, ActorSocketEffect, ActorSocketInvocation},
sockets::{SocketRegistry, browser::SocketDispatcher},
};
use anyhow::{Context, Result, ensure};
use async_trait::async_trait;
use axum::http::StatusCode;
use serde::{Deserialize, Serialize};
use std::{sync::Arc, time::Duration};
use tokio::time::Instant;
use tokio_util::sync::CancellationToken;
pub(crate) use crate::sockets::operations::{SocketOperation, SocketOperationReply};
#[derive(Serialize, Deserialize)]
pub(crate) struct SocketEventRequest {
#[serde(default)]
pub host_state: HostState,
pub routing_ms: f64,
pub owner_epoch: u64,
pub invocation: ActorSocketInvocation,
}
#[derive(Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub(crate) enum SocketEventReply {
Completed { effects: Vec<ActorSocketEffect> },
NotExecuted,
Failed { message: String },
}
type ActorRoute = Arc<tokio::sync::Mutex<Option<ActorTarget>>>;
pub(super) struct SocketGateway {
pub registry: SocketRegistry,
pub stop: CancellationToken,
pub(super) bus: Arc<dyn SocketBus>,
pub(super) accepts_sockets: bool,
routes: moka::future::Cache<ActorKey, ActorRoute>,
pub(super) http: reqwest::Client,
}
impl SocketGateway {
pub(super) fn start(
bus: Arc<dyn SocketBus>,
accepts_sockets: bool,
stop: CancellationToken,
http: reqwest::Client,
) -> Result<Arc<Self>> {
Ok(Arc::new(Self {
registry: bus.registry().clone(),
bus,
accepts_sockets,
stop,
routes: moka::future::Cache::builder()
.max_capacity(100_000)
.time_to_idle(Duration::from_secs(300))
.build(),
http,
}))
}
pub(super) fn generation(&self) -> Option<u64> {
if self.stop.is_cancelled() {
return None;
}
self.bus.generation()
}
pub(super) fn ensure_delivery(&self, generation: u64) -> Result<()> {
ensure!(
self.generation() == Some(generation),
"socket gateway lost its subscription"
);
Ok(())
}
pub(super) async fn operation(
&self,
actor: &ActorKey,
operation: SocketOperation,
) -> Result<SocketOperationReply> {
ensure!(!self.stop.is_cancelled(), "socket gateway stopping");
match operation {
SocketOperation::Publish { effects } => {
crate::actor::validate_socket_effects(&effects)?;
self.bus.publish(actor, effects).await?;
Ok(SocketOperationReply::Published)
}
SocketOperation::Connections { tag } => Ok(SocketOperationReply::Connections {
connections: self.bus.connections(actor, tag.as_deref()).await?,
}),
SocketOperation::Count => Ok(SocketOperationReply::Count {
count: self.bus.count(actor).await?,
}),
}
}
pub(super) async fn inventory(&self, project: &str) -> Result<SocketInventory> {
ensure!(!self.stop.is_cancelled(), "socket gateway stopping");
self.bus.inventory(project).await
}
async fn target(
&self,
service: &ControlPlaneService,
ticket: &SocketTicket,
) -> Result<ActorTarget> {
service
.ensure_project_access(&ticket.actor.project_id)
.await?;
let route = self.route(&ticket.actor).await;
let mut cached = route.lock().await;
let now = crate::clock::Clock::now_ms(&crate::clock::SystemClock)? as i64;
if let Some(target) = cached.as_ref().filter(|target| target.expires_at_ms > now)
&& service
.lifecycle
.routing(
&ticket.actor.storage_key(),
target.host_id.as_str(),
&target.session_id,
target.owner_epoch,
)
.await
.is_none()
{
return Ok(target.clone());
}
*cached = None;
let target = service
.resolve_actor_route(&ticket.actor, ticket.home_region.as_deref(), None, None)
.await?;
let mut cached_target = target.clone();
cached_target.host_state = HostState::Warm;
*cached = Some(cached_target);
Ok(target)
}
async fn route(&self, actor: &ActorKey) -> ActorRoute {
self.routes
.get_with(actor.clone(), async { Arc::default() })
.await
}
async fn invalidate(&self, actor: &ActorKey, target: &ActorTarget) {
let route = self.route(actor).await;
let mut cached = route.lock().await;
if cached.as_ref().is_some_and(|current| {
current.owner_epoch == target.owner_epoch && current.route == target.route
}) {
*cached = None;
}
}
async fn dispatch(
&self,
service: &ControlPlaneService,
ticket: &SocketTicket,
invocation: ActorSocketInvocation,
) -> Result<Vec<ActorSocketEffect>> {
let started = Instant::now();
let retry_deadline = started + super::CONTROL_PLANE_REQUEST_TIMEOUT;
let mut host_state = HostState::Warm;
loop {
ensure!(!self.stop.is_cancelled(), "socket gateway stopping");
ensure!(
Instant::now() < retry_deadline,
"actor remained unavailable during socket handoff"
);
let target = tokio::time::timeout_at(retry_deadline, self.target(service, ticket))
.await
.context("actor remained unavailable during socket handoff")?;
let target = match target {
Ok(target) => target,
Err(error) if error.is::<crate::sandbox::HostNotReady>() => {
self.wait_to_retry().await?;
continue;
}
Err(error) => return Err(error),
};
if target.host_state == HostState::Cold {
host_state = HostState::Cold;
}
let actor = &ticket.actor;
let url = format!(
"{}/v1/projects/{}/actors/{}/{}/socket-events",
target.route.trim_end_matches('/'),
actor.project_id,
actor.actor_name,
actor.actor_id
);
let reply = self
.http
.post(url)
.bearer_auth(&target.token)
.json(&SocketEventRequest {
host_state,
routing_ms: started.elapsed().as_secs_f64() * 1000.0,
owner_epoch: target.owner_epoch,
invocation: invocation.clone(),
})
.send()
.await;
match reply {
Err(error) if error.is_connect() => {
self.invalidate(actor, &target).await;
}
Err(error) => {
return Err(error)
.context("socket event outcome unknown; event was not replayed");
}
Ok(reply)
if matches!(
reply.status(),
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN
) =>
{
self.invalidate(actor, &target).await;
}
Ok(reply) => match reply.error_for_status()?.json::<SocketEventReply>().await? {
SocketEventReply::Completed { effects } => return Ok(effects),
SocketEventReply::NotExecuted => {
self.invalidate(actor, &target).await;
}
SocketEventReply::Failed { message } => {
anyhow::bail!("socket handler failed: {message}")
}
},
}
self.wait_to_retry().await?;
}
}
async fn wait_to_retry(&self) -> Result<()> {
tokio::select! {
_ = self.stop.cancelled() => anyhow::bail!("socket gateway stopped"),
_ = tokio::time::sleep(Duration::from_millis(100)) => Ok(()),
}
}
}
pub(super) struct GatewaySocketDispatcher {
pub gateway: Arc<SocketGateway>,
pub service: ControlPlaneService,
pub generation: u64,
}
#[async_trait]
impl SocketDispatcher for GatewaySocketDispatcher {
fn ensure_authority(&self) -> Result<()> {
self.gateway.ensure_delivery(self.generation)
}
async fn dispatch(
&self,
ticket: &SocketTicket,
invocation: ActorSocketInvocation,
) -> Result<Vec<ActorSocketEffect>> {
let result = self
.gateway
.dispatch(&self.service, ticket, invocation)
.await;
if let Err(error) = &result {
tracing::warn!(actor = %ticket.actor.storage_key().as_str(), error = %format!("{error:#}"), "socket event delivery failed");
}
result
}
async fn deliver(
&self,
_: &SocketRegistry,
actor: &ActorKey,
effects: Vec<ActorSocketEffect>,
) -> Result<()> {
self.gateway.bus.publish(actor, effects).await
}
fn notify(&self, ticket: &SocketTicket, event: &crate::actor::ActorSocketEvent) {
self.service
.deliver_socket_message_event(&ticket.actor, None, event);
}
}