use std::{
sync::{Arc, RwLock},
time::Duration,
};
use anyhow::{Context, Result, ensure};
use async_trait::async_trait;
use tonic::{
Request,
metadata::MetadataValue,
transport::{Channel, Endpoint},
};
use crate::{
actor::{ActorKey, ActorSocketEffect, ActorSocketPublisher},
grpc::proto::actor_control_plane_service_client::ActorControlPlaneServiceClient,
host::HostId,
host_leases::{HostLease, HostLeaseRegistry, HostLeaseRequest},
storage_urls::StateWriteTicket,
};
use super::{
CONTROL_PLANE_REQUEST_TIMEOUT, MAX_CONTROL_PLANE_MESSAGE_BYTES,
protocol::{ControlPlaneCommand, ControlPlaneCommandReply, decode_reply, encode_command},
};
const CONTROL_PLANE_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
#[derive(Clone)]
pub struct ControlPlaneClient {
client: ActorControlPlaneServiceClient<Channel>,
authorization: Arc<RwLock<MetadataValue<tonic::metadata::Ascii>>>,
socket_gateway: String,
http: reqwest::Client,
}
impl ControlPlaneClient {
pub async fn connect(endpoint: impl Into<String>, token: impl AsRef<str>) -> Result<Self> {
let endpoint = endpoint.into();
let channel = Endpoint::new(endpoint.clone())
.context("parse actor control-plane endpoint")?
.connect_timeout(CONTROL_PLANE_CONNECT_TIMEOUT)
.timeout(CONTROL_PLANE_REQUEST_TIMEOUT)
.connect()
.await
.context("connect to actor control plane")?;
Ok(Self {
socket_gateway: endpoint,
http: reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.timeout(CONTROL_PLANE_REQUEST_TIMEOUT)
.build()?,
client: ActorControlPlaneServiceClient::new(channel)
.max_decoding_message_size(MAX_CONTROL_PLANE_MESSAGE_BYTES)
.max_encoding_message_size(MAX_CONTROL_PLANE_MESSAGE_BYTES),
authorization: Arc::new(RwLock::new(bearer_authorization(token.as_ref())?)),
})
}
pub(crate) fn with_socket_gateway(mut self, endpoint: &str) -> Self {
self.socket_gateway = endpoint.to_owned();
self
}
pub async fn prepare_state_write(
&self,
actor: &ActorKey,
host_id: &HostId,
owner_epoch: u64,
expected_version: u64,
) -> Result<StateWriteTicket> {
match self
.execute(ControlPlaneCommand::PrepareStateWrite {
actor: actor.clone(),
host_id: host_id.clone(),
owner_epoch,
expected_version,
})
.await?
{
ControlPlaneCommandReply::StateWriteTicket { ticket } => Ok(ticket),
reply => anyhow::bail!("unexpected prepare-state-write reply: {reply:?}"),
}
}
#[allow(clippy::too_many_arguments)]
pub async fn commit_state(
&self,
actor: &ActorKey,
host_id: &HostId,
owner_epoch: u64,
expected_version: u64,
state_object: &str,
request_id: &str,
) -> Result<(u64, Option<StateWriteTicket>)> {
match self
.execute(ControlPlaneCommand::CommitState {
actor: actor.clone(),
host_id: host_id.clone(),
owner_epoch,
expected_version,
state_object: state_object.to_owned(),
request_id: request_id.to_owned(),
})
.await?
{
ControlPlaneCommandReply::StateCommitted {
state_version,
next_write,
} => Ok((state_version, next_write)),
reply => anyhow::bail!("unexpected commit-state reply: {reply:?}"),
}
}
}
#[async_trait]
impl ActorSocketPublisher for ControlPlaneClient {
async fn publish(&self, actor: &ActorKey, effects: Vec<ActorSocketEffect>) -> Result<()> {
actor.validate()?;
let authorization = self
.authorization
.read()
.map_err(|_| anyhow::anyhow!("actor authorization lock poisoned"))?
.to_str()?
.to_owned();
let url = format!(
"{}/v1/namespaces/{}/actors/{}/{}/socket-effects",
self.socket_gateway.trim_end_matches('/'),
actor.namespace_id,
actor.actor_type,
actor.actor_id
);
let response = self
.http
.post(url)
.header("authorization", authorization)
.json(&serde_json::json!({ "effects": effects }))
.send()
.await
.context("publish actor socket output")?;
ensure!(
response.status().is_success(),
"socket gateway rejected actor output with HTTP {}",
response.status()
);
Ok(())
}
}
#[async_trait]
impl HostLeaseRegistry for ControlPlaneClient {
async fn register(&self, request: &HostLeaseRequest) -> Result<HostLease> {
match self
.execute(ControlPlaneCommand::RegisterLease {
request: request.clone(),
})
.await?
{
ControlPlaneCommandReply::Lease {
lease,
replacement_token,
} => {
if let Some(token) = replacement_token {
self.replace_token(&token)?;
}
Ok(lease)
}
reply => anyhow::bail!("unexpected register-lease reply: {reply:?}"),
}
}
async fn unregister(&self, id: &HostId, _session_id: &str) -> Result<()> {
match self
.execute(ControlPlaneCommand::UnregisterLease {
host_id: id.clone(),
})
.await?
{
ControlPlaneCommandReply::Unit => Ok(()),
reply => anyhow::bail!("unexpected unregister-lease reply: {reply:?}"),
}
}
}
impl ControlPlaneClient {
async fn execute(&self, command: ControlPlaneCommand) -> Result<ControlPlaneCommandReply> {
let mut request = Request::new(encode_command(command)?);
request.set_timeout(CONTROL_PLANE_REQUEST_TIMEOUT);
request.metadata_mut().insert(
"authorization",
self.authorization
.read()
.map_err(|_| anyhow::anyhow!("actor authorization lock poisoned"))?
.clone(),
);
let reply = self
.client
.clone()
.execute(request)
.await
.context("execute actor control-plane command")?
.into_inner();
decode_reply(reply).context("decode actor control-plane reply")
}
fn replace_token(&self, token: &str) -> Result<()> {
*self
.authorization
.write()
.map_err(|_| anyhow::anyhow!("actor authorization lock poisoned"))? =
bearer_authorization(token)?;
Ok(())
}
}
fn bearer_authorization(token: &str) -> Result<MetadataValue<tonic::metadata::Ascii>> {
ensure!(
!token.is_empty() && token.trim() == token,
"actor token is invalid"
);
format!("Bearer {token}")
.parse()
.context("actor token is not valid gRPC metadata")
}