1use std::{
2 collections::{BTreeMap, BTreeSet, HashMap},
3 fmt,
4 net::SocketAddr,
5 sync::{Arc, Mutex},
6 time::{Duration, SystemTime},
7};
8
9use datum_cluster::{ClusterConfig, ClusterNode, ClusterState, Member, MemberState, Signal};
10use datum_net::quic::quinn;
11use tokio::{
12 sync::{mpsc, oneshot, watch},
13 task::JoinHandle,
14 time::{Instant, sleep},
15};
16
17use crate::{
18 Agent, AgentConfig, AgentHandle, ClusterJobMetadata, ClusterPlacementHistory,
19 JobRegistryHandle, PlacementSpec as RegistryPlacementSpec,
20 PlacementStrategy as RegistryPlacementStrategy,
21 dcp::{
22 ClientKind, ClusterJobList, ClusterJobNode, ClusterJobStart, ClusterNodeError,
23 ClusterNodeList, ClusterNodeStatus, ClusterViewProvider, CompleteShardingAsk, DcpClient,
24 DcpError, DcpJobFactories, DcpServer, DcpServerConfig, DcpServerHandle,
25 ForwardShardEnvelopes, Hello, PlacementSpec, PlacementStrategy, RememberClusterAssignment,
26 ResponseStatus, ShardAllocation, ShardAllocationTable, ShardEnvelopeBatchResult,
27 ShardPipeClient, ShardPipeFrame, SubmitClusterJob,
28 server::{
29 cluster_metadata_from_wire, placement_spec_from_wire, wire_cluster_job_start,
30 wire_job_status,
31 },
32 },
33};
34
35pub const AGENT_ROLE: &str = "agent";
37
38pub type ClusterAgentResult<T> = Result<T, ClusterAgentError>;
40
41#[derive(Debug, thiserror::Error)]
43pub enum ClusterAgentError {
44 #[error("invalid cluster-agent config: {0}")]
45 InvalidConfig(String),
46 #[error(transparent)]
47 Agent(#[from] crate::AgentError),
48 #[error(transparent)]
49 Dcp(#[from] DcpError),
50 #[error(transparent)]
51 Cluster(#[from] datum_cluster::ClusterError),
52}
53
54#[derive(Clone)]
56pub enum NodeSessionTransport {
57 TcpLoopback,
59 QuicMtls {
61 server_name: String,
62 client_config: quinn::ClientConfig,
63 },
64}
65
66impl fmt::Debug for NodeSessionTransport {
67 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
68 match self {
69 Self::TcpLoopback => formatter.write_str("TcpLoopback"),
70 Self::QuicMtls { server_name, .. } => formatter
71 .debug_struct("QuicMtls")
72 .field("server_name", server_name)
73 .finish_non_exhaustive(),
74 }
75 }
76}
77
78#[derive(Clone, Debug)]
80pub struct NodeSessionConfig {
81 pub agent_role: String,
83 pub transport: NodeSessionTransport,
85 pub reconnect_min_backoff: Duration,
87 pub reconnect_max_backoff: Duration,
89 pub request_timeout: Duration,
91 pub command_buffer: usize,
93}
94
95impl Default for NodeSessionConfig {
96 fn default() -> Self {
97 Self {
98 agent_role: AGENT_ROLE.to_owned(),
99 transport: NodeSessionTransport::TcpLoopback,
100 reconnect_min_backoff: Duration::from_millis(50),
101 reconnect_max_backoff: Duration::from_secs(2),
102 request_timeout: Duration::from_millis(750),
103 command_buffer: 32,
104 }
105 }
106}
107
108#[derive(Clone, Default)]
110pub struct ClusterAgentConfig {
111 pub agent: AgentConfig,
113 pub cluster: ClusterConfig,
115 pub dcp: DcpServerConfig,
117 pub sessions: NodeSessionConfig,
119}
120
121impl fmt::Debug for ClusterAgentConfig {
122 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
123 formatter
124 .debug_struct("ClusterAgentConfig")
125 .field("agent", &self.agent)
126 .field("cluster", &self.cluster)
127 .field("dcp", &"<DcpServerConfig>")
128 .field("sessions", &self.sessions)
129 .finish()
130 }
131}
132
133pub struct ClusterAgent;
136
137impl ClusterAgent {
138 pub async fn start(
139 config: ClusterAgentConfig,
140 factories: DcpJobFactories,
141 ) -> ClusterAgentResult<ClusterAgentHandle> {
142 let mut config = config;
143 validate_cluster_agent_config(&config)?;
144 ensure_agent_role(&mut config.cluster, &config.sessions.agent_role);
145 config.dcp.node_id = config.cluster.node_id.clone();
146
147 let agent = Agent::start_with_config(config.agent.clone())?;
148 let server = DcpServer::from_agent(&agent, factories.clone(), config.dcp.clone());
149 let server_handle = server.start().await?;
150 let agent_addr = advertised_agent_addr(&config.sessions.transport, &server_handle)?;
151 config.cluster.agent_addr = Some(agent_addr);
152
153 let cluster = ClusterNode::start(config.cluster.clone()).await?;
154 let sessions = NodeSessionManagerHandle::start(
155 config.sessions.clone(),
156 cluster.node_id().to_owned(),
157 cluster.state(),
158 )?;
159 let placement = PlacementCoordinatorHandle::start(
160 cluster.node_id().to_owned(),
161 config.sessions.agent_role.clone(),
162 cluster.state(),
163 agent.registry().clone(),
164 sessions.clone(),
165 factories.clone(),
166 config.sessions.request_timeout,
167 )?;
168 let provider = Arc::new(ClusterView {
169 registry: agent.registry().clone(),
170 state: cluster.state(),
171 sessions: sessions.clone(),
172 placement: placement.clone(),
173 self_node: cluster.node_id().to_owned(),
174 agent_role: config.sessions.agent_role.clone(),
175 });
176 server.set_cluster_view(provider);
177
178 Ok(ClusterAgentHandle {
179 agent,
180 cluster,
181 server,
182 server_handle: Some(server_handle),
183 sessions,
184 placement,
185 })
186 }
187}
188
189pub struct ClusterAgentHandle {
191 agent: AgentHandle,
192 cluster: ClusterNode,
193 server: DcpServer,
194 server_handle: Option<DcpServerHandle>,
195 sessions: NodeSessionManagerHandle,
196 placement: PlacementCoordinatorHandle,
197}
198
199impl ClusterAgentHandle {
200 #[must_use]
201 pub fn agent(&self) -> &AgentHandle {
202 &self.agent
203 }
204
205 #[must_use]
206 pub fn cluster(&self) -> &ClusterNode {
207 &self.cluster
208 }
209
210 #[must_use]
211 pub fn sessions(&self) -> &NodeSessionManagerHandle {
212 &self.sessions
213 }
214
215 #[must_use]
216 pub fn server(&self) -> &DcpServer {
217 &self.server
218 }
219
220 #[must_use]
221 pub fn tcp_addr(&self) -> Option<SocketAddr> {
222 self.server_handle
223 .as_ref()
224 .and_then(DcpServerHandle::tcp_addr)
225 }
226
227 #[must_use]
228 pub fn quic_addr(&self) -> Option<SocketAddr> {
229 self.server_handle
230 .as_ref()
231 .and_then(DcpServerHandle::quic_addr)
232 }
233
234 pub async fn shutdown(mut self) -> ClusterAgentResult<()> {
235 self.server.clear_cluster_view();
236 self.server.clear_sharding_view();
237 self.placement.shutdown().await;
238 self.sessions.shutdown().await;
239 let _ = self.cluster.abort().await;
240 if let Some(handle) = self.server_handle.take() {
241 handle.shutdown().await;
242 }
243 self.agent.registry().shutdown()?;
244 Ok(())
245 }
246}
247
248impl Drop for ClusterAgentHandle {
249 fn drop(&mut self) {
250 self.server.clear_cluster_view();
251 self.server.clear_sharding_view();
252 }
253}
254
255#[derive(Clone)]
256pub struct NodeSessionManagerHandle {
257 inner: Arc<NodeSessionManagerInner>,
258}
259
260struct NodeSessionManagerInner {
261 config: NodeSessionConfig,
262 self_node: String,
263 state: Signal<ClusterState>,
264 sessions: tokio::sync::Mutex<BTreeMap<String, PeerSession>>,
265 tasks: Mutex<Vec<JoinHandle<()>>>,
266}
267
268struct PeerSession {
269 member: Member,
270 commands: mpsc::Sender<SessionCommand>,
271 pipe_commands: mpsc::Sender<ShardPipeCommand>,
272 stop: watch::Sender<bool>,
273 state: watch::Receiver<PeerSessionState>,
274 task: JoinHandle<()>,
275}
276
277#[derive(Debug, Clone, PartialEq, Eq)]
278enum PeerSessionState {
279 Connecting,
280 Connected,
281 BackingOff(String),
282 Closed,
283}
284
285impl PeerSessionState {
286 fn as_text(&self) -> String {
287 match self {
288 Self::Connecting => "connecting".to_owned(),
289 Self::Connected => "connected".to_owned(),
290 Self::BackingOff(error) => format!("backing_off:{error}"),
291 Self::Closed => "closed".to_owned(),
292 }
293 }
294}
295
296enum SessionCommand {
297 ListJobs {
298 timeout: Duration,
299 reply: oneshot::Sender<Result<Vec<crate::dcp::proto::JobStatus>, String>>,
300 },
301 SubmitClusterJob {
302 request: SubmitClusterJob,
303 timeout: Duration,
304 reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>,
305 },
306 StartClusterJob {
307 factory_name: String,
308 instance_name: String,
309 params: HashMap<String, String>,
310 assignment: ClusterJobStart,
311 timeout: Duration,
312 reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>,
313 },
314 ClusterJobStatus {
315 name: String,
316 cluster: bool,
317 timeout: Duration,
318 reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>,
319 },
320 DrainClusterJob {
321 name: String,
322 cluster: bool,
323 timeout: Duration,
324 reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>,
325 },
326 StopClusterJob {
327 name: String,
328 cluster: bool,
329 timeout: Duration,
330 reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>,
331 },
332 ListClusterJobs {
333 timeout: Duration,
334 reply: oneshot::Sender<Result<ClusterJobList, String>>,
335 },
336 RememberClusterAssignment {
337 instance_name: String,
338 assignment: ClusterJobStart,
339 timeout: Duration,
340 reply: oneshot::Sender<Result<(), String>>,
341 },
342 AllocateShard {
343 type_name: String,
344 shard_id: String,
345 timeout: Duration,
346 reply: oneshot::Sender<Result<ShardAllocation, String>>,
347 },
348 RememberShardAllocations {
349 table: ShardAllocationTable,
350 timeout: Duration,
351 reply: oneshot::Sender<Result<(), String>>,
352 },
353 GetShardAllocations {
354 type_name: String,
355 timeout: Duration,
356 reply: oneshot::Sender<Result<ShardAllocationTable, String>>,
357 },
358 ForwardShardEnvelopes {
359 batch: ForwardShardEnvelopes,
360 timeout: Duration,
361 reply: oneshot::Sender<Result<ShardEnvelopeBatchResult, String>>,
362 },
363 CompleteShardingAsk {
364 response: CompleteShardingAsk,
365 timeout: Duration,
366 reply: oneshot::Sender<Result<(), String>>,
367 },
368}
369
370enum ShardPipeCommand {
371 Forward { batch: ForwardShardEnvelopes },
372 Reply { response: CompleteShardingAsk },
373}
374
375impl NodeSessionManagerHandle {
376 fn start(
377 config: NodeSessionConfig,
378 self_node: String,
379 state: Signal<ClusterState>,
380 ) -> ClusterAgentResult<Self> {
381 let inner = Arc::new(NodeSessionManagerInner {
382 config,
383 self_node,
384 state,
385 sessions: tokio::sync::Mutex::new(BTreeMap::new()),
386 tasks: Mutex::new(Vec::new()),
387 });
388
389 let manager_task = {
390 let inner = Arc::clone(&inner);
391 tokio::spawn(async move {
392 run_node_session_manager(inner).await;
393 })
394 };
395 {
396 let mut tasks = inner.tasks.lock().expect("node-session tasks poisoned");
397 tasks.push(manager_task);
398 }
399
400 Ok(Self { inner })
401 }
402
403 pub async fn shutdown(&self) {
404 let sessions = {
405 let mut locked = self.inner.sessions.lock().await;
406 std::mem::take(&mut *locked)
407 };
408 for (_, session) in sessions {
409 let _ = session.stop.send(true);
410 session.task.abort();
411 }
412 let tasks = {
413 let mut locked = self
414 .inner
415 .tasks
416 .lock()
417 .expect("node-session tasks poisoned");
418 std::mem::take(&mut *locked)
419 };
420 for task in tasks {
421 task.abort();
422 }
423 }
424
425 async fn list_jobs(
426 &self,
427 node_id: &str,
428 timeout: Duration,
429 ) -> Result<Vec<crate::dcp::proto::JobStatus>, String> {
430 let sender = {
431 let locked = self.inner.sessions.lock().await;
432 locked.get(node_id).map(|session| session.commands.clone())
433 }
434 .ok_or_else(|| "node session unavailable".to_owned())?;
435 let (reply, receiver) = oneshot::channel();
436 sender
437 .try_send(SessionCommand::ListJobs { timeout, reply })
438 .map_err(|error| format!("node session command queue unavailable: {error}"))?;
439 tokio::time::timeout(timeout, receiver)
440 .await
441 .map_err(|_| "node session request timed out".to_owned())?
442 .map_err(|_| "node session closed".to_owned())?
443 }
444
445 async fn submit_cluster_job(
446 &self,
447 node_id: &str,
448 request: SubmitClusterJob,
449 timeout: Duration,
450 ) -> Result<crate::dcp::proto::JobStatus, String> {
451 self.request_status(node_id, timeout, |reply| SessionCommand::SubmitClusterJob {
452 request,
453 timeout,
454 reply,
455 })
456 .await
457 }
458
459 async fn start_cluster_job(
460 &self,
461 node_id: &str,
462 factory_name: String,
463 instance_name: String,
464 params: HashMap<String, String>,
465 assignment: ClusterJobStart,
466 timeout: Duration,
467 ) -> Result<crate::dcp::proto::JobStatus, String> {
468 self.request_status(node_id, timeout, |reply| SessionCommand::StartClusterJob {
469 factory_name,
470 instance_name,
471 params,
472 assignment,
473 timeout,
474 reply,
475 })
476 .await
477 }
478
479 async fn cluster_job_status(
480 &self,
481 node_id: &str,
482 name: String,
483 timeout: Duration,
484 ) -> Result<crate::dcp::proto::JobStatus, String> {
485 self.request_status(node_id, timeout, |reply| SessionCommand::ClusterJobStatus {
486 name,
487 cluster: false,
488 timeout,
489 reply,
490 })
491 .await
492 }
493
494 async fn drain_cluster_job(
495 &self,
496 node_id: &str,
497 name: String,
498 timeout: Duration,
499 ) -> Result<crate::dcp::proto::JobStatus, String> {
500 self.request_status(node_id, timeout, |reply| SessionCommand::DrainClusterJob {
501 name,
502 cluster: false,
503 timeout,
504 reply,
505 })
506 .await
507 }
508
509 async fn stop_cluster_job(
510 &self,
511 node_id: &str,
512 name: String,
513 timeout: Duration,
514 ) -> Result<crate::dcp::proto::JobStatus, String> {
515 self.request_status(node_id, timeout, |reply| SessionCommand::StopClusterJob {
516 name,
517 cluster: false,
518 timeout,
519 reply,
520 })
521 .await
522 }
523
524 async fn forward_cluster_job_status(
525 &self,
526 node_id: &str,
527 name: String,
528 timeout: Duration,
529 ) -> Result<crate::dcp::proto::JobStatus, String> {
530 self.request_status(node_id, timeout, |reply| SessionCommand::ClusterJobStatus {
531 name,
532 cluster: true,
533 timeout,
534 reply,
535 })
536 .await
537 }
538
539 async fn forward_drain_cluster_job(
540 &self,
541 node_id: &str,
542 name: String,
543 timeout: Duration,
544 ) -> Result<crate::dcp::proto::JobStatus, String> {
545 self.request_status(node_id, timeout, |reply| SessionCommand::DrainClusterJob {
546 name,
547 cluster: true,
548 timeout,
549 reply,
550 })
551 .await
552 }
553
554 async fn forward_stop_cluster_job(
555 &self,
556 node_id: &str,
557 name: String,
558 timeout: Duration,
559 ) -> Result<crate::dcp::proto::JobStatus, String> {
560 self.request_status(node_id, timeout, |reply| SessionCommand::StopClusterJob {
561 name,
562 cluster: true,
563 timeout,
564 reply,
565 })
566 .await
567 }
568
569 async fn list_cluster_jobs(
570 &self,
571 node_id: &str,
572 timeout: Duration,
573 ) -> Result<ClusterJobList, String> {
574 let sender = self.session_sender(node_id).await?;
575 let (reply, receiver) = oneshot::channel();
576 sender
577 .try_send(SessionCommand::ListClusterJobs { timeout, reply })
578 .map_err(|error| format!("node session command queue unavailable: {error}"))?;
579 wait_session_reply(timeout, receiver).await
580 }
581
582 async fn remember_cluster_assignment(
583 &self,
584 node_id: &str,
585 instance_name: String,
586 assignment: ClusterJobStart,
587 timeout: Duration,
588 ) -> Result<(), String> {
589 let sender = self.session_sender(node_id).await?;
590 let (reply, receiver) = oneshot::channel();
591 sender
592 .try_send(SessionCommand::RememberClusterAssignment {
593 instance_name,
594 assignment,
595 timeout,
596 reply,
597 })
598 .map_err(|error| format!("node session command queue unavailable: {error}"))?;
599 wait_session_reply(timeout, receiver).await
600 }
601
602 pub async fn allocate_shard(
603 &self,
604 node_id: &str,
605 type_name: String,
606 shard_id: String,
607 timeout: Duration,
608 ) -> Result<ShardAllocation, String> {
609 let sender = self.session_sender(node_id).await?;
610 let (reply, receiver) = oneshot::channel();
611 sender
612 .try_send(SessionCommand::AllocateShard {
613 type_name,
614 shard_id,
615 timeout,
616 reply,
617 })
618 .map_err(|error| format!("node session command queue unavailable: {error}"))?;
619 wait_session_reply(timeout, receiver).await
620 }
621
622 pub async fn remember_shard_allocations(
623 &self,
624 node_id: &str,
625 table: ShardAllocationTable,
626 timeout: Duration,
627 ) -> Result<(), String> {
628 let sender = self.session_sender(node_id).await?;
629 let (reply, receiver) = oneshot::channel();
630 sender
631 .try_send(SessionCommand::RememberShardAllocations {
632 table,
633 timeout,
634 reply,
635 })
636 .map_err(|error| format!("node session command queue unavailable: {error}"))?;
637 wait_session_reply(timeout, receiver).await
638 }
639
640 pub async fn get_shard_allocations(
641 &self,
642 node_id: &str,
643 type_name: String,
644 timeout: Duration,
645 ) -> Result<ShardAllocationTable, String> {
646 let sender = self.session_sender(node_id).await?;
647 let (reply, receiver) = oneshot::channel();
648 sender
649 .try_send(SessionCommand::GetShardAllocations {
650 type_name,
651 timeout,
652 reply,
653 })
654 .map_err(|error| format!("node session command queue unavailable: {error}"))?;
655 wait_session_reply(timeout, receiver).await
656 }
657
658 pub async fn forward_shard_envelopes(
659 &self,
660 node_id: &str,
661 batch: ForwardShardEnvelopes,
662 timeout: Duration,
663 ) -> Result<ShardEnvelopeBatchResult, String> {
664 let sender = self.session_sender(node_id).await?;
665 let (reply, receiver) = oneshot::channel();
666 sender
667 .try_send(SessionCommand::ForwardShardEnvelopes {
668 batch,
669 timeout,
670 reply,
671 })
672 .map_err(|error| format!("node session command queue unavailable: {error}"))?;
673 wait_session_reply(timeout, receiver).await
674 }
675
676 pub async fn complete_sharding_ask(
677 &self,
678 node_id: &str,
679 response: CompleteShardingAsk,
680 timeout: Duration,
681 ) -> Result<(), String> {
682 let sender = self.session_sender(node_id).await?;
683 let (reply, receiver) = oneshot::channel();
684 sender
685 .try_send(SessionCommand::CompleteShardingAsk {
686 response,
687 timeout,
688 reply,
689 })
690 .map_err(|error| format!("node session command queue unavailable: {error}"))?;
691 wait_session_reply(timeout, receiver).await
692 }
693
694 pub async fn forward_shard_pipe_envelopes(
695 &self,
696 node_id: &str,
697 batch: ForwardShardEnvelopes,
698 ) -> Result<(), String> {
699 let sender = self.pipe_sender(node_id).await?;
700 sender
701 .try_send(ShardPipeCommand::Forward { batch })
702 .map_err(|error| format!("node session shard pipe queue unavailable: {error}"))
703 }
704
705 pub async fn complete_sharding_ask_pipe(
706 &self,
707 node_id: &str,
708 response: CompleteShardingAsk,
709 ) -> Result<(), String> {
710 let sender = self.pipe_sender(node_id).await?;
711 sender
712 .try_send(ShardPipeCommand::Reply { response })
713 .map_err(|error| format!("node session shard pipe queue unavailable: {error}"))
714 }
715
716 pub fn try_complete_sharding_ask_pipe(
717 &self,
718 node_id: &str,
719 response: CompleteShardingAsk,
720 ) -> Result<(), String> {
721 let locked = self
722 .inner
723 .sessions
724 .try_lock()
725 .map_err(|_| "node session table busy".to_owned())?;
726 let sender = locked
727 .get(node_id)
728 .map(|session| session.pipe_commands.clone())
729 .ok_or_else(|| "node session unavailable".to_owned())?;
730 sender
731 .try_send(ShardPipeCommand::Reply { response })
732 .map_err(|error| format!("node session shard pipe queue unavailable: {error}"))
733 }
734
735 async fn request_status<F>(
736 &self,
737 node_id: &str,
738 timeout: Duration,
739 make_command: F,
740 ) -> Result<crate::dcp::proto::JobStatus, String>
741 where
742 F: FnOnce(oneshot::Sender<Result<crate::dcp::proto::JobStatus, String>>) -> SessionCommand,
743 {
744 let sender = self.session_sender(node_id).await?;
745 let (reply, receiver) = oneshot::channel();
746 sender
747 .try_send(make_command(reply))
748 .map_err(|error| format!("node session command queue unavailable: {error}"))?;
749 wait_session_reply(timeout, receiver).await
750 }
751
752 async fn session_sender(&self, node_id: &str) -> Result<mpsc::Sender<SessionCommand>, String> {
753 let locked = self.inner.sessions.lock().await;
754 locked
755 .get(node_id)
756 .map(|session| session.commands.clone())
757 .ok_or_else(|| "node session unavailable".to_owned())
758 }
759
760 async fn pipe_sender(&self, node_id: &str) -> Result<mpsc::Sender<ShardPipeCommand>, String> {
761 let locked = self.inner.sessions.lock().await;
762 locked
763 .get(node_id)
764 .map(|session| session.pipe_commands.clone())
765 .ok_or_else(|| "node session unavailable".to_owned())
766 }
767
768 async fn session_states(&self) -> BTreeMap<String, String> {
769 let locked = self.inner.sessions.lock().await;
770 locked
771 .iter()
772 .map(|(node_id, session)| (node_id.clone(), session.state.borrow().clone().as_text()))
773 .collect()
774 }
775}
776
777async fn wait_session_reply<T>(
778 timeout: Duration,
779 receiver: oneshot::Receiver<Result<T, String>>,
780) -> Result<T, String> {
781 tokio::time::timeout(timeout, receiver)
782 .await
783 .map_err(|_| "node session request timed out".to_owned())?
784 .map_err(|_| "node session closed".to_owned())?
785}
786
787async fn run_node_session_manager(inner: Arc<NodeSessionManagerInner>) {
788 let interval = inner
789 .config
790 .reconnect_min_backoff
791 .min(Duration::from_millis(100))
792 .max(Duration::from_millis(10));
793 loop {
794 reconcile_sessions(&inner).await;
795 tokio::time::sleep(interval).await;
796 }
797}
798
799async fn reconcile_sessions(inner: &Arc<NodeSessionManagerInner>) {
800 let state = inner.state.get();
801 let mut wanted = BTreeSet::new();
802 for member in state.members.values() {
803 if eligible_member(inner, member) {
804 wanted.insert(member.node_id.clone());
805 ensure_session(inner, member.clone()).await;
806 }
807 }
808 close_unwanted_sessions(inner, &wanted).await;
809}
810
811fn eligible_member(inner: &NodeSessionManagerInner, member: &Member) -> bool {
812 member.node_id != inner.self_node
813 && member.state == MemberState::Up
814 && member.has_role(&inner.config.agent_role)
815 && member.agent_addr.is_some()
816}
817
818async fn ensure_session(inner: &Arc<NodeSessionManagerInner>, member: Member) {
819 let mut locked = inner.sessions.lock().await;
820 if let Some(existing) = locked.get(&member.node_id)
821 && existing.member.agent_addr == member.agent_addr
822 && existing.member.incarnation == member.incarnation
823 {
824 return;
825 }
826
827 if let Some(existing) = locked.remove(&member.node_id) {
828 let _ = existing.stop.send(true);
829 existing.task.abort();
830 }
831
832 let (commands, receiver) = mpsc::channel(inner.config.command_buffer.max(1));
833 let pipe_buffer = inner.config.command_buffer.saturating_mul(4).max(1);
834 let (pipe_commands, pipe_receiver) = mpsc::channel(pipe_buffer);
835 let (stop, stop_receiver) = watch::channel(false);
836 let (state_sender, state_receiver) = watch::channel(PeerSessionState::Connecting);
837 let config = inner.config.clone();
838 let local_node = inner.self_node.clone();
839 let task_member = member.clone();
840 let task = tokio::spawn(async move {
841 run_peer_session(
842 config,
843 local_node,
844 task_member,
845 receiver,
846 pipe_receiver,
847 stop_receiver,
848 state_sender,
849 )
850 .await;
851 });
852
853 locked.insert(
854 member.node_id.clone(),
855 PeerSession {
856 member,
857 commands,
858 pipe_commands,
859 stop,
860 state: state_receiver,
861 task,
862 },
863 );
864}
865
866async fn close_unwanted_sessions(inner: &Arc<NodeSessionManagerInner>, wanted: &BTreeSet<String>) {
867 let to_close = {
868 let locked = inner.sessions.lock().await;
869 locked
870 .keys()
871 .filter(|node_id| !wanted.contains(*node_id))
872 .cloned()
873 .collect::<Vec<_>>()
874 };
875 for node_id in to_close {
876 close_session(inner, &node_id).await;
877 }
878}
879
880async fn close_session(inner: &Arc<NodeSessionManagerInner>, node_id: &str) {
881 let removed = {
882 let mut locked = inner.sessions.lock().await;
883 locked.remove(node_id)
884 };
885 if let Some(session) = removed {
886 let _ = session.stop.send(true);
887 session.task.abort();
888 }
889}
890
891async fn run_peer_session(
892 config: NodeSessionConfig,
893 local_node: String,
894 member: Member,
895 mut commands: mpsc::Receiver<SessionCommand>,
896 mut pipe_commands: mpsc::Receiver<ShardPipeCommand>,
897 mut stop: watch::Receiver<bool>,
898 state: watch::Sender<PeerSessionState>,
899) {
900 let mut backoff = config.reconnect_min_backoff;
901 loop {
902 if *stop.borrow() {
903 break;
904 }
905 let _ = state.send(PeerSessionState::Connecting);
906 match connect_peer(&config.transport, &local_node, &member).await {
907 Ok(client) => match connect_shard_pipe(&config.transport, &local_node, &member).await {
908 Ok(pipe) => {
909 let _ = state.send(PeerSessionState::Connected);
910 backoff = config.reconnect_min_backoff;
911 if !run_connected_peer(
912 client,
913 pipe,
914 &mut commands,
915 &mut pipe_commands,
916 &mut stop,
917 )
918 .await
919 {
920 break;
921 }
922 }
923 Err(error) => {
924 let _ = state.send(PeerSessionState::BackingOff(error.clone()));
925 if !backoff_or_stop(
926 backoff,
927 error,
928 &mut commands,
929 &mut pipe_commands,
930 &mut stop,
931 )
932 .await
933 {
934 break;
935 }
936 backoff = next_backoff(backoff, config.reconnect_max_backoff);
937 }
938 },
939 Err(error) => {
940 let _ = state.send(PeerSessionState::BackingOff(error.clone()));
941 if !backoff_or_stop(backoff, error, &mut commands, &mut pipe_commands, &mut stop)
942 .await
943 {
944 break;
945 }
946 backoff = next_backoff(backoff, config.reconnect_max_backoff);
947 }
948 }
949 }
950 let _ = state.send(PeerSessionState::Closed);
951}
952
953async fn connect_peer(
954 transport: &NodeSessionTransport,
955 local_node: &str,
956 member: &Member,
957) -> Result<DcpClient, String> {
958 let addr = member
959 .agent_addr
960 .ok_or_else(|| "member did not advertise a DCP agent address".to_owned())?;
961 let mut hello = Hello::new(local_node.to_owned(), ClientKind::ClusterNode);
962 hello.capabilities.push("cluster-node-sessions".to_owned());
963 match transport {
964 NodeSessionTransport::TcpLoopback => {
965 if !addr.ip().is_loopback() {
966 return Err(format!(
967 "plaintext node-session TCP requires loopback, got {addr}"
968 ));
969 }
970 DcpClient::connect_tcp(addr, hello)
971 .await
972 .map_err(|error| error.to_string())
973 }
974 NodeSessionTransport::QuicMtls {
975 server_name,
976 client_config,
977 } => DcpClient::connect_quic(addr, server_name, client_config.clone(), hello)
978 .await
979 .map_err(|error| error.to_string()),
980 }
981}
982
983async fn connect_shard_pipe(
984 transport: &NodeSessionTransport,
985 local_node: &str,
986 member: &Member,
987) -> Result<ShardPipeClient, String> {
988 let addr = member
989 .agent_addr
990 .ok_or_else(|| "member did not advertise a DCP agent address".to_owned())?;
991 let mut hello = Hello::new(local_node.to_owned(), ClientKind::ClusterNode);
992 hello.capabilities.push("cluster-node-sessions".to_owned());
993 hello.capabilities.push("cluster-sharding-pipe".to_owned());
994 match transport {
995 NodeSessionTransport::TcpLoopback => {
996 if !addr.ip().is_loopback() {
997 return Err(format!(
998 "plaintext node-session TCP requires loopback, got {addr}"
999 ));
1000 }
1001 ShardPipeClient::connect_tcp(addr, hello)
1002 .await
1003 .map_err(|error| error.to_string())
1004 }
1005 NodeSessionTransport::QuicMtls {
1006 server_name,
1007 client_config,
1008 } => ShardPipeClient::connect_quic(addr, server_name, client_config.clone(), hello)
1009 .await
1010 .map_err(|error| error.to_string()),
1011 }
1012}
1013
1014async fn run_connected_peer(
1015 client: DcpClient,
1016 pipe: ShardPipeClient,
1017 commands: &mut mpsc::Receiver<SessionCommand>,
1018 pipe_commands: &mut mpsc::Receiver<ShardPipeCommand>,
1019 stop: &mut watch::Receiver<bool>,
1020) -> bool {
1021 loop {
1022 tokio::select! {
1023 changed = stop.changed() => {
1024 return changed.is_ok() && !*stop.borrow();
1025 }
1026 command = pipe_commands.recv() => {
1027 let Some(command) = command else {
1028 return false;
1029 };
1030 let frame = build_shard_pipe_frame(command, pipe_commands);
1031 if pipe.send(frame).await.is_err() {
1032 return true;
1033 }
1034 }
1035 command = commands.recv() => {
1036 let Some(command) = command else {
1037 return false;
1038 };
1039 match command {
1040 SessionCommand::ListJobs { timeout, reply } => {
1041 let result = tokio::time::timeout(timeout, client.list_jobs()).await
1042 .map_err(|_| "peer ListJobs timed out".to_owned())
1043 .and_then(|result| result.map_err(|error| error.to_string()));
1044 let reconnect = result.is_err();
1045 let _ = reply.send(result);
1046 if reconnect {
1047 return true;
1048 }
1049 }
1050 SessionCommand::SubmitClusterJob {
1051 request,
1052 timeout,
1053 reply,
1054 } => {
1055 let timeout_ms = duration_millis_u64(timeout);
1056 let result = tokio::time::timeout(
1057 timeout,
1058 client.submit_cluster_job(
1059 request.factory_name,
1060 request.instance_name,
1061 request.params,
1062 request.placement.unwrap_or_else(default_wire_placement),
1063 timeout_ms,
1064 ),
1065 )
1066 .await
1067 .map_err(|_| "peer SubmitClusterJob timed out".to_owned())
1068 .and_then(|result| result.map_err(|error| error.to_string()));
1069 let reconnect = result.is_err();
1070 let _ = reply.send(result);
1071 if reconnect {
1072 return true;
1073 }
1074 }
1075 SessionCommand::StartClusterJob {
1076 factory_name,
1077 instance_name,
1078 params,
1079 assignment,
1080 timeout,
1081 reply,
1082 } => {
1083 let result = tokio::time::timeout(
1084 timeout,
1085 client.start_cluster_job_on_node(
1086 factory_name,
1087 instance_name,
1088 params,
1089 assignment,
1090 ),
1091 )
1092 .await
1093 .map_err(|_| "peer StartJob timed out".to_owned())
1094 .and_then(|result| result.map_err(|error| error.to_string()));
1095 let reconnect = result.is_err();
1096 let _ = reply.send(result);
1097 if reconnect {
1098 return true;
1099 }
1100 }
1101 SessionCommand::ClusterJobStatus {
1102 name,
1103 cluster,
1104 timeout,
1105 reply,
1106 } => {
1107 let request = async {
1108 if cluster {
1109 client
1110 .cluster_job_status(name, duration_millis_u64(timeout))
1111 .await
1112 } else {
1113 client.job_status(name).await
1114 }
1115 };
1116 let result = tokio::time::timeout(timeout, request).await
1117 .map_err(|_| "peer JobStatus timed out".to_owned())
1118 .and_then(|result| result.map_err(|error| error.to_string()));
1119 let reconnect = result.is_err();
1120 let _ = reply.send(result);
1121 if reconnect {
1122 return true;
1123 }
1124 }
1125 SessionCommand::DrainClusterJob {
1126 name,
1127 cluster,
1128 timeout,
1129 reply,
1130 } => {
1131 let request = async {
1132 if cluster {
1133 client
1134 .drain_cluster_job(name, duration_millis_u64(timeout))
1135 .await
1136 } else {
1137 client.drain_job(name).await
1138 }
1139 };
1140 let result = tokio::time::timeout(timeout, request).await
1141 .map_err(|_| "peer DrainJob timed out".to_owned())
1142 .and_then(|result| result.map_err(|error| error.to_string()));
1143 let reconnect = result.is_err();
1144 let _ = reply.send(result);
1145 if reconnect {
1146 return true;
1147 }
1148 }
1149 SessionCommand::StopClusterJob {
1150 name,
1151 cluster,
1152 timeout,
1153 reply,
1154 } => {
1155 let request = async {
1156 if cluster {
1157 client
1158 .stop_cluster_job(name, duration_millis_u64(timeout))
1159 .await
1160 } else {
1161 client.stop_job(name).await
1162 }
1163 };
1164 let result = tokio::time::timeout(timeout, request).await
1165 .map_err(|_| "peer StopJob timed out".to_owned())
1166 .and_then(|result| result.map_err(|error| error.to_string()));
1167 let reconnect = result.is_err();
1168 let _ = reply.send(result);
1169 if reconnect {
1170 return true;
1171 }
1172 }
1173 SessionCommand::ListClusterJobs { timeout, reply } => {
1174 let result = tokio::time::timeout(
1175 timeout,
1176 client.list_cluster_jobs(duration_millis_u64(timeout)),
1177 )
1178 .await
1179 .map_err(|_| "peer ListClusterJobs timed out".to_owned())
1180 .and_then(|result| result.map_err(|error| error.to_string()));
1181 let reconnect = result.is_err();
1182 let _ = reply.send(result);
1183 if reconnect {
1184 return true;
1185 }
1186 }
1187 SessionCommand::RememberClusterAssignment {
1188 instance_name,
1189 assignment,
1190 timeout,
1191 reply,
1192 } => {
1193 let result = tokio::time::timeout(
1194 timeout,
1195 client.remember_cluster_assignment(instance_name, assignment),
1196 )
1197 .await
1198 .map_err(|_| "peer RememberClusterAssignment timed out".to_owned())
1199 .and_then(|result| result.map_err(|error| error.to_string()));
1200 let reconnect = result.is_err();
1201 let _ = reply.send(result);
1202 if reconnect {
1203 return true;
1204 }
1205 }
1206 SessionCommand::AllocateShard {
1207 type_name,
1208 shard_id,
1209 timeout,
1210 reply,
1211 } => {
1212 let timeout_ms = duration_millis_u64(timeout);
1213 let result = tokio::time::timeout(
1214 timeout,
1215 client.allocate_shard(type_name, shard_id, timeout_ms),
1216 )
1217 .await
1218 .map_err(|_| "peer AllocateShard timed out".to_owned())
1219 .and_then(|result| result.map_err(|error| error.to_string()));
1220 let reconnect = result.is_err();
1221 let _ = reply.send(result);
1222 if reconnect {
1223 return true;
1224 }
1225 }
1226 SessionCommand::RememberShardAllocations {
1227 table,
1228 timeout,
1229 reply,
1230 } => {
1231 let result = tokio::time::timeout(
1232 timeout,
1233 client.remember_shard_allocations(table),
1234 )
1235 .await
1236 .map_err(|_| "peer RememberShardAllocations timed out".to_owned())
1237 .and_then(|result| result.map_err(|error| error.to_string()));
1238 let reconnect = result.is_err();
1239 let _ = reply.send(result);
1240 if reconnect {
1241 return true;
1242 }
1243 }
1244 SessionCommand::GetShardAllocations {
1245 type_name,
1246 timeout,
1247 reply,
1248 } => {
1249 let timeout_ms = duration_millis_u64(timeout);
1250 let result = tokio::time::timeout(
1251 timeout,
1252 client.get_shard_allocations(type_name, timeout_ms),
1253 )
1254 .await
1255 .map_err(|_| "peer GetShardAllocations timed out".to_owned())
1256 .and_then(|result| result.map_err(|error| error.to_string()));
1257 let reconnect = result.is_err();
1258 let _ = reply.send(result);
1259 if reconnect {
1260 return true;
1261 }
1262 }
1263 SessionCommand::ForwardShardEnvelopes {
1264 batch,
1265 timeout,
1266 reply,
1267 } => {
1268 let timeout_ms = duration_millis_u64(timeout);
1269 let result = tokio::time::timeout(
1270 timeout,
1271 client.forward_shard_envelopes(batch, timeout_ms),
1272 )
1273 .await
1274 .map_err(|_| "peer ForwardShardEnvelopes timed out".to_owned())
1275 .and_then(|result| result.map_err(|error| error.to_string()));
1276 let reconnect = result.is_err();
1277 let _ = reply.send(result);
1278 if reconnect {
1279 return true;
1280 }
1281 }
1282 SessionCommand::CompleteShardingAsk {
1283 response,
1284 timeout,
1285 reply,
1286 } => {
1287 let result =
1288 tokio::time::timeout(timeout, client.complete_sharding_ask(response))
1289 .await
1290 .map_err(|_| "peer CompleteShardingAsk timed out".to_owned())
1291 .and_then(|result| result.map_err(|error| error.to_string()));
1292 let reconnect = result.is_err();
1293 let _ = reply.send(result);
1294 if reconnect {
1295 return true;
1296 }
1297 }
1298 }
1299 }
1300 }
1301 }
1302}
1303
1304const SHARD_PIPE_MAX_BATCH_COMMANDS: usize = 1024;
1305
1306fn build_shard_pipe_frame(
1307 first: ShardPipeCommand,
1308 pipe_commands: &mut mpsc::Receiver<ShardPipeCommand>,
1309) -> ShardPipeFrame {
1310 let mut frame = ShardPipeFrame {
1311 forwards: Vec::new(),
1312 replies: Vec::new(),
1313 };
1314 push_shard_pipe_command(&mut frame, first);
1315 for _ in 1..SHARD_PIPE_MAX_BATCH_COMMANDS {
1316 match pipe_commands.try_recv() {
1317 Ok(command) => push_shard_pipe_command(&mut frame, command),
1318 Err(_) => break,
1319 }
1320 }
1321 frame
1322}
1323
1324fn push_shard_pipe_command(frame: &mut ShardPipeFrame, command: ShardPipeCommand) {
1325 match command {
1326 ShardPipeCommand::Forward { batch } => {
1327 if let Some(existing) = frame
1328 .forwards
1329 .iter_mut()
1330 .find(|existing| existing.type_name == batch.type_name)
1331 {
1332 existing.envelopes.extend(batch.envelopes);
1333 } else {
1334 frame.forwards.push(batch);
1335 }
1336 }
1337 ShardPipeCommand::Reply { response } => {
1338 frame.replies.push(response);
1339 }
1340 }
1341}
1342
1343async fn backoff_or_stop(
1344 backoff: Duration,
1345 error: String,
1346 commands: &mut mpsc::Receiver<SessionCommand>,
1347 pipe_commands: &mut mpsc::Receiver<ShardPipeCommand>,
1348 stop: &mut watch::Receiver<bool>,
1349) -> bool {
1350 let sleep = sleep(backoff);
1351 tokio::pin!(sleep);
1352 loop {
1353 tokio::select! {
1354 () = &mut sleep => return true,
1355 changed = stop.changed() => {
1356 return changed.is_ok() && !*stop.borrow();
1357 }
1358 command = pipe_commands.recv() => {
1359 if command.is_none() {
1360 return false;
1361 }
1362 }
1363 command = commands.recv() => {
1364 let Some(command) = command else {
1365 return false;
1366 };
1367 match command {
1368 SessionCommand::ListJobs { reply, .. } => {
1369 let _ = reply.send(Err(format!("node session unreachable: {error}")));
1370 }
1371 SessionCommand::SubmitClusterJob { reply, .. }
1372 | SessionCommand::StartClusterJob { reply, .. }
1373 | SessionCommand::ClusterJobStatus { reply, .. }
1374 | SessionCommand::DrainClusterJob { reply, .. }
1375 | SessionCommand::StopClusterJob { reply, .. } => {
1376 let _ = reply.send(Err(format!("node session unreachable: {error}")));
1377 }
1378 SessionCommand::ListClusterJobs { reply, .. } => {
1379 let _ = reply.send(Err(format!("node session unreachable: {error}")));
1380 }
1381 SessionCommand::RememberClusterAssignment { reply, .. } => {
1382 let _ = reply.send(Err(format!("node session unreachable: {error}")));
1383 }
1384 SessionCommand::AllocateShard { reply, .. } => {
1385 let _ = reply.send(Err(format!("node session unreachable: {error}")));
1386 }
1387 SessionCommand::RememberShardAllocations { reply, .. } => {
1388 let _ = reply.send(Err(format!("node session unreachable: {error}")));
1389 }
1390 SessionCommand::GetShardAllocations { reply, .. } => {
1391 let _ = reply.send(Err(format!("node session unreachable: {error}")));
1392 }
1393 SessionCommand::ForwardShardEnvelopes { reply, .. } => {
1394 let _ = reply.send(Err(format!("node session unreachable: {error}")));
1395 }
1396 SessionCommand::CompleteShardingAsk { reply, .. } => {
1397 let _ = reply.send(Err(format!("node session unreachable: {error}")));
1398 }
1399 }
1400 }
1401 }
1402 }
1403}
1404
1405fn default_wire_placement() -> PlacementSpec {
1406 PlacementSpec {
1407 role_constraint: String::new(),
1408 strategy: PlacementStrategy::LeastJobs as i32,
1409 pinned_node_id: String::new(),
1410 }
1411}
1412
1413fn duration_millis_u64(duration: Duration) -> u64 {
1414 duration.as_millis().min(u128::from(u64::MAX)) as u64
1415}
1416
1417fn next_backoff(current: Duration, max: Duration) -> Duration {
1418 current.saturating_mul(2).min(max)
1419}
1420
1421#[derive(Clone)]
1422struct PlacementCoordinatorHandle {
1423 commands: mpsc::Sender<PlacementCommand>,
1424 tasks: Arc<Mutex<Vec<JoinHandle<()>>>>,
1425}
1426
1427enum PlacementCommand {
1428 Submit {
1429 request: SubmitClusterJob,
1430 timeout: Duration,
1431 reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, DcpError>>,
1432 },
1433 Status {
1434 name: String,
1435 timeout: Duration,
1436 reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, DcpError>>,
1437 },
1438 Drain {
1439 name: String,
1440 timeout: Duration,
1441 reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, DcpError>>,
1442 },
1443 Stop {
1444 name: String,
1445 timeout: Duration,
1446 reply: oneshot::Sender<Result<crate::dcp::proto::JobStatus, DcpError>>,
1447 },
1448 Remember {
1449 instance_name: String,
1450 assignment: ClusterJobStart,
1451 reply: oneshot::Sender<Result<(), DcpError>>,
1452 },
1453 Tick,
1454 Shutdown,
1455}
1456
1457struct PlacementActor {
1458 self_node: String,
1459 agent_role: String,
1460 state: Signal<ClusterState>,
1461 registry: JobRegistryHandle,
1462 sessions: NodeSessionManagerHandle,
1463 factories: DcpJobFactories,
1464 request_timeout: Duration,
1465 active_coordinator: bool,
1466 assignments: BTreeMap<String, ClusterJobMetadata>,
1467}
1468
1469impl PlacementCoordinatorHandle {
1470 fn start(
1471 self_node: String,
1472 agent_role: String,
1473 state: Signal<ClusterState>,
1474 registry: JobRegistryHandle,
1475 sessions: NodeSessionManagerHandle,
1476 factories: DcpJobFactories,
1477 request_timeout: Duration,
1478 ) -> ClusterAgentResult<Self> {
1479 let (commands, receiver) = mpsc::channel(128);
1480 let tasks = Arc::new(Mutex::new(Vec::new()));
1481 let actor = PlacementActor {
1482 self_node,
1483 agent_role,
1484 state,
1485 registry,
1486 sessions,
1487 factories,
1488 request_timeout,
1489 active_coordinator: false,
1490 assignments: BTreeMap::new(),
1491 };
1492 let actor_task = tokio::spawn(run_placement_actor(actor, receiver));
1493 let tick_commands = commands.clone();
1494 let tick_task = tokio::spawn(async move {
1495 let mut interval = tokio::time::interval(Duration::from_millis(50));
1496 loop {
1497 interval.tick().await;
1498 if tick_commands.send(PlacementCommand::Tick).await.is_err() {
1499 break;
1500 }
1501 }
1502 });
1503 {
1504 let mut locked = tasks.lock().expect("placement tasks poisoned");
1505 locked.push(actor_task);
1506 locked.push(tick_task);
1507 }
1508 Ok(Self { commands, tasks })
1509 }
1510
1511 async fn submit_cluster_job(
1512 &self,
1513 request: SubmitClusterJob,
1514 timeout: Duration,
1515 ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
1516 self.request(timeout, |reply| PlacementCommand::Submit {
1517 request,
1518 timeout,
1519 reply,
1520 })
1521 .await
1522 }
1523
1524 async fn cluster_job_status(
1525 &self,
1526 name: String,
1527 timeout: Duration,
1528 ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
1529 self.request(timeout, |reply| PlacementCommand::Status {
1530 name,
1531 timeout,
1532 reply,
1533 })
1534 .await
1535 }
1536
1537 async fn drain_cluster_job(
1538 &self,
1539 name: String,
1540 timeout: Duration,
1541 ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
1542 self.request(timeout, |reply| PlacementCommand::Drain {
1543 name,
1544 timeout,
1545 reply,
1546 })
1547 .await
1548 }
1549
1550 async fn stop_cluster_job(
1551 &self,
1552 name: String,
1553 timeout: Duration,
1554 ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
1555 self.request(timeout, |reply| PlacementCommand::Stop {
1556 name,
1557 timeout,
1558 reply,
1559 })
1560 .await
1561 }
1562
1563 async fn remember_cluster_assignment(
1564 &self,
1565 instance_name: String,
1566 assignment: ClusterJobStart,
1567 ) -> Result<(), DcpError> {
1568 self.request(Duration::from_millis(250), |reply| {
1569 PlacementCommand::Remember {
1570 instance_name,
1571 assignment,
1572 reply,
1573 }
1574 })
1575 .await
1576 }
1577
1578 async fn shutdown(&self) {
1579 let _ = self.commands.send(PlacementCommand::Shutdown).await;
1580 let tasks = {
1581 let mut locked = self.tasks.lock().expect("placement tasks poisoned");
1582 std::mem::take(&mut *locked)
1583 };
1584 for task in tasks {
1585 task.abort();
1586 }
1587 }
1588
1589 async fn request<T, F>(&self, timeout: Duration, make: F) -> Result<T, DcpError>
1590 where
1591 T: Send + 'static,
1592 F: FnOnce(oneshot::Sender<Result<T, DcpError>>) -> PlacementCommand,
1593 {
1594 let (reply, receiver) = oneshot::channel();
1595 self.commands
1596 .send(make(reply))
1597 .await
1598 .map_err(|_| DcpError::response(ResponseStatus::Failed, "placement actor stopped"))?;
1599 tokio::time::timeout(timeout.saturating_add(Duration::from_millis(250)), receiver)
1600 .await
1601 .map_err(|_| {
1602 DcpError::response(
1603 ResponseStatus::DeadlineExceeded,
1604 "placement request timed out",
1605 )
1606 })?
1607 .map_err(|_| DcpError::response(ResponseStatus::Failed, "placement actor stopped"))?
1608 }
1609}
1610
1611async fn run_placement_actor(
1612 mut actor: PlacementActor,
1613 mut receiver: mpsc::Receiver<PlacementCommand>,
1614) {
1615 while let Some(command) = receiver.recv().await {
1616 match command {
1617 PlacementCommand::Submit {
1618 request,
1619 timeout,
1620 reply,
1621 } => {
1622 let result = actor.submit(request, timeout).await;
1623 let _ = reply.send(result);
1624 }
1625 PlacementCommand::Status {
1626 name,
1627 timeout,
1628 reply,
1629 } => {
1630 let result = actor.cluster_job_status(name, timeout).await;
1631 let _ = reply.send(result);
1632 }
1633 PlacementCommand::Drain {
1634 name,
1635 timeout,
1636 reply,
1637 } => {
1638 let result = actor.drain_cluster_job(name, timeout).await;
1639 let _ = reply.send(result);
1640 }
1641 PlacementCommand::Stop {
1642 name,
1643 timeout,
1644 reply,
1645 } => {
1646 let result = actor.stop_cluster_job(name, timeout).await;
1647 let _ = reply.send(result);
1648 }
1649 PlacementCommand::Remember {
1650 instance_name,
1651 assignment,
1652 reply,
1653 } => {
1654 let result = actor.remember(instance_name, assignment);
1655 let _ = reply.send(result);
1656 }
1657 PlacementCommand::Tick => {
1658 let _ = actor.reconcile().await;
1659 }
1660 PlacementCommand::Shutdown => break,
1661 }
1662 }
1663}
1664
1665impl PlacementActor {
1666 async fn submit(
1667 &mut self,
1668 request: SubmitClusterJob,
1669 timeout: Duration,
1670 ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
1671 self.ensure_active_coordinator(timeout).await?;
1672 validate_submit(&request)?;
1673 if self.assignments.contains_key(&request.instance_name) {
1674 return Err(DcpError::response(
1675 ResponseStatus::Conflict,
1676 format!("cluster job already exists: {}", request.instance_name),
1677 ));
1678 }
1679
1680 let placement = placement_spec_from_wire(request.placement.clone())?;
1681 let target = self.choose_submit_target(&placement)?;
1682 let history = ClusterPlacementHistory {
1683 generation: 1,
1684 from_node: None,
1685 to_node: target.clone(),
1686 reason: "submitted".to_owned(),
1687 timestamp: SystemTime::now(),
1688 };
1689 let metadata = ClusterJobMetadata {
1690 factory_name: request.factory_name.clone(),
1691 params: request.params.clone().into_iter().collect(),
1692 placement,
1693 coordinator_node: self.self_node.clone(),
1694 assigned_node: target.clone(),
1695 placement_generation: 1,
1696 history: vec![history],
1697 };
1698 let status = self
1699 .start_on_node(
1700 &target,
1701 request.factory_name,
1702 request.instance_name.clone(),
1703 request.params,
1704 metadata.clone(),
1705 timeout,
1706 )
1707 .await?;
1708 self.assignments
1709 .insert(request.instance_name.clone(), metadata.clone());
1710 self.replicate_assignment(request.instance_name, metadata, timeout)
1711 .await;
1712 Ok(status)
1713 }
1714
1715 async fn cluster_job_status(
1716 &mut self,
1717 name: String,
1718 timeout: Duration,
1719 ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
1720 self.ensure_active_coordinator(timeout).await?;
1721 let metadata = self.assignment_or_rebuild(&name, timeout).await?;
1722 self.status_on_node(&metadata.assigned_node, name, timeout)
1723 .await
1724 }
1725
1726 async fn drain_cluster_job(
1727 &mut self,
1728 name: String,
1729 timeout: Duration,
1730 ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
1731 self.ensure_active_coordinator(timeout).await?;
1732 let metadata = self.assignment_or_rebuild(&name, timeout).await?;
1733 let status = self
1734 .drain_on_node(&metadata.assigned_node, name.clone(), timeout)
1735 .await?;
1736 self.assignments.remove(&name);
1737 Ok(status)
1738 }
1739
1740 async fn stop_cluster_job(
1741 &mut self,
1742 name: String,
1743 timeout: Duration,
1744 ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
1745 self.ensure_active_coordinator(timeout).await?;
1746 let metadata = self.assignment_or_rebuild(&name, timeout).await?;
1747 let status = self
1748 .stop_on_node(&metadata.assigned_node, name.clone(), timeout)
1749 .await?;
1750 self.assignments.remove(&name);
1751 Ok(status)
1752 }
1753
1754 fn remember(
1755 &mut self,
1756 instance_name: String,
1757 assignment: ClusterJobStart,
1758 ) -> Result<(), DcpError> {
1759 let metadata = cluster_metadata_from_wire(assignment)?;
1760 self.assignments.insert(instance_name, metadata);
1761 Ok(())
1762 }
1763
1764 async fn reconcile(&mut self) -> Result<(), DcpError> {
1765 let snapshot = self.state.get();
1766 let is_coordinator = snapshot.is_placement_coordinator(&self.self_node);
1767 drop(snapshot);
1768 if !is_coordinator {
1769 self.active_coordinator = false;
1770 return Ok(());
1771 }
1772 self.ensure_active_coordinator(self.request_timeout).await?;
1773 self.replace_down_members(self.request_timeout).await
1774 }
1775
1776 async fn ensure_active_coordinator(&mut self, timeout: Duration) -> Result<(), DcpError> {
1777 let snapshot = self.state.get();
1778 if !snapshot.is_placement_coordinator(&self.self_node) {
1779 return Err(DcpError::response(
1780 ResponseStatus::Failed,
1781 "this node is not the placement coordinator",
1782 ));
1783 }
1784 drop(snapshot);
1785 if !self.active_coordinator {
1786 self.rebuild_from_registries(timeout).await?;
1787 self.active_coordinator = true;
1788 }
1789 Ok(())
1790 }
1791
1792 async fn assignment_or_rebuild(
1793 &mut self,
1794 name: &str,
1795 timeout: Duration,
1796 ) -> Result<ClusterJobMetadata, DcpError> {
1797 if let Some(metadata) = self.assignments.get(name).cloned() {
1798 return Ok(metadata);
1799 }
1800 self.rebuild_from_registries(timeout).await?;
1801 self.assignments.get(name).cloned().ok_or_else(|| {
1802 DcpError::response(
1803 ResponseStatus::NotFound,
1804 format!("cluster job not found: {name}"),
1805 )
1806 })
1807 }
1808
1809 async fn rebuild_from_registries(&mut self, timeout: Duration) -> Result<(), DcpError> {
1810 let statuses = self.collect_job_statuses(timeout).await?;
1811 for status in statuses {
1812 if status.cluster_job
1813 && status.desired_state != "Stopped"
1814 && let Some(metadata) = metadata_from_status(&status)?
1815 {
1816 let replace = self
1817 .assignments
1818 .get(&status.name)
1819 .map(|existing| {
1820 metadata.placement_generation > existing.placement_generation
1821 || (metadata.placement_generation == existing.placement_generation
1822 && metadata.assigned_node < existing.assigned_node)
1823 })
1824 .unwrap_or(true);
1825 if replace {
1826 self.assignments.insert(status.name.clone(), metadata);
1827 }
1828 }
1829 }
1830 Ok(())
1831 }
1832
1833 async fn replace_down_members(&mut self, timeout: Duration) -> Result<(), DcpError> {
1834 let snapshot = self.state.get();
1835 let down_nodes = snapshot
1836 .members
1837 .values()
1838 .filter(|member| member.state == MemberState::Down)
1839 .map(|member| member.node_id.clone())
1840 .collect::<BTreeSet<_>>();
1841 drop(snapshot);
1842 if down_nodes.is_empty() {
1843 return Ok(());
1844 }
1845
1846 let to_replace = self
1847 .assignments
1848 .iter()
1849 .filter(|(_name, metadata)| down_nodes.contains(&metadata.assigned_node))
1850 .map(|(name, metadata)| (name.clone(), metadata.clone()))
1851 .collect::<Vec<_>>();
1852
1853 for (name, mut metadata) in to_replace {
1854 let old_node = metadata.assigned_node.clone();
1855 let target = self.choose_replacement_target(&metadata.placement, &old_node)?;
1856 metadata.placement_generation = metadata.placement_generation.saturating_add(1);
1857 metadata.assigned_node = target.clone();
1858 metadata.coordinator_node = self.self_node.clone();
1859 metadata.history.push(ClusterPlacementHistory {
1860 generation: metadata.placement_generation,
1861 from_node: Some(old_node.clone()),
1862 to_node: target.clone(),
1863 reason: format!("node_down:{old_node}"),
1864 timestamp: SystemTime::now(),
1865 });
1866 let params = metadata.params.clone().into_iter().collect();
1867 let factory_name = metadata.factory_name.clone();
1868 self.start_on_node(
1869 &target,
1870 factory_name,
1871 name.clone(),
1872 params,
1873 metadata.clone(),
1874 timeout,
1875 )
1876 .await?;
1877 self.assignments.insert(name.clone(), metadata.clone());
1878 self.replicate_assignment(name, metadata, timeout).await;
1879 }
1880 Ok(())
1881 }
1882
1883 fn choose_submit_target(&self, placement: &RegistryPlacementSpec) -> Result<String, DcpError> {
1884 let snapshot = self.state.get();
1885 let candidates = eligible_members(
1886 &snapshot,
1887 &self.agent_role,
1888 placement.role_constraint.as_deref(),
1889 );
1890 if candidates.is_empty() {
1891 return Err(no_eligible_node_error(placement));
1892 }
1893 match &placement.strategy {
1894 RegistryPlacementStrategy::LeastJobs => self.least_loaded(candidates),
1895 RegistryPlacementStrategy::Pinned { node_id } => candidates
1896 .into_iter()
1897 .find(|member| member.node_id == *node_id)
1898 .map(|member| member.node_id.clone())
1899 .ok_or_else(|| {
1900 DcpError::response(
1901 ResponseStatus::NotFound,
1902 format!("pinned node is not eligible for placement: {node_id}"),
1903 )
1904 }),
1905 }
1906 }
1907
1908 fn choose_replacement_target(
1909 &self,
1910 placement: &RegistryPlacementSpec,
1911 old_node: &str,
1912 ) -> Result<String, DcpError> {
1913 let snapshot = self.state.get();
1914 let candidates = eligible_members(
1915 &snapshot,
1916 &self.agent_role,
1917 placement.role_constraint.as_deref(),
1918 )
1919 .into_iter()
1920 .filter(|member| member.node_id != old_node)
1921 .collect::<Vec<_>>();
1922 if candidates.is_empty() {
1923 return Err(no_eligible_node_error(placement));
1924 }
1925 self.least_loaded(candidates)
1926 }
1927
1928 fn least_loaded(&self, candidates: Vec<&Member>) -> Result<String, DcpError> {
1929 candidates
1930 .into_iter()
1931 .min_by(|left, right| {
1932 let left_count = self.jobs_on_node(&left.node_id);
1933 let right_count = self.jobs_on_node(&right.node_id);
1934 left_count
1935 .cmp(&right_count)
1936 .then_with(|| left.node_id.cmp(&right.node_id))
1937 })
1938 .map(|member| member.node_id.clone())
1939 .ok_or_else(|| DcpError::response(ResponseStatus::NotFound, "no eligible node"))
1940 }
1941
1942 fn jobs_on_node(&self, node_id: &str) -> usize {
1943 self.assignments
1944 .values()
1945 .filter(|metadata| metadata.assigned_node == node_id)
1946 .count()
1947 }
1948
1949 async fn start_on_node(
1950 &self,
1951 node_id: &str,
1952 factory_name: String,
1953 instance_name: String,
1954 params: HashMap<String, String>,
1955 metadata: ClusterJobMetadata,
1956 timeout: Duration,
1957 ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
1958 if node_id == self.self_node {
1959 let mut spec = self
1960 .factories
1961 .build(&factory_name, instance_name.clone(), params)?;
1962 spec = spec.with_cluster_metadata(metadata);
1963 let registry = self.registry.clone();
1964 let status = tokio::task::spawn_blocking(move || {
1965 registry.submit(spec)?;
1966 registry.start(instance_name)
1967 })
1968 .await??;
1969 return Ok(wire_job_status(&status));
1970 }
1971
1972 self.sessions
1973 .start_cluster_job(
1974 node_id,
1975 factory_name,
1976 instance_name,
1977 params,
1978 wire_cluster_job_start(&metadata),
1979 timeout,
1980 )
1981 .await
1982 .map_err(session_error)
1983 }
1984
1985 async fn status_on_node(
1986 &self,
1987 node_id: &str,
1988 name: String,
1989 timeout: Duration,
1990 ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
1991 if node_id == self.self_node {
1992 let registry = self.registry.clone();
1993 let status = tokio::task::spawn_blocking(move || registry.status(name)).await??;
1994 return Ok(wire_job_status(&status));
1995 }
1996 self.sessions
1997 .cluster_job_status(node_id, name, timeout)
1998 .await
1999 .map_err(session_error)
2000 }
2001
2002 async fn drain_on_node(
2003 &self,
2004 node_id: &str,
2005 name: String,
2006 timeout: Duration,
2007 ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
2008 if node_id == self.self_node {
2009 let registry = self.registry.clone();
2010 let status = tokio::task::spawn_blocking(move || registry.drain(name)).await??;
2011 return Ok(wire_job_status(&status));
2012 }
2013 self.sessions
2014 .drain_cluster_job(node_id, name, timeout)
2015 .await
2016 .map_err(session_error)
2017 }
2018
2019 async fn stop_on_node(
2020 &self,
2021 node_id: &str,
2022 name: String,
2023 timeout: Duration,
2024 ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
2025 if node_id == self.self_node {
2026 let registry = self.registry.clone();
2027 let status = tokio::task::spawn_blocking(move || registry.stop(name)).await??;
2028 return Ok(wire_job_status(&status));
2029 }
2030 self.sessions
2031 .stop_cluster_job(node_id, name, timeout)
2032 .await
2033 .map_err(session_error)
2034 }
2035
2036 async fn collect_job_statuses(
2037 &self,
2038 timeout: Duration,
2039 ) -> Result<Vec<crate::dcp::proto::JobStatus>, DcpError> {
2040 let snapshot = self.state.get();
2041 let mut statuses = Vec::new();
2042 let registry = self.registry.clone();
2043 let local_jobs = tokio::task::spawn_blocking(move || registry.list()).await??;
2044 statuses.extend(local_jobs.iter().map(wire_job_status));
2045
2046 let mut pending = Vec::new();
2047 for member in snapshot.members.values() {
2048 if member.node_id == self.self_node
2049 || !member.has_role(&self.agent_role)
2050 || member.state != MemberState::Up
2051 || member.unreachable
2052 {
2053 continue;
2054 }
2055 let node_id = member.node_id.clone();
2056 let sessions = self.sessions.clone();
2057 pending.push(tokio::spawn(async move {
2058 sessions.list_jobs(&node_id, timeout).await
2059 }));
2060 }
2061 drop(snapshot);
2062
2063 let deadline = Instant::now() + timeout;
2064 for task in pending {
2065 let remaining = deadline.saturating_duration_since(Instant::now());
2066 if remaining.is_zero() {
2067 break;
2068 }
2069 if let Ok(Ok(Ok(mut peer_statuses))) = tokio::time::timeout(remaining, task).await {
2070 statuses.append(&mut peer_statuses);
2071 }
2072 }
2073 Ok(statuses)
2074 }
2075
2076 async fn replicate_assignment(
2077 &self,
2078 instance_name: String,
2079 metadata: ClusterJobMetadata,
2080 timeout: Duration,
2081 ) {
2082 let snapshot = self.state.get();
2083 let peers = snapshot
2084 .members
2085 .values()
2086 .filter(|member| {
2087 member.node_id != self.self_node
2088 && member.has_role(&self.agent_role)
2089 && member.state == MemberState::Up
2090 && !member.unreachable
2091 })
2092 .map(|member| member.node_id.clone())
2093 .collect::<Vec<_>>();
2094 drop(snapshot);
2095 let assignment = wire_cluster_job_start(&metadata);
2096 let timeout = timeout.min(Duration::from_millis(250));
2097 for node_id in peers {
2098 let _ = self
2099 .sessions
2100 .remember_cluster_assignment(
2101 &node_id,
2102 instance_name.clone(),
2103 assignment.clone(),
2104 timeout,
2105 )
2106 .await;
2107 }
2108 }
2109}
2110
2111fn validate_submit(request: &SubmitClusterJob) -> Result<(), DcpError> {
2112 if request.factory_name.trim().is_empty() || request.instance_name.trim().is_empty() {
2113 return Err(DcpError::response(
2114 ResponseStatus::BadRequest,
2115 "SubmitClusterJob requires factory_name and instance_name",
2116 ));
2117 }
2118 Ok(())
2119}
2120
2121fn eligible_members<'a>(
2122 snapshot: &'a ClusterState,
2123 agent_role: &str,
2124 role_constraint: Option<&str>,
2125) -> Vec<&'a Member> {
2126 snapshot
2127 .members
2128 .values()
2129 .filter(|member| member.state == MemberState::Up && !member.unreachable)
2130 .filter(|member| member.has_role(agent_role))
2131 .filter(|member| role_constraint.is_none_or(|role| member.has_role(role)))
2132 .collect()
2133}
2134
2135fn no_eligible_node_error(placement: &RegistryPlacementSpec) -> DcpError {
2136 let role = placement
2137 .role_constraint
2138 .as_deref()
2139 .map(|role| format!(" with role '{role}'"))
2140 .unwrap_or_default();
2141 DcpError::response(
2142 ResponseStatus::NotFound,
2143 format!("no eligible Up placement nodes{role}"),
2144 )
2145}
2146
2147fn metadata_from_status(
2148 status: &crate::dcp::proto::JobStatus,
2149) -> Result<Option<ClusterJobMetadata>, DcpError> {
2150 if !status.cluster_job {
2151 return Ok(None);
2152 }
2153 let placement = placement_spec_from_wire(status.placement.clone())?;
2154 Ok(Some(ClusterJobMetadata {
2155 factory_name: status.factory_name.clone(),
2156 params: status.params.clone().into_iter().collect(),
2157 placement,
2158 coordinator_node: status.coordinator_node_id.clone(),
2159 assigned_node: status.placement_node_id.clone(),
2160 placement_generation: status.placement_generation,
2161 history: status
2162 .placement_history
2163 .iter()
2164 .map(|history| ClusterPlacementHistory {
2165 generation: history.generation,
2166 from_node: if history.from_node_id.is_empty() {
2167 None
2168 } else {
2169 Some(history.from_node_id.clone())
2170 },
2171 to_node: history.to_node_id.clone(),
2172 reason: history.reason.clone(),
2173 timestamp: std::time::UNIX_EPOCH + Duration::from_millis(history.timestamp_ms),
2174 })
2175 .collect(),
2176 }))
2177}
2178
2179fn session_error(message: String) -> DcpError {
2180 let status = if message.contains("DCP response NotFound") {
2181 ResponseStatus::NotFound
2182 } else if message.contains("DCP response Conflict") {
2183 ResponseStatus::Conflict
2184 } else if message.contains("DCP response BadRequest") {
2185 ResponseStatus::BadRequest
2186 } else if message.contains("DCP response DeadlineExceeded") {
2187 ResponseStatus::DeadlineExceeded
2188 } else {
2189 ResponseStatus::Failed
2190 };
2191 DcpError::response(status, message)
2192}
2193
2194struct ClusterView {
2195 registry: JobRegistryHandle,
2196 state: Signal<ClusterState>,
2197 sessions: NodeSessionManagerHandle,
2198 placement: PlacementCoordinatorHandle,
2199 self_node: String,
2200 agent_role: String,
2201}
2202
2203impl ClusterViewProvider for ClusterView {
2204 fn submit_cluster_job(
2205 &self,
2206 request: SubmitClusterJob,
2207 timeout: Duration,
2208 ) -> crate::dcp::server::ClusterViewFuture<'_, crate::dcp::proto::JobStatus> {
2209 Box::pin(async move { self.submit_cluster_job_inner(request, timeout).await })
2210 }
2211
2212 fn list_cluster_jobs(
2213 &self,
2214 timeout: Duration,
2215 ) -> crate::dcp::server::ClusterViewFuture<'_, ClusterJobList> {
2216 Box::pin(async move { self.list_cluster_jobs_inner(timeout).await })
2217 }
2218
2219 fn cluster_node_info(
2220 &self,
2221 timeout: Duration,
2222 ) -> crate::dcp::server::ClusterViewFuture<'_, ClusterNodeList> {
2223 Box::pin(async move { self.cluster_node_info_inner(timeout).await })
2224 }
2225
2226 fn cluster_job_status(
2227 &self,
2228 name: String,
2229 timeout: Duration,
2230 ) -> crate::dcp::server::ClusterViewFuture<'_, crate::dcp::proto::JobStatus> {
2231 Box::pin(async move { self.cluster_job_status_inner(name, timeout).await })
2232 }
2233
2234 fn drain_cluster_job(
2235 &self,
2236 name: String,
2237 timeout: Duration,
2238 ) -> crate::dcp::server::ClusterViewFuture<'_, crate::dcp::proto::JobStatus> {
2239 Box::pin(async move { self.drain_cluster_job_inner(name, timeout).await })
2240 }
2241
2242 fn stop_cluster_job(
2243 &self,
2244 name: String,
2245 timeout: Duration,
2246 ) -> crate::dcp::server::ClusterViewFuture<'_, crate::dcp::proto::JobStatus> {
2247 Box::pin(async move { self.stop_cluster_job_inner(name, timeout).await })
2248 }
2249
2250 fn remember_cluster_assignment(
2251 &self,
2252 request: RememberClusterAssignment,
2253 ) -> crate::dcp::server::ClusterViewFuture<'_, ()> {
2254 Box::pin(async move { self.remember_cluster_assignment_inner(request).await })
2255 }
2256}
2257
2258impl ClusterView {
2259 async fn submit_cluster_job_inner(
2260 &self,
2261 request: SubmitClusterJob,
2262 timeout: Duration,
2263 ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
2264 if self.is_local_coordinator() {
2265 return self.placement.submit_cluster_job(request, timeout).await;
2266 }
2267 let coordinator = self.coordinator_node_id()?;
2268 self.sessions
2269 .submit_cluster_job(&coordinator, request, timeout)
2270 .await
2271 .map_err(session_error)
2272 }
2273
2274 async fn list_cluster_jobs_inner(&self, timeout: Duration) -> Result<ClusterJobList, DcpError> {
2275 if !self.is_local_coordinator() {
2276 let coordinator = self.coordinator_node_id()?;
2277 return self
2278 .sessions
2279 .list_cluster_jobs(&coordinator, timeout)
2280 .await
2281 .map_err(session_error);
2282 }
2283
2284 let timeout = if timeout.is_zero() {
2285 self.sessions.inner.config.request_timeout
2286 } else {
2287 timeout
2288 };
2289 let snapshot = self.state.get();
2290 let local_member = snapshot.member(&self.self_node).cloned();
2291 let local_jobs = list_local_jobs(self.registry.clone()).await?;
2292 let mut nodes = Vec::new();
2293 nodes.push(ClusterJobNode {
2294 node_id: self.self_node.clone(),
2295 address: local_member
2296 .as_ref()
2297 .map(|member| member.address.to_string())
2298 .unwrap_or_default(),
2299 local: true,
2300 jobs: local_jobs,
2301 });
2302
2303 let mut errors = Vec::new();
2304 let mut pending = Vec::new();
2305 for member in snapshot.members.values() {
2306 if member.node_id == self.self_node || !member.has_role(&self.agent_role) {
2307 continue;
2308 }
2309 if member.state == MemberState::Removed {
2310 continue;
2311 }
2312 if member.state != MemberState::Up || member.unreachable {
2313 errors.push(ClusterNodeError {
2314 node_id: member.node_id.clone(),
2315 message: if member.unreachable {
2316 "member is unreachable".to_owned()
2317 } else {
2318 format!("member is {:?}", member.state)
2319 },
2320 });
2321 continue;
2322 }
2323 let node_id = member.node_id.clone();
2324 let address = member.address.to_string();
2325 let sessions = self.sessions.clone();
2326 pending.push(tokio::spawn(async move {
2327 let result = sessions.list_jobs(&node_id, timeout).await;
2328 (node_id, address, result)
2329 }));
2330 }
2331
2332 let deadline = Instant::now() + timeout;
2333 for task in pending {
2334 let remaining = deadline.saturating_duration_since(Instant::now());
2335 match tokio::time::timeout(remaining, task).await {
2336 Ok(Ok((node_id, address, Ok(jobs)))) => nodes.push(ClusterJobNode {
2337 node_id,
2338 address,
2339 local: false,
2340 jobs,
2341 }),
2342 Ok(Ok((node_id, _address, Err(message)))) => {
2343 errors.push(ClusterNodeError { node_id, message });
2344 }
2345 Ok(Err(error)) => errors.push(ClusterNodeError {
2346 node_id: "unknown".to_owned(),
2347 message: format!("cluster fan-out task failed: {error}"),
2348 }),
2349 Err(_) => errors.push(ClusterNodeError {
2350 node_id: "unknown".to_owned(),
2351 message: "cluster fan-out timed out".to_owned(),
2352 }),
2353 }
2354 }
2355
2356 nodes.sort_by(|left, right| left.node_id.cmp(&right.node_id));
2357 errors.sort_by(|left, right| left.node_id.cmp(&right.node_id));
2358 Ok(ClusterJobList {
2359 partial: !errors.is_empty(),
2360 nodes,
2361 errors,
2362 })
2363 }
2364
2365 async fn cluster_node_info_inner(
2366 &self,
2367 _timeout: Duration,
2368 ) -> Result<ClusterNodeList, DcpError> {
2369 let snapshot = self.state.get();
2370 let session_states = self.sessions.session_states().await;
2371 let mut nodes = snapshot
2372 .members
2373 .values()
2374 .map(|member| ClusterNodeStatus {
2375 node_id: member.node_id.clone(),
2376 member_state: format!("{:?}", member.state),
2377 address: member.address.to_string(),
2378 agent_addr: member
2379 .agent_addr
2380 .map(|addr| addr.to_string())
2381 .unwrap_or_default(),
2382 roles: member.roles.clone(),
2383 unreachable: member.unreachable,
2384 local: member.node_id == self.self_node,
2385 session_state: if member.node_id == self.self_node {
2386 "local".to_owned()
2387 } else {
2388 session_states
2389 .get(&member.node_id)
2390 .cloned()
2391 .unwrap_or_else(|| "not_connected".to_owned())
2392 },
2393 })
2394 .collect::<Vec<_>>();
2395 nodes.sort_by(|left, right| left.node_id.cmp(&right.node_id));
2396
2397 let mut errors = nodes
2398 .iter()
2399 .filter(|node| !node.local)
2400 .filter(|node| node.roles.iter().any(|role| role == &self.agent_role))
2401 .filter(|node| {
2402 node.member_state == "Up" && !node.unreachable && node.session_state != "connected"
2403 })
2404 .map(|node| ClusterNodeError {
2405 node_id: node.node_id.clone(),
2406 message: format!("node session {}", node.session_state),
2407 })
2408 .collect::<Vec<_>>();
2409 errors.sort_by(|left, right| left.node_id.cmp(&right.node_id));
2410
2411 Ok(ClusterNodeList {
2412 partial: !errors.is_empty(),
2413 nodes,
2414 errors,
2415 })
2416 }
2417
2418 async fn cluster_job_status_inner(
2419 &self,
2420 name: String,
2421 timeout: Duration,
2422 ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
2423 if self.is_local_coordinator() {
2424 return self.placement.cluster_job_status(name, timeout).await;
2425 }
2426 let coordinator = self.coordinator_node_id()?;
2427 self.sessions
2428 .forward_cluster_job_status(&coordinator, name, timeout)
2429 .await
2430 .map_err(session_error)
2431 }
2432
2433 async fn drain_cluster_job_inner(
2434 &self,
2435 name: String,
2436 timeout: Duration,
2437 ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
2438 if self.is_local_coordinator() {
2439 return self.placement.drain_cluster_job(name, timeout).await;
2440 }
2441 let coordinator = self.coordinator_node_id()?;
2442 self.sessions
2443 .forward_drain_cluster_job(&coordinator, name, timeout)
2444 .await
2445 .map_err(session_error)
2446 }
2447
2448 async fn stop_cluster_job_inner(
2449 &self,
2450 name: String,
2451 timeout: Duration,
2452 ) -> Result<crate::dcp::proto::JobStatus, DcpError> {
2453 if self.is_local_coordinator() {
2454 return self.placement.stop_cluster_job(name, timeout).await;
2455 }
2456 let coordinator = self.coordinator_node_id()?;
2457 self.sessions
2458 .forward_stop_cluster_job(&coordinator, name, timeout)
2459 .await
2460 .map_err(session_error)
2461 }
2462
2463 async fn remember_cluster_assignment_inner(
2464 &self,
2465 request: RememberClusterAssignment,
2466 ) -> Result<(), DcpError> {
2467 let assignment = request.assignment.ok_or_else(|| {
2468 DcpError::response(
2469 ResponseStatus::BadRequest,
2470 "RememberClusterAssignment requires assignment",
2471 )
2472 })?;
2473 self.placement
2474 .remember_cluster_assignment(request.instance_name, assignment)
2475 .await
2476 }
2477
2478 fn is_local_coordinator(&self) -> bool {
2479 self.state.get().is_placement_coordinator(&self.self_node)
2480 }
2481
2482 fn coordinator_node_id(&self) -> Result<String, DcpError> {
2483 self.state
2484 .get()
2485 .placement_coordinator()
2486 .map(|member| member.node_id.clone())
2487 .ok_or_else(|| {
2488 DcpError::response(
2489 ResponseStatus::Failed,
2490 "no placement coordinator is available",
2491 )
2492 })
2493 }
2494}
2495
2496async fn list_local_jobs(
2497 registry: JobRegistryHandle,
2498) -> Result<Vec<crate::dcp::proto::JobStatus>, DcpError> {
2499 let statuses = tokio::task::spawn_blocking(move || registry.list()).await??;
2500 Ok(statuses.iter().map(wire_job_status).collect())
2501}
2502
2503fn validate_cluster_agent_config(config: &ClusterAgentConfig) -> ClusterAgentResult<()> {
2504 if config.sessions.agent_role.trim().is_empty() {
2505 return Err(ClusterAgentError::InvalidConfig(
2506 "sessions.agent_role must not be empty".to_owned(),
2507 ));
2508 }
2509 if config.sessions.reconnect_min_backoff.is_zero()
2510 || config.sessions.reconnect_max_backoff < config.sessions.reconnect_min_backoff
2511 {
2512 return Err(ClusterAgentError::InvalidConfig(
2513 "node-session reconnect backoff must be non-zero and ordered".to_owned(),
2514 ));
2515 }
2516 match config.sessions.transport {
2517 NodeSessionTransport::TcpLoopback if config.dcp.tcp.is_none() => {
2518 Err(ClusterAgentError::InvalidConfig(
2519 "node-session TCP transport requires a DCP TCP listener".to_owned(),
2520 ))
2521 }
2522 NodeSessionTransport::QuicMtls { .. } if config.dcp.quic.is_none() => {
2523 Err(ClusterAgentError::InvalidConfig(
2524 "node-session QUIC transport requires a DCP QUIC listener".to_owned(),
2525 ))
2526 }
2527 _ => Ok(()),
2528 }
2529}
2530
2531fn advertised_agent_addr(
2532 transport: &NodeSessionTransport,
2533 handle: &DcpServerHandle,
2534) -> ClusterAgentResult<SocketAddr> {
2535 match transport {
2536 NodeSessionTransport::TcpLoopback => handle.tcp_addr().ok_or_else(|| {
2537 ClusterAgentError::InvalidConfig("DCP TCP listener did not start".to_owned())
2538 }),
2539 NodeSessionTransport::QuicMtls { .. } => handle.quic_addr().ok_or_else(|| {
2540 ClusterAgentError::InvalidConfig("DCP QUIC listener did not start".to_owned())
2541 }),
2542 }
2543}
2544
2545fn ensure_agent_role(config: &mut ClusterConfig, role: &str) {
2546 if !config.roles.iter().any(|candidate| candidate == role) {
2547 config.roles.push(role.to_owned());
2548 }
2549}