use std::collections::HashMap;
use crate::prelude::*;
use serde::{Deserialize, Serialize};
use crate::cluster::framing::SpawnRequest;
use crate::app::election::LeaderElection;
use crate::app::error::OrchestratorError;
use crate::app::node_info::{ClusterView, NodeInfo};
use crate::app::placement::{self, PlacementDecision, PlacementStrategy};
use crate::app::spawn_sender::SpawnSender;
use crate::app::spec::{ActorSpec, CrashStrategy};
#[derive(Debug)]
pub struct Coordinator;
pub struct CoordinatorState {
pub cluster_view: ClusterView,
pub specs: HashMap<String, ActorSpec>,
pub placement_strategy: Box<dyn PlacementStrategy>,
pub election: Box<dyn LeaderElection>,
pub local_node_id: String,
pub pending_spawns: HashMap<u64, PendingSpawn>,
pub waiting_for_return: HashMap<String, WaitingSpec>,
next_request_id: u64,
spawn_sender: Option<SpawnSender>,
}
#[derive(Debug, Clone)]
pub struct PendingSpawn {
pub spec_label: String,
pub target_node_id: String,
}
#[derive(Debug, Clone)]
pub struct WaitingSpec {
pub spec: ActorSpec,
pub failed_node_id: String,
}
impl Actor for Coordinator {
type State = CoordinatorState;
}
#[derive(Debug, Clone, Serialize, Deserialize, Message)]
#[message(result = Result<PlacementDecision, String>, remote = "coordinator::SubmitSpec")]
pub struct SubmitSpec {
pub spec: ActorSpec,
}
#[derive(Debug, Clone, Serialize, Deserialize, Message)]
#[message(result = Result<(), String>, remote = "coordinator::RemoveSpec")]
pub struct RemoveSpec {
pub label: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, Message)]
#[message(result = ClusterViewSnapshot, remote = "coordinator::GetClusterView")]
pub struct GetClusterView;
#[derive(Debug, Clone, Serialize, Deserialize, Message)]
#[message(result = Vec<SpecStatus>, remote = "coordinator::GetSpecs")]
pub struct GetSpecs;
#[derive(Debug, Clone, Serialize, Deserialize, Message)]
#[message(result = (), remote = "coordinator::NotifyNodeJoined")]
pub struct NotifyNodeJoined {
pub node_id: String,
pub info: SerializableNodeInfo,
}
#[derive(Debug, Clone, Serialize, Deserialize, Message)]
#[message(result = (), remote = "coordinator::NotifyNodeFailed")]
pub struct NotifyNodeFailed {
pub node_id: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, Message)]
#[message(result = (), remote = "coordinator::NotifyNodeLeft")]
pub struct NotifyNodeLeft {
pub node_id: String,
}
#[derive(Debug, Clone, Serialize, Deserialize, Message)]
#[message(result = (), remote = "coordinator::NotifySpawnAck")]
pub struct NotifySpawnAck {
pub request_id: u64,
pub success: bool,
pub error: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Message)]
#[message(result = (), remote = "coordinator::WaitForReturnTimeout")]
pub struct WaitForReturnTimeout {
pub label: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClusterViewSnapshot {
pub nodes: Vec<NodeSnapshot>,
pub alive_count: usize,
pub total_count: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeSnapshot {
pub node_id: String,
pub name: String,
pub class: String,
pub actor_count: usize,
pub is_alive: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SpecStatus {
pub label: String,
pub actor_type: String,
pub placed_on: Option<String>,
pub state: SpecState,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum SpecState {
Running,
Pending,
WaitingForReturn,
Unplaced,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SerializableNodeInfo {
pub name: String,
pub host: String,
pub port: u16,
pub incarnation: u64,
pub class: crate::cluster::config::NodeClass,
pub metadata: HashMap<String, String>,
}
#[handlers]
impl Coordinator {
#[handler]
fn submit_spec(
&mut self,
_ctx: &ActorContext<Self>,
state: &mut CoordinatorState,
msg: SubmitSpec,
) -> Result<PlacementDecision, String> {
if !state.is_leader() {
return Err(OrchestratorError::NotLeader.to_string());
}
let label = msg.spec.label.clone();
if state.specs.contains_key(&label) {
return Err(OrchestratorError::SpecAlreadyExists { label }.to_string());
}
let decision = placement::select_node(
state.placement_strategy.as_ref(),
&msg.spec,
&state.cluster_view,
)
.ok_or_else(|| {
OrchestratorError::NoEligibleNodes {
reason: format!("no node satisfies constraints for {label}"),
}
.to_string()
})?;
tracing::info!(
"Placing {} on {} ({})",
label,
decision.node_id,
decision.reason
);
state.specs.insert(label.clone(), msg.spec.clone());
let request_id = state.next_request_id();
state.pending_spawns.insert(
request_id,
PendingSpawn {
spec_label: label,
target_node_id: decision.node_id.clone(),
},
);
state.send_spawn(&decision.node_id, request_id, &msg.spec);
Ok(decision)
}
#[handler]
fn remove_spec(
&mut self,
_ctx: &ActorContext<Self>,
state: &mut CoordinatorState,
msg: RemoveSpec,
) -> Result<(), String> {
if !state.is_leader() {
return Err(OrchestratorError::NotLeader.to_string());
}
if state.specs.remove(&msg.label).is_none() {
return Err(OrchestratorError::SpecNotFound {
label: msg.label.clone(),
}
.to_string());
}
state
.pending_spawns
.retain(|_, ps| ps.spec_label != msg.label);
state.waiting_for_return.remove(&msg.label);
state.cluster_view.remove_actor_anywhere(&msg.label);
tracing::info!("Removed spec: {}", msg.label);
Ok(())
}
#[handler]
fn get_cluster_view(
&mut self,
_ctx: &ActorContext<Self>,
state: &mut CoordinatorState,
_msg: GetClusterView,
) -> ClusterViewSnapshot {
ClusterViewSnapshot {
alive_count: state.cluster_view.alive_count(),
total_count: state.cluster_view.total_count(),
nodes: state
.cluster_view
.nodes
.values()
.map(|n| NodeSnapshot {
node_id: n.node_id(),
name: n.identity.name.clone(),
class: n.class.to_string(),
actor_count: n.actor_count(),
is_alive: n.is_alive,
})
.collect(),
}
}
#[handler]
fn get_specs(
&mut self,
_ctx: &ActorContext<Self>,
state: &mut CoordinatorState,
_msg: GetSpecs,
) -> Vec<SpecStatus> {
state
.specs
.iter()
.map(|(label, spec)| {
let placed_on = state.cluster_view.find_actor(label).map(|s| s.to_string());
let spec_state = if state.waiting_for_return.contains_key(label) {
SpecState::WaitingForReturn
} else if state
.pending_spawns
.values()
.any(|ps| ps.spec_label == *label)
{
SpecState::Pending
} else if placed_on.is_some() {
SpecState::Running
} else {
SpecState::Unplaced
};
SpecStatus {
label: label.clone(),
actor_type: spec.actor_type_name.clone(),
placed_on,
state: spec_state,
}
})
.collect()
}
#[handler]
fn notify_node_joined(
&mut self,
_ctx: &ActorContext<Self>,
state: &mut CoordinatorState,
msg: NotifyNodeJoined,
) {
let info = NodeInfo::new(
crate::cluster::config::NodeIdentity {
name: msg.info.name,
host: msg.info.host,
port: msg.info.port,
incarnation: msg.info.incarnation,
},
msg.info.class,
msg.info.metadata,
);
tracing::info!("Node joined: {}", msg.node_id);
state.cluster_view.upsert_node(info);
let rejoined: Vec<String> = state
.waiting_for_return
.iter()
.filter(|(_, ws)| ws.failed_node_id == msg.node_id)
.map(|(label, _)| label.clone())
.collect();
for label in rejoined {
if let Some(ws) = state.waiting_for_return.remove(&label) {
tracing::info!(
"Node {} returned — re-spawning spec {} on it",
msg.node_id,
ws.spec.label
);
let request_id = state.next_request_id();
state.pending_spawns.insert(
request_id,
PendingSpawn {
spec_label: ws.spec.label.clone(),
target_node_id: msg.node_id.clone(),
},
);
state.send_spawn(&msg.node_id, request_id, &ws.spec);
}
}
}
#[handler]
fn notify_node_failed(
&mut self,
ctx: &ActorContext<Self>,
state: &mut CoordinatorState,
msg: NotifyNodeFailed,
) {
tracing::warn!("Node failed: {}", msg.node_id);
state.cluster_view.mark_failed(&msg.node_id);
let timers = state.handle_node_departure(&msg.node_id, false);
for (label, duration) in timers {
let endpoint = ctx.endpoint();
tokio::spawn(async move {
tokio::time::sleep(duration).await;
let _ = endpoint.send(WaitForReturnTimeout { label }).await;
});
}
}
#[handler]
fn notify_node_left(
&mut self,
_ctx: &ActorContext<Self>,
state: &mut CoordinatorState,
msg: NotifyNodeLeft,
) {
tracing::info!("Node left gracefully: {}", msg.node_id);
let _timers = state.handle_node_departure(&msg.node_id, true);
state.cluster_view.remove_node(&msg.node_id);
}
#[handler]
fn notify_spawn_ack(
&mut self,
_ctx: &ActorContext<Self>,
state: &mut CoordinatorState,
msg: NotifySpawnAck,
) {
if let Some(pending) = state.pending_spawns.remove(&msg.request_id) {
if msg.success {
tracing::info!(
"Spawn confirmed: {} on {}",
pending.spec_label,
pending.target_node_id
);
state
.cluster_view
.add_actor(&pending.target_node_id, &pending.spec_label);
} else {
tracing::warn!(
"Spawn failed for {}: {}",
pending.spec_label,
msg.error.as_deref().unwrap_or("unknown error")
);
}
}
}
#[handler]
fn wait_for_return_timeout(
&mut self,
_ctx: &ActorContext<Self>,
state: &mut CoordinatorState,
msg: WaitForReturnTimeout,
) {
if let Some(ws) = state.waiting_for_return.remove(&msg.label) {
tracing::warn!(
"WaitForReturn timeout for {} (was on {}) — falling back to Redistribute",
msg.label,
ws.failed_node_id,
);
if let Some(decision) = placement::select_node(
state.placement_strategy.as_ref(),
&ws.spec,
&state.cluster_view,
) {
let request_id = state.next_request_id();
state.pending_spawns.insert(
request_id,
PendingSpawn {
spec_label: msg.label.clone(),
target_node_id: decision.node_id.clone(),
},
);
state.send_spawn(&decision.node_id, request_id, &ws.spec);
tracing::info!(
"Re-placed {} on {} after timeout ({})",
msg.label,
decision.node_id,
decision.reason
);
} else {
tracing::warn!(
"No eligible node for {} after WaitForReturn timeout — spec remains unplaced",
msg.label
);
}
}
}
}
impl CoordinatorState {
pub fn new(
local_node_id: impl Into<String>,
placement_strategy: Box<dyn PlacementStrategy>,
election: Box<dyn LeaderElection>,
) -> Self {
Self {
cluster_view: ClusterView::new(),
specs: HashMap::new(),
placement_strategy,
election,
local_node_id: local_node_id.into(),
pending_spawns: HashMap::new(),
waiting_for_return: HashMap::new(),
next_request_id: 0,
spawn_sender: None,
}
}
pub(crate) fn with_spawn_sender(mut self, sender: SpawnSender) -> Self {
self.spawn_sender = Some(sender);
self
}
fn next_request_id(&mut self) -> u64 {
let id = self.next_request_id;
self.next_request_id += 1;
id
}
pub fn is_leader(&self) -> bool {
self.election
.elect(&self.cluster_view)
.is_some_and(|leader| leader == self.local_node_id)
}
fn send_spawn(&self, target_node_id: &str, request_id: u64, spec: &ActorSpec) {
if let Some(sender) = &self.spawn_sender {
sender.send_spawn(
target_node_id,
SpawnRequest {
request_id,
label: spec.label.clone(),
actor_type_name: spec.actor_type_name.clone(),
initial_state: spec.initial_state.clone(),
},
);
} else {
tracing::warn!(
"No spawn sender configured — spawn request for {} dropped",
spec.label
);
}
}
fn handle_node_departure(
&mut self,
node_id: &str,
graceful: bool,
) -> Vec<(String, std::time::Duration)> {
let mut timers_needed = Vec::new();
let affected_labels: Vec<String> = self
.specs
.keys()
.filter(|label| {
let running_on = self
.cluster_view
.find_actor(label)
.is_some_and(|n| n == node_id);
let pending_on = self
.pending_spawns
.values()
.any(|ps| ps.spec_label == **label && ps.target_node_id == node_id);
running_on || pending_on
})
.cloned()
.collect();
self.pending_spawns
.retain(|_, ps| ps.target_node_id != node_id);
for label in affected_labels {
let spec = match self.specs.get(&label) {
Some(s) => s.clone(),
None => continue,
};
self.cluster_view.remove_actor(node_id, &label);
match &spec.crash_strategy {
CrashStrategy::Abandon => {
tracing::info!("Abandoning {label} (node {node_id} departed)");
self.specs.remove(&label);
}
CrashStrategy::WaitForReturn(duration) if !graceful => {
tracing::warn!(
"Waiting {:?} for node {} to return (spec: {label})",
duration,
node_id,
);
self.waiting_for_return.insert(
label.clone(),
WaitingSpec {
spec: spec.clone(),
failed_node_id: node_id.to_string(),
},
);
timers_needed.push((label, *duration));
}
_ => {
let reason = if graceful {
"graceful departure"
} else {
"failure"
};
tracing::info!("Redistributing {label} (node {node_id} {reason})");
if let Some(decision) = placement::select_node(
self.placement_strategy.as_ref(),
&spec,
&self.cluster_view,
) {
let request_id = self.next_request_id();
self.pending_spawns.insert(
request_id,
PendingSpawn {
spec_label: label.clone(),
target_node_id: decision.node_id.clone(),
},
);
self.send_spawn(&decision.node_id, request_id, &spec);
tracing::info!(
"Re-placed {label} on {} ({})",
decision.node_id,
decision.reason
);
} else {
tracing::warn!("No eligible node for {label} — spec remains unplaced");
}
}
}
}
timers_needed
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::app::election::OldestNode;
use crate::app::placement::LeastLoaded;
use crate::cluster::config::{NodeClass, NodeIdentity};
fn make_system_and_coordinator() -> (crate::System, Endpoint<Coordinator>) {
let system = crate::System::local();
let alpha_identity = NodeIdentity {
name: "alpha".into(),
host: "127.0.0.1".into(),
port: 7100,
incarnation: 1,
};
let local_node_id = alpha_identity.node_id_string();
let mut state = CoordinatorState::new(
&local_node_id,
Box::new(LeastLoaded),
Box::new(OldestNode::any()),
);
state.cluster_view.upsert_node(NodeInfo::new(
alpha_identity,
NodeClass::Worker,
HashMap::new(),
));
state.cluster_view.upsert_node(NodeInfo::new(
NodeIdentity {
name: "beta".into(),
host: "127.0.0.1".into(),
port: 7200,
incarnation: 2,
},
NodeClass::Worker,
HashMap::new(),
));
let ep = system.start("coordinator", Coordinator, state);
(system, ep)
}
#[tokio::test]
async fn test_submit_spec() {
let (_system, coordinator) = make_system_and_coordinator();
let result = coordinator
.send(SubmitSpec {
spec: ActorSpec::new("worker/0", "app::Worker"),
})
.await
.unwrap();
assert!(result.is_ok());
let decision = result.unwrap();
assert!(!decision.node_id.is_empty());
}
#[tokio::test]
async fn test_submit_duplicate_spec() {
let (_system, coordinator) = make_system_and_coordinator();
let _ = coordinator
.send(SubmitSpec {
spec: ActorSpec::new("worker/0", "app::Worker"),
})
.await
.unwrap();
let result = coordinator
.send(SubmitSpec {
spec: ActorSpec::new("worker/0", "app::Worker"),
})
.await
.unwrap();
assert!(result.is_err());
}
#[tokio::test]
async fn test_remove_spec() {
let (_system, coordinator) = make_system_and_coordinator();
let _ = coordinator
.send(SubmitSpec {
spec: ActorSpec::new("worker/0", "app::Worker"),
})
.await
.unwrap();
let result = coordinator
.send(RemoveSpec {
label: "worker/0".into(),
})
.await
.unwrap();
assert!(result.is_ok());
}
#[tokio::test]
async fn test_get_specs() {
let (_system, coordinator) = make_system_and_coordinator();
let _ = coordinator
.send(SubmitSpec {
spec: ActorSpec::new("worker/0", "app::Worker"),
})
.await
.unwrap();
let _ = coordinator
.send(SubmitSpec {
spec: ActorSpec::new("worker/1", "app::Worker"),
})
.await
.unwrap();
let specs = coordinator.send(GetSpecs).await.unwrap();
assert_eq!(specs.len(), 2);
}
#[tokio::test]
async fn test_node_failure_redistributes() {
let (_system, coordinator) = make_system_and_coordinator();
let result = coordinator
.send(SubmitSpec {
spec: ActorSpec::new("worker/0", "app::Worker")
.with_crash_strategy(CrashStrategy::Redistribute),
})
.await
.unwrap()
.unwrap();
let original_node = result.node_id.clone();
coordinator
.send(NotifyNodeFailed {
node_id: original_node.clone(),
})
.await
.unwrap();
let specs = coordinator.send(GetSpecs).await.unwrap();
assert_eq!(specs.len(), 1);
assert!(matches!(specs[0].state, SpecState::Pending));
coordinator
.send(NotifySpawnAck {
request_id: 1,
success: true,
error: None,
})
.await
.unwrap();
let specs = coordinator.send(GetSpecs).await.unwrap();
assert_eq!(specs.len(), 1);
assert!(matches!(specs[0].state, SpecState::Running));
assert_ne!(specs[0].placed_on.as_deref(), Some(original_node.as_str()));
}
#[tokio::test]
async fn test_node_failure_abandon() {
let (_system, coordinator) = make_system_and_coordinator();
let result = coordinator
.send(SubmitSpec {
spec: ActorSpec::new("worker/0", "app::Worker")
.with_crash_strategy(CrashStrategy::Abandon),
})
.await
.unwrap()
.unwrap();
coordinator
.send(NotifyNodeFailed {
node_id: result.node_id,
})
.await
.unwrap();
let specs = coordinator.send(GetSpecs).await.unwrap();
assert_eq!(specs.len(), 0);
}
#[tokio::test]
async fn test_get_cluster_view() {
let (_system, coordinator) = make_system_and_coordinator();
let view = coordinator.send(GetClusterView).await.unwrap();
assert_eq!(view.alive_count, 2);
assert_eq!(view.total_count, 2);
assert_eq!(view.nodes.len(), 2);
}
#[tokio::test]
async fn test_wait_for_return_timeout_redistributes() {
let (_system, coordinator) = make_system_and_coordinator();
let result = coordinator
.send(SubmitSpec {
spec: ActorSpec::new("worker/0", "app::Worker").with_crash_strategy(
CrashStrategy::WaitForReturn(std::time::Duration::from_millis(50)),
),
})
.await
.unwrap()
.unwrap();
let original_node = result.node_id.clone();
coordinator
.send(NotifyNodeFailed {
node_id: original_node.clone(),
})
.await
.unwrap();
let specs = coordinator.send(GetSpecs).await.unwrap();
assert!(matches!(specs[0].state, SpecState::WaitingForReturn));
coordinator
.send(WaitForReturnTimeout {
label: "worker/0".into(),
})
.await
.unwrap();
let specs = coordinator.send(GetSpecs).await.unwrap();
assert_eq!(specs.len(), 1);
assert!(!matches!(specs[0].state, SpecState::WaitingForReturn));
}
#[tokio::test]
async fn test_wait_for_return_node_rejoins() {
let (_system, coordinator) = make_system_and_coordinator();
let result = coordinator
.send(SubmitSpec {
spec: ActorSpec::new("worker/0", "app::Worker").with_crash_strategy(
CrashStrategy::WaitForReturn(std::time::Duration::from_secs(60)),
),
})
.await
.unwrap()
.unwrap();
let original_node = result.node_id.clone();
coordinator
.send(NotifyNodeFailed {
node_id: original_node.clone(),
})
.await
.unwrap();
coordinator
.send(NotifyNodeJoined {
node_id: original_node.clone(),
info: SerializableNodeInfo {
name: "alpha".into(),
host: "127.0.0.1".into(),
port: 7100,
incarnation: 1,
class: NodeClass::Worker,
metadata: HashMap::new(),
},
})
.await
.unwrap();
let specs = coordinator.send(GetSpecs).await.unwrap();
assert_eq!(specs.len(), 1);
assert!(matches!(specs[0].state, SpecState::Pending));
}
}