use std::{
collections::{BTreeMap, BTreeSet, HashMap},
fmt,
net::SocketAddr,
sync::{Arc, Mutex},
time::{Duration, SystemTime},
};
use datum_cluster::{ClusterConfig, ClusterNode, ClusterState, Member, MemberState, Signal};
use datum_net::quic::quinn;
use tokio::{
sync::{mpsc, oneshot, watch},
task::JoinHandle,
time::{Instant, sleep},
};
use crate::{
Agent, AgentConfig, AgentHandle, ClusterJobMetadata, ClusterPlacementHistory,
JobRegistryHandle, PlacementSpec as RegistryPlacementSpec,
PlacementStrategy as RegistryPlacementStrategy,
dcp::{
ClientKind, ClusterJobList, ClusterJobNode, ClusterJobStart, ClusterNodeError,
ClusterNodeList, ClusterNodeStatus, ClusterViewProvider, CompleteShardingAsk, DcpClient,
DcpError, DcpJobFactories, DcpServer, DcpServerConfig, DcpServerHandle,
ForwardShardEnvelopes, Hello, PlacementSpec, PlacementStrategy, RememberClusterAssignment,
ResponseStatus, ShardAllocation, ShardAllocationTable, ShardEnvelopeBatchResult,
ShardPipeClient, ShardPipeFrame, SubmitClusterJob,
server::{
cluster_metadata_from_wire, placement_spec_from_wire, wire_cluster_job_start,
wire_job_status,
},
},
};
pub const AGENT_ROLE: &str = "agent";
pub type ClusterAgentResult<T> = Result<T, ClusterAgentError>;
#[derive(Debug, thiserror::Error)]
pub enum ClusterAgentError {
#[error("invalid cluster-agent config: {0}")]
InvalidConfig(String),
#[error(transparent)]
Agent(#[from] crate::AgentError),
#[error(transparent)]
Dcp(#[from] DcpError),
#[error(transparent)]
Cluster(#[from] datum_cluster::ClusterError),
}
#[derive(Clone)]
pub enum NodeSessionTransport {
TcpLoopback,
QuicMtls {
server_name: String,
client_config: quinn::ClientConfig,
},
}
impl fmt::Debug for NodeSessionTransport {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::TcpLoopback => formatter.write_str("TcpLoopback"),
Self::QuicMtls { server_name, .. } => formatter
.debug_struct("QuicMtls")
.field("server_name", server_name)
.finish_non_exhaustive(),
}
}
}
#[derive(Clone, Debug)]
pub struct NodeSessionConfig {
pub agent_role: String,
pub transport: NodeSessionTransport,
pub reconnect_min_backoff: Duration,
pub reconnect_max_backoff: Duration,
pub request_timeout: Duration,
pub command_buffer: usize,
}
impl Default for NodeSessionConfig {
fn default() -> Self {
Self {
agent_role: AGENT_ROLE.to_owned(),
transport: NodeSessionTransport::TcpLoopback,
reconnect_min_backoff: Duration::from_millis(50),
reconnect_max_backoff: Duration::from_secs(2),
request_timeout: Duration::from_millis(750),
command_buffer: 32,
}
}
}
#[derive(Clone, Default)]
pub struct ClusterAgentConfig {
pub agent: AgentConfig,
pub cluster: ClusterConfig,
pub dcp: DcpServerConfig,
pub sessions: NodeSessionConfig,
}
impl fmt::Debug for ClusterAgentConfig {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("ClusterAgentConfig")
.field("agent", &self.agent)
.field("cluster", &self.cluster)
.field("dcp", &"<DcpServerConfig>")
.field("sessions", &self.sessions)
.finish()
}
}
pub struct ClusterAgent;
impl ClusterAgent {
pub async fn start(
config: ClusterAgentConfig,
factories: DcpJobFactories,
) -> ClusterAgentResult<ClusterAgentHandle> {
let mut config = config;
validate_cluster_agent_config(&config)?;
ensure_agent_role(&mut config.cluster, &config.sessions.agent_role);
config.dcp.node_id = config.cluster.node_id.clone();
let agent = Agent::start_with_config(config.agent.clone())?;
let server = DcpServer::from_agent(&agent, factories.clone(), config.dcp.clone());
let server_handle = server.start().await?;
let agent_addr = advertised_agent_addr(&config.sessions.transport, &server_handle)?;
config.cluster.agent_addr = Some(agent_addr);
let cluster = ClusterNode::start(config.cluster.clone()).await?;
let sessions = NodeSessionManagerHandle::start(
config.sessions.clone(),
cluster.node_id().to_owned(),
cluster.state(),
)?;
let placement = PlacementCoordinatorHandle::start(
cluster.node_id().to_owned(),
config.sessions.agent_role.clone(),
cluster.state(),
agent.registry().clone(),
sessions.clone(),
factories.clone(),
config.sessions.request_timeout,
)?;
let provider = Arc::new(ClusterView {
registry: agent.registry().clone(),
state: cluster.state(),
sessions: sessions.clone(),
placement: placement.clone(),
self_node: cluster.node_id().to_owned(),
agent_role: config.sessions.agent_role.clone(),
});
server.set_cluster_view(provider);
Ok(ClusterAgentHandle {
agent,
cluster,
server,
server_handle: Some(server_handle),
sessions,
placement,
})
}
}
pub struct ClusterAgentHandle {
agent: AgentHandle,
cluster: ClusterNode,
server: DcpServer,
server_handle: Option<DcpServerHandle>,
sessions: NodeSessionManagerHandle,
placement: PlacementCoordinatorHandle,
}
impl ClusterAgentHandle {
#[must_use]
pub fn agent(&self) -> &AgentHandle {
&self.agent
}
#[must_use]
pub fn cluster(&self) -> &ClusterNode {
&self.cluster
}
#[must_use]
pub fn sessions(&self) -> &NodeSessionManagerHandle {
&self.sessions
}
#[must_use]
pub fn server(&self) -> &DcpServer {
&self.server
}
#[must_use]
pub fn tcp_addr(&self) -> Option<SocketAddr> {
self.server_handle
.as_ref()
.and_then(DcpServerHandle::tcp_addr)
}
#[must_use]
pub fn quic_addr(&self) -> Option<SocketAddr> {
self.server_handle
.as_ref()
.and_then(DcpServerHandle::quic_addr)
}
pub async fn shutdown(mut self) -> ClusterAgentResult<()> {
self.server.clear_cluster_view();
self.server.clear_sharding_view();
self.placement.shutdown().await;
self.sessions.shutdown().await;
let _ = self.cluster.abort().await;
if let Some(handle) = self.server_handle.take() {
handle.shutdown().await;
}
self.agent.registry().shutdown()?;
Ok(())
}
}
impl Drop for ClusterAgentHandle {
fn drop(&mut self) {
self.server.clear_cluster_view();
self.server.clear_sharding_view();
}
}
#[derive(Clone)]
pub struct NodeSessionManagerHandle {
inner: Arc<NodeSessionManagerInner>,
}
struct NodeSessionManagerInner {
config: NodeSessionConfig,
self_node: String,
state: Signal<ClusterState>,
sessions: tokio::sync::Mutex<BTreeMap<String, PeerSession>>,
tasks: Mutex<Vec<JoinHandle<()>>>,
}
struct PeerSession {
member: Member,
commands: mpsc::Sender<SessionCommand>,
pipe_commands: mpsc::Sender<ShardPipeCommand>,
stop: watch::Sender<bool>,
state: watch::Receiver<PeerSessionState>,
task: JoinHandle<()>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum PeerSessionState {
Connecting,
Connected,
BackingOff(String),
Closed,
}
impl PeerSessionState {
fn as_text(&self) -> String {
match self {
Self::Connecting => "connecting".to_owned(),
Self::Connected => "connected".to_owned(),
Self::BackingOff(error) => format!("backing_off:{error}"),
Self::Closed => "closed".to_owned(),
}
}
}
enum SessionCommand {
ListJobs {
timeout: Duration,
reply: oneshot::Sender<Result<Vec<crate::dcp::proto::JobStatus>, String>>,
},
SubmitClusterJob {
request: SubmitClusterJob,
timeout: Duration,
reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>,
},
StartClusterJob {
factory_name: String,
instance_name: String,
params: HashMap<String, String>,
assignment: ClusterJobStart,
timeout: Duration,
reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>,
},
ClusterJobStatus {
name: String,
cluster: bool,
timeout: Duration,
reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>,
},
DrainClusterJob {
name: String,
cluster: bool,
timeout: Duration,
reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>,
},
StopClusterJob {
name: String,
cluster: bool,
timeout: Duration,
reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>,
},
ListClusterJobs {
timeout: Duration,
reply: oneshot::Sender<Result<ClusterJobList, String>>,
},
RememberClusterAssignment {
instance_name: String,
assignment: ClusterJobStart,
timeout: Duration,
reply: oneshot::Sender<Result<(), String>>,
},
AllocateShard {
type_name: String,
shard_id: String,
timeout: Duration,
reply: oneshot::Sender<Result<ShardAllocation, String>>,
},
RememberShardAllocations {
table: ShardAllocationTable,
timeout: Duration,
reply: oneshot::Sender<Result<(), String>>,
},
GetShardAllocations {
type_name: String,
timeout: Duration,
reply: oneshot::Sender<Result<ShardAllocationTable, String>>,
},
ForwardShardEnvelopes {
batch: ForwardShardEnvelopes,
timeout: Duration,
reply: oneshot::Sender<Result<ShardEnvelopeBatchResult, String>>,
},
CompleteShardingAsk {
response: CompleteShardingAsk,
timeout: Duration,
reply: oneshot::Sender<Result<(), String>>,
},
}
enum ShardPipeCommand {
Forward { batch: ForwardShardEnvelopes },
Reply { response: CompleteShardingAsk },
}
impl NodeSessionManagerHandle {
fn start(
config: NodeSessionConfig,
self_node: String,
state: Signal<ClusterState>,
) -> ClusterAgentResult<Self> {
let inner = Arc::new(NodeSessionManagerInner {
config,
self_node,
state,
sessions: tokio::sync::Mutex::new(BTreeMap::new()),
tasks: Mutex::new(Vec::new()),
});
let manager_task = {
let inner = Arc::clone(&inner);
tokio::spawn(async move {
run_node_session_manager(inner).await;
})
};
{
let mut tasks = inner.tasks.lock().expect("node-session tasks poisoned");
tasks.push(manager_task);
}
Ok(Self { inner })
}
pub async fn shutdown(&self) {
let sessions = {
let mut locked = self.inner.sessions.lock().await;
std::mem::take(&mut *locked)
};
for (_, session) in sessions {
let _ = session.stop.send(true);
session.task.abort();
}
let tasks = {
let mut locked = self
.inner
.tasks
.lock()
.expect("node-session tasks poisoned");
std::mem::take(&mut *locked)
};
for task in tasks {
task.abort();
}
}
async fn list_jobs(
&self,
node_id: &str,
timeout: Duration,
) -> Result<Vec<crate::dcp::proto::JobStatus>, String> {
let sender = {
let locked = self.inner.sessions.lock().await;
locked.get(node_id).map(|session| session.commands.clone())
}
.ok_or_else(|| "node session unavailable".to_owned())?;
let (reply, receiver) = oneshot::channel();
sender
.try_send(SessionCommand::ListJobs { timeout, reply })
.map_err(|error| format!("node session command queue unavailable: {error}"))?;
tokio::time::timeout(timeout, receiver)
.await
.map_err(|_| "node session request timed out".to_owned())?
.map_err(|_| "node session closed".to_owned())?
}
async fn submit_cluster_job(
&self,
node_id: &str,
request: SubmitClusterJob,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, String> {
self.request_status(node_id, timeout, |reply| SessionCommand::SubmitClusterJob {
request,
timeout,
reply,
})
.await
}
async fn start_cluster_job(
&self,
node_id: &str,
factory_name: String,
instance_name: String,
params: HashMap<String, String>,
assignment: ClusterJobStart,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, String> {
self.request_status(node_id, timeout, |reply| SessionCommand::StartClusterJob {
factory_name,
instance_name,
params,
assignment,
timeout,
reply,
})
.await
}
async fn cluster_job_status(
&self,
node_id: &str,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, String> {
self.request_status(node_id, timeout, |reply| SessionCommand::ClusterJobStatus {
name,
cluster: false,
timeout,
reply,
})
.await
}
async fn drain_cluster_job(
&self,
node_id: &str,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, String> {
self.request_status(node_id, timeout, |reply| SessionCommand::DrainClusterJob {
name,
cluster: false,
timeout,
reply,
})
.await
}
async fn stop_cluster_job(
&self,
node_id: &str,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, String> {
self.request_status(node_id, timeout, |reply| SessionCommand::StopClusterJob {
name,
cluster: false,
timeout,
reply,
})
.await
}
async fn forward_cluster_job_status(
&self,
node_id: &str,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, String> {
self.request_status(node_id, timeout, |reply| SessionCommand::ClusterJobStatus {
name,
cluster: true,
timeout,
reply,
})
.await
}
async fn forward_drain_cluster_job(
&self,
node_id: &str,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, String> {
self.request_status(node_id, timeout, |reply| SessionCommand::DrainClusterJob {
name,
cluster: true,
timeout,
reply,
})
.await
}
async fn forward_stop_cluster_job(
&self,
node_id: &str,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, String> {
self.request_status(node_id, timeout, |reply| SessionCommand::StopClusterJob {
name,
cluster: true,
timeout,
reply,
})
.await
}
async fn list_cluster_jobs(
&self,
node_id: &str,
timeout: Duration,
) -> Result<ClusterJobList, String> {
let sender = self.session_sender(node_id).await?;
let (reply, receiver) = oneshot::channel();
sender
.try_send(SessionCommand::ListClusterJobs { timeout, reply })
.map_err(|error| format!("node session command queue unavailable: {error}"))?;
wait_session_reply(timeout, receiver).await
}
async fn remember_cluster_assignment(
&self,
node_id: &str,
instance_name: String,
assignment: ClusterJobStart,
timeout: Duration,
) -> Result<(), String> {
let sender = self.session_sender(node_id).await?;
let (reply, receiver) = oneshot::channel();
sender
.try_send(SessionCommand::RememberClusterAssignment {
instance_name,
assignment,
timeout,
reply,
})
.map_err(|error| format!("node session command queue unavailable: {error}"))?;
wait_session_reply(timeout, receiver).await
}
pub async fn allocate_shard(
&self,
node_id: &str,
type_name: String,
shard_id: String,
timeout: Duration,
) -> Result<ShardAllocation, String> {
let sender = self.session_sender(node_id).await?;
let (reply, receiver) = oneshot::channel();
sender
.try_send(SessionCommand::AllocateShard {
type_name,
shard_id,
timeout,
reply,
})
.map_err(|error| format!("node session command queue unavailable: {error}"))?;
wait_session_reply(timeout, receiver).await
}
pub async fn remember_shard_allocations(
&self,
node_id: &str,
table: ShardAllocationTable,
timeout: Duration,
) -> Result<(), String> {
let sender = self.session_sender(node_id).await?;
let (reply, receiver) = oneshot::channel();
sender
.try_send(SessionCommand::RememberShardAllocations {
table,
timeout,
reply,
})
.map_err(|error| format!("node session command queue unavailable: {error}"))?;
wait_session_reply(timeout, receiver).await
}
pub async fn get_shard_allocations(
&self,
node_id: &str,
type_name: String,
timeout: Duration,
) -> Result<ShardAllocationTable, String> {
let sender = self.session_sender(node_id).await?;
let (reply, receiver) = oneshot::channel();
sender
.try_send(SessionCommand::GetShardAllocations {
type_name,
timeout,
reply,
})
.map_err(|error| format!("node session command queue unavailable: {error}"))?;
wait_session_reply(timeout, receiver).await
}
pub async fn forward_shard_envelopes(
&self,
node_id: &str,
batch: ForwardShardEnvelopes,
timeout: Duration,
) -> Result<ShardEnvelopeBatchResult, String> {
let sender = self.session_sender(node_id).await?;
let (reply, receiver) = oneshot::channel();
sender
.try_send(SessionCommand::ForwardShardEnvelopes {
batch,
timeout,
reply,
})
.map_err(|error| format!("node session command queue unavailable: {error}"))?;
wait_session_reply(timeout, receiver).await
}
pub async fn complete_sharding_ask(
&self,
node_id: &str,
response: CompleteShardingAsk,
timeout: Duration,
) -> Result<(), String> {
let sender = self.session_sender(node_id).await?;
let (reply, receiver) = oneshot::channel();
sender
.try_send(SessionCommand::CompleteShardingAsk {
response,
timeout,
reply,
})
.map_err(|error| format!("node session command queue unavailable: {error}"))?;
wait_session_reply(timeout, receiver).await
}
pub async fn forward_shard_pipe_envelopes(
&self,
node_id: &str,
batch: ForwardShardEnvelopes,
) -> Result<(), String> {
let sender = self.pipe_sender(node_id).await?;
sender
.try_send(ShardPipeCommand::Forward { batch })
.map_err(|error| format!("node session shard pipe queue unavailable: {error}"))
}
pub async fn complete_sharding_ask_pipe(
&self,
node_id: &str,
response: CompleteShardingAsk,
) -> Result<(), String> {
let sender = self.pipe_sender(node_id).await?;
sender
.try_send(ShardPipeCommand::Reply { response })
.map_err(|error| format!("node session shard pipe queue unavailable: {error}"))
}
pub fn try_complete_sharding_ask_pipe(
&self,
node_id: &str,
response: CompleteShardingAsk,
) -> Result<(), String> {
let locked = self
.inner
.sessions
.try_lock()
.map_err(|_| "node session table busy".to_owned())?;
let sender = locked
.get(node_id)
.map(|session| session.pipe_commands.clone())
.ok_or_else(|| "node session unavailable".to_owned())?;
sender
.try_send(ShardPipeCommand::Reply { response })
.map_err(|error| format!("node session shard pipe queue unavailable: {error}"))
}
async fn request_status<F>(
&self,
node_id: &str,
timeout: Duration,
make_command: F,
) -> Result<crate::dcp::proto::JobStatus, String>
where
F: FnOnce(oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>) -> SessionCommand,
{
let sender = self.session_sender(node_id).await?;
let (reply, receiver) = oneshot::channel();
sender
.try_send(make_command(reply))
.map_err(|error| format!("node session command queue unavailable: {error}"))?;
wait_session_reply(timeout, receiver).await
}
async fn session_sender(&self, node_id: &str) -> Result<mpsc::Sender<SessionCommand>, String> {
let locked = self.inner.sessions.lock().await;
locked
.get(node_id)
.map(|session| session.commands.clone())
.ok_or_else(|| "node session unavailable".to_owned())
}
async fn pipe_sender(&self, node_id: &str) -> Result<mpsc::Sender<ShardPipeCommand>, String> {
let locked = self.inner.sessions.lock().await;
locked
.get(node_id)
.map(|session| session.pipe_commands.clone())
.ok_or_else(|| "node session unavailable".to_owned())
}
async fn session_states(&self) -> BTreeMap<String, String> {
let locked = self.inner.sessions.lock().await;
locked
.iter()
.map(|(node_id, session)| (node_id.clone(), session.state.borrow().clone().as_text()))
.collect()
}
}
async fn wait_session_reply<T>(
timeout: Duration,
receiver: oneshot::Receiver<Result<T, String>>,
) -> Result<T, String> {
tokio::time::timeout(timeout, receiver)
.await
.map_err(|_| "node session request timed out".to_owned())?
.map_err(|_| "node session closed".to_owned())?
}
async fn run_node_session_manager(inner: Arc<NodeSessionManagerInner>) {
let interval = inner
.config
.reconnect_min_backoff
.min(Duration::from_millis(100))
.max(Duration::from_millis(10));
loop {
reconcile_sessions(&inner).await;
tokio::time::sleep(interval).await;
}
}
async fn reconcile_sessions(inner: &Arc<NodeSessionManagerInner>) {
let state = inner.state.get();
let mut wanted = BTreeSet::new();
for member in state.members.values() {
if eligible_member(inner, member) {
wanted.insert(member.node_id.clone());
ensure_session(inner, member.clone()).await;
}
}
close_unwanted_sessions(inner, &wanted).await;
}
fn eligible_member(inner: &NodeSessionManagerInner, member: &Member) -> bool {
member.node_id != inner.self_node
&& member.state == MemberState::Up
&& member.has_role(&inner.config.agent_role)
&& member.agent_addr.is_some()
}
async fn ensure_session(inner: &Arc<NodeSessionManagerInner>, member: Member) {
let mut locked = inner.sessions.lock().await;
if let Some(existing) = locked.get(&member.node_id)
&& existing.member.agent_addr == member.agent_addr
&& existing.member.incarnation == member.incarnation
{
return;
}
if let Some(existing) = locked.remove(&member.node_id) {
let _ = existing.stop.send(true);
existing.task.abort();
}
let (commands, receiver) = mpsc::channel(inner.config.command_buffer.max(1));
let pipe_buffer = inner.config.command_buffer.saturating_mul(4).max(1);
let (pipe_commands, pipe_receiver) = mpsc::channel(pipe_buffer);
let (stop, stop_receiver) = watch::channel(false);
let (state_sender, state_receiver) = watch::channel(PeerSessionState::Connecting);
let config = inner.config.clone();
let local_node = inner.self_node.clone();
let task_member = member.clone();
let task = tokio::spawn(async move {
run_peer_session(
config,
local_node,
task_member,
receiver,
pipe_receiver,
stop_receiver,
state_sender,
)
.await;
});
locked.insert(
member.node_id.clone(),
PeerSession {
member,
commands,
pipe_commands,
stop,
state: state_receiver,
task,
},
);
}
async fn close_unwanted_sessions(inner: &Arc<NodeSessionManagerInner>, wanted: &BTreeSet<String>) {
let to_close = {
let locked = inner.sessions.lock().await;
locked
.keys()
.filter(|node_id| !wanted.contains(*node_id))
.cloned()
.collect::<Vec<_>>()
};
for node_id in to_close {
close_session(inner, &node_id).await;
}
}
async fn close_session(inner: &Arc<NodeSessionManagerInner>, node_id: &str) {
let removed = {
let mut locked = inner.sessions.lock().await;
locked.remove(node_id)
};
if let Some(session) = removed {
let _ = session.stop.send(true);
session.task.abort();
}
}
async fn run_peer_session(
config: NodeSessionConfig,
local_node: String,
member: Member,
mut commands: mpsc::Receiver<SessionCommand>,
mut pipe_commands: mpsc::Receiver<ShardPipeCommand>,
mut stop: watch::Receiver<bool>,
state: watch::Sender<PeerSessionState>,
) {
let mut backoff = config.reconnect_min_backoff;
loop {
if *stop.borrow() {
break;
}
let _ = state.send(PeerSessionState::Connecting);
match connect_peer(&config.transport, &local_node, &member).await {
Ok(client) => match connect_shard_pipe(&config.transport, &local_node, &member).await {
Ok(pipe) => {
let _ = state.send(PeerSessionState::Connected);
backoff = config.reconnect_min_backoff;
if !run_connected_peer(
client,
pipe,
&mut commands,
&mut pipe_commands,
&mut stop,
)
.await
{
break;
}
}
Err(error) => {
let _ = state.send(PeerSessionState::BackingOff(error.clone()));
if !backoff_or_stop(
backoff,
error,
&mut commands,
&mut pipe_commands,
&mut stop,
)
.await
{
break;
}
backoff = next_backoff(backoff, config.reconnect_max_backoff);
}
},
Err(error) => {
let _ = state.send(PeerSessionState::BackingOff(error.clone()));
if !backoff_or_stop(backoff, error, &mut commands, &mut pipe_commands, &mut stop)
.await
{
break;
}
backoff = next_backoff(backoff, config.reconnect_max_backoff);
}
}
}
let _ = state.send(PeerSessionState::Closed);
}
async fn connect_peer(
transport: &NodeSessionTransport,
local_node: &str,
member: &Member,
) -> Result<DcpClient, String> {
let addr = member
.agent_addr
.ok_or_else(|| "member did not advertise a DCP agent address".to_owned())?;
let mut hello = Hello::new(local_node.to_owned(), ClientKind::ClusterNode);
hello.capabilities.push("cluster-node-sessions".to_owned());
match transport {
NodeSessionTransport::TcpLoopback => {
if !addr.ip().is_loopback() {
return Err(format!(
"plaintext node-session TCP requires loopback, got {addr}"
));
}
DcpClient::connect_tcp(addr, hello)
.await
.map_err(|error| error.to_string())
}
NodeSessionTransport::QuicMtls {
server_name,
client_config,
} => DcpClient::connect_quic(addr, server_name, client_config.clone(), hello)
.await
.map_err(|error| error.to_string()),
}
}
async fn connect_shard_pipe(
transport: &NodeSessionTransport,
local_node: &str,
member: &Member,
) -> Result<ShardPipeClient, String> {
let addr = member
.agent_addr
.ok_or_else(|| "member did not advertise a DCP agent address".to_owned())?;
let mut hello = Hello::new(local_node.to_owned(), ClientKind::ClusterNode);
hello.capabilities.push("cluster-node-sessions".to_owned());
hello.capabilities.push("cluster-sharding-pipe".to_owned());
match transport {
NodeSessionTransport::TcpLoopback => {
if !addr.ip().is_loopback() {
return Err(format!(
"plaintext node-session TCP requires loopback, got {addr}"
));
}
ShardPipeClient::connect_tcp(addr, hello)
.await
.map_err(|error| error.to_string())
}
NodeSessionTransport::QuicMtls {
server_name,
client_config,
} => ShardPipeClient::connect_quic(addr, server_name, client_config.clone(), hello)
.await
.map_err(|error| error.to_string()),
}
}
async fn run_connected_peer(
client: DcpClient,
pipe: ShardPipeClient,
commands: &mut mpsc::Receiver<SessionCommand>,
pipe_commands: &mut mpsc::Receiver<ShardPipeCommand>,
stop: &mut watch::Receiver<bool>,
) -> bool {
loop {
tokio::select! {
changed = stop.changed() => {
return changed.is_ok() && !*stop.borrow();
}
command = pipe_commands.recv() => {
let Some(command) = command else {
return false;
};
let frame = build_shard_pipe_frame(command, pipe_commands);
if pipe.send(frame).await.is_err() {
return true;
}
}
command = commands.recv() => {
let Some(command) = command else {
return false;
};
match command {
SessionCommand::ListJobs { timeout, reply } => {
let result = tokio::time::timeout(timeout, client.list_jobs()).await
.map_err(|_| "peer ListJobs timed out".to_owned())
.and_then(|result| result.map_err(|error| error.to_string()));
let reconnect = result.is_err();
let _ = reply.send(result);
if reconnect {
return true;
}
}
SessionCommand::SubmitClusterJob {
request,
timeout,
reply,
} => {
let timeout_ms = duration_millis_u64(timeout);
let result = tokio::time::timeout(
timeout,
client.submit_cluster_job(
request.factory_name,
request.instance_name,
request.params,
request.placement.unwrap_or_else(default_wire_placement),
timeout_ms,
),
)
.await
.map_err(|_| "peer SubmitClusterJob timed out".to_owned())
.and_then(|result| result.map_err(|error| error.to_string()));
let reconnect = result.is_err();
let _ = reply.send(result);
if reconnect {
return true;
}
}
SessionCommand::StartClusterJob {
factory_name,
instance_name,
params,
assignment,
timeout,
reply,
} => {
let result = tokio::time::timeout(
timeout,
client.start_cluster_job_on_node(
factory_name,
instance_name,
params,
assignment,
),
)
.await
.map_err(|_| "peer StartJob timed out".to_owned())
.and_then(|result| result.map_err(|error| error.to_string()));
let reconnect = result.is_err();
let _ = reply.send(result);
if reconnect {
return true;
}
}
SessionCommand::ClusterJobStatus {
name,
cluster,
timeout,
reply,
} => {
let request = async {
if cluster {
client
.cluster_job_status(name, duration_millis_u64(timeout))
.await
} else {
client.job_status(name).await
}
};
let result = tokio::time::timeout(timeout, request).await
.map_err(|_| "peer JobStatus timed out".to_owned())
.and_then(|result| result.map_err(|error| error.to_string()));
let reconnect = result.is_err();
let _ = reply.send(result);
if reconnect {
return true;
}
}
SessionCommand::DrainClusterJob {
name,
cluster,
timeout,
reply,
} => {
let request = async {
if cluster {
client
.drain_cluster_job(name, duration_millis_u64(timeout))
.await
} else {
client.drain_job(name).await
}
};
let result = tokio::time::timeout(timeout, request).await
.map_err(|_| "peer DrainJob timed out".to_owned())
.and_then(|result| result.map_err(|error| error.to_string()));
let reconnect = result.is_err();
let _ = reply.send(result);
if reconnect {
return true;
}
}
SessionCommand::StopClusterJob {
name,
cluster,
timeout,
reply,
} => {
let request = async {
if cluster {
client
.stop_cluster_job(name, duration_millis_u64(timeout))
.await
} else {
client.stop_job(name).await
}
};
let result = tokio::time::timeout(timeout, request).await
.map_err(|_| "peer StopJob timed out".to_owned())
.and_then(|result| result.map_err(|error| error.to_string()));
let reconnect = result.is_err();
let _ = reply.send(result);
if reconnect {
return true;
}
}
SessionCommand::ListClusterJobs { timeout, reply } => {
let result = tokio::time::timeout(
timeout,
client.list_cluster_jobs(duration_millis_u64(timeout)),
)
.await
.map_err(|_| "peer ListClusterJobs timed out".to_owned())
.and_then(|result| result.map_err(|error| error.to_string()));
let reconnect = result.is_err();
let _ = reply.send(result);
if reconnect {
return true;
}
}
SessionCommand::RememberClusterAssignment {
instance_name,
assignment,
timeout,
reply,
} => {
let result = tokio::time::timeout(
timeout,
client.remember_cluster_assignment(instance_name, assignment),
)
.await
.map_err(|_| "peer RememberClusterAssignment timed out".to_owned())
.and_then(|result| result.map_err(|error| error.to_string()));
let reconnect = result.is_err();
let _ = reply.send(result);
if reconnect {
return true;
}
}
SessionCommand::AllocateShard {
type_name,
shard_id,
timeout,
reply,
} => {
let timeout_ms = duration_millis_u64(timeout);
let result = tokio::time::timeout(
timeout,
client.allocate_shard(type_name, shard_id, timeout_ms),
)
.await
.map_err(|_| "peer AllocateShard timed out".to_owned())
.and_then(|result| result.map_err(|error| error.to_string()));
let reconnect = result.is_err();
let _ = reply.send(result);
if reconnect {
return true;
}
}
SessionCommand::RememberShardAllocations {
table,
timeout,
reply,
} => {
let result = tokio::time::timeout(
timeout,
client.remember_shard_allocations(table),
)
.await
.map_err(|_| "peer RememberShardAllocations timed out".to_owned())
.and_then(|result| result.map_err(|error| error.to_string()));
let reconnect = result.is_err();
let _ = reply.send(result);
if reconnect {
return true;
}
}
SessionCommand::GetShardAllocations {
type_name,
timeout,
reply,
} => {
let timeout_ms = duration_millis_u64(timeout);
let result = tokio::time::timeout(
timeout,
client.get_shard_allocations(type_name, timeout_ms),
)
.await
.map_err(|_| "peer GetShardAllocations timed out".to_owned())
.and_then(|result| result.map_err(|error| error.to_string()));
let reconnect = result.is_err();
let _ = reply.send(result);
if reconnect {
return true;
}
}
SessionCommand::ForwardShardEnvelopes {
batch,
timeout,
reply,
} => {
let timeout_ms = duration_millis_u64(timeout);
let result = tokio::time::timeout(
timeout,
client.forward_shard_envelopes(batch, timeout_ms),
)
.await
.map_err(|_| "peer ForwardShardEnvelopes timed out".to_owned())
.and_then(|result| result.map_err(|error| error.to_string()));
let reconnect = result.is_err();
let _ = reply.send(result);
if reconnect {
return true;
}
}
SessionCommand::CompleteShardingAsk {
response,
timeout,
reply,
} => {
let result =
tokio::time::timeout(timeout, client.complete_sharding_ask(response))
.await
.map_err(|_| "peer CompleteShardingAsk timed out".to_owned())
.and_then(|result| result.map_err(|error| error.to_string()));
let reconnect = result.is_err();
let _ = reply.send(result);
if reconnect {
return true;
}
}
}
}
}
}
}
const SHARD_PIPE_MAX_BATCH_COMMANDS: usize = 1024;
fn build_shard_pipe_frame(
first: ShardPipeCommand,
pipe_commands: &mut mpsc::Receiver<ShardPipeCommand>,
) -> ShardPipeFrame {
let mut frame = ShardPipeFrame {
forwards: Vec::new(),
replies: Vec::new(),
};
push_shard_pipe_command(&mut frame, first);
for _ in 1..SHARD_PIPE_MAX_BATCH_COMMANDS {
match pipe_commands.try_recv() {
Ok(command) => push_shard_pipe_command(&mut frame, command),
Err(_) => break,
}
}
frame
}
fn push_shard_pipe_command(frame: &mut ShardPipeFrame, command: ShardPipeCommand) {
match command {
ShardPipeCommand::Forward { batch } => {
if let Some(existing) = frame
.forwards
.iter_mut()
.find(|existing| existing.type_name == batch.type_name)
{
existing.envelopes.extend(batch.envelopes);
} else {
frame.forwards.push(batch);
}
}
ShardPipeCommand::Reply { response } => {
frame.replies.push(response);
}
}
}
async fn backoff_or_stop(
backoff: Duration,
error: String,
commands: &mut mpsc::Receiver<SessionCommand>,
pipe_commands: &mut mpsc::Receiver<ShardPipeCommand>,
stop: &mut watch::Receiver<bool>,
) -> bool {
let sleep = sleep(backoff);
tokio::pin!(sleep);
loop {
tokio::select! {
() = &mut sleep => return true,
changed = stop.changed() => {
return changed.is_ok() && !*stop.borrow();
}
command = pipe_commands.recv() => {
if command.is_none() {
return false;
}
}
command = commands.recv() => {
let Some(command) = command else {
return false;
};
match command {
SessionCommand::ListJobs { reply, .. } => {
let _ = reply.send(Err(format!("node session unreachable: {error}")));
}
SessionCommand::SubmitClusterJob { reply, .. }
| SessionCommand::StartClusterJob { reply, .. }
| SessionCommand::ClusterJobStatus { reply, .. }
| SessionCommand::DrainClusterJob { reply, .. }
| SessionCommand::StopClusterJob { reply, .. } => {
let _ = reply.send(Err(format!("node session unreachable: {error}")));
}
SessionCommand::ListClusterJobs { reply, .. } => {
let _ = reply.send(Err(format!("node session unreachable: {error}")));
}
SessionCommand::RememberClusterAssignment { reply, .. } => {
let _ = reply.send(Err(format!("node session unreachable: {error}")));
}
SessionCommand::AllocateShard { reply, .. } => {
let _ = reply.send(Err(format!("node session unreachable: {error}")));
}
SessionCommand::RememberShardAllocations { reply, .. } => {
let _ = reply.send(Err(format!("node session unreachable: {error}")));
}
SessionCommand::GetShardAllocations { reply, .. } => {
let _ = reply.send(Err(format!("node session unreachable: {error}")));
}
SessionCommand::ForwardShardEnvelopes { reply, .. } => {
let _ = reply.send(Err(format!("node session unreachable: {error}")));
}
SessionCommand::CompleteShardingAsk { reply, .. } => {
let _ = reply.send(Err(format!("node session unreachable: {error}")));
}
}
}
}
}
}
fn default_wire_placement() -> PlacementSpec {
PlacementSpec {
role_constraint: String::new(),
strategy: PlacementStrategy::LeastJobs as i32,
pinned_node_id: String::new(),
}
}
fn duration_millis_u64(duration: Duration) -> u64 {
duration.as_millis().min(u128::from(u64::MAX)) as u64
}
fn next_backoff(current: Duration, max: Duration) -> Duration {
current.saturating_mul(2).min(max)
}
#[derive(Clone)]
struct PlacementCoordinatorHandle {
commands: mpsc::Sender<PlacementCommand>,
tasks: Arc<Mutex<Vec<JoinHandle<()>>>>,
}
enum PlacementCommand {
Submit {
request: SubmitClusterJob,
timeout: Duration,
reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, DcpError>>,
},
Status {
name: String,
timeout: Duration,
reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, DcpError>>,
},
Drain {
name: String,
timeout: Duration,
reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, DcpError>>,
},
Stop {
name: String,
timeout: Duration,
reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, DcpError>>,
},
Remember {
instance_name: String,
assignment: ClusterJobStart,
reply: oneshot::Sender<Result<(), DcpError>>,
},
Tick,
Shutdown,
}
struct PlacementActor {
self_node: String,
agent_role: String,
state: Signal<ClusterState>,
registry: JobRegistryHandle,
sessions: NodeSessionManagerHandle,
factories: DcpJobFactories,
request_timeout: Duration,
active_coordinator: bool,
assignments: BTreeMap<String, ClusterJobMetadata>,
}
impl PlacementCoordinatorHandle {
fn start(
self_node: String,
agent_role: String,
state: Signal<ClusterState>,
registry: JobRegistryHandle,
sessions: NodeSessionManagerHandle,
factories: DcpJobFactories,
request_timeout: Duration,
) -> ClusterAgentResult<Self> {
let (commands, receiver) = mpsc::channel(128);
let tasks = Arc::new(Mutex::new(Vec::new()));
let actor = PlacementActor {
self_node,
agent_role,
state,
registry,
sessions,
factories,
request_timeout,
active_coordinator: false,
assignments: BTreeMap::new(),
};
let actor_task = tokio::spawn(run_placement_actor(actor, receiver));
let tick_commands = commands.clone();
let tick_task = tokio::spawn(async move {
let mut interval = tokio::time::interval(Duration::from_millis(50));
loop {
interval.tick().await;
if tick_commands.send(PlacementCommand::Tick).await.is_err() {
break;
}
}
});
{
let mut locked = tasks.lock().expect("placement tasks poisoned");
locked.push(actor_task);
locked.push(tick_task);
}
Ok(Self { commands, tasks })
}
async fn submit_cluster_job(
&self,
request: SubmitClusterJob,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, DcpError> {
self.request(timeout, |reply| PlacementCommand::Submit {
request,
timeout,
reply,
})
.await
}
async fn cluster_job_status(
&self,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, DcpError> {
self.request(timeout, |reply| PlacementCommand::Status {
name,
timeout,
reply,
})
.await
}
async fn drain_cluster_job(
&self,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, DcpError> {
self.request(timeout, |reply| PlacementCommand::Drain {
name,
timeout,
reply,
})
.await
}
async fn stop_cluster_job(
&self,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, DcpError> {
self.request(timeout, |reply| PlacementCommand::Stop {
name,
timeout,
reply,
})
.await
}
async fn remember_cluster_assignment(
&self,
instance_name: String,
assignment: ClusterJobStart,
) -> Result<(), DcpError> {
self.request(Duration::from_millis(250), |reply| {
PlacementCommand::Remember {
instance_name,
assignment,
reply,
}
})
.await
}
async fn shutdown(&self) {
let _ = self.commands.send(PlacementCommand::Shutdown).await;
let tasks = {
let mut locked = self.tasks.lock().expect("placement tasks poisoned");
std::mem::take(&mut *locked)
};
for task in tasks {
task.abort();
}
}
async fn request<T, F>(&self, timeout: Duration, make: F) -> Result<T, DcpError>
where
T: Send + 'static,
F: FnOnce(oneshot::Sender<Result<T, DcpError>>) -> PlacementCommand,
{
let (reply, receiver) = oneshot::channel();
self.commands
.send(make(reply))
.await
.map_err(|_| DcpError::response(ResponseStatus::Failed, "placement actor stopped"))?;
tokio::time::timeout(timeout.saturating_add(Duration::from_millis(250)), receiver)
.await
.map_err(|_| {
DcpError::response(
ResponseStatus::DeadlineExceeded,
"placement request timed out",
)
})?
.map_err(|_| DcpError::response(ResponseStatus::Failed, "placement actor stopped"))?
}
}
async fn run_placement_actor(
mut actor: PlacementActor,
mut receiver: mpsc::Receiver<PlacementCommand>,
) {
while let Some(command) = receiver.recv().await {
match command {
PlacementCommand::Submit {
request,
timeout,
reply,
} => {
let result = actor.submit(request, timeout).await;
let _ = reply.send(result);
}
PlacementCommand::Status {
name,
timeout,
reply,
} => {
let result = actor.cluster_job_status(name, timeout).await;
let _ = reply.send(result);
}
PlacementCommand::Drain {
name,
timeout,
reply,
} => {
let result = actor.drain_cluster_job(name, timeout).await;
let _ = reply.send(result);
}
PlacementCommand::Stop {
name,
timeout,
reply,
} => {
let result = actor.stop_cluster_job(name, timeout).await;
let _ = reply.send(result);
}
PlacementCommand::Remember {
instance_name,
assignment,
reply,
} => {
let result = actor.remember(instance_name, assignment);
let _ = reply.send(result);
}
PlacementCommand::Tick => {
let _ = actor.reconcile().await;
}
PlacementCommand::Shutdown => break,
}
}
}
impl PlacementActor {
async fn submit(
&mut self,
request: SubmitClusterJob,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, DcpError> {
self.ensure_active_coordinator(timeout).await?;
validate_submit(&request)?;
if self.assignments.contains_key(&request.instance_name) {
return Err(DcpError::response(
ResponseStatus::Conflict,
format!("cluster job already exists: {}", request.instance_name),
));
}
let placement = placement_spec_from_wire(request.placement.clone())?;
let target = self.choose_submit_target(&placement)?;
let history = ClusterPlacementHistory {
generation: 1,
from_node: None,
to_node: target.clone(),
reason: "submitted".to_owned(),
timestamp: SystemTime::now(),
};
let metadata = ClusterJobMetadata {
factory_name: request.factory_name.clone(),
params: request.params.clone().into_iter().collect(),
placement,
coordinator_node: self.self_node.clone(),
assigned_node: target.clone(),
placement_generation: 1,
history: vec![history],
};
let status = self
.start_on_node(
&target,
request.factory_name,
request.instance_name.clone(),
request.params,
metadata.clone(),
timeout,
)
.await?;
self.assignments
.insert(request.instance_name.clone(), metadata.clone());
self.replicate_assignment(request.instance_name, metadata, timeout)
.await;
Ok(status)
}
async fn cluster_job_status(
&mut self,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, DcpError> {
self.ensure_active_coordinator(timeout).await?;
let metadata = self.assignment_or_rebuild(&name, timeout).await?;
self.status_on_node(&metadata.assigned_node, name, timeout)
.await
}
async fn drain_cluster_job(
&mut self,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, DcpError> {
self.ensure_active_coordinator(timeout).await?;
let metadata = self.assignment_or_rebuild(&name, timeout).await?;
let status = self
.drain_on_node(&metadata.assigned_node, name.clone(), timeout)
.await?;
self.assignments.remove(&name);
Ok(status)
}
async fn stop_cluster_job(
&mut self,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, DcpError> {
self.ensure_active_coordinator(timeout).await?;
let metadata = self.assignment_or_rebuild(&name, timeout).await?;
let status = self
.stop_on_node(&metadata.assigned_node, name.clone(), timeout)
.await?;
self.assignments.remove(&name);
Ok(status)
}
fn remember(
&mut self,
instance_name: String,
assignment: ClusterJobStart,
) -> Result<(), DcpError> {
let metadata = cluster_metadata_from_wire(assignment)?;
self.assignments.insert(instance_name, metadata);
Ok(())
}
async fn reconcile(&mut self) -> Result<(), DcpError> {
let snapshot = self.state.get();
let is_coordinator = snapshot.is_placement_coordinator(&self.self_node);
drop(snapshot);
if !is_coordinator {
self.active_coordinator = false;
return Ok(());
}
self.ensure_active_coordinator(self.request_timeout).await?;
self.replace_down_members(self.request_timeout).await
}
async fn ensure_active_coordinator(&mut self, timeout: Duration) -> Result<(), DcpError> {
let snapshot = self.state.get();
if !snapshot.is_placement_coordinator(&self.self_node) {
return Err(DcpError::response(
ResponseStatus::Failed,
"this node is not the placement coordinator",
));
}
drop(snapshot);
if !self.active_coordinator {
self.rebuild_from_registries(timeout).await?;
self.active_coordinator = true;
}
Ok(())
}
async fn assignment_or_rebuild(
&mut self,
name: &str,
timeout: Duration,
) -> Result<ClusterJobMetadata, DcpError> {
if let Some(metadata) = self.assignments.get(name).cloned() {
return Ok(metadata);
}
self.rebuild_from_registries(timeout).await?;
self.assignments.get(name).cloned().ok_or_else(|| {
DcpError::response(
ResponseStatus::NotFound,
format!("cluster job not found: {name}"),
)
})
}
async fn rebuild_from_registries(&mut self, timeout: Duration) -> Result<(), DcpError> {
let statuses = self.collect_job_statuses(timeout).await?;
for status in statuses {
if status.cluster_job
&& status.desired_state != "Stopped"
&& let Some(metadata) = metadata_from_status(&status)?
{
let replace = self
.assignments
.get(&status.name)
.map(|existing| {
metadata.placement_generation > existing.placement_generation
|| (metadata.placement_generation == existing.placement_generation
&& metadata.assigned_node < existing.assigned_node)
})
.unwrap_or(true);
if replace {
self.assignments.insert(status.name.clone(), metadata);
}
}
}
Ok(())
}
async fn replace_down_members(&mut self, timeout: Duration) -> Result<(), DcpError> {
let snapshot = self.state.get();
let down_nodes = snapshot
.members
.values()
.filter(|member| member.state == MemberState::Down)
.map(|member| member.node_id.clone())
.collect::<BTreeSet<_>>();
drop(snapshot);
if down_nodes.is_empty() {
return Ok(());
}
let to_replace = self
.assignments
.iter()
.filter(|(_name, metadata)| down_nodes.contains(&metadata.assigned_node))
.map(|(name, metadata)| (name.clone(), metadata.clone()))
.collect::<Vec<_>>();
for (name, mut metadata) in to_replace {
let old_node = metadata.assigned_node.clone();
let target = self.choose_replacement_target(&metadata.placement, &old_node)?;
metadata.placement_generation = metadata.placement_generation.saturating_add(1);
metadata.assigned_node = target.clone();
metadata.coordinator_node = self.self_node.clone();
metadata.history.push(ClusterPlacementHistory {
generation: metadata.placement_generation,
from_node: Some(old_node.clone()),
to_node: target.clone(),
reason: format!("node_down:{old_node}"),
timestamp: SystemTime::now(),
});
let params = metadata.params.clone().into_iter().collect();
let factory_name = metadata.factory_name.clone();
self.start_on_node(
&target,
factory_name,
name.clone(),
params,
metadata.clone(),
timeout,
)
.await?;
self.assignments.insert(name.clone(), metadata.clone());
self.replicate_assignment(name, metadata, timeout).await;
}
Ok(())
}
fn choose_submit_target(&self, placement: &RegistryPlacementSpec) -> Result<String, DcpError> {
let snapshot = self.state.get();
let candidates = eligible_members(
&snapshot,
&self.agent_role,
placement.role_constraint.as_deref(),
);
if candidates.is_empty() {
return Err(no_eligible_node_error(placement));
}
match &placement.strategy {
RegistryPlacementStrategy::LeastJobs => self.least_loaded(candidates),
RegistryPlacementStrategy::Pinned { node_id } => candidates
.into_iter()
.find(|member| member.node_id == *node_id)
.map(|member| member.node_id.clone())
.ok_or_else(|| {
DcpError::response(
ResponseStatus::NotFound,
format!("pinned node is not eligible for placement: {node_id}"),
)
}),
}
}
fn choose_replacement_target(
&self,
placement: &RegistryPlacementSpec,
old_node: &str,
) -> Result<String, DcpError> {
let snapshot = self.state.get();
let candidates = eligible_members(
&snapshot,
&self.agent_role,
placement.role_constraint.as_deref(),
)
.into_iter()
.filter(|member| member.node_id != old_node)
.collect::<Vec<_>>();
if candidates.is_empty() {
return Err(no_eligible_node_error(placement));
}
self.least_loaded(candidates)
}
fn least_loaded(&self, candidates: Vec<&Member>) -> Result<String, DcpError> {
candidates
.into_iter()
.min_by(|left, right| {
let left_count = self.jobs_on_node(&left.node_id);
let right_count = self.jobs_on_node(&right.node_id);
left_count
.cmp(&right_count)
.then_with(|| left.node_id.cmp(&right.node_id))
})
.map(|member| member.node_id.clone())
.ok_or_else(|| DcpError::response(ResponseStatus::NotFound, "no eligible node"))
}
fn jobs_on_node(&self, node_id: &str) -> usize {
self.assignments
.values()
.filter(|metadata| metadata.assigned_node == node_id)
.count()
}
async fn start_on_node(
&self,
node_id: &str,
factory_name: String,
instance_name: String,
params: HashMap<String, String>,
metadata: ClusterJobMetadata,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, DcpError> {
if node_id == self.self_node {
let mut spec = self
.factories
.build(&factory_name, instance_name.clone(), params)?;
spec = spec.with_cluster_metadata(metadata);
let registry = self.registry.clone();
let status = tokio::task::spawn_blocking(move || {
registry.submit(spec)?;
registry.start(instance_name)
})
.await??;
return Ok(wire_job_status(&status));
}
self.sessions
.start_cluster_job(
node_id,
factory_name,
instance_name,
params,
wire_cluster_job_start(&metadata),
timeout,
)
.await
.map_err(session_error)
}
async fn status_on_node(
&self,
node_id: &str,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, DcpError> {
if node_id == self.self_node {
let registry = self.registry.clone();
let status = tokio::task::spawn_blocking(move || registry.status(name)).await??;
return Ok(wire_job_status(&status));
}
self.sessions
.cluster_job_status(node_id, name, timeout)
.await
.map_err(session_error)
}
async fn drain_on_node(
&self,
node_id: &str,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, DcpError> {
if node_id == self.self_node {
let registry = self.registry.clone();
let status = tokio::task::spawn_blocking(move || registry.drain(name)).await??;
return Ok(wire_job_status(&status));
}
self.sessions
.drain_cluster_job(node_id, name, timeout)
.await
.map_err(session_error)
}
async fn stop_on_node(
&self,
node_id: &str,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, DcpError> {
if node_id == self.self_node {
let registry = self.registry.clone();
let status = tokio::task::spawn_blocking(move || registry.stop(name)).await??;
return Ok(wire_job_status(&status));
}
self.sessions
.stop_cluster_job(node_id, name, timeout)
.await
.map_err(session_error)
}
async fn collect_job_statuses(
&self,
timeout: Duration,
) -> Result<Vec<crate::dcp::proto::JobStatus>, DcpError> {
let snapshot = self.state.get();
let mut statuses = Vec::new();
let registry = self.registry.clone();
let local_jobs = tokio::task::spawn_blocking(move || registry.list()).await??;
statuses.extend(local_jobs.iter().map(wire_job_status));
let mut pending = Vec::new();
for member in snapshot.members.values() {
if member.node_id == self.self_node
|| !member.has_role(&self.agent_role)
|| member.state != MemberState::Up
|| member.unreachable
{
continue;
}
let node_id = member.node_id.clone();
let sessions = self.sessions.clone();
pending.push(tokio::spawn(async move {
sessions.list_jobs(&node_id, timeout).await
}));
}
drop(snapshot);
let deadline = Instant::now() + timeout;
for task in pending {
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
break;
}
if let Ok(Ok(Ok(mut peer_statuses))) = tokio::time::timeout(remaining, task).await {
statuses.append(&mut peer_statuses);
}
}
Ok(statuses)
}
async fn replicate_assignment(
&self,
instance_name: String,
metadata: ClusterJobMetadata,
timeout: Duration,
) {
let snapshot = self.state.get();
let peers = snapshot
.members
.values()
.filter(|member| {
member.node_id != self.self_node
&& member.has_role(&self.agent_role)
&& member.state == MemberState::Up
&& !member.unreachable
})
.map(|member| member.node_id.clone())
.collect::<Vec<_>>();
drop(snapshot);
let assignment = wire_cluster_job_start(&metadata);
let timeout = timeout.min(Duration::from_millis(250));
for node_id in peers {
let _ = self
.sessions
.remember_cluster_assignment(
&node_id,
instance_name.clone(),
assignment.clone(),
timeout,
)
.await;
}
}
}
fn validate_submit(request: &SubmitClusterJob) -> Result<(), DcpError> {
if request.factory_name.trim().is_empty() || request.instance_name.trim().is_empty() {
return Err(DcpError::response(
ResponseStatus::BadRequest,
"SubmitClusterJob requires factory_name and instance_name",
));
}
Ok(())
}
fn eligible_members<'a>(
snapshot: &'a ClusterState,
agent_role: &str,
role_constraint: Option<&str>,
) -> Vec<&'a Member> {
snapshot
.members
.values()
.filter(|member| member.state == MemberState::Up && !member.unreachable)
.filter(|member| member.has_role(agent_role))
.filter(|member| role_constraint.is_none_or(|role| member.has_role(role)))
.collect()
}
fn no_eligible_node_error(placement: &RegistryPlacementSpec) -> DcpError {
let role = placement
.role_constraint
.as_deref()
.map(|role| format!(" with role '{role}'"))
.unwrap_or_default();
DcpError::response(
ResponseStatus::NotFound,
format!("no eligible Up placement nodes{role}"),
)
}
fn metadata_from_status(
status: &crate::dcp::proto::JobStatus,
) -> Result<Option<ClusterJobMetadata>, DcpError> {
if !status.cluster_job {
return Ok(None);
}
let placement = placement_spec_from_wire(status.placement.clone())?;
Ok(Some(ClusterJobMetadata {
factory_name: status.factory_name.clone(),
params: status.params.clone().into_iter().collect(),
placement,
coordinator_node: status.coordinator_node_id.clone(),
assigned_node: status.placement_node_id.clone(),
placement_generation: status.placement_generation,
history: status
.placement_history
.iter()
.map(|history| ClusterPlacementHistory {
generation: history.generation,
from_node: if history.from_node_id.is_empty() {
None
} else {
Some(history.from_node_id.clone())
},
to_node: history.to_node_id.clone(),
reason: history.reason.clone(),
timestamp: std::time::UNIX_EPOCH + Duration::from_millis(history.timestamp_ms),
})
.collect(),
}))
}
fn session_error(message: String) -> DcpError {
let status = if message.contains("DCP response NotFound") {
ResponseStatus::NotFound
} else if message.contains("DCP response Conflict") {
ResponseStatus::Conflict
} else if message.contains("DCP response BadRequest") {
ResponseStatus::BadRequest
} else if message.contains("DCP response DeadlineExceeded") {
ResponseStatus::DeadlineExceeded
} else {
ResponseStatus::Failed
};
DcpError::response(status, message)
}
struct ClusterView {
registry: JobRegistryHandle,
state: Signal<ClusterState>,
sessions: NodeSessionManagerHandle,
placement: PlacementCoordinatorHandle,
self_node: String,
agent_role: String,
}
impl ClusterViewProvider for ClusterView {
fn submit_cluster_job(
&self,
request: SubmitClusterJob,
timeout: Duration,
) -> crate::dcp::server::ClusterViewFuture<'_, crate::dcp::proto::JobStatus> {
Box::pin(async move { self.submit_cluster_job_inner(request, timeout).await })
}
fn list_cluster_jobs(
&self,
timeout: Duration,
) -> crate::dcp::server::ClusterViewFuture<'_, ClusterJobList> {
Box::pin(async move { self.list_cluster_jobs_inner(timeout).await })
}
fn cluster_node_info(
&self,
timeout: Duration,
) -> crate::dcp::server::ClusterViewFuture<'_, ClusterNodeList> {
Box::pin(async move { self.cluster_node_info_inner(timeout).await })
}
fn cluster_job_status(
&self,
name: String,
timeout: Duration,
) -> crate::dcp::server::ClusterViewFuture<'_, crate::dcp::proto::JobStatus> {
Box::pin(async move { self.cluster_job_status_inner(name, timeout).await })
}
fn drain_cluster_job(
&self,
name: String,
timeout: Duration,
) -> crate::dcp::server::ClusterViewFuture<'_, crate::dcp::proto::JobStatus> {
Box::pin(async move { self.drain_cluster_job_inner(name, timeout).await })
}
fn stop_cluster_job(
&self,
name: String,
timeout: Duration,
) -> crate::dcp::server::ClusterViewFuture<'_, crate::dcp::proto::JobStatus> {
Box::pin(async move { self.stop_cluster_job_inner(name, timeout).await })
}
fn remember_cluster_assignment(
&self,
request: RememberClusterAssignment,
) -> crate::dcp::server::ClusterViewFuture<'_, ()> {
Box::pin(async move { self.remember_cluster_assignment_inner(request).await })
}
}
impl ClusterView {
async fn submit_cluster_job_inner(
&self,
request: SubmitClusterJob,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, DcpError> {
if self.is_local_coordinator() {
return self.placement.submit_cluster_job(request, timeout).await;
}
let coordinator = self.coordinator_node_id()?;
self.sessions
.submit_cluster_job(&coordinator, request, timeout)
.await
.map_err(session_error)
}
async fn list_cluster_jobs_inner(&self, timeout: Duration) -> Result<ClusterJobList, DcpError> {
if !self.is_local_coordinator() {
let coordinator = self.coordinator_node_id()?;
return self
.sessions
.list_cluster_jobs(&coordinator, timeout)
.await
.map_err(session_error);
}
let timeout = if timeout.is_zero() {
self.sessions.inner.config.request_timeout
} else {
timeout
};
let snapshot = self.state.get();
let local_member = snapshot.member(&self.self_node).cloned();
let local_jobs = list_local_jobs(self.registry.clone()).await?;
let mut nodes = Vec::new();
nodes.push(ClusterJobNode {
node_id: self.self_node.clone(),
address: local_member
.as_ref()
.map(|member| member.address.to_string())
.unwrap_or_default(),
local: true,
jobs: local_jobs,
});
let mut errors = Vec::new();
let mut pending = Vec::new();
for member in snapshot.members.values() {
if member.node_id == self.self_node || !member.has_role(&self.agent_role) {
continue;
}
if member.state == MemberState::Removed {
continue;
}
if member.state != MemberState::Up || member.unreachable {
errors.push(ClusterNodeError {
node_id: member.node_id.clone(),
message: if member.unreachable {
"member is unreachable".to_owned()
} else {
format!("member is {:?}", member.state)
},
});
continue;
}
let node_id = member.node_id.clone();
let address = member.address.to_string();
let sessions = self.sessions.clone();
pending.push(tokio::spawn(async move {
let result = sessions.list_jobs(&node_id, timeout).await;
(node_id, address, result)
}));
}
let deadline = Instant::now() + timeout;
for task in pending {
let remaining = deadline.saturating_duration_since(Instant::now());
match tokio::time::timeout(remaining, task).await {
Ok(Ok((node_id, address, Ok(jobs)))) => nodes.push(ClusterJobNode {
node_id,
address,
local: false,
jobs,
}),
Ok(Ok((node_id, _address, Err(message)))) => {
errors.push(ClusterNodeError { node_id, message });
}
Ok(Err(error)) => errors.push(ClusterNodeError {
node_id: "unknown".to_owned(),
message: format!("cluster fan-out task failed: {error}"),
}),
Err(_) => errors.push(ClusterNodeError {
node_id: "unknown".to_owned(),
message: "cluster fan-out timed out".to_owned(),
}),
}
}
nodes.sort_by(|left, right| left.node_id.cmp(&right.node_id));
errors.sort_by(|left, right| left.node_id.cmp(&right.node_id));
Ok(ClusterJobList {
partial: !errors.is_empty(),
nodes,
errors,
})
}
async fn cluster_node_info_inner(
&self,
_timeout: Duration,
) -> Result<ClusterNodeList, DcpError> {
let snapshot = self.state.get();
let session_states = self.sessions.session_states().await;
let mut nodes = snapshot
.members
.values()
.map(|member| ClusterNodeStatus {
node_id: member.node_id.clone(),
member_state: format!("{:?}", member.state),
address: member.address.to_string(),
agent_addr: member
.agent_addr
.map(|addr| addr.to_string())
.unwrap_or_default(),
roles: member.roles.clone(),
unreachable: member.unreachable,
local: member.node_id == self.self_node,
session_state: if member.node_id == self.self_node {
"local".to_owned()
} else {
session_states
.get(&member.node_id)
.cloned()
.unwrap_or_else(|| "not_connected".to_owned())
},
})
.collect::<Vec<_>>();
nodes.sort_by(|left, right| left.node_id.cmp(&right.node_id));
let mut errors = nodes
.iter()
.filter(|node| !node.local)
.filter(|node| node.roles.iter().any(|role| role == &self.agent_role))
.filter(|node| {
node.member_state == "Up" && !node.unreachable && node.session_state != "connected"
})
.map(|node| ClusterNodeError {
node_id: node.node_id.clone(),
message: format!("node session {}", node.session_state),
})
.collect::<Vec<_>>();
errors.sort_by(|left, right| left.node_id.cmp(&right.node_id));
Ok(ClusterNodeList {
partial: !errors.is_empty(),
nodes,
errors,
})
}
async fn cluster_job_status_inner(
&self,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, DcpError> {
if self.is_local_coordinator() {
return self.placement.cluster_job_status(name, timeout).await;
}
let coordinator = self.coordinator_node_id()?;
self.sessions
.forward_cluster_job_status(&coordinator, name, timeout)
.await
.map_err(session_error)
}
async fn drain_cluster_job_inner(
&self,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, DcpError> {
if self.is_local_coordinator() {
return self.placement.drain_cluster_job(name, timeout).await;
}
let coordinator = self.coordinator_node_id()?;
self.sessions
.forward_drain_cluster_job(&coordinator, name, timeout)
.await
.map_err(session_error)
}
async fn stop_cluster_job_inner(
&self,
name: String,
timeout: Duration,
) -> Result<crate::dcp::proto::JobStatus, DcpError> {
if self.is_local_coordinator() {
return self.placement.stop_cluster_job(name, timeout).await;
}
let coordinator = self.coordinator_node_id()?;
self.sessions
.forward_stop_cluster_job(&coordinator, name, timeout)
.await
.map_err(session_error)
}
async fn remember_cluster_assignment_inner(
&self,
request: RememberClusterAssignment,
) -> Result<(), DcpError> {
let assignment = request.assignment.ok_or_else(|| {
DcpError::response(
ResponseStatus::BadRequest,
"RememberClusterAssignment requires assignment",
)
})?;
self.placement
.remember_cluster_assignment(request.instance_name, assignment)
.await
}
fn is_local_coordinator(&self) -> bool {
self.state.get().is_placement_coordinator(&self.self_node)
}
fn coordinator_node_id(&self) -> Result<String, DcpError> {
self.state
.get()
.placement_coordinator()
.map(|member| member.node_id.clone())
.ok_or_else(|| {
DcpError::response(
ResponseStatus::Failed,
"no placement coordinator is available",
)
})
}
}
async fn list_local_jobs(
registry: JobRegistryHandle,
) -> Result<Vec<crate::dcp::proto::JobStatus>, DcpError> {
let statuses = tokio::task::spawn_blocking(move || registry.list()).await??;
Ok(statuses.iter().map(wire_job_status).collect())
}
fn validate_cluster_agent_config(config: &ClusterAgentConfig) -> ClusterAgentResult<()> {
if config.sessions.agent_role.trim().is_empty() {
return Err(ClusterAgentError::InvalidConfig(
"sessions.agent_role must not be empty".to_owned(),
));
}
if config.sessions.reconnect_min_backoff.is_zero()
|| config.sessions.reconnect_max_backoff < config.sessions.reconnect_min_backoff
{
return Err(ClusterAgentError::InvalidConfig(
"node-session reconnect backoff must be non-zero and ordered".to_owned(),
));
}
match config.sessions.transport {
NodeSessionTransport::TcpLoopback if config.dcp.tcp.is_none() => {
Err(ClusterAgentError::InvalidConfig(
"node-session TCP transport requires a DCP TCP listener".to_owned(),
))
}
NodeSessionTransport::QuicMtls { .. } if config.dcp.quic.is_none() => {
Err(ClusterAgentError::InvalidConfig(
"node-session QUIC transport requires a DCP QUIC listener".to_owned(),
))
}
_ => Ok(()),
}
}
fn advertised_agent_addr(
transport: &NodeSessionTransport,
handle: &DcpServerHandle,
) -> ClusterAgentResult<SocketAddr> {
match transport {
NodeSessionTransport::TcpLoopback => handle.tcp_addr().ok_or_else(|| {
ClusterAgentError::InvalidConfig("DCP TCP listener did not start".to_owned())
}),
NodeSessionTransport::QuicMtls { .. } => handle.quic_addr().ok_or_else(|| {
ClusterAgentError::InvalidConfig("DCP QUIC listener did not start".to_owned())
}),
}
}
fn ensure_agent_role(config: &mut ClusterConfig, role: &str) {
if !config.roles.iter().any(|candidate| candidate == role) {
config.roles.push(role.to_owned());
}
}