use std::{
sync::Arc,
time::{Duration, Instant},
};
use anyhow::{Context, Result, ensure};
use async_trait::async_trait;
use moka::future::Cache;
use tonic::transport::{Channel, Endpoint};
use tonic::{Request, Response, Status};
use tracing::{info, warn};
use crate::{
actor::{ActorKey, ActorSocketEffect, ActorSocketEvent, ActorSocketInvocation},
grpc::proto::{
ControlPlaneReply, ControlPlaneRequest, HostSocketEventRequest,
actor_control_plane_service_server::{
ActorControlPlaneService, ActorControlPlaneServiceServer,
},
actor_host_service_client::ActorHostServiceClient,
},
host::{ActorProcessRole, HostId},
host_leases::{HostLease, HostLeaseStore},
placement::{
ObjectPlacement, ObjectPlacementStore, PlacementClaim, StateCommit, StateCommitRequest,
},
sandbox::{
EnsureHostRequest, HostSandboxRuntimeConfig, HostTermination, ImageWarmup,
ProviderCommandFailure, PublicHostRouteRequest, SandboxProvider, TerminateHostsRequest,
WarmImageRequest,
},
storage_urls::{StateWriteTicket, StorageUrlSigner, validate_snapshot_object_name},
};
use super::{
MAX_CONTROL_PLANE_MESSAGE_BYTES,
admin::{AdminRegistry, AdminService, HostLaunchSpec},
auth::{ActorJwtVerifier, ActorPrincipal},
issuer::ActorJwtIssuer,
protocol::{ControlPlaneCommand, ControlPlaneCommandReply, decode_command, encode_reply},
};
const HOST_TOKEN_RENEWAL_WINDOW_SECONDS: i64 = 10 * 60;
#[derive(Clone)]
pub struct ControlPlaneService {
leases: Arc<dyn HostLeaseStore>,
placements: Arc<dyn ObjectPlacementStore>,
storage_urls: Arc<dyn StorageUrlSigner>,
auth: ActorJwtVerifier,
host_token_issuer: ActorJwtIssuer,
registry: Arc<dyn AdminRegistry>,
provisioner: Option<Arc<dyn HostProvisioner>>,
public_routes: Cache<(HostId, String), String>,
socket_targets: Cache<(ActorKey, HostId, String, i64), Arc<WorkflowActorTarget>>,
host_channels: Cache<String, Channel>,
socket_events: Option<Arc<dyn super::event_sink::SocketMessageEventSink>>,
socket_authenticator: Option<Arc<dyn super::socket_auth::SocketAuthenticator>>,
}
impl ControlPlaneService {
pub(crate) fn new(
leases: Arc<dyn HostLeaseStore>,
placements: Arc<dyn ObjectPlacementStore>,
storage_urls: Arc<dyn StorageUrlSigner>,
auth: ActorJwtVerifier,
registry: Arc<dyn AdminRegistry>,
issuer: ActorJwtIssuer,
provisioner: Option<Arc<dyn HostProvisioner>>,
) -> Self {
Self {
leases,
placements,
storage_urls,
auth,
host_token_issuer: issuer,
registry,
provisioner,
public_routes: Cache::builder()
.max_capacity(1024)
.time_to_live(Duration::from_secs(60))
.build(),
socket_targets: Cache::builder()
.max_capacity(4096)
.time_to_live(Duration::from_secs(55))
.build(),
host_channels: Cache::builder()
.max_capacity(256)
.time_to_idle(Duration::from_secs(300))
.build(),
socket_events: None,
socket_authenticator: None,
}
}
pub(crate) fn with_socket_event_sink(
mut self,
sink: Option<Arc<dyn super::event_sink::SocketMessageEventSink>>,
) -> Self {
self.socket_events = sink;
self
}
pub(crate) fn with_socket_authenticator(
mut self,
authenticator: Option<Arc<dyn super::socket_auth::SocketAuthenticator>>,
) -> Self {
self.socket_authenticator = authenticator;
self
}
pub fn into_internal_service(self) -> ActorControlPlaneServiceServer<Self> {
ActorControlPlaneServiceServer::new(self)
.max_decoding_message_size(MAX_CONTROL_PLANE_MESSAGE_BYTES)
.max_encoding_message_size(MAX_CONTROL_PLANE_MESSAGE_BYTES)
}
pub(super) fn authenticate_workflow(&self, authorization: &str) -> Result<ActorPrincipal> {
self.auth.authenticate_authorization(authorization)
}
pub(super) async fn authorize_socket(
&self,
request: super::socket_auth::SocketAuthorizationRequest,
) -> std::result::Result<
super::socket_auth::SocketAuthorization,
super::socket_auth::SocketAuthorizationError,
> {
self.socket_authenticator
.as_ref()
.ok_or_else(|| {
super::socket_auth::SocketAuthorizationError::Unavailable(anyhow::anyhow!(
"socket authorization is not configured"
))
})?
.authorize(request)
.await
}
pub(super) async fn register_deployment(
&self,
admin: &AdminService,
spec: &HostLaunchSpec,
) -> Result<bool> {
let previous = admin.current_deployment(&spec.namespace_id).await?;
let changed = admin.ensure_namespace_and_register_deployment(spec).await?;
if changed && let Some(previous) = previous {
self.terminate_deployment_hosts(&previous).await;
}
Ok(changed)
}
pub(super) fn warm_deployment_image(&self, spec: HostLaunchSpec, region: String) {
let Some(provisioner) = self.provisioner.clone() else {
return;
};
if !self.storage_urls.regions().contains(®ion) {
warn!(
event = "actor_image_warmup",
namespace_id = %spec.namespace_id,
code_revision = %spec.code_revision,
region,
outcome = "invalid_region",
"actor image warmup skipped"
);
return;
}
tokio::spawn(async move {
let started_at = Instant::now();
match provisioner.warm_image(&spec, ®ion).await {
Ok(warmup) => info!(
event = "actor_image_warmup",
namespace_id = %spec.namespace_id,
code_revision = %spec.code_revision,
region,
provider = %warmup.provider,
provider_resource_id = %warmup.resource_id,
provider_total_ms = warmup.total_ms,
total_ms = elapsed_ms(started_at),
outcome = "warmed",
"actor image warmup completed"
),
Err(error) => warn!(
event = "actor_image_warmup",
namespace_id = %spec.namespace_id,
code_revision = %spec.code_revision,
region,
total_ms = elapsed_ms(started_at),
outcome = "failed",
error = %format!("{error:#}"),
"actor image warmup failed"
),
}
});
}
async fn terminate_deployment_hosts(&self, spec: &HostLaunchSpec) {
let Some(provisioner) = self.provisioner.as_ref() else {
return;
};
let started_at = Instant::now();
match provisioner
.terminate_hosts(spec, &self.storage_urls.regions())
.await
{
Ok(termination) => info!(
event = "actor_hosts_terminated",
namespace_id = %spec.namespace_id,
code_revision = %spec.code_revision,
provider = %termination.provider,
resource_count = termination.resource_ids.len(),
total_ms = elapsed_ms(started_at),
outcome = "terminated",
"replaced deployment hosts terminated"
),
Err(error) => warn!(
event = "actor_hosts_terminated",
namespace_id = %spec.namespace_id,
code_revision = %spec.code_revision,
total_ms = elapsed_ms(started_at),
outcome = "failed",
error = %format!("{error:#}"),
"replaced deployment hosts could not be terminated"
),
}
}
#[cfg(test)]
pub(super) async fn resolve_workflow_target(
&self,
principal: &ActorPrincipal,
actor: &ActorKey,
) -> Result<WorkflowActorTarget> {
self.resolve_workflow_route(principal, actor, WorkflowRoute::Invocation, None)
.await
}
pub(super) async fn resolve_workflow_target_timed(
&self,
principal: &ActorPrincipal,
actor: &ActorKey,
timings: &mut TargetResolutionTimings,
) -> Result<WorkflowActorTarget> {
self.resolve_workflow_route(principal, actor, WorkflowRoute::Invocation, Some(timings))
.await
}
pub(super) async fn dispatch_socket_event(
&self,
principal: &ActorPrincipal,
invocation: ActorSocketInvocation,
) -> Result<Vec<ActorSocketEffect>> {
let key = (
invocation.actor.clone(),
principal.host_id.clone(),
principal.session_id.clone(),
principal.expires_at,
);
for attempt in 0..2 {
let target = self.socket_target(principal, &invocation.actor).await?;
match self.send_socket_event(&target, &invocation).await {
Ok(crate::actor::ActorExecutionResult::Completed { effects, .. }) => {
return Ok(effects);
}
Ok(
crate::actor::ActorExecutionResult::Reroute
| crate::actor::ActorExecutionResult::HostUnavailable,
) => {
self.socket_targets.invalidate(&key).await;
if attempt == 1 {
anyhow::bail!("actor host remained unavailable after socket reroute");
}
}
Ok(crate::actor::ActorExecutionResult::Failed { failure }) => {
anyhow::bail!("{}: {}", failure.code, failure.message)
}
Err(error) => {
self.socket_targets.invalidate(&key).await;
self.host_channels.invalidate(&target.route).await;
return Err(error);
}
}
}
unreachable!("socket dispatch returns after its final attempt")
}
async fn send_socket_event(
&self,
target: &WorkflowActorTarget,
invocation: &ActorSocketInvocation,
) -> Result<crate::actor::ActorExecutionResult> {
let channel = self
.host_channels
.try_get_with(target.route.clone(), async {
Endpoint::new(target.route.clone())?
.connect_timeout(Duration::from_secs(5))
.timeout(super::CONTROL_PLANE_REQUEST_TIMEOUT)
.connect()
.await
.context("connect to actor host")
})
.await
.map_err(|error| anyhow::anyhow!("{error:#}"))?;
let mut client = ActorHostServiceClient::new(channel)
.max_decoding_message_size(crate::actor::MAX_ACTOR_EXECUTOR_MESSAGE_BYTES)
.max_encoding_message_size(crate::actor::MAX_ACTOR_EXECUTOR_MESSAGE_BYTES);
let mut request = Request::new(HostSocketEventRequest {
request_id: invocation.request_id.clone(),
actor: Some(invocation.actor.clone().into()),
event_json: serde_json::to_vec(&invocation.event)?,
connections_json: serde_json::to_vec(&invocation.connections)?,
owner_epoch: target.owner_epoch,
state_read_url: target.state_read_url.clone(),
state_version: target.state_version,
});
request
.metadata_mut()
.insert("authorization", format!("Bearer {}", target.token).parse()?);
let reply = client.handle_socket(request).await?.into_inner();
crate::actor::ActorExecutionResult::try_from(reply)
}
async fn socket_target(
&self,
principal: &ActorPrincipal,
actor: &ActorKey,
) -> Result<Arc<WorkflowActorTarget>> {
actor.validate()?;
ensure!(
principal.process_role == ActorProcessRole::Workflow && principal.scope.contains(actor),
"workflow token cannot resolve this actor"
);
let now = unix_seconds()?;
ensure!(principal.expires_at > now, "workflow token has expired");
let key = (
actor.clone(),
principal.host_id.clone(),
principal.session_id.clone(),
principal.expires_at,
);
if let Some(target) = self.socket_targets.get(&key).await {
if target.expires_at_ms > (now + 5) * 1000 {
return Ok(target);
}
self.socket_targets.invalidate(&key).await;
}
self.socket_targets
.try_get_with(key, async {
self.resolve_workflow_route(principal, actor, WorkflowRoute::ControlPlane, None)
.await
.map(Arc::new)
})
.await
.map_err(|error| anyhow::anyhow!("{error:#}"))
}
pub(super) fn deliver_socket_message_event(
&self,
actor: &ActorKey,
trigger_id: Option<String>,
event: &ActorSocketEvent,
) {
let ActorSocketEvent::Message {
connection_id,
message,
} = event
else {
return;
};
if let Some(sink) = self.socket_events.clone() {
let event = super::event_sink::SocketMessageEvent::new(
actor,
trigger_id,
connection_id,
message,
);
tokio::spawn(async move {
if let Err(error) = sink.deliver(event).await {
warn!(error = %format!("{error:#}"), "actor socket message trigger delivery failed");
}
});
}
info!(
event = "actor_socket_message_committed",
namespace_id = %actor.namespace_id,
actor_type = %actor.actor_type,
actor_id = %actor.actor_id,
connection_id,
message_kind = match message {
crate::actor::ActorSocketMessage::Text { .. } => "text",
crate::actor::ActorSocketMessage::Binary { .. } => "binary",
},
"actor socket message committed"
);
}
async fn resolve_workflow_route(
&self,
principal: &ActorPrincipal,
actor: &ActorKey,
route_kind: WorkflowRoute,
mut timings: Option<&mut TargetResolutionTimings>,
) -> Result<WorkflowActorTarget> {
actor.validate()?;
ensure!(
principal.process_role == ActorProcessRole::Workflow && principal.scope.contains(actor),
"workflow token cannot resolve this actor"
);
let target = self
.route_actor(actor, &principal.region, timings.as_deref_mut())
.await?;
let state_read_url = match &target.placement.state_object {
Some(object_name) => {
self.storage_urls
.read_url(&target.placement.home_region, object_name)
.await?
}
None => String::new(),
};
if let Some(timings) = timings.as_deref_mut() {
timings.state_url_signed_at_ms = Some(timings.elapsed_ms());
}
let issued = self.host_token_issuer.issue_invocation_target(
actor,
&target.lease.id,
&target.lease.session_id,
&target.spec.code_revision,
&target.placement.home_region,
target.placement.owner_epoch,
target.placement.state_version,
&state_read_url,
principal.expires_at,
)?;
if let Some(timings) = timings.as_deref_mut() {
timings.invocation_token_issued_at_ms = Some(timings.elapsed_ms());
}
let route = match route_kind {
WorkflowRoute::Invocation
if principal.private_routing
&& principal.region == target.placement.home_region =>
{
target.lease.route.clone()
}
WorkflowRoute::Invocation | WorkflowRoute::ControlPlane => {
self.public_route(&target).await?
}
};
if let Some(timings) = timings {
timings.route_selected_at_ms = Some(timings.elapsed_ms());
}
Ok(WorkflowActorTarget {
route,
token: issued.token,
owner_epoch: target.placement.owner_epoch,
state_version: target.placement.state_version,
state_read_url,
expires_at_ms: issued.expires_at_ms,
})
}
fn provisioner(&self) -> Result<&Arc<dyn HostProvisioner>> {
self.provisioner
.as_ref()
.context("sandbox provider is not configured")
}
async fn public_route(&self, target: &RoutedActor) -> Result<String> {
if target.lease.route.starts_with("https://") {
return Ok(target.lease.route.clone());
}
let key = (target.lease.id.clone(), target.lease.session_id.clone());
self.public_routes
.try_get_with(key, async {
self.provisioner()?
.public_host_route(&target.spec, &target.placement.home_region)
.await
})
.await
.map_err(|error| anyhow::anyhow!("{error:#}"))
}
}
#[derive(Clone, Copy)]
enum WorkflowRoute {
Invocation,
ControlPlane,
}
pub(super) struct WorkflowActorTarget {
pub route: String,
pub token: String,
pub owner_epoch: u64,
pub state_version: u64,
pub state_read_url: String,
pub expires_at_ms: i64,
}
pub(super) struct TargetResolutionTimings {
started_at: Instant,
pub request_validated_at_ms: Option<f64>,
pub workflow_authenticated_at_ms: Option<f64>,
pub deployment_loaded_at_ms: Option<f64>,
pub placement_loaded_at_ms: Option<f64>,
pub lease_checked_at_ms: Option<f64>,
pub host_ensured_at_ms: Option<f64>,
pub placement_claimed_at_ms: Option<f64>,
pub state_url_signed_at_ms: Option<f64>,
pub invocation_token_issued_at_ms: Option<f64>,
pub route_selected_at_ms: Option<f64>,
}
impl TargetResolutionTimings {
pub fn new() -> Self {
Self {
started_at: Instant::now(),
request_validated_at_ms: None,
workflow_authenticated_at_ms: None,
deployment_loaded_at_ms: None,
placement_loaded_at_ms: None,
lease_checked_at_ms: None,
host_ensured_at_ms: None,
placement_claimed_at_ms: None,
state_url_signed_at_ms: None,
invocation_token_issued_at_ms: None,
route_selected_at_ms: None,
}
}
pub fn elapsed_ms(&self) -> f64 {
elapsed_ms(self.started_at)
}
}
#[tonic::async_trait]
impl ActorControlPlaneService for ControlPlaneService {
async fn execute(
&self,
request: Request<ControlPlaneRequest>,
) -> std::result::Result<Response<ControlPlaneReply>, Status> {
let principal = self.auth.authenticate(&request).await?;
let command = decode_command(request.into_inner())
.map_err(|error| Status::invalid_argument(format!("invalid command: {error:#}")))?;
let reply = self
.execute_command(&principal, command)
.await
.map_err(failed_precondition)?;
Ok(Response::new(encode_reply(reply).map_err(internal)?))
}
}
impl ControlPlaneService {
async fn execute_command(
&self,
principal: &ActorPrincipal,
command: ControlPlaneCommand,
) -> Result<ControlPlaneCommandReply> {
ensure!(
principal.process_role == ActorProcessRole::Host,
"only hosts may use the internal control-plane API"
);
match command {
ControlPlaneCommand::RegisterLease { request } => {
self.register_lease(principal, request).await
}
ControlPlaneCommand::UnregisterLease { host_id } => {
self.unregister_lease(principal, host_id).await
}
ControlPlaneCommand::PrepareStateWrite {
actor,
host_id,
owner_epoch,
expected_version,
} => {
self.prepare_state_write(principal, actor, host_id, owner_epoch, expected_version)
.await
}
ControlPlaneCommand::CommitState {
actor,
host_id,
owner_epoch,
expected_version,
state_object,
request_id,
} => {
self.commit_state(
principal,
actor,
host_id,
owner_epoch,
expected_version,
state_object,
request_id,
)
.await
}
}
}
async fn register_lease(
&self,
principal: &ActorPrincipal,
request: crate::host_leases::HostLeaseRequest,
) -> Result<ControlPlaneCommandReply> {
principal.validate_host_id(request.id.as_str())?;
ensure!(
request.session_id == principal.session_id,
"lease session does not match host token"
);
validate_host_route(&request.route)?;
let lease = self.leases.register(&request).await?;
Ok(ControlPlaneCommandReply::Lease {
lease,
replacement_token: self.replacement_host_token(principal)?,
})
}
async fn unregister_lease(
&self,
principal: &ActorPrincipal,
host_id: HostId,
) -> Result<ControlPlaneCommandReply> {
principal.validate_host_id(host_id.as_str())?;
self.leases
.unregister(&host_id, &principal.session_id)
.await?;
Ok(ControlPlaneCommandReply::Unit)
}
async fn prepare_state_write(
&self,
principal: &ActorPrincipal,
actor: ActorKey,
host_id: HostId,
owner_epoch: u64,
expected_version: u64,
) -> Result<ControlPlaneCommandReply> {
actor.validate()?;
ensure!(
principal.scope.contains(&actor),
"actor crossed the host namespace"
);
principal.validate_host_id(host_id.as_str())?;
self.require_active_host(principal).await?;
let placement = self.current_placement(&actor).await?;
validate_state_owner(principal, &host_id, owner_epoch, &placement)?;
ensure!(
placement.state_version == expected_version,
"actor state version changed before write preparation"
);
Ok(ControlPlaneCommandReply::StateWriteTicket {
ticket: self
.state_write_ticket(&placement.home_region, &actor, expected_version)
.await?,
})
}
#[allow(clippy::too_many_arguments)]
async fn commit_state(
&self,
principal: &ActorPrincipal,
actor: ActorKey,
host_id: HostId,
owner_epoch: u64,
expected_version: u64,
state_object: String,
request_id: String,
) -> Result<ControlPlaneCommandReply> {
actor.validate()?;
ensure!(
principal.scope.contains(&actor),
"actor crossed the host namespace"
);
principal.validate_host_id(host_id.as_str())?;
self.require_active_host(principal).await?;
validate_snapshot_object_name(
&actor,
expected_version
.checked_add(1)
.context("actor state version overflow")?,
&state_object,
)?;
let committed = self
.placements
.commit_state(&StateCommitRequest {
object: actor.storage_key(),
owner: host_id,
session_id: principal.session_id.clone(),
owner_epoch,
expected_version,
state_object,
request_id,
})
.await?;
let placement = match committed {
StateCommit::Committed(placement) => placement,
StateCommit::Current(_) => {
anyhow::bail!("actor ownership or state version changed before commit")
}
};
let next_write = self
.state_write_ticket(&placement.home_region, &actor, placement.state_version)
.await
.map_err(|error| {
warn!(
event = "actor_state_write_ticket",
actor = %actor.storage_key(),
state_version = placement.state_version,
error = %format!("{error:#}"),
"next actor state write ticket could not be prepared"
);
error
})
.ok();
Ok(ControlPlaneCommandReply::StateCommitted {
state_version: placement.state_version,
next_write,
})
}
async fn state_write_ticket(
&self,
region: &str,
actor: &ActorKey,
expected_version: u64,
) -> Result<StateWriteTicket> {
self.storage_urls
.write_ticket(
region,
actor,
expected_version
.checked_add(1)
.context("actor state version overflow")?,
)
.await
}
async fn route_actor(
&self,
actor: &ActorKey,
storage_region: &str,
mut timings: Option<&mut TargetResolutionTimings>,
) -> Result<RoutedActor> {
let spec = self
.registry
.launch_spec(&actor.namespace_id)
.await?
.context("project has no registered actor code")?;
if let Some(timings) = timings.as_deref_mut() {
timings.deployment_loaded_at_ms = Some(timings.elapsed_ms());
}
loop {
let current = self.placements.get(&actor.storage_key()).await?;
if let Some(timings) = timings.as_deref_mut() {
timings.placement_loaded_at_ms = Some(timings.elapsed_ms());
}
let lease_checked = current.as_ref().is_some_and(|placement| {
host_matches_revision(&placement.owner, &spec.namespace_id, &spec.code_revision)
});
let active = self.active_target(¤t, &spec).await?;
if lease_checked && let Some(timings) = timings.as_deref_mut() {
timings.lease_checked_at_ms = Some(timings.elapsed_ms());
}
if let Some(target) = active {
return Ok(target);
}
let region = self.target_region(current.as_ref(), storage_region)?;
let provisioner = self.provisioner()?;
let lease = provisioner.ensure_host(&spec, ®ion).await?;
if let Some(timings) = timings.as_deref_mut() {
timings.host_ensured_at_ms = Some(timings.elapsed_ms());
}
match self
.placements
.claim(&actor.storage_key(), current.as_ref(), &lease.id, ®ion)
.await?
{
PlacementClaim::Acquired(placement) | PlacementClaim::Current(placement)
if placement.owner == lease.id =>
{
if let Some(timings) = timings.as_deref_mut() {
timings.placement_claimed_at_ms = Some(timings.elapsed_ms());
}
return Ok(RoutedActor {
placement,
lease,
spec,
});
}
PlacementClaim::Acquired(_) | PlacementClaim::Current(_) => continue,
}
}
}
fn replacement_host_token(&self, principal: &ActorPrincipal) -> Result<Option<String>> {
if principal.expires_at.saturating_sub(unix_seconds()?) > HOST_TOKEN_RENEWAL_WINDOW_SECONDS
{
return Ok(None);
}
Ok(Some(
self.host_token_issuer
.issue_host(
&principal.scope.namespace_id,
&principal.host_id,
&principal.session_id,
principal
.code_revision
.as_deref()
.context("host token has no code revision")?,
&principal.region,
)?
.token,
))
}
async fn require_active_host(&self, principal: &ActorPrincipal) -> Result<HostLease> {
let status = self.leases.lease_status(&principal.host_id).await?;
ensure!(status.is_active(), "host lease is not active");
let lease = status.lease.context("active host lease is missing")?;
ensure!(
lease.session_id == principal.session_id,
"host lease belongs to another session"
);
Ok(lease)
}
async fn current_placement(&self, actor: &ActorKey) -> Result<ObjectPlacement> {
self.placements
.get(&actor.storage_key())
.await?
.context("actor has no current placement")
}
async fn active_target(
&self,
current: &Option<ObjectPlacement>,
spec: &HostLaunchSpec,
) -> Result<Option<RoutedActor>> {
let Some(placement) = current else {
return Ok(None);
};
if !host_matches_revision(&placement.owner, &spec.namespace_id, &spec.code_revision) {
return Ok(None);
}
let status = self.leases.lease_status(&placement.owner).await?;
if !status.is_active() {
return Ok(None);
}
Ok(Some(RoutedActor {
placement: placement.clone(),
lease: status.lease.context("active lease is missing")?,
spec: spec.clone(),
}))
}
fn target_region(
&self,
current: Option<&ObjectPlacement>,
storage_region: &str,
) -> Result<String> {
select_target_region(current, storage_region, &self.storage_urls.regions())
}
}
fn select_target_region(
current: Option<&ObjectPlacement>,
storage_region: &str,
configured_regions: &[String],
) -> Result<String> {
if let Some(placement) = current {
return Ok(placement.home_region.clone());
}
ensure!(
configured_regions
.iter()
.any(|region| region == storage_region),
"workflow storage region has no Standard bucket"
);
Ok(storage_region.to_owned())
}
#[async_trait]
pub(crate) trait HostProvisioner: Send + Sync {
async fn ensure_host(&self, spec: &HostLaunchSpec, region: &str) -> Result<HostLease>;
async fn public_host_route(&self, spec: &HostLaunchSpec, region: &str) -> Result<String>;
async fn warm_image(&self, spec: &HostLaunchSpec, region: &str) -> Result<ImageWarmup>;
async fn terminate_hosts(
&self,
spec: &HostLaunchSpec,
regions: &[String],
) -> Result<HostTermination>;
}
pub(crate) struct SandboxHostProvisioner {
provider: Arc<dyn SandboxProvider>,
runtime: HostSandboxRuntimeConfig,
issuer: ActorJwtIssuer,
leases: Arc<dyn HostLeaseStore>,
}
impl SandboxHostProvisioner {
pub(crate) fn new(
provider: Arc<dyn SandboxProvider>,
runtime: HostSandboxRuntimeConfig,
issuer: ActorJwtIssuer,
leases: Arc<dyn HostLeaseStore>,
) -> Self {
Self {
provider,
runtime,
issuer,
leases,
}
}
}
#[async_trait]
impl HostProvisioner for SandboxHostProvisioner {
async fn ensure_host(&self, spec: &HostLaunchSpec, region: &str) -> Result<HostLease> {
let started_at = Instant::now();
let request = self.request(spec, region)?;
let handle = match self.provider.ensure_host(&request).await {
Ok(handle) => handle,
Err(error) => {
let command = error.downcast_ref::<ProviderCommandFailure>();
warn!(
event = "actor_host_provisioning",
namespace_id = %spec.namespace_id,
code_revision = %spec.code_revision,
region,
host_id = %request.host_id,
started_at_ms = 0,
provider_process_spawned_at_ms = command.and_then(ProviderCommandFailure::spawned_at_ms),
provider_request_written_at_ms = command.and_then(ProviderCommandFailure::request_written_at_ms),
provider_process_completed_at_ms = command.and_then(ProviderCommandFailure::process_completed_at_ms),
provider_response_decoded_at_ms = command.and_then(ProviderCommandFailure::response_decoded_at_ms),
completed_at_ms = elapsed_ms(started_at),
outcome = "provider_failed",
error = %format!("{error:#}"),
"actor host provisioning failed"
);
return Err(error);
}
};
let provisioning = handle.provisioning.clone();
let lease = self.active_lease(handle, region).await?;
let lease_validated_at_ms = elapsed_ms(started_at);
info!(
event = "actor_host_provisioning",
namespace_id = %spec.namespace_id,
code_revision = %spec.code_revision,
region,
host_id = %lease.id,
provider = provisioning.as_ref().map(|value| value.provider.as_str()).unwrap_or("unknown"),
provider_resource_id = provisioning.as_ref().map(|value| value.resource_id.as_str()).unwrap_or(""),
provider_reused = provisioning.as_ref().is_some_and(|value| value.reused),
started_at_ms = 0,
provider_process_spawned_at_ms = provisioning.as_ref().and_then(|value| value.command_spawned_at_ms),
provider_request_written_at_ms = provisioning.as_ref().and_then(|value| value.request_written_at_ms),
provider_process_completed_at_ms = provisioning.as_ref().and_then(|value| value.process_completed_at_ms),
provider_response_decoded_at_ms = provisioning.as_ref().and_then(|value| value.response_decoded_at_ms),
modal_provider_started_at_ms = provisioning.as_ref().map(|value| value.started_at_ms),
modal_input_parsed_at_ms = provisioning.as_ref().and_then(|value| value.input_parsed_at_ms),
modal_sdk_loaded_at_ms = provisioning.as_ref().and_then(|value| value.sdk_loaded_at_ms),
modal_resources_resolved_at_ms = provisioning.as_ref().and_then(|value| value.resources_resolved_at_ms),
modal_existing_host_checked_at_ms = provisioning.as_ref().and_then(|value| value.existing_host_checked_at_ms),
modal_sandbox_scheduled_at_ms = provisioning.as_ref().and_then(|value| value.sandbox_scheduled_at_ms),
modal_host_ready_observed_at_ms = provisioning.as_ref().and_then(|value| value.host_ready_observed_at_ms),
modal_route_read_at_ms = provisioning.as_ref().and_then(|value| value.route_read_at_ms),
modal_metadata_written_at_ms = provisioning.as_ref().and_then(|value| value.metadata_written_at_ms),
modal_provider_completed_at_ms = provisioning.as_ref().map(|value| value.completed_at_ms),
lease_validated_at_ms,
completed_at_ms = elapsed_ms(started_at),
outcome = "ready",
"actor host provisioning completed"
);
Ok(lease)
}
async fn public_host_route(&self, spec: &HostLaunchSpec, region: &str) -> Result<String> {
let response = self
.provider
.public_host_route(&PublicHostRouteRequest {
namespace_id: spec.namespace_id.clone(),
code_revision: spec.code_revision.clone(),
canonical_region: region.to_owned(),
})
.await?;
validate_host_route(&response.route)?;
Ok(response.route)
}
async fn warm_image(&self, spec: &HostLaunchSpec, region: &str) -> Result<ImageWarmup> {
self.provider
.warm_image(&WarmImageRequest {
namespace_id: spec.namespace_id.clone(),
code_revision: spec.code_revision.clone(),
canonical_region: region.to_owned(),
image_ref: spec.image_ref.clone(),
})
.await
}
async fn terminate_hosts(
&self,
spec: &HostLaunchSpec,
regions: &[String],
) -> Result<HostTermination> {
self.provider
.terminate_hosts(&TerminateHostsRequest {
namespace_id: spec.namespace_id.clone(),
code_revision: spec.code_revision.clone(),
canonical_regions: regions.to_vec(),
})
.await
}
}
impl SandboxHostProvisioner {
fn request(&self, spec: &HostLaunchSpec, region: &str) -> Result<EnsureHostRequest> {
let host_id = HostId::new(format!(
"host.v1.{}.{}.{}",
spec.namespace_id,
spec.code_revision,
uuid::Uuid::new_v4()
));
let session_id = uuid::Uuid::new_v4().to_string();
let host_token = self
.issuer
.issue_host(
&spec.namespace_id,
&host_id,
&session_id,
&spec.code_revision,
region,
)?
.token;
Ok(EnsureHostRequest {
namespace_id: spec.namespace_id.clone(),
code_revision: spec.code_revision.clone(),
canonical_region: region.to_owned(),
host_id,
session_id,
host_token,
jwt_public_keys: self.issuer.verifier_keys_json()?,
control_plane_url: self.runtime.control_plane_url.clone(),
jwt_issuer: self.runtime.jwt_issuer.clone(),
invocation_jwt_audience: self.runtime.invocation_jwt_audience.clone(),
image_ref: spec.image_ref.clone(),
working_directory: spec.working_directory.clone(),
actor_entrypoint: spec.actor_entrypoint.clone(),
actor_idle_timeout_ms: self.runtime.actor_idle_timeout_ms,
host_idle_timeout_ms: self.runtime.host_idle_timeout_ms,
})
}
async fn active_lease(
&self,
handle: crate::sandbox::ActorHostHandle,
region: &str,
) -> Result<HostLease> {
ensure!(
handle.canonical_region == region,
"sandbox provider returned the wrong region"
);
validate_host_route(&handle.route)?;
let status = self.leases.lease_status(&handle.host_id).await?;
ensure!(status.is_active(), "sandbox host lease is not active");
let lease = status.lease.context("active sandbox lease is missing")?;
ensure!(
lease.route == handle.route,
"sandbox route does not match its lease"
);
Ok(lease)
}
}
struct RoutedActor {
placement: ObjectPlacement,
lease: HostLease,
spec: HostLaunchSpec,
}
fn elapsed_ms(started_at: Instant) -> f64 {
started_at.elapsed().as_secs_f64() * 1_000.0
}
fn validate_state_owner(
principal: &ActorPrincipal,
host_id: &HostId,
owner_epoch: u64,
placement: &ObjectPlacement,
) -> Result<()> {
ensure!(placement.owner == *host_id, "host does not own this actor");
ensure!(
placement.owner_epoch == owner_epoch,
"actor owner epoch is stale"
);
ensure!(
placement.home_region == principal.region,
"host is outside the actor home region"
);
Ok(())
}
fn host_matches_revision(host: &HostId, namespace: &str, revision: &str) -> bool {
host.as_str()
.starts_with(&format!("host.v1.{namespace}.{revision}."))
}
fn validate_host_route(route: &str) -> Result<()> {
let route = reqwest::Url::parse(route)?;
ensure!(
matches!(route.scheme(), "http" | "https") && route.host_str().is_some(),
"host route must be an HTTP origin"
);
ensure!(
route.username().is_empty() && route.password().is_none(),
"host route must not contain credentials"
);
ensure!(
route.path() == "/" && route.query().is_none() && route.fragment().is_none(),
"host route must not contain a path, query, or fragment"
);
Ok(())
}
fn unix_seconds() -> Result<i64> {
Ok(i64::try_from(
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)?
.as_secs(),
)?)
}
fn failed_precondition(error: impl std::fmt::Display) -> Status {
Status::failed_precondition(error.to_string())
}
fn internal(error: impl std::fmt::Display) -> Status {
Status::internal(error.to_string())
}
#[cfg(test)]
mod tests {
use std::{collections::HashMap, sync::Mutex, time::Duration};
use super::super::{ActorTokenPurpose, admin::LocalAdminRegistry};
use super::*;
use crate::{
actor::ActorScope,
actor_state::ActorStorageKey,
host_leases::{HostLeaseRegistry, HostLeaseRequest, HostLeaseStatus},
placement::testing::LocalObjectPlacementStore,
};
use aws_lc_rs::{rand::SystemRandom, signature::Ed25519KeyPair};
use base64::{Engine, engine::general_purpose::STANDARD};
struct FakeLeaseStore {
leases: Mutex<HashMap<HostId, HostLease>>,
}
struct FakeSocketEventSink {
delivered: tokio::sync::mpsc::UnboundedSender<serde_json::Value>,
}
#[async_trait]
impl super::super::event_sink::SocketMessageEventSink for FakeSocketEventSink {
async fn deliver(&self, event: super::super::event_sink::SocketMessageEvent) -> Result<()> {
self.delivered.send(serde_json::to_value(event)?)?;
Ok(())
}
}
#[async_trait]
impl HostLeaseRegistry for FakeLeaseStore {
async fn register(&self, request: &HostLeaseRequest) -> Result<HostLease> {
let lease = HostLease {
id: request.id.clone(),
session_id: request.session_id.clone(),
route: request.route.clone(),
expires_at_ms: 10_000,
};
self.leases
.lock()
.unwrap()
.insert(lease.id.clone(), lease.clone());
Ok(lease)
}
async fn unregister(&self, id: &HostId, _session_id: &str) -> Result<()> {
self.leases.lock().unwrap().remove(id);
Ok(())
}
}
#[async_trait]
impl HostLeaseStore for FakeLeaseStore {
async fn lease_status(&self, id: &HostId) -> Result<HostLeaseStatus> {
Ok(HostLeaseStatus {
lease: self.leases.lock().unwrap().get(id).cloned(),
store_now_ms: 0,
})
}
}
struct FakeStorageUrls;
#[async_trait]
impl StorageUrlSigner for FakeStorageUrls {
async fn read_url(&self, _region: &str, _object_name: &str) -> Result<String> {
Ok("https://storage.example.com/state".into())
}
async fn write_ticket(
&self,
_region: &str,
_actor: &ActorKey,
state_version: u64,
) -> Result<StateWriteTicket> {
Ok(StateWriteTicket {
state_version,
object_name: format!(
"snapshots/00/00000000000000000000000000000000/project-1/Counter/counter-1/{state_version}.json"
),
url: "https://storage.example.com/write".into(),
expires_at_ms: i64::MAX,
})
}
fn regions(&self) -> Vec<String> {
vec!["us-east".into()]
}
}
struct FakeWarmProvisioner {
warmed: tokio::sync::mpsc::UnboundedSender<(HostLaunchSpec, String)>,
}
#[async_trait]
impl HostProvisioner for FakeWarmProvisioner {
async fn ensure_host(&self, _spec: &HostLaunchSpec, _region: &str) -> Result<HostLease> {
anyhow::bail!("host creation is outside this test")
}
async fn public_host_route(&self, _spec: &HostLaunchSpec, _region: &str) -> Result<String> {
anyhow::bail!("public host routing is outside this test")
}
async fn warm_image(&self, spec: &HostLaunchSpec, region: &str) -> Result<ImageWarmup> {
self.warmed.send((spec.clone(), region.to_owned()))?;
Ok(ImageWarmup {
provider: "test".into(),
resource_id: "sandbox-1".into(),
total_ms: 1,
})
}
async fn terminate_hosts(
&self,
_spec: &HostLaunchSpec,
_regions: &[String],
) -> Result<HostTermination> {
anyhow::bail!("host termination is outside this test")
}
}
struct FakeRetiringProvisioner {
retired: tokio::sync::mpsc::UnboundedSender<(HostLaunchSpec, Vec<String>)>,
}
struct FakeRouteProvisioner {
routes: Mutex<Vec<(HostLaunchSpec, String)>>,
}
#[async_trait]
impl HostProvisioner for FakeRouteProvisioner {
async fn ensure_host(&self, _spec: &HostLaunchSpec, _region: &str) -> Result<HostLease> {
anyhow::bail!("host creation is outside this test")
}
async fn public_host_route(&self, spec: &HostLaunchSpec, region: &str) -> Result<String> {
self.routes
.lock()
.unwrap()
.push((spec.clone(), region.to_owned()));
Ok("https://actor.example.com/".into())
}
async fn warm_image(&self, _spec: &HostLaunchSpec, _region: &str) -> Result<ImageWarmup> {
anyhow::bail!("image warmup is outside this test")
}
async fn terminate_hosts(
&self,
_spec: &HostLaunchSpec,
_regions: &[String],
) -> Result<HostTermination> {
anyhow::bail!("host termination is outside this test")
}
}
#[async_trait]
impl HostProvisioner for FakeRetiringProvisioner {
async fn ensure_host(&self, _spec: &HostLaunchSpec, _region: &str) -> Result<HostLease> {
anyhow::bail!("host creation is outside this test")
}
async fn public_host_route(&self, _spec: &HostLaunchSpec, _region: &str) -> Result<String> {
anyhow::bail!("public host routing is outside this test")
}
async fn warm_image(&self, _spec: &HostLaunchSpec, _region: &str) -> Result<ImageWarmup> {
anyhow::bail!("image warmup is outside this test")
}
async fn terminate_hosts(
&self,
spec: &HostLaunchSpec,
regions: &[String],
) -> Result<HostTermination> {
self.retired.send((spec.clone(), regions.to_vec()))?;
Ok(HostTermination {
provider: "test".into(),
resource_ids: vec!["sandbox-1".into()],
})
}
}
#[tokio::test]
async fn replacing_a_deployment_terminates_the_previous_revision_hosts() -> Result<()> {
let issuer = test_issuer()?;
let auth = ActorJwtVerifier::for_scope(
issuer.verifier_keys_json()?,
"issuer",
"invocation",
ActorTokenPurpose::Invocation,
Duration::from_secs(60),
)?;
let registry = Arc::new(LocalAdminRegistry::default());
let admin = super::super::admin::AdminService::new(
"admin-token".into(),
registry.clone(),
issuer.clone(),
)?;
let (retired_tx, mut retired_rx) = tokio::sync::mpsc::unbounded_channel();
let service = ControlPlaneService::new(
Arc::new(FakeLeaseStore {
leases: Mutex::new(HashMap::new()),
}),
Arc::new(LocalObjectPlacementStore::default()),
Arc::new(FakeStorageUrls),
auth,
registry,
issuer,
Some(Arc::new(FakeRetiringProvisioner {
retired: retired_tx,
})),
);
let first = HostLaunchSpec {
namespace_id: "project-1".into(),
code_revision: "revision-1".into(),
image_ref: "image-1".into(),
working_directory: "/workspace".into(),
actor_entrypoint: None,
};
let mut replacement = first.clone();
replacement.code_revision = "revision-2".into();
replacement.image_ref = "image-2".into();
assert!(service.register_deployment(&admin, &first).await?);
assert!(retired_rx.try_recv().is_err());
assert!(service.register_deployment(&admin, &replacement).await?);
assert_eq!(
retired_rx.recv().await,
Some((first, vec!["us-east".into()]))
);
Ok(())
}
#[tokio::test]
async fn deployment_image_warmup_runs_in_the_background_without_creating_an_actor() -> Result<()>
{
let issuer = test_issuer()?;
let auth = ActorJwtVerifier::for_scope(
issuer.verifier_keys_json()?,
"issuer",
"invocation",
ActorTokenPurpose::Invocation,
Duration::from_secs(60),
)?;
let leases = Arc::new(FakeLeaseStore {
leases: Mutex::new(HashMap::new()),
});
let placements = Arc::new(LocalObjectPlacementStore::default());
let registry = Arc::new(LocalAdminRegistry::default());
let (warmed_tx, mut warmed_rx) = tokio::sync::mpsc::unbounded_channel();
let service = ControlPlaneService::new(
leases,
placements,
Arc::new(FakeStorageUrls),
auth,
registry,
issuer,
Some(Arc::new(FakeWarmProvisioner { warmed: warmed_tx })),
);
let spec = HostLaunchSpec {
namespace_id: "project-1".into(),
code_revision: "revision-1".into(),
image_ref: "image-1".into(),
working_directory: "/workspace".into(),
actor_entrypoint: None,
};
service.warm_deployment_image(spec.clone(), "us-east".into());
let warmed = tokio::time::timeout(Duration::from_secs(1), warmed_rx.recv())
.await?
.context("warmup task stopped")?;
assert_eq!(warmed, (spec, "us-east".into()));
Ok(())
}
#[tokio::test]
async fn accepted_socket_messages_are_delivered_to_the_configured_event_sink() -> Result<()> {
let issuer = test_issuer()?;
let auth = ActorJwtVerifier::for_scope(
issuer.verifier_keys_json()?,
"issuer",
"invocation",
ActorTokenPurpose::Invocation,
Duration::from_secs(60),
)?;
let (delivered_tx, mut delivered_rx) = tokio::sync::mpsc::unbounded_channel();
let service = ControlPlaneService::new(
Arc::new(FakeLeaseStore {
leases: Mutex::new(HashMap::new()),
}),
Arc::new(LocalObjectPlacementStore::default()),
Arc::new(FakeStorageUrls),
auth,
Arc::new(LocalAdminRegistry::default()),
issuer,
None,
)
.with_socket_event_sink(Some(Arc::new(FakeSocketEventSink {
delivered: delivered_tx,
})));
let actor = ActorKey {
namespace_id: "project-1".into(),
actor_type: "ChatRoom".into(),
actor_id: "room-1".into(),
};
service.deliver_socket_message_event(
&actor,
Some("trigger-1".into()),
&ActorSocketEvent::Message {
connection_id: "socket-1".into(),
message: crate::actor::ActorSocketMessage::Text {
data: "hello".into(),
},
},
);
let delivered = tokio::time::timeout(Duration::from_secs(1), delivered_rx.recv())
.await?
.context("socket event delivery task stopped")?;
assert_eq!(delivered["namespaceId"], "project-1");
assert_eq!(delivered["actorType"], "ChatRoom");
assert_eq!(delivered["actorId"], "room-1");
assert_eq!(delivered["triggerId"], "trigger-1");
assert_eq!(delivered["connectionId"], "socket-1");
assert_eq!(delivered["message"]["type"], "text");
assert_eq!(delivered["message"]["data"], "hello");
Ok(())
}
#[tokio::test]
async fn workflow_target_uses_private_routing_only_for_capable_same_region_callers()
-> Result<()> {
let issuer = test_issuer()?;
let auth = ActorJwtVerifier::for_scope(
issuer.verifier_keys_json()?,
"issuer",
"invocation",
ActorTokenPurpose::Invocation,
Duration::from_secs(60),
)?;
let host_id = HostId::new("host.v1.project-1.revision-1.host-1");
let leases = Arc::new(FakeLeaseStore {
leases: Mutex::new(HashMap::from([(
host_id.clone(),
HostLease {
id: host_id.clone(),
session_id: "00000000-0000-4000-8000-000000000001".into(),
route: "http://[fd00:cafe::1234]:7101/".into(),
expires_at_ms: 10_000,
},
)])),
});
let placements = Arc::new(LocalObjectPlacementStore::default());
let registry = Arc::new(LocalAdminRegistry::default());
let actor = ActorKey {
namespace_id: "project-1".into(),
actor_type: "Counter".into(),
actor_id: "counter-1".into(),
};
registry
.ensure_namespace_and_register_deployment(&HostLaunchSpec {
namespace_id: "project-1".into(),
code_revision: "revision-1".into(),
image_ref: "image-1".into(),
working_directory: "/workspace".into(),
actor_entrypoint: None,
})
.await?;
placements
.claim(&actor.storage_key(), None, &host_id, "us-east")
.await?;
let provisioner = Arc::new(FakeRouteProvisioner {
routes: Mutex::new(Vec::new()),
});
let service = ControlPlaneService::new(
leases,
placements,
Arc::new(FakeStorageUrls),
auth,
registry,
issuer.clone(),
Some(provisioner.clone()),
);
let target = service
.resolve_workflow_target(
&ActorPrincipal {
scope: ActorScope {
namespace_id: "project-1".into(),
},
host_id: HostId::new("workflow.v1.project-1.execution-1"),
session_id: uuid::Uuid::new_v4().to_string(),
process_role: ActorProcessRole::Workflow,
region: "us-east".into(),
private_routing: true,
code_revision: None,
expires_at: unix_seconds()? + 30,
invocation: None,
},
&actor,
)
.await?;
assert_eq!(target.route, "http://[fd00:cafe::1234]:7101/");
assert!(provisioner.routes.lock().unwrap().is_empty());
assert_eq!(target.owner_epoch, 1);
assert_eq!(target.state_version, 0);
assert!(target.state_read_url.is_empty());
let verifier = ActorJwtVerifier::for_scope(
issuer.verifier_keys_json()?,
"issuer",
"invocation",
ActorTokenPurpose::Invocation,
Duration::from_secs(60),
)?;
let principal = verifier.authenticate_authorization(&format!("Bearer {}", target.token))?;
assert_eq!(
principal
.invocation
.expect("direct invocation capability")
.actor,
actor
);
let local_target = service
.resolve_workflow_target(
&ActorPrincipal {
scope: ActorScope {
namespace_id: "project-1".into(),
},
host_id: HostId::new("workflow.v1.project-1.execution-local"),
session_id: uuid::Uuid::new_v4().to_string(),
process_role: ActorProcessRole::Workflow,
region: "us-east".into(),
private_routing: false,
code_revision: None,
expires_at: unix_seconds()? + 30,
invocation: None,
},
&actor,
)
.await?;
assert_eq!(local_target.route, "https://actor.example.com/");
let cross_region_target = service
.resolve_workflow_target(
&ActorPrincipal {
scope: ActorScope {
namespace_id: "project-1".into(),
},
host_id: HostId::new("workflow.v1.project-1.execution-2"),
session_id: uuid::Uuid::new_v4().to_string(),
process_role: ActorProcessRole::Workflow,
region: "us-west".into(),
private_routing: true,
code_revision: None,
expires_at: unix_seconds()? + 30,
invocation: None,
},
&actor,
)
.await?;
assert_eq!(cross_region_target.route, "https://actor.example.com/");
let principal = ActorPrincipal {
scope: ActorScope {
namespace_id: "project-1".into(),
},
host_id: HostId::new("workflow.v1.project-1.socket-test"),
session_id: uuid::Uuid::new_v4().to_string(),
process_role: ActorProcessRole::Workflow,
region: "us-east".into(),
private_routing: false,
code_revision: None,
expires_at: unix_seconds()? + 30,
invocation: None,
};
let first = service.socket_target(&principal, &actor).await?;
let second = service.socket_target(&principal, &actor).await?;
assert!(Arc::ptr_eq(&first, &second));
let cached_key = (
actor.clone(),
principal.host_id.clone(),
principal.session_id.clone(),
principal.expires_at,
);
service
.socket_targets
.insert(
cached_key,
Arc::new(WorkflowActorTarget {
route: first.route.clone(),
token: first.token.clone(),
owner_epoch: first.owner_epoch,
state_version: first.state_version,
state_read_url: first.state_read_url.clone(),
expires_at_ms: (unix_seconds()? + 4) * 1000,
}),
)
.await;
let refreshed = service.socket_target(&principal, &actor).await?;
assert!(refreshed.expires_at_ms > (unix_seconds()? + 5) * 1000);
let mut shorter = principal.clone();
shorter.expires_at -= 10;
let shorter_target = service.socket_target(&shorter, &actor).await?;
assert!(shorter_target.expires_at_ms <= shorter.expires_at * 1000);
let mut other_session = principal.clone();
other_session.session_id = uuid::Uuid::new_v4().to_string();
assert!(!Arc::ptr_eq(
&first,
&service.socket_target(&other_session, &actor).await?
));
let mut expired = principal.clone();
expired.expires_at = unix_seconds()?;
assert!(service.socket_target(&expired, &actor).await.is_err());
assert_eq!(
provisioner.routes.lock().unwrap().as_slice(),
&[(
HostLaunchSpec {
namespace_id: "project-1".into(),
code_revision: "revision-1".into(),
image_ref: "image-1".into(),
working_directory: "/workspace".into(),
actor_entrypoint: None,
},
"us-east".into()
)]
);
let mut routed = service.route_actor(&actor, "us-east", None).await?;
routed.lease.route = "https://already-public.example.com/".into();
assert_eq!(service.public_route(&routed).await?, routed.lease.route);
assert_eq!(provisioner.routes.lock().unwrap().len(), 1);
routed.lease.route = "http://[fd00:cafe::5678]:7101/".into();
routed.lease.session_id = uuid::Uuid::new_v4().to_string();
service.public_route(&routed).await?;
assert_eq!(provisioner.routes.lock().unwrap().len(), 2);
Ok(())
}
#[tokio::test]
async fn socket_dispatch_reuses_a_connection_and_does_not_replay_transport_failures()
-> Result<()> {
use crate::grpc::proto::{
self,
actor_host_service_server::{ActorHostService, ActorHostServiceServer},
};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use tokio_stream::{StreamExt, wrappers::TcpListenerStream};
struct Host {
calls: Arc<AtomicUsize>,
fail: Arc<AtomicBool>,
}
#[tonic::async_trait]
impl ActorHostService for Host {
async fn invoke(
&self,
_: Request<proto::HostInvokeActorRequest>,
) -> Result<tonic::Response<proto::InvokeActorReply>, Status> {
Err(Status::unimplemented(
"method invocation is outside this test",
))
}
async fn handle_socket(
&self,
request: Request<HostSocketEventRequest>,
) -> Result<tonic::Response<proto::InvokeActorReply>, Status> {
assert_eq!(
request.metadata().get("authorization").unwrap(),
"Bearer test-token"
);
self.calls.fetch_add(1, Ordering::SeqCst);
if self.fail.load(Ordering::SeqCst) {
return Err(Status::unavailable("reply lost after execution"));
}
Ok(tonic::Response::new(
crate::actor::ActorExecutionResult::Completed {
result: serde_json::Value::Null,
effects: vec![],
}
.into(),
))
}
}
let connections = Arc::new(AtomicUsize::new(0));
let calls = Arc::new(AtomicUsize::new(0));
let fail = Arc::new(AtomicBool::new(false));
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
let route = format!("http://{}", listener.local_addr()?);
let accepted = connections.clone();
let incoming = TcpListenerStream::new(listener).map(move |stream| {
if stream.is_ok() {
accepted.fetch_add(1, Ordering::SeqCst);
}
stream
});
let stop = tokio_util::sync::CancellationToken::new();
let shutdown = stop.clone();
let host = Host {
calls: calls.clone(),
fail: fail.clone(),
};
let server = tokio::spawn(
tonic::transport::Server::builder()
.add_service(ActorHostServiceServer::new(host))
.serve_with_incoming_shutdown(incoming, async move { shutdown.cancelled().await }),
);
let issuer = test_issuer()?;
let auth = ActorJwtVerifier::for_scope(
issuer.verifier_keys_json()?,
"issuer",
"invocation",
ActorTokenPurpose::Invocation,
Duration::from_secs(60),
)?;
let service = ControlPlaneService::new(
Arc::new(FakeLeaseStore {
leases: Mutex::new(HashMap::new()),
}),
Arc::new(LocalObjectPlacementStore::default()),
Arc::new(FakeStorageUrls),
auth,
Arc::new(LocalAdminRegistry::default()),
issuer,
None,
);
let actor = ActorKey {
namespace_id: "project-1".into(),
actor_type: "Counter".into(),
actor_id: "one".into(),
};
let principal = ActorPrincipal {
scope: ActorScope {
namespace_id: actor.namespace_id.clone(),
},
host_id: HostId::new("workflow.v1.project-1.test"),
session_id: uuid::Uuid::new_v4().to_string(),
process_role: ActorProcessRole::Workflow,
region: "us-east".into(),
private_routing: false,
code_revision: None,
expires_at: unix_seconds()? + 60,
invocation: None,
};
let key = (
actor.clone(),
principal.host_id.clone(),
principal.session_id.clone(),
principal.expires_at,
);
service
.socket_targets
.insert(
key.clone(),
Arc::new(WorkflowActorTarget {
route: route.clone(),
token: "test-token".into(),
owner_epoch: 1,
state_version: 0,
state_read_url: String::new(),
expires_at_ms: (unix_seconds()? + 30) * 1000,
}),
)
.await;
let invocation = |id: &str| ActorSocketInvocation {
request_id: id.into(),
actor: actor.clone(),
event: ActorSocketEvent::Message {
connection_id: "one".into(),
message: crate::actor::ActorSocketMessage::Text {
data: "hello".into(),
},
},
connections: vec![],
};
service
.dispatch_socket_event(&principal, invocation("first"))
.await?;
service
.dispatch_socket_event(&principal, invocation("second"))
.await?;
assert_eq!(connections.load(Ordering::SeqCst), 1);
fail.store(true, Ordering::SeqCst);
assert!(
service
.dispatch_socket_event(&principal, invocation("third"))
.await
.is_err()
);
assert_eq!(calls.load(Ordering::SeqCst), 3);
assert!(service.socket_targets.get(&key).await.is_none());
assert!(service.host_channels.get(&route).await.is_none());
stop.cancel();
server.await??;
Ok(())
}
#[test]
fn new_actors_use_the_workflow_region_and_existing_actors_stay_pinned() -> Result<()> {
let regions = vec![
"north-america-east".into(),
"north-america-central".into(),
"north-america-west".into(),
];
let actor = ActorStorageKey::new("object.v1.project.Counter.one");
let current = ObjectPlacement {
object: actor,
owner: HostId::new("host.v1.project.revision.host"),
owner_epoch: 1,
home_region: "north-america-east".into(),
state_version: 0,
state_object: None,
last_request_id: None,
};
assert_eq!(
select_target_region(None, "north-america-central", ®ions)?,
"north-america-central"
);
assert_eq!(
select_target_region(Some(¤t), "north-america-west", ®ions)?,
"north-america-east"
);
assert!(select_target_region(None, "europe-west", ®ions).is_err());
Ok(())
}
fn test_issuer() -> Result<ActorJwtIssuer> {
let pkcs8 = Ed25519KeyPair::generate_pkcs8(&SystemRandom::new())?;
ActorJwtIssuer::from_base64_pkcs8(
&STANDARD.encode(pkcs8.as_ref()),
"test-key",
"issuer",
"authority",
"invocation",
Duration::from_secs(60),
)
}
}