use std::{
collections::{BTreeMap, BTreeSet, HashMap},
fmt,
net::SocketAddr,
sync::{Arc, Mutex},
time::{Duration, SystemTime, UNIX_EPOCH},
};
use datum::{NotUsed, Sink, StreamCompletion, Topic, TopicOverflow, TopicTryPublishError};
use datum_cluster::{
ClusterConfig, ClusterNode, ClusterState, Member, MemberEvent, MemberEventKind, MemberState,
Signal,
};
use datum_net::quic::quinn;
use tokio::{
sync::{mpsc, oneshot, watch},
task::JoinHandle,
time::{Instant, sleep},
};
use crate::{
Agent, AgentConfig, AgentError, AgentHandle, ClusterJobMetadata, ClusterPlacementHistory,
DesiredJobState, JobRegistryHandle, JobState, PlacementSpec as RegistryPlacementSpec,
PlacementStrategy as RegistryPlacementStrategy,
dcp::{
ClientKind, ClusterEvent, 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,
client::MetricSubscription,
proto::MetricSample,
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 cluster_events = ClusterEventPublisher::new(config.agent.event_buffer)?;
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 membership_events = cluster
.events()
.changes()
.run_with(Sink::foreach({
let cluster_events = cluster_events.clone();
move |event| cluster_events.publish_member(&event)
}))
.map_err(crate::AgentError::from)?;
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(),
PlacementDependencies {
registry: agent.registry().clone(),
sessions: sessions.clone(),
factories: factories.clone(),
cluster_events: cluster_events.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(),
cluster_events: cluster_events.clone(),
});
server.set_cluster_view(provider);
Ok(ClusterAgentHandle {
agent,
cluster,
server,
server_handle: Some(server_handle),
sessions,
placement,
cluster_events,
_membership_events: membership_events,
})
}
}
pub struct ClusterAgentHandle {
agent: AgentHandle,
cluster: ClusterNode,
server: DcpServer,
server_handle: Option<DcpServerHandle>,
sessions: NodeSessionManagerHandle,
placement: PlacementCoordinatorHandle,
cluster_events: ClusterEventPublisher,
_membership_events: StreamCompletion<NotUsed>,
}
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<()> {
let node_id = self.cluster.node_id().to_owned();
self.server.clear_cluster_view();
self.server.clear_sharding_view();
self.placement.shutdown().await;
eprintln!(
"datum-agent INFO shutdown_component_stopped node_id={node_id} component=placement"
);
self.sessions.shutdown().await;
eprintln!(
"datum-agent INFO shutdown_component_stopped node_id={node_id} component=sessions"
);
let _ = self.cluster.abort().await;
eprintln!(
"datum-agent INFO shutdown_component_stopped node_id={node_id} component=cluster"
);
let _ = self.cluster_events.close();
if let Some(handle) = self.server_handle.take() {
handle.shutdown().await;
}
eprintln!(
"datum-agent INFO shutdown_component_stopped node_id={node_id} component=dcp_server"
);
self.agent.registry().shutdown()?;
eprintln!(
"datum-agent INFO shutdown_component_stopped node_id={node_id} component=registry"
);
eprintln!("datum-agent INFO shutdown_done node_id={node_id}");
Ok(())
}
}
impl Drop for ClusterAgentHandle {
fn drop(&mut self) {
self.server.clear_cluster_view();
self.server.clear_sharding_view();
let _ = self.cluster_events.close();
}
}
#[derive(Clone)]
struct ClusterEventPublisher {
topic: Topic<ClusterEvent>,
next_sequence: Arc<Mutex<u64>>,
subscriber_buffer: usize,
}
impl ClusterEventPublisher {
fn new(buffer: usize) -> crate::AgentResult<Self> {
let subscriber_buffer = buffer.max(1);
Ok(Self {
topic: Topic::new(subscriber_buffer, TopicOverflow::Sliding)?,
next_sequence: Arc::new(Mutex::new(1)),
subscriber_buffer,
})
}
fn subscribe(&self) -> crate::dcp::DcpResult<mpsc::Receiver<ClusterEvent>> {
let (sender, receiver) = mpsc::channel(self.subscriber_buffer);
let completion = self
.topic
.subscribe()
.run_with(Sink::foreach_result(move |event| {
sender
.blocking_send(event)
.map_err(|_| datum::StreamError::Cancelled)
}))?;
tokio::task::spawn_blocking(move || {
let _ = completion.wait();
});
Ok(receiver)
}
fn publish_member(&self, event: &MemberEvent) {
self.publish(
system_time_ms(event.at),
member_event_kind(event.kind),
event.member.node_id.clone(),
format!(
"state={:?} address={} unreachable={} incarnation={}",
event.member.state,
event.member.address,
event.member.unreachable,
event.member.incarnation
),
Some(event.member.incarnation),
);
}
fn publish_coordinator_changed(
&self,
previous: Option<&str>,
current: Option<&str>,
observer: &str,
) {
self.publish(
system_time_ms(SystemTime::now()),
"CoordinatorChanged",
current.unwrap_or(observer).to_owned(),
format!(
"previous={} current={}",
previous.unwrap_or("none"),
current.unwrap_or("none")
),
None,
);
}
fn publish_job_replaced(
&self,
at: SystemTime,
name: &str,
old_node: &str,
target: &str,
reason: &str,
generation: u64,
) {
self.publish(
system_time_ms(at),
"JobReplaced",
target.to_owned(),
format!("job={name} from={old_node} to={target} reason={reason}"),
Some(generation),
);
}
fn publish(
&self,
timestamp_ms: u64,
kind: &str,
node_id: String,
detail: String,
generation: Option<u64>,
) {
let mut next_sequence = self
.next_sequence
.lock()
.unwrap_or_else(|poison| poison.into_inner());
let event = ClusterEvent {
sequence: *next_sequence,
timestamp_ms,
kind: kind.to_owned(),
node_id,
detail,
generation,
};
*next_sequence = next_sequence.wrapping_add(1).max(1);
match self.topic.try_publish(event) {
Ok(())
| Err(TopicTryPublishError::Full(_))
| Err(TopicTryPublishError::Busy(_))
| Err(TopicTryPublishError::Closed(_)) => {}
}
}
fn close(&self) -> datum::StreamResult<()> {
if self.topic.is_closed() {
Ok(())
} else {
self.topic.close()
}
}
}
const fn member_event_kind(kind: MemberEventKind) -> &'static str {
match kind {
MemberEventKind::Initialized => "Initialized",
MemberEventKind::MemberJoining => "MemberJoining",
MemberEventKind::MemberUp => "MemberUp",
MemberEventKind::MemberUpdated => "MemberUpdated",
MemberEventKind::MemberUnreachable => "MemberUnreachable",
MemberEventKind::MemberReachable => "MemberReachable",
MemberEventKind::MemberLeaving => "MemberLeaving",
MemberEventKind::MemberExiting => "MemberExiting",
MemberEventKind::MemberDown => "MemberDown",
MemberEventKind::MemberRemoved => "MemberRemoved",
MemberEventKind::MemberRejoined => "MemberRejoined",
}
}
fn system_time_ms(time: SystemTime) -> u64 {
time.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_millis().min(u128::from(u64::MAX)) as u64)
.unwrap_or(0)
}
#[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(u64),
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(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct PeerSessionSnapshot {
member_incarnation: u64,
agent_addr: Option<SocketAddr>,
state: PeerSessionState,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct MetricRouteTarget {
member_incarnation: u64,
agent_addr: Option<SocketAddr>,
connection_generation: u64,
}
impl PeerSessionSnapshot {
fn metric_route_target(&self) -> Option<MetricRouteTarget> {
match &self.state {
PeerSessionState::Connected(connection_generation) => Some(MetricRouteTarget {
member_incarnation: self.member_incarnation,
agent_addr: self.agent_addr,
connection_generation: *connection_generation,
}),
PeerSessionState::Connecting
| PeerSessionState::BackingOff(_)
| PeerSessionState::Closed => None,
}
}
}
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>>,
},
SubscribeMetrics {
interval: Duration,
job_names: Vec<String>,
timeout: Duration,
reply: oneshot::Sender<Result<MetricSubscription, String>>,
},
RememberClusterAssignment {
instance_name: String,
assignment: Option<ClusterJobStart>,
tombstone: Option<AssignmentTombstone>,
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 subscribe_metrics(
&self,
node_id: &str,
interval: Duration,
job_names: Vec<String>,
timeout: Duration,
) -> Result<MetricSubscription, String> {
let sender = self.session_sender(node_id).await?;
let (reply, receiver) = oneshot::channel();
sender
.try_send(SessionCommand::SubscribeMetrics {
interval,
job_names,
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: Option<ClusterJobStart>,
tombstone: Option<AssignmentTombstone>,
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,
tombstone,
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_snapshots(&self) -> BTreeMap<String, PeerSessionSnapshot> {
let locked = self.inner.sessions.lock().await;
locked
.iter()
.map(|(node_id, session)| {
(
node_id.clone(),
PeerSessionSnapshot {
member_incarnation: session.member.incarnation,
agent_addr: session.member.agent_addr,
state: session.state.borrow().clone(),
},
)
})
.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;
let mut connection_generation = 0_u64;
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) => {
connection_generation = connection_generation.saturating_add(1).max(1);
let _ = state.send(PeerSessionState::Connected(connection_generation));
eprintln!(
"datum-agent INFO peer_session_connected node_id={local_node} peer={}",
member.node_id
);
backoff = config.reconnect_min_backoff;
if !run_connected_peer(
client,
pipe,
&mut commands,
&mut pipe_commands,
&mut stop,
)
.await
{
break;
}
eprintln!(
"datum-agent WARN peer_session_lost node_id={local_node} peer={}",
member.node_id
);
}
Err(error) => {
log_peer_connect_failed(&local_node, &member.node_id, &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);
log_peer_reconnecting(&local_node, &member.node_id, backoff);
}
},
Err(error) => {
log_peer_connect_failed(&local_node, &member.node_id, &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);
log_peer_reconnecting(&local_node, &member.node_id, backoff);
}
}
}
let _ = state.send(PeerSessionState::Closed);
eprintln!(
"datum-agent INFO peer_session_closed node_id={local_node} peer={}",
member.node_id
);
}
fn log_peer_connect_failed(local_node: &str, peer: &str, error: &str) {
eprintln!(
"datum-agent WARN peer_session_connect_failed node_id={local_node} peer={peer} error={error}"
);
}
fn log_peer_reconnecting(local_node: &str, peer: &str, backoff: Duration) {
eprintln!(
"datum-agent INFO peer_session_reconnecting node_id={local_node} peer={peer} backoff_ms={}",
backoff.as_millis()
);
}
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, reconnect) = classify_peer_request(
tokio::time::timeout(timeout, client.list_jobs()).await,
"ListJobs",
);
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;
let (result, reconnect) =
classify_peer_request(result, "SubmitClusterJob");
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;
let (result, reconnect) = classify_peer_request(result, "StartJob");
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, reconnect) = classify_peer_request(
tokio::time::timeout(timeout, request).await,
"JobStatus",
);
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, reconnect) = classify_peer_request(
tokio::time::timeout(timeout, request).await,
"DrainJob",
);
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, reconnect) = classify_peer_request(
tokio::time::timeout(timeout, request).await,
"StopJob",
);
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;
let (result, reconnect) =
classify_peer_request(result, "ListClusterJobs");
let _ = reply.send(result);
if reconnect {
return true;
}
}
SessionCommand::SubscribeMetrics {
interval,
job_names,
timeout,
reply,
} => {
let result = tokio::time::timeout(
timeout,
client.subscribe_local_metrics(
duration_millis_u64(interval),
job_names,
),
)
.await;
let (result, reconnect) =
classify_peer_request(result, "SubscribeMetrics");
let _ = reply.send(result);
if reconnect {
return true;
}
}
SessionCommand::RememberClusterAssignment {
instance_name,
assignment,
tombstone,
timeout,
reply,
} => {
let (tombstone_generation, tombstone_coordinator) = tombstone
.map(|tombstone| {
(tombstone.placement_generation, tombstone.coordinator_node)
})
.unwrap_or_default();
let result = tokio::time::timeout(
timeout,
client.remember_cluster_assignment(
instance_name,
assignment,
tombstone_generation,
tombstone_coordinator,
),
)
.await;
let (result, reconnect) =
classify_peer_request(result, "RememberClusterAssignment");
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;
let (result, reconnect) = classify_peer_request(result, "AllocateShard");
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;
let (result, reconnect) =
classify_peer_request(result, "RememberShardAllocations");
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;
let (result, reconnect) =
classify_peer_request(result, "GetShardAllocations");
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;
let (result, reconnect) =
classify_peer_request(result, "ForwardShardEnvelopes");
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;
let (result, reconnect) =
classify_peer_request(result, "CompleteShardingAsk");
let _ = reply.send(result);
if reconnect {
return true;
}
}
}
}
}
}
}
fn classify_peer_request<T>(
result: Result<Result<T, DcpError>, tokio::time::error::Elapsed>,
operation: &str,
) -> (Result<T, String>, bool) {
match result {
Ok(Ok(value)) => (Ok(value), false),
Err(_) => (Err(format!("peer {operation} timed out")), false),
Ok(Err(error)) => {
let reconnect = dcp_error_requires_reconnect(&error);
(Err(error.to_string()), reconnect)
}
}
}
fn dcp_error_requires_reconnect(error: &DcpError) -> bool {
match error {
DcpError::Closed | DcpError::Protocol(_) | DcpError::Io(_) | DcpError::Decode(_) => true,
DcpError::Response { .. }
| DcpError::Encode(_)
| DcpError::Agent(_)
| DcpError::Stream(_)
| DcpError::Join(_) => false,
}
}
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::SubscribeMetrics { 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: Option<ClusterJobStart>,
tombstone: Option<AssignmentTombstone>,
reply: oneshot::Sender<Result<(), DcpError>>,
},
RegisterRestarted {
instance_name: String,
assignment: ClusterJobStart,
timeout: Duration,
reply: oneshot::Sender<Result<(), DcpError>>,
},
AssignmentSyncCompleted {
node_id: String,
incarnation: u64,
instance_name: String,
sync_mark: AssignmentSyncMark,
ok: bool,
},
Tick,
Shutdown,
}
struct PlacementActor {
commands: mpsc::Sender<PlacementCommand>,
self_node: String,
agent_role: String,
state: Signal<ClusterState>,
registry: JobRegistryHandle,
sessions: NodeSessionManagerHandle,
factories: DcpJobFactories,
request_timeout: Duration,
active_coordinator: bool,
last_known_coordinator: Option<String>,
assignments: BTreeMap<String, ClusterJobMetadata>,
assignment_tombstones: BTreeMap<String, AssignmentTombstone>,
replicated_assignments: BTreeMap<String, PeerAssignmentSync>,
inflight_assignment_syncs: BTreeSet<AssignmentSyncKey>,
pending_replacements: BTreeMap<String, ClusterJobMetadata>,
assignment_updates: watch::Sender<Arc<BTreeMap<String, String>>>,
cluster_events: ClusterEventPublisher,
}
#[derive(Clone, Debug, PartialEq, Eq)]
struct AssignmentTombstone {
placement_generation: u64,
coordinator_node: String,
}
impl AssignmentTombstone {
fn from_metadata(metadata: &ClusterJobMetadata) -> Self {
Self {
placement_generation: metadata.placement_generation,
coordinator_node: metadata.coordinator_node.clone(),
}
}
fn from_wire(request: &RememberClusterAssignment) -> Self {
Self {
placement_generation: request.tombstone_placement_generation,
coordinator_node: request.tombstone_coordinator_node_id.clone(),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
enum AssignmentSyncMark {
Active {
placement_generation: u64,
assigned_node: String,
coordinator_node: String,
},
Tombstone {
placement_generation: u64,
coordinator_node: String,
},
}
impl AssignmentSyncMark {
fn active(metadata: &ClusterJobMetadata) -> Self {
Self::Active {
placement_generation: metadata.placement_generation,
assigned_node: metadata.assigned_node.clone(),
coordinator_node: metadata.coordinator_node.clone(),
}
}
fn tombstone(tombstone: &AssignmentTombstone) -> Self {
Self::Tombstone {
placement_generation: tombstone.placement_generation,
coordinator_node: tombstone.coordinator_node.clone(),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
struct PeerAssignmentSync {
incarnation: u64,
assignments: BTreeMap<String, AssignmentSyncMark>,
}
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
struct AssignmentSyncKey {
node_id: String,
incarnation: u64,
instance_name: String,
sync_mark: AssignmentSyncMark,
}
struct PlacementDependencies {
registry: JobRegistryHandle,
sessions: NodeSessionManagerHandle,
factories: DcpJobFactories,
cluster_events: ClusterEventPublisher,
}
impl PlacementCoordinatorHandle {
fn start(
self_node: String,
agent_role: String,
state: Signal<ClusterState>,
dependencies: PlacementDependencies,
request_timeout: Duration,
) -> ClusterAgentResult<Self> {
let (commands, receiver) = mpsc::channel(128);
let (assignment_updates, _assignments) = watch::channel(Arc::new(BTreeMap::new()));
let tasks = Arc::new(Mutex::new(Vec::new()));
let actor = PlacementActor {
commands: commands.clone(),
self_node,
agent_role,
state,
registry: dependencies.registry,
sessions: dependencies.sessions,
factories: dependencies.factories,
request_timeout,
active_coordinator: false,
last_known_coordinator: None,
assignments: BTreeMap::new(),
assignment_tombstones: BTreeMap::new(),
replicated_assignments: BTreeMap::new(),
inflight_assignment_syncs: BTreeSet::new(),
pending_replacements: BTreeMap::new(),
assignment_updates,
cluster_events: dependencies.cluster_events,
};
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: Option<ClusterJobStart>,
tombstone: Option<AssignmentTombstone>,
) -> Result<(), DcpError> {
self.request(Duration::from_millis(250), |reply| {
PlacementCommand::Remember {
instance_name,
assignment,
tombstone,
reply,
}
})
.await
}
async fn register_restarted_cluster_assignment(
&self,
instance_name: String,
assignment: ClusterJobStart,
timeout: Duration,
) -> Result<(), DcpError> {
self.request(timeout, |reply| PlacementCommand::RegisterRestarted {
instance_name,
assignment,
timeout,
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,
tombstone,
reply,
} => {
let result = actor.remember(instance_name, assignment, tombstone).await;
let _ = reply.send(result);
}
PlacementCommand::RegisterRestarted {
instance_name,
assignment,
timeout,
reply,
} => {
let result = actor
.register_restarted(instance_name, assignment, timeout)
.await;
let _ = reply.send(result);
}
PlacementCommand::AssignmentSyncCompleted {
node_id,
incarnation,
instance_name,
sync_mark,
ok,
} => {
actor.assignment_sync_completed(node_id, incarnation, instance_name, sync_mark, ok);
}
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.assignment_tombstones.remove(&request.instance_name);
self.pending_replacements.remove(&request.instance_name);
self.publish_assignments();
self.publish_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_inactive(&name, timeout).await?;
let mut status = self
.status_on_node(&metadata.assigned_node, name, timeout)
.await?;
status.coordinator_node_id.clone_from(&self.self_node);
Ok(status)
}
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?;
if let Some(removed) = self.assignments.remove(&name) {
self.record_assignment_tombstone(&name, &removed);
}
self.pending_replacements.remove(&name);
self.publish_assignments();
self.replicate_assignment_mutation(name, None, timeout);
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?;
if let Some(removed) = self.assignments.remove(&name) {
self.record_assignment_tombstone(&name, &removed);
}
self.pending_replacements.remove(&name);
self.publish_assignments();
self.replicate_assignment_mutation(name, None, timeout);
Ok(status)
}
async fn remember(
&mut self,
instance_name: String,
assignment: Option<ClusterJobStart>,
tombstone: Option<AssignmentTombstone>,
) -> Result<(), DcpError> {
let current_coordinator = self.current_placement_coordinator();
let Some(assignment) = assignment else {
let tombstone = tombstone.unwrap_or(AssignmentTombstone {
placement_generation: 0,
coordinator_node: String::new(),
});
ensure_mutation_from_current_coordinator(
&tombstone.coordinator_node,
current_coordinator.as_deref(),
)?;
if let Some(existing) = self.assignments.get(&instance_name)
&& !tombstone_dominates_assignment(
&tombstone,
existing,
current_coordinator.as_deref(),
)
{
return Ok(());
}
self.record_assignment_tombstone_value(
&instance_name,
tombstone.clone(),
current_coordinator.as_deref(),
);
self.stop_local_cluster_assignment(
&instance_name,
&tombstone,
current_coordinator.as_deref(),
)
.await?;
let removed = self.assignments.remove(&instance_name);
if removed.is_some() {
self.publish_assignments();
}
self.pending_replacements.remove(&instance_name);
return Ok(());
};
let metadata = cluster_metadata_from_wire(assignment)?;
ensure_mutation_from_current_coordinator(
&metadata.coordinator_node,
current_coordinator.as_deref(),
)?;
if let Some(tombstone) = self.assignment_tombstones.get(&instance_name)
&& !active_assignment_supersedes_tombstone(&metadata, tombstone)
{
return Err(stale_assignment_conflict(
&instance_name,
metadata.placement_generation,
tombstone,
));
}
if !remember_assignment(
&mut self.assignments,
instance_name.clone(),
metadata.clone(),
current_coordinator.as_deref(),
) {
return Ok(());
}
self.assignment_tombstones.remove(&instance_name);
self.pending_replacements.remove(&instance_name);
self.publish_assignments();
self.reconcile_local_assignment(&instance_name, &metadata)
.await?;
Ok(())
}
async fn register_restarted(
&mut self,
instance_name: String,
assignment: ClusterJobStart,
timeout: Duration,
) -> Result<(), DcpError> {
let mut metadata = cluster_metadata_from_wire(assignment)?;
if metadata.assigned_node != self.self_node {
return Err(DcpError::response(
ResponseStatus::Conflict,
"only the owning node can re-register restarted cluster intent",
));
}
if let Some(existing) = self.assignments.get(&instance_name).cloned()
&& active_assignment_blocks_local_restart(&metadata, &existing)
{
self.stop_stale_restarted_local_assignment(&instance_name, &existing)
.await?;
return Err(stale_restart_conflict(&instance_name, &metadata, &existing));
}
if let Some(tombstone) = self.assignment_tombstones.get(&instance_name)
&& metadata.placement_generation <= tombstone.placement_generation
{
metadata.placement_generation = tombstone.placement_generation.saturating_add(1);
metadata.history.push(ClusterPlacementHistory {
generation: metadata.placement_generation,
from_node: Some(self.self_node.clone()),
to_node: self.self_node.clone(),
reason: "local_restart_after_terminal_intent".to_owned(),
timestamp: SystemTime::now(),
});
}
let snapshot = self.state.get();
metadata.coordinator_node = snapshot
.placement_coordinator()
.map(|member| member.node_id.clone())
.ok_or_else(|| {
DcpError::response(
ResponseStatus::Failed,
"no placement coordinator is available",
)
})?;
drop(snapshot);
self.assignments
.insert(instance_name.clone(), metadata.clone());
self.assignment_tombstones.remove(&instance_name);
self.pending_replacements.remove(&instance_name);
self.publish_assignments();
self.publish_assignment(instance_name, metadata, timeout)
.await
}
async fn reconcile(&mut self) -> Result<(), DcpError> {
let snapshot = self.state.get();
let coordinator_id = snapshot
.placement_coordinator()
.map(|member| member.node_id.clone());
let is_coordinator = snapshot.is_placement_coordinator(&self.self_node);
drop(snapshot);
if coordinator_id != self.last_known_coordinator {
self.cluster_events.publish_coordinator_changed(
self.last_known_coordinator.as_deref(),
coordinator_id.as_deref(),
&self.self_node,
);
if let Some(coordinator_id) = &coordinator_id {
eprintln!(
"datum-agent INFO coordinator_changed node_id={} new_coordinator={coordinator_id}",
self.self_node
);
}
self.last_known_coordinator = coordinator_id;
}
if !is_coordinator {
if self.active_coordinator {
eprintln!(
"datum-agent INFO placement_coordinator_inactive node_id={}",
self.self_node
);
}
self.active_coordinator = false;
return Ok(());
}
self.ensure_active_coordinator(self.request_timeout).await?;
self.replace_down_members(self.request_timeout).await?;
self.sync_assignments_to_peers(self.request_timeout);
Ok(())
}
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;
eprintln!(
"datum-agent INFO placement_coordinator_active node_id={}",
self.self_node
);
}
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 assignment_or_inactive(
&mut self,
name: &str,
timeout: Duration,
) -> Result<ClusterJobMetadata, DcpError> {
if let Some(metadata) = self.assignments.get(name).cloned() {
return Ok(metadata);
}
let inactive = self.rebuild_from_registries(timeout).await?;
self.assignments
.get(name)
.or_else(|| inactive.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<BTreeMap<String, ClusterJobMetadata>, DcpError> {
let statuses = self.collect_job_statuses(timeout).await?;
let mut running = BTreeSet::new();
let mut inactive = BTreeSet::new();
let mut inactive_metadata = BTreeMap::new();
for status in statuses {
if !status.cluster_job {
continue;
}
if status.desired_state == "Running"
&& let Some(mut metadata) = metadata_from_status(&status)?
{
metadata.coordinator_node.clone_from(&self.self_node);
if let Some(tombstone) = self.assignment_tombstones.get(&status.name)
&& registry_observation_is_tombstoned(tombstone, &metadata)
{
let retained_assignment_is_newer =
self.assignments.get(&status.name).is_some_and(|existing| {
active_assignment_supersedes_tombstone(existing, tombstone)
});
if !retained_assignment_is_newer && !running.contains(&status.name) {
inactive.insert(status.name.clone());
}
continue;
}
running.insert(status.name.clone());
inactive.remove(&status.name);
inactive_metadata.remove(&status.name);
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.assignment_tombstones.remove(&status.name);
self.pending_replacements.remove(&status.name);
self.assignments.insert(status.name.clone(), metadata);
self.publish_assignments();
}
} else if !running.contains(&status.name) {
if let Some(mut metadata) = metadata_from_status(&status)? {
metadata.coordinator_node.clone_from(&self.self_node);
inactive_metadata.insert(status.name.clone(), metadata);
}
inactive.insert(status.name);
}
}
for name in &inactive {
let removed = self.assignments.remove(name);
if let Some(metadata) = inactive_metadata.get(name) {
self.record_assignment_tombstone(name, metadata);
} else if let Some(removed) = removed.as_ref() {
self.record_assignment_tombstone(name, removed);
}
if removed.is_some() {
self.pending_replacements.remove(name);
self.publish_assignments();
}
}
for metadata in self.assignments.values_mut() {
metadata.coordinator_node.clone_from(&self.self_node);
}
self.publish_assignments();
let assignments = self
.assignments
.iter()
.map(|(name, metadata)| (name.clone(), metadata.clone()))
.collect::<Vec<_>>();
for (name, metadata) in assignments {
self.publish_assignment(name, metadata, timeout).await?;
}
for name in inactive {
if let Some(metadata) = inactive_metadata.get(&name).cloned() {
self.publish_assignment(name.clone(), metadata, timeout)
.await?;
}
self.replicate_assignment_mutation(name, None, timeout);
}
Ok(inactive_metadata)
}
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, current) in to_replace {
let metadata = if let Some(pending) = self.pending_replacements.get(&name)
&& pending
.history
.last()
.and_then(|event| event.from_node.as_ref())
.is_some_and(|node| down_nodes.contains(node))
&& self.is_replacement_target_eligible(pending)
{
pending.clone()
} else {
let mut metadata = current.clone();
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,
reason: format!("node_down:{old_node}"),
timestamp: SystemTime::now(),
});
self.pending_replacements
.insert(name.clone(), metadata.clone());
metadata
};
let params = metadata.params.clone().into_iter().collect();
let factory_name = metadata.factory_name.clone();
self.ensure_started_on_node(
&metadata.assigned_node,
factory_name,
name.clone(),
params,
metadata.clone(),
timeout,
)
.await?;
self.assignments.insert(name.clone(), metadata.clone());
self.assignment_tombstones.remove(&name);
self.pending_replacements.remove(&name);
self.publish_assignments();
self.publish_committed_replacement(&name, &metadata);
self.publish_assignment(name, metadata, timeout).await?;
}
Ok(())
}
fn is_replacement_target_eligible(&self, metadata: &ClusterJobMetadata) -> bool {
let snapshot = self.state.get();
snapshot
.member(&metadata.assigned_node)
.is_some_and(|member| member.state == MemberState::Up && !member.unreachable)
}
fn current_placement_coordinator(&self) -> Option<String> {
self.state
.get()
.placement_coordinator()
.map(|member| member.node_id.clone())
}
fn publish_committed_replacement(&self, name: &str, metadata: &ClusterJobMetadata) {
let Some(history) = metadata.history.last() else {
return;
};
let Some(old_node) = history.from_node.as_deref() else {
return;
};
eprintln!(
"datum-agent INFO job_replaced job={name} from_node={old_node} to_node={} reason={} generation={}",
metadata.assigned_node, history.reason, metadata.placement_generation
);
self.cluster_events.publish_job_replaced(
history.timestamp,
name,
old_node,
&metadata.assigned_node,
&history.reason,
metadata.placement_generation,
);
}
fn publish_assignments(&self) {
let assignments = self
.assignments
.iter()
.map(|(name, metadata)| (name.clone(), metadata.assigned_node.clone()))
.collect();
self.assignment_updates.send_replace(Arc::new(assignments));
}
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 ensure_started_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> {
match self
.start_on_node(
node_id,
factory_name,
instance_name.clone(),
params,
metadata.clone(),
timeout,
)
.await
{
Ok(status) => Ok(status),
Err(error) if placement_start_may_have_committed(&error) => {
match self.status_on_node(node_id, instance_name, timeout).await {
Ok(status) if status_matches_assignment(&status, &metadata) => Ok(status),
_ => Err(error),
}
}
Err(error) => Err(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() {
return Err(DcpError::response(
ResponseStatus::DeadlineExceeded,
"registry reconciliation timed out",
));
}
match tokio::time::timeout(remaining, task).await {
Ok(Ok(Ok(mut peer_statuses))) => statuses.append(&mut peer_statuses),
Ok(Ok(Err(message))) => return Err(session_error(message)),
Ok(Err(error)) => {
return Err(DcpError::response(
ResponseStatus::Failed,
format!("registry reconciliation task failed: {error}"),
));
}
Err(_) => {
return Err(DcpError::response(
ResponseStatus::DeadlineExceeded,
"registry reconciliation timed out",
));
}
}
}
Ok(statuses)
}
async fn publish_assignment(
&mut self,
instance_name: String,
metadata: ClusterJobMetadata,
timeout: Duration,
) -> Result<(), DcpError> {
self.reconcile_local_assignment(&instance_name, &metadata)
.await?;
self.replicate_assignment_mutation(
instance_name,
Some(wire_cluster_job_start(&metadata)),
timeout,
);
Ok(())
}
async fn reconcile_local_assignment(
&self,
instance_name: &str,
metadata: &ClusterJobMetadata,
) -> Result<(), DcpError> {
if metadata.assigned_node == self.self_node {
return self.update_local_assignment(instance_name, metadata).await;
}
self.stop_stale_local_assignment(instance_name, metadata)
.await
}
async fn update_local_assignment(
&self,
instance_name: &str,
metadata: &ClusterJobMetadata,
) -> Result<(), DcpError> {
if metadata.assigned_node != self.self_node {
return Ok(());
}
let registry = self.registry.clone();
let instance_name = instance_name.to_owned();
let metadata = metadata.clone();
tokio::task::spawn_blocking(move || {
registry.update_cluster_metadata(instance_name, metadata)
})
.await??;
Ok(())
}
async fn stop_stale_local_assignment(
&self,
instance_name: &str,
metadata: &ClusterJobMetadata,
) -> Result<(), DcpError> {
let registry = self.registry.clone();
let instance_name = instance_name.to_owned();
let self_node = self.self_node.clone();
let placement_generation = metadata.placement_generation;
tokio::task::spawn_blocking(move || -> Result<(), AgentError> {
let status = match registry.status(instance_name.clone()) {
Ok(status) => status,
Err(AgentError::JobNotFound(_)) => return Ok(()),
Err(error) => return Err(error),
};
let Some(local_metadata) = status.cluster.as_ref() else {
return Ok(());
};
if local_metadata.assigned_node == self_node
&& local_metadata.placement_generation < placement_generation
&& should_stop_stale_local_job(status.state, status.desired_state)
{
let _ = registry.stop(instance_name)?;
}
Ok(())
})
.await??;
Ok(())
}
async fn stop_local_cluster_assignment(
&self,
instance_name: &str,
tombstone: &AssignmentTombstone,
current_coordinator: Option<&str>,
) -> Result<(), DcpError> {
let registry = self.registry.clone();
let instance_name = instance_name.to_owned();
let tombstone = tombstone.clone();
let current_coordinator = current_coordinator.map(str::to_owned);
tokio::task::spawn_blocking(move || -> Result<(), AgentError> {
let status = match registry.status(instance_name.clone()) {
Ok(status) => status,
Err(AgentError::JobNotFound(_)) => return Ok(()),
Err(error) => return Err(error),
};
let Some(mut local_metadata) = status.cluster.clone() else {
return Ok(());
};
if !tombstone_dominates_assignment(
&tombstone,
&local_metadata,
current_coordinator.as_deref(),
) {
return Ok(());
}
local_metadata
.coordinator_node
.clone_from(&tombstone.coordinator_node);
if local_metadata.placement_generation < tombstone.placement_generation {
local_metadata.placement_generation = tombstone.placement_generation;
}
registry.update_cluster_metadata(instance_name.clone(), local_metadata)?;
if should_stop_stale_local_job(status.state, status.desired_state) {
let _ = registry.stop(instance_name)?;
}
Ok(())
})
.await??;
Ok(())
}
async fn stop_stale_restarted_local_assignment(
&self,
instance_name: &str,
known_assignment: &ClusterJobMetadata,
) -> Result<(), DcpError> {
let registry = self.registry.clone();
let instance_name = instance_name.to_owned();
let self_node = self.self_node.clone();
let known_assignment = known_assignment.clone();
tokio::task::spawn_blocking(move || -> Result<(), AgentError> {
let status = match registry.status(instance_name.clone()) {
Ok(status) => status,
Err(AgentError::JobNotFound(_)) => return Ok(()),
Err(error) => return Err(error),
};
let Some(local_metadata) = status.cluster.as_ref() else {
return Ok(());
};
let stale_local_owner = local_metadata.assigned_node == self_node
&& (local_metadata.placement_generation < known_assignment.placement_generation
|| (local_metadata.placement_generation
== known_assignment.placement_generation
&& known_assignment.assigned_node != self_node));
if stale_local_owner && should_stop_stale_local_job(status.state, status.desired_state)
{
let _ = registry.stop(instance_name)?;
}
Ok(())
})
.await??;
Ok(())
}
fn replicate_assignment_mutation(
&mut self,
instance_name: String,
assignment: Option<ClusterJobStart>,
timeout: Duration,
) {
let peers = self.assignment_sync_peers();
let sync_mark = self.assignment_sync_mark(&instance_name, assignment.as_ref());
let timeout = timeout.min(Duration::from_millis(250));
for (node_id, incarnation) in peers {
self.sync_assignment_to_peer(
node_id,
incarnation,
instance_name.clone(),
assignment.clone(),
sync_mark.clone(),
timeout,
)
}
}
fn sync_assignments_to_peers(&mut self, timeout: Duration) {
let peers = self.assignment_sync_peers();
if peers.is_empty() {
return;
}
let timeout = timeout.min(Duration::from_millis(250));
let assignments = self
.assignments
.iter()
.map(|(name, metadata)| {
(
name.clone(),
Some(wire_cluster_job_start(metadata)),
AssignmentSyncMark::active(metadata),
)
})
.collect::<Vec<_>>();
let tombstones = self
.assignment_tombstones
.iter()
.filter(|(name, _tombstone)| !self.assignments.contains_key(*name))
.map(|(name, tombstone)| (name.clone(), None, AssignmentSyncMark::tombstone(tombstone)))
.collect::<Vec<_>>();
for (node_id, incarnation) in peers {
for (name, assignment, sync_mark) in assignments.iter().chain(tombstones.iter()) {
self.sync_assignment_to_peer(
node_id.clone(),
incarnation,
name.clone(),
assignment.clone(),
sync_mark.clone(),
timeout,
)
}
}
}
fn assignment_sync_peers(&mut self) -> Vec<(String, u64)> {
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(), member.incarnation))
.collect::<BTreeMap<_, _>>();
drop(snapshot);
self.replicated_assignments.retain(|node_id, sync| {
peers
.get(node_id)
.is_some_and(|incarnation| *incarnation == sync.incarnation)
});
self.inflight_assignment_syncs.retain(|key| {
peers
.get(&key.node_id)
.is_some_and(|incarnation| *incarnation == key.incarnation)
});
peers.into_iter().collect()
}
fn sync_assignment_to_peer(
&mut self,
node_id: String,
incarnation: u64,
instance_name: String,
assignment: Option<ClusterJobStart>,
sync_mark: AssignmentSyncMark,
timeout: Duration,
) {
if self
.replicated_assignments
.get(&node_id)
.filter(|sync| sync.incarnation == incarnation)
.and_then(|sync| sync.assignments.get(&instance_name))
== Some(&sync_mark)
{
return;
}
let sync_key = AssignmentSyncKey {
node_id: node_id.clone(),
incarnation,
instance_name: instance_name.clone(),
sync_mark: sync_mark.clone(),
};
if !self.inflight_assignment_syncs.insert(sync_key) {
return;
}
let sessions = self.sessions.clone();
let commands = self.commands.clone();
let tombstone = self
.assignment_tombstone_for_sync(&instance_name, assignment.as_ref())
.cloned();
tokio::spawn(async move {
let ok = sessions
.remember_cluster_assignment(
&node_id,
instance_name.clone(),
assignment,
tombstone,
timeout,
)
.await
.is_ok();
let _ = commands
.send(PlacementCommand::AssignmentSyncCompleted {
node_id,
incarnation,
instance_name,
sync_mark,
ok,
})
.await;
});
}
fn assignment_sync_completed(
&mut self,
node_id: String,
incarnation: u64,
instance_name: String,
sync_mark: AssignmentSyncMark,
ok: bool,
) {
self.inflight_assignment_syncs.remove(&AssignmentSyncKey {
node_id: node_id.clone(),
incarnation,
instance_name: instance_name.clone(),
sync_mark: sync_mark.clone(),
});
if ok {
if self.current_assignment_sync_mark(&instance_name) != sync_mark {
return;
}
self.replicated_assignments
.entry(node_id)
.and_modify(|sync| {
if sync.incarnation != incarnation {
sync.incarnation = incarnation;
sync.assignments.clear();
}
sync.assignments
.insert(instance_name.clone(), sync_mark.clone());
})
.or_insert_with(|| PeerAssignmentSync {
incarnation,
assignments: BTreeMap::from([(instance_name, sync_mark)]),
});
}
}
fn current_assignment_sync_mark(&self, instance_name: &str) -> AssignmentSyncMark {
self.assignments
.get(instance_name)
.map(AssignmentSyncMark::active)
.unwrap_or_else(|| {
self.assignment_tombstones
.get(instance_name)
.map(AssignmentSyncMark::tombstone)
.unwrap_or_else(|| {
AssignmentSyncMark::tombstone(&AssignmentTombstone {
placement_generation: 0,
coordinator_node: String::new(),
})
})
})
}
fn assignment_tombstone_for_sync(
&self,
instance_name: &str,
assignment: Option<&ClusterJobStart>,
) -> Option<&AssignmentTombstone> {
assignment
.is_none()
.then(|| self.assignment_tombstones.get(instance_name))
.flatten()
}
fn assignment_sync_mark(
&self,
instance_name: &str,
assignment: Option<&ClusterJobStart>,
) -> AssignmentSyncMark {
match assignment {
Some(assignment) => AssignmentSyncMark::Active {
placement_generation: assignment.placement_generation,
assigned_node: assignment.assigned_node_id.clone(),
coordinator_node: assignment.coordinator_node_id.clone(),
},
None => self
.assignment_tombstones
.get(instance_name)
.map(AssignmentSyncMark::tombstone)
.unwrap_or_else(|| {
AssignmentSyncMark::tombstone(&AssignmentTombstone {
placement_generation: 0,
coordinator_node: String::new(),
})
}),
}
}
fn record_assignment_tombstone(&mut self, instance_name: &str, metadata: &ClusterJobMetadata) {
let current_coordinator = self.current_placement_coordinator();
self.record_assignment_tombstone_value(
instance_name,
AssignmentTombstone::from_metadata(metadata),
current_coordinator.as_deref(),
);
}
fn record_assignment_tombstone_value(
&mut self,
instance_name: &str,
tombstone: AssignmentTombstone,
current_coordinator: Option<&str>,
) {
self.assignment_tombstones
.entry(instance_name.to_owned())
.and_modify(|existing| {
if tombstone_supersedes(&tombstone, existing, current_coordinator) {
*existing = tombstone.clone();
}
})
.or_insert(tombstone);
}
}
fn remember_assignment(
assignments: &mut BTreeMap<String, ClusterJobMetadata>,
instance_name: String,
metadata: ClusterJobMetadata,
current_coordinator: Option<&str>,
) -> bool {
let accept = assignments.get(&instance_name).is_none_or(|existing| {
metadata.placement_generation > existing.placement_generation
|| (metadata.placement_generation == existing.placement_generation
&& metadata.assigned_node == existing.assigned_node
&& (metadata.coordinator_node == existing.coordinator_node
|| current_coordinator
.is_some_and(|coordinator| coordinator == metadata.coordinator_node)))
});
if accept {
assignments.insert(instance_name, metadata);
}
accept
}
fn ensure_mutation_from_current_coordinator(
mutation_coordinator: &str,
current_coordinator: Option<&str>,
) -> Result<(), DcpError> {
if mutation_coordinator.is_empty() {
return Err(DcpError::response(
ResponseStatus::Conflict,
"assignment mutation is missing coordinator fence",
));
}
if current_coordinator == Some(mutation_coordinator) {
return Ok(());
}
Err(DcpError::response(
ResponseStatus::Conflict,
format!(
"assignment mutation from non-current coordinator: mutation={mutation_coordinator} current={}",
current_coordinator.unwrap_or("<none>")
),
))
}
fn tombstone_dominates_assignment(
tombstone: &AssignmentTombstone,
metadata: &ClusterJobMetadata,
current_coordinator: Option<&str>,
) -> bool {
tombstone.placement_generation > metadata.placement_generation
|| (tombstone.placement_generation == metadata.placement_generation
&& current_coordinator == Some(tombstone.coordinator_node.as_str()))
}
fn active_assignment_supersedes_tombstone(
metadata: &ClusterJobMetadata,
tombstone: &AssignmentTombstone,
) -> bool {
metadata.placement_generation > tombstone.placement_generation
}
fn registry_observation_is_tombstoned(
tombstone: &AssignmentTombstone,
metadata: &ClusterJobMetadata,
) -> bool {
!active_assignment_supersedes_tombstone(metadata, tombstone)
}
fn active_assignment_blocks_local_restart(
restarted: &ClusterJobMetadata,
existing: &ClusterJobMetadata,
) -> bool {
existing.placement_generation > restarted.placement_generation
|| (existing.placement_generation == restarted.placement_generation
&& existing.assigned_node != restarted.assigned_node)
}
fn stale_assignment_conflict(
instance_name: &str,
placement_generation: u64,
tombstone: &AssignmentTombstone,
) -> DcpError {
DcpError::response(
ResponseStatus::Conflict,
format!(
"cluster assignment {instance_name} generation {placement_generation} does not supersede tombstone generation {} from coordinator {}",
tombstone.placement_generation, tombstone.coordinator_node
),
)
}
fn stale_restart_conflict(
instance_name: &str,
restarted: &ClusterJobMetadata,
existing: &ClusterJobMetadata,
) -> DcpError {
DcpError::response(
ResponseStatus::Conflict,
format!(
"local restart for cluster assignment {instance_name} generation {} on {} is superseded by active generation {} on {}",
restarted.placement_generation,
restarted.assigned_node,
existing.placement_generation,
existing.assigned_node
),
)
}
fn placement_start_may_have_committed(error: &DcpError) -> bool {
match error {
DcpError::Closed | DcpError::Protocol(_) | DcpError::Io(_) | DcpError::Decode(_) => true,
DcpError::Response { status, message } => {
matches!(
status,
ResponseStatus::Conflict
| ResponseStatus::DeadlineExceeded
| ResponseStatus::Failed
) || message.contains("peer StartJob timed out")
|| message.contains("DCP connection closed")
}
DcpError::Agent(AgentError::JobAlreadyExists(_))
| DcpError::Agent(AgentError::JobAlreadyRunning(_)) => true,
DcpError::Encode(_) | DcpError::Agent(_) | DcpError::Stream(_) | DcpError::Join(_) => false,
}
}
fn status_matches_assignment(
status: &crate::dcp::proto::JobStatus,
metadata: &ClusterJobMetadata,
) -> bool {
status.cluster_job
&& status.desired_state == "Running"
&& status.placement_node_id == metadata.assigned_node
&& status.placement_generation == metadata.placement_generation
&& status.coordinator_node_id == metadata.coordinator_node
}
const fn should_stop_stale_local_job(state: JobState, desired_state: DesiredJobState) -> bool {
matches!(desired_state, DesiredJobState::Running)
&& matches!(
state,
JobState::Starting | JobState::Running | JobState::BackingOff
)
}
fn tombstone_supersedes(
candidate: &AssignmentTombstone,
existing: &AssignmentTombstone,
current_coordinator: Option<&str>,
) -> bool {
candidate.placement_generation > existing.placement_generation
|| (candidate.placement_generation == existing.placement_generation
&& !candidate.coordinator_node.is_empty()
&& (candidate.coordinator_node == existing.coordinator_node
|| current_coordinator == Some(candidate.coordinator_node.as_str())))
}
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,
cluster_events: ClusterEventPublisher,
}
impl ClusterViewProvider for ClusterView {
fn subscribe_cluster_events(&self) -> crate::dcp::DcpResult<mpsc::Receiver<ClusterEvent>> {
self.cluster_events.subscribe()
}
fn subscribe_cluster_metrics(
&self,
interval: Duration,
) -> Option<mpsc::Receiver<MetricSample>> {
Some(spawn_cluster_metrics_proxy(
self.self_node.clone(),
self.sessions.clone(),
interval,
))
}
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 })
}
fn register_restarted_cluster_assignment(
&self,
instance_name: String,
assignment: ClusterJobStart,
timeout: Duration,
) -> crate::dcp::server::ClusterViewFuture<'_, ()> {
Box::pin(async move {
self.placement
.register_restarted_cluster_assignment(instance_name, assignment, timeout)
.await
})
}
}
struct RemoteMetricRoute {
target: MetricRouteTarget,
cancel: oneshot::Sender<()>,
task: JoinHandle<()>,
}
struct RemoteMetricUpdate {
node_id: String,
sample: MetricSample,
}
struct RemoteMetricSubscriptionAttempt {
node_id: String,
target: MetricRouteTarget,
result: Result<MetricSubscription, String>,
}
#[derive(Default)]
struct RemoteMetricProxyState {
routes: BTreeMap<String, RemoteMetricRoute>,
latest: BTreeMap<String, MetricSample>,
pending: BTreeMap<String, MetricRouteTarget>,
wanted: BTreeMap<String, MetricRouteTarget>,
}
fn spawn_cluster_metrics_proxy(
self_node: String,
sessions: NodeSessionManagerHandle,
interval: Duration,
) -> mpsc::Receiver<MetricSample> {
let (output, receiver) = mpsc::channel(8);
tokio::spawn(run_cluster_metrics_proxy(
self_node, sessions, interval, output,
));
receiver
}
async fn run_cluster_metrics_proxy(
self_node: String,
sessions: NodeSessionManagerHandle,
interval: Duration,
output: mpsc::Sender<MetricSample>,
) {
let (updates, mut update_receiver) = mpsc::channel(32);
let (attempts, mut attempt_receiver) = mpsc::channel(32);
let mut proxy = RemoteMetricProxyState::default();
let retry_interval = interval
.max(Duration::from_millis(50))
.min(Duration::from_secs(1));
let mut retry = tokio::time::interval(retry_interval);
let mut emit = tokio::time::interval(interval);
reconcile_metric_routes(&self_node, &sessions, interval, &attempts, &mut proxy).await;
loop {
tokio::select! {
_ = output.closed() => break,
_ = retry.tick() => {
reconcile_metric_routes(
&self_node,
&sessions,
interval,
&attempts,
&mut proxy,
).await;
}
attempt = attempt_receiver.recv() => {
if let Some(attempt) = attempt {
install_metric_route(
attempt,
&mut proxy,
&updates,
).await;
}
}
update = update_receiver.recv() => {
if let Some(update) = update
&& proxy.routes.contains_key(&update.node_id)
{
proxy.latest.insert(update.node_id, update.sample);
}
}
_ = emit.tick() => {
let streams = proxy.latest
.values()
.flat_map(|sample| sample.streams.iter().cloned())
.collect();
let nodes = proxy.latest
.values()
.flat_map(|sample| sample.nodes.iter().cloned())
.collect();
let timestamp_ms = proxy.latest
.values()
.map(|sample| sample.timestamp_ms)
.max()
.unwrap_or(0);
match output.try_send(MetricSample {
timestamp_ms,
streams,
nodes,
}) {
Ok(()) | Err(mpsc::error::TrySendError::Full(_)) => {}
Err(mpsc::error::TrySendError::Closed(_)) => break,
}
}
}
}
drop(attempt_receiver);
drop(attempts);
cancel_metric_routes(proxy.routes).await;
}
async fn reconcile_metric_routes(
self_node: &str,
sessions: &NodeSessionManagerHandle,
interval: Duration,
attempts: &mpsc::Sender<RemoteMetricSubscriptionAttempt>,
proxy: &mut RemoteMetricProxyState,
) {
let next_wanted = sessions
.session_snapshots()
.await
.into_iter()
.filter_map(|(node_id, snapshot)| {
if node_id == self_node {
return None;
}
snapshot
.metric_route_target()
.map(|target| (node_id, target))
})
.collect::<BTreeMap<_, _>>();
proxy.wanted = next_wanted;
let stale = proxy
.routes
.iter()
.filter(|(node_id, route)| {
proxy.wanted.get(*node_id) != Some(&route.target) || route.task.is_finished()
})
.map(|(node_id, _)| node_id.clone())
.collect::<Vec<_>>();
for node_id in stale {
proxy.latest.remove(&node_id);
if let Some(route) = proxy.routes.remove(&node_id) {
cancel_metric_route(route).await;
}
}
for (node_id, target) in &proxy.wanted {
if proxy.routes.contains_key(node_id) || proxy.pending.get(node_id) == Some(target) {
continue;
}
proxy.pending.insert(node_id.clone(), target.clone());
let timeout = sessions.inner.config.request_timeout;
let sessions = sessions.clone();
let attempts = attempts.clone();
let node_id = node_id.clone();
let target = target.clone();
let _task = tokio::spawn(async move {
let result = sessions
.subscribe_metrics(&node_id, interval, Vec::new(), timeout)
.await;
let attempt = RemoteMetricSubscriptionAttempt {
node_id,
target,
result,
};
if let Err(error) = attempts.send(attempt).await
&& let Ok(subscription) = error.0.result
{
let _ = subscription.cancel().await;
}
});
}
}
async fn install_metric_route(
attempt: RemoteMetricSubscriptionAttempt,
proxy: &mut RemoteMetricProxyState,
updates: &mpsc::Sender<RemoteMetricUpdate>,
) {
let RemoteMetricSubscriptionAttempt {
node_id,
target,
result,
} = attempt;
if proxy.pending.get(&node_id) == Some(&target) {
proxy.pending.remove(&node_id);
}
let is_current =
proxy.wanted.get(&node_id) == Some(&target) && !proxy.routes.contains_key(&node_id);
let Ok(subscription) = result else {
return;
};
if !is_current {
let _ = subscription.cancel().await;
return;
}
let (cancel, cancel_receiver) = oneshot::channel();
let task = tokio::spawn(pump_remote_metrics(
node_id.clone(),
subscription,
updates.clone(),
cancel_receiver,
));
proxy.routes.insert(
node_id,
RemoteMetricRoute {
target,
cancel,
task,
},
);
}
async fn pump_remote_metrics(
node_id: String,
mut subscription: MetricSubscription,
updates: mpsc::Sender<RemoteMetricUpdate>,
mut cancel: oneshot::Receiver<()>,
) {
loop {
tokio::select! {
_ = &mut cancel => {
let _ = subscription.cancel().await;
break;
}
sample = subscription.recv() => {
let Some(sample) = sample else {
break;
};
match updates.try_send(RemoteMetricUpdate {
node_id: node_id.clone(),
sample,
}) {
Ok(()) | Err(mpsc::error::TrySendError::Full(_)) => {}
Err(mpsc::error::TrySendError::Closed(_)) => {
let _ = subscription.cancel().await;
break;
}
}
}
}
}
}
async fn cancel_metric_route(route: RemoteMetricRoute) {
let _ = route.cancel.send(());
let _ = route.task.await;
}
async fn cancel_metric_routes(routes: BTreeMap<String, RemoteMetricRoute>) {
for (_, route) in routes {
cancel_metric_route(route).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 coordinator_node_id = snapshot
.placement_coordinator()
.map(|member| member.node_id.clone())
.unwrap_or_default();
let session_states = self.sessions.session_snapshots().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)
.map(|snapshot| snapshot.state.as_text())
.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,
coordinator_node_id,
})
}
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 tombstone = request
.assignment
.is_none()
.then(|| AssignmentTombstone::from_wire(&request));
self.placement
.remember_cluster_assignment(request.instance_name, request.assignment, tombstone)
.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());
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{JobMat, JobRegistry, JobSpec};
use datum::{Keep, Source};
use prost::Message as _;
use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _, DuplexStream};
#[tokio::test]
async fn connected_peer_keeps_session_after_request_timeout() {
let (client_io, mut peer_io) = tokio::io::duplex(4096);
let (client_reader, client_writer) = tokio::io::split(client_io);
let client = DcpClient::from_test_parts(client_reader, client_writer);
let peer_task = tokio::spawn(async move {
let first_request = read_test_request(&mut peer_io).await;
tokio::time::sleep(Duration::from_millis(100)).await;
write_empty_job_list(&mut peer_io, first_request).await;
let second_request = read_test_request(&mut peer_io).await;
write_empty_job_list(&mut peer_io, second_request).await;
});
let (pipe_io, _pipe_peer_io) = tokio::io::duplex(4096);
let (pipe_reader, pipe_writer) = tokio::io::split(pipe_io);
let pipe = ShardPipeClient::from_test_parts(pipe_reader, pipe_writer);
let (command_sender, mut commands) = mpsc::channel(2);
let (_pipe_sender, mut pipe_commands) = mpsc::channel(1);
let (stop_sender, mut stop) = watch::channel(false);
let session_task = tokio::spawn(async move {
run_connected_peer(client, pipe, &mut commands, &mut pipe_commands, &mut stop).await
});
let (first_reply, first_result) = oneshot::channel();
command_sender
.send(SessionCommand::ListJobs {
timeout: Duration::from_millis(20),
reply: first_reply,
})
.await
.expect("first command accepted");
assert_eq!(
first_result.await.expect("first reply sent"),
Err("peer ListJobs timed out".to_owned())
);
let (second_reply, second_result) = oneshot::channel();
command_sender
.send(SessionCommand::ListJobs {
timeout: Duration::from_secs(1),
reply: second_reply,
})
.await
.expect("same session accepts second command");
assert_eq!(
second_result.await.expect("second reply sent"),
Ok(Vec::new())
);
stop_sender.send(true).expect("session still running");
assert!(
!session_task.await.expect("session task joins"),
"orderly stop must not be reported as a lost peer session"
);
peer_task.await.expect("test peer joins");
}
#[tokio::test]
async fn connected_peer_reconnects_after_transport_error() {
let (client_io, mut peer_io) = tokio::io::duplex(4096);
let (client_reader, client_writer) = tokio::io::split(client_io);
let client = DcpClient::from_test_parts(client_reader, client_writer);
let peer_task = tokio::spawn(async move {
let _request_id = read_test_request(&mut peer_io).await;
drop(peer_io);
});
let (pipe_io, _pipe_peer_io) = tokio::io::duplex(4096);
let (pipe_reader, pipe_writer) = tokio::io::split(pipe_io);
let pipe = ShardPipeClient::from_test_parts(pipe_reader, pipe_writer);
let (command_sender, mut commands) = mpsc::channel(1);
let (_pipe_sender, mut pipe_commands) = mpsc::channel(1);
let (_stop_sender, mut stop) = watch::channel(false);
let session_task = tokio::spawn(async move {
run_connected_peer(client, pipe, &mut commands, &mut pipe_commands, &mut stop).await
});
let (reply, result) = oneshot::channel();
command_sender
.send(SessionCommand::ListJobs {
timeout: Duration::from_secs(1),
reply,
})
.await
.expect("command accepted");
let error = result
.await
.expect("transport error reply sent")
.expect_err("closed transport fails request");
assert!(!error.contains("timed out"), "unexpected error: {error}");
assert!(
tokio::time::timeout(Duration::from_secs(1), session_task)
.await
.expect("connected loop exits promptly")
.expect("session task joins"),
"transport error must ask the outer peer-session loop to reconnect"
);
peer_task.await.expect("test peer joins");
}
#[test]
fn peer_response_errors_do_not_request_reconnect() {
for status in [
ResponseStatus::Conflict,
ResponseStatus::NotFound,
ResponseStatus::DeadlineExceeded,
ResponseStatus::BadRequest,
] {
let (result, reconnect) = classify_peer_request::<()>(
Ok(Err(DcpError::response(status, "application response"))),
"StartJob",
);
assert!(result.is_err());
assert!(
!reconnect,
"DCP response {status:?} must not tear down the peer session"
);
}
let (_result, reconnect) =
classify_peer_request::<()>(Ok(Err(DcpError::Closed)), "StartJob");
assert!(reconnect, "closed transports still reconnect");
}
async fn read_test_request(stream: &mut DuplexStream) -> u64 {
let frame = read_test_frame(stream).await;
let request = match frame.frame {
Some(crate::dcp::proto::dcp_frame::Frame::Request(request)) => request,
_ => panic!("expected DCP request frame"),
};
assert!(matches!(
request.command,
Some(crate::dcp::proto::request::Command::ListJobs(_))
));
request.request_id
}
async fn read_test_frame(stream: &mut DuplexStream) -> crate::dcp::proto::DcpFrame {
let mut header = [0_u8; 4];
stream
.read_exact(&mut header)
.await
.expect("DCP frame header");
let mut payload = vec![0_u8; u32::from_be_bytes(header) as usize];
stream
.read_exact(&mut payload)
.await
.expect("DCP frame payload");
crate::dcp::proto::DcpFrame::decode(payload.as_slice()).expect("valid DCP frame")
}
async fn write_empty_job_list(stream: &mut DuplexStream, request_id: u64) {
let payload = crate::dcp::proto::JobList { jobs: Vec::new() }.encode_to_vec();
let frame = crate::dcp::proto::DcpFrame::response(crate::dcp::proto::Response::ok(
request_id, payload,
));
let payload = frame.encode_to_vec();
let len = u32::try_from(payload.len()).expect("test DCP frame length");
stream
.write_all(&len.to_be_bytes())
.await
.expect("write DCP frame header");
stream
.write_all(&payload)
.await
.expect("write DCP frame payload");
stream.flush().await.expect("flush DCP frame");
}
#[test]
fn remembered_assignment_generation_is_monotonic() {
let mut assignments = BTreeMap::new();
assert!(remember_assignment(
&mut assignments,
"job".to_owned(),
assignment(2, "current-owner"),
Some("coordinator"),
));
assert!(!remember_assignment(
&mut assignments,
"job".to_owned(),
assignment(1, "stale-owner"),
Some("coordinator"),
));
assert_eq!(assignments["job"].assigned_node, "current-owner");
assert_eq!(assignments["job"].placement_generation, 2);
assert!(!remember_assignment(
&mut assignments,
"job".to_owned(),
assignment(2, "equal-generation-owner"),
Some("coordinator"),
));
assert_eq!(assignments["job"].assigned_node, "current-owner");
let mut stale_same_owner = assignment(2, "current-owner");
stale_same_owner.coordinator_node = "old-coordinator".to_owned();
assert!(!remember_assignment(
&mut assignments,
"job".to_owned(),
stale_same_owner,
Some("new-coordinator"),
));
assert_eq!(assignments["job"].coordinator_node, "coordinator");
let mut same_owner_refresh = assignment(2, "current-owner");
same_owner_refresh.coordinator_node = "new-coordinator".to_owned();
assert!(remember_assignment(
&mut assignments,
"job".to_owned(),
same_owner_refresh,
Some("new-coordinator"),
));
assert_eq!(assignments["job"].assigned_node, "current-owner");
assert_eq!(assignments["job"].coordinator_node, "new-coordinator");
assert!(remember_assignment(
&mut assignments,
"job".to_owned(),
assignment(3, "new-owner"),
Some("new-coordinator"),
));
assert_eq!(assignments["job"].assigned_node, "new-owner");
assert_eq!(assignments["job"].placement_generation, 3);
}
#[test]
fn registry_rebuild_preserves_tombstone_against_stale_running_observations() {
let tombstone = AssignmentTombstone {
placement_generation: 2,
coordinator_node: "new-coordinator".to_owned(),
};
let older = assignment_with_coordinator(1, "stale-owner", "old-coordinator");
assert!(registry_observation_is_tombstoned(&tombstone, &older));
let equal = assignment_with_coordinator(2, "stale-owner", "new-coordinator");
assert!(registry_observation_is_tombstoned(&tombstone, &equal));
let newer = assignment_with_coordinator(3, "new-owner", "new-coordinator");
assert!(!registry_observation_is_tombstoned(&tombstone, &newer));
}
#[tokio::test]
async fn stale_tombstone_does_not_stop_newer_local_owner() {
let registry = JobRegistry::start(AgentConfig::default()).expect("registry starts");
let sessions = test_sessions("node-a");
let mut actor = test_placement_actor_with_coordinator(
"node-a",
"new-coordinator",
registry.clone(),
sessions.clone(),
);
let mut current = assignment(2, "node-a");
current.coordinator_node = "new-coordinator".to_owned();
actor
.assignments
.insert("local-job".to_owned(), current.clone());
submit_running_cluster_job(®istry, "local-job", current.clone());
let stale = actor
.remember(
"local-job".to_owned(),
None,
Some(AssignmentTombstone {
placement_generation: 1,
coordinator_node: "old-coordinator".to_owned(),
}),
)
.await;
assert!(
matches!(
stale,
Err(DcpError::Response {
status: ResponseStatus::Conflict,
..
})
),
"non-current coordinator tombstone must not be falsely acknowledged: {stale:?}"
);
let status = registry.status("local-job").expect("local job status");
assert_eq!(status.state, JobState::Running);
assert!(actor.assignments.contains_key("local-job"));
actor
.remember(
"local-job".to_owned(),
None,
Some(AssignmentTombstone {
placement_generation: 2,
coordinator_node: "new-coordinator".to_owned(),
}),
)
.await
.expect("current tombstone applies");
let status = registry.status("local-job").expect("local job status");
assert_eq!(status.state, JobState::Stopped);
assert!(!actor.assignments.contains_key("local-job"));
sessions.shutdown().await;
registry.shutdown().expect("registry shuts down");
}
#[tokio::test]
async fn delayed_active_after_recorded_tombstone_is_rejected() {
let registry = JobRegistry::start(AgentConfig::default()).expect("registry starts");
let sessions = test_sessions("node-a");
let mut actor = test_placement_actor_with_coordinator(
"node-a",
"new-coordinator",
registry.clone(),
sessions.clone(),
);
actor
.remember(
"local-job".to_owned(),
None,
Some(AssignmentTombstone {
placement_generation: 2,
coordinator_node: "new-coordinator".to_owned(),
}),
)
.await
.expect("valid no-existing-assignment tombstone is retained");
assert_eq!(
actor
.assignment_tombstones
.get("local-job")
.map(|tombstone| (
tombstone.placement_generation,
tombstone.coordinator_node.as_str()
)),
Some((2, "new-coordinator"))
);
let mut stale = assignment(2, "node-b");
stale.coordinator_node = "new-coordinator".to_owned();
let result = actor
.remember(
"local-job".to_owned(),
Some(wire_cluster_job_start(&stale)),
None,
)
.await;
assert!(matches!(
result,
Err(DcpError::Response {
status: ResponseStatus::Conflict,
..
})
));
assert!(!actor.assignments.contains_key("local-job"));
assert!(actor.assignment_tombstones.contains_key("local-job"));
let mut newer = assignment(3, "node-b");
newer.coordinator_node = "new-coordinator".to_owned();
actor
.remember(
"local-job".to_owned(),
Some(wire_cluster_job_start(&newer)),
None,
)
.await
.expect("strictly newer active assignment supersedes tombstone");
assert_eq!(actor.assignments["local-job"].placement_generation, 3);
assert!(!actor.assignment_tombstones.contains_key("local-job"));
sessions.shutdown().await;
registry.shutdown().expect("registry shuts down");
}
#[tokio::test]
async fn new_coordinator_tombstone_applies_before_same_generation_refresh() {
let registry = JobRegistry::start(AgentConfig::default()).expect("registry starts");
let sessions = test_sessions("node-a");
let mut actor = test_placement_actor_with_coordinator(
"node-a",
"new-coordinator",
registry.clone(),
sessions.clone(),
);
let mut local = assignment(2, "node-a");
local.coordinator_node = "old-coordinator".to_owned();
actor
.assignments
.insert("local-job".to_owned(), local.clone());
submit_running_cluster_job(®istry, "local-job", local);
actor
.remember(
"local-job".to_owned(),
None,
Some(AssignmentTombstone {
placement_generation: 2,
coordinator_node: "new-coordinator".to_owned(),
}),
)
.await
.expect("current coordinator tombstone applies across handoff");
let status = registry.status("local-job").expect("local job status");
assert_eq!(status.state, JobState::Stopped);
assert!(!actor.assignments.contains_key("local-job"));
assert!(actor.assignment_tombstones.contains_key("local-job"));
let mut same_generation_refresh = assignment(2, "node-a");
same_generation_refresh.coordinator_node = "new-coordinator".to_owned();
let result = actor
.remember(
"local-job".to_owned(),
Some(wire_cluster_job_start(&same_generation_refresh)),
None,
)
.await;
assert!(matches!(
result,
Err(DcpError::Response {
status: ResponseStatus::Conflict,
..
})
));
assert!(!actor.assignments.contains_key("local-job"));
assert!(actor.assignment_tombstones.contains_key("local-job"));
sessions.shutdown().await;
registry.shutdown().expect("registry shuts down");
}
#[tokio::test]
async fn current_tombstone_refreshes_terminal_local_metadata() {
let registry = JobRegistry::start(AgentConfig::default()).expect("registry starts");
let sessions = test_sessions("node-a");
let mut actor = test_placement_actor_with_coordinator(
"node-a",
"new-coordinator",
registry.clone(),
sessions.clone(),
);
let mut local = assignment(2, "node-a");
local.coordinator_node = "old-coordinator".to_owned();
submit_running_cluster_job(®istry, "local-job", local);
registry.stop("local-job").expect("local job stopped");
actor
.remember(
"local-job".to_owned(),
None,
Some(AssignmentTombstone {
placement_generation: 2,
coordinator_node: "new-coordinator".to_owned(),
}),
)
.await
.expect("current coordinator tombstone refreshes terminal metadata");
let status = registry.status("local-job").expect("local job status");
assert_eq!(status.state, JobState::Stopped);
let metadata = status.cluster.expect("cluster metadata retained");
assert_eq!(metadata.placement_generation, 2);
assert_eq!(metadata.coordinator_node, "new-coordinator");
assert!(!actor.assignments.contains_key("local-job"));
assert!(actor.assignment_tombstones.contains_key("local-job"));
sessions.shutdown().await;
registry.shutdown().expect("registry shuts down");
}
#[tokio::test]
async fn delayed_active_replacement_stops_stale_local_owner() {
let registry = JobRegistry::start(AgentConfig::default()).expect("registry starts");
let sessions = test_sessions("node-a");
let mut actor = test_placement_actor_with_coordinator(
"node-a",
"new-coordinator",
registry.clone(),
sessions.clone(),
);
let mut stale = assignment(1, "node-a");
stale.coordinator_node = "old-coordinator".to_owned();
actor
.assignments
.insert("local-job".to_owned(), stale.clone());
submit_running_cluster_job(®istry, "local-job", stale);
let mut replacement = assignment(2, "node-b");
replacement.coordinator_node = "new-coordinator".to_owned();
actor
.remember(
"local-job".to_owned(),
Some(wire_cluster_job_start(&replacement)),
None,
)
.await
.expect("delayed active replacement is accepted");
let status = registry.status("local-job").expect("local job status");
assert_eq!(status.state, JobState::Stopped);
assert_eq!(actor.assignments["local-job"].assigned_node, "node-b");
assert_eq!(actor.assignments["local-job"].placement_generation, 2);
sessions.shutdown().await;
registry.shutdown().expect("registry shuts down");
}
#[tokio::test]
async fn register_restarted_rejects_stale_local_restart_over_active_assignment() {
let registry = JobRegistry::start(AgentConfig::default()).expect("registry starts");
let sessions = test_sessions("node-a");
let mut actor = test_placement_actor_with_coordinator(
"node-a",
"new-coordinator",
registry.clone(),
sessions.clone(),
);
let mut stale_higher = assignment(1, "node-a");
stale_higher.coordinator_node = "new-coordinator".to_owned();
actor.assignments.insert(
"higher-job".to_owned(),
assignment_with_coordinator(2, "node-b", "new-coordinator"),
);
submit_running_cluster_job(®istry, "higher-job", stale_higher.clone());
let result = actor
.register_restarted(
"higher-job".to_owned(),
wire_cluster_job_start(&stale_higher),
Duration::from_millis(50),
)
.await;
assert!(matches!(
result,
Err(DcpError::Response {
status: ResponseStatus::Conflict,
..
})
));
assert_eq!(actor.assignments["higher-job"].assigned_node, "node-b");
assert_eq!(actor.assignments["higher-job"].placement_generation, 2);
assert_eq!(
registry.status("higher-job").expect("higher status").state,
JobState::Stopped
);
let mut stale_equal = assignment(2, "node-a");
stale_equal.coordinator_node = "new-coordinator".to_owned();
actor.assignments.insert(
"equal-job".to_owned(),
assignment_with_coordinator(2, "node-b", "new-coordinator"),
);
submit_running_cluster_job(®istry, "equal-job", stale_equal.clone());
let result = actor
.register_restarted(
"equal-job".to_owned(),
wire_cluster_job_start(&stale_equal),
Duration::from_millis(50),
)
.await;
assert!(matches!(
result,
Err(DcpError::Response {
status: ResponseStatus::Conflict,
..
})
));
assert_eq!(actor.assignments["equal-job"].assigned_node, "node-b");
assert_eq!(actor.assignments["equal-job"].placement_generation, 2);
assert_eq!(
registry.status("equal-job").expect("equal status").state,
JobState::Stopped
);
sessions.shutdown().await;
registry.shutdown().expect("registry shuts down");
}
#[tokio::test]
async fn tombstone_stop_error_keeps_assignment_for_retry() {
let registry = JobRegistry::start(AgentConfig::default()).expect("registry starts");
let sessions = test_sessions("node-a");
let mut actor = test_placement_actor_with_coordinator(
"node-a",
"new-coordinator",
registry.clone(),
sessions.clone(),
);
let mut current = assignment(2, "node-a");
current.coordinator_node = "new-coordinator".to_owned();
actor
.assignments
.insert("local-job".to_owned(), current.clone());
registry
.shutdown()
.expect("registry shuts down before cleanup");
let result = actor
.remember(
"local-job".to_owned(),
None,
Some(AssignmentTombstone {
placement_generation: 2,
coordinator_node: "new-coordinator".to_owned(),
}),
)
.await;
assert!(result.is_err());
assert!(
actor.assignments.contains_key("local-job"),
"failed cleanup must leave assignment for the next retry"
);
sessions.shutdown().await;
}
#[tokio::test]
async fn membership_coordinator_and_placement_sources_publish_cluster_events() {
let publisher = ClusterEventPublisher::new(8).expect("cluster event topic");
let mut events = publisher.subscribe().expect("cluster event subscription");
let at = UNIX_EPOCH + Duration::from_millis(1_720_000_000_000);
publisher.publish_member(&MemberEvent {
sequence: 11,
kind: MemberEventKind::MemberDown,
member: Member {
node_id: "node-1".to_owned(),
address: "127.0.0.1:2551".parse().expect("member address"),
agent_addr: Some("127.0.0.1:9555".parse().expect("agent address")),
roles: vec![AGENT_ROLE.to_owned()],
state: MemberState::Down,
unreachable: true,
incarnation: 4,
seen_at: at,
unreachable_since: Some(at),
},
at,
});
publisher.publish_coordinator_changed(Some("node-1"), Some("node-2"), "node-2");
publisher.publish_job_replaced(at, "ingest", "node-1", "node-2", "node_down:node-1", 3);
let mut received = Vec::new();
for _ in 0..3 {
received.push(
tokio::time::timeout(Duration::from_secs(1), events.recv())
.await
.expect("cluster event arrives")
.expect("cluster event subscription remains open"),
);
}
assert_eq!(
received
.iter()
.map(|event| (event.sequence, event.kind.as_str(), event.node_id.as_str()))
.collect::<Vec<_>>(),
vec![
(1, "MemberDown", "node-1"),
(2, "CoordinatorChanged", "node-2"),
(3, "JobReplaced", "node-2"),
]
);
assert_eq!(received[0].generation, Some(4));
assert_eq!(received[2].generation, Some(3));
assert!(received[2].detail.contains("job=ingest"));
}
fn assignment(placement_generation: u64, assigned_node: &str) -> ClusterJobMetadata {
ClusterJobMetadata {
factory_name: "ticker".to_owned(),
params: BTreeMap::new(),
placement: RegistryPlacementSpec::least_jobs(None),
coordinator_node: "coordinator".to_owned(),
assigned_node: assigned_node.to_owned(),
placement_generation,
history: Vec::new(),
}
}
fn assignment_with_coordinator(
placement_generation: u64,
assigned_node: &str,
coordinator_node: &str,
) -> ClusterJobMetadata {
let mut metadata = assignment(placement_generation, assigned_node);
metadata.coordinator_node = coordinator_node.to_owned();
metadata
}
fn submit_running_cluster_job(
registry: &JobRegistryHandle,
name: &str,
metadata: ClusterJobMetadata,
) {
let spec = JobSpec::new(name, |context| {
let control = context.control();
Ok(Source::tick(Duration::ZERO, Duration::from_secs(60), 1_u64)
.via_mat(context.drain_flow(), Keep::right)
.to_mat(Sink::ignore(), move |_switch, completion| {
JobMat::new(completion, control.clone())
}))
})
.with_cluster_metadata(metadata);
registry.submit(spec).expect("job submitted");
registry.start(name).expect("job starts");
}
fn test_sessions(node_id: &str) -> NodeSessionManagerHandle {
let state = Signal::new(test_cluster_state(node_id)).expect("cluster state");
NodeSessionManagerHandle::start(NodeSessionConfig::default(), node_id.to_owned(), state)
.expect("sessions start")
}
fn test_placement_actor_with_coordinator(
node_id: &str,
coordinator_node: &str,
registry: JobRegistryHandle,
sessions: NodeSessionManagerHandle,
) -> PlacementActor {
let state = Signal::new(test_cluster_state_with_coordinator(
node_id,
coordinator_node,
))
.expect("cluster state");
let (commands, _receiver) = mpsc::channel(8);
let (assignment_updates, _assignments) = watch::channel(Arc::new(BTreeMap::new()));
PlacementActor {
commands,
self_node: node_id.to_owned(),
agent_role: AGENT_ROLE.to_owned(),
state,
registry,
sessions,
factories: DcpJobFactories::new(),
request_timeout: Duration::from_millis(50),
active_coordinator: false,
last_known_coordinator: None,
assignments: BTreeMap::new(),
assignment_tombstones: BTreeMap::new(),
replicated_assignments: BTreeMap::new(),
inflight_assignment_syncs: BTreeSet::new(),
pending_replacements: BTreeMap::new(),
assignment_updates,
cluster_events: ClusterEventPublisher::new(8).expect("cluster events"),
}
}
fn test_cluster_state(node_id: &str) -> ClusterState {
test_cluster_state_with_coordinator(node_id, node_id)
}
fn test_cluster_state_with_coordinator(node_id: &str, coordinator_node: &str) -> ClusterState {
let mut members = BTreeMap::new();
members.insert(node_id.to_owned(), test_member_with_incarnation(node_id, 1));
members.insert(
coordinator_node.to_owned(),
test_member_with_incarnation(coordinator_node, 0),
);
ClusterState {
self_node: node_id.to_owned(),
members,
}
}
fn test_member_with_incarnation(node_id: &str, incarnation: u64) -> Member {
Member {
node_id: node_id.to_owned(),
address: "127.0.0.1:2551".parse().expect("member address"),
agent_addr: Some("127.0.0.1:9555".parse().expect("agent address")),
roles: vec![AGENT_ROLE.to_owned()],
state: MemberState::Up,
unreachable: false,
incarnation,
seen_at: UNIX_EPOCH,
unreachable_since: None,
}
}
}