Skip to main content

datum_agent/
cluster.rs

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
35/// Role used by cluster-aware agents to discover peers.
36pub const AGENT_ROLE: &str = "agent";
37
38/// Result type for cluster-aware agent startup and shutdown.
39pub type ClusterAgentResult<T> = Result<T, ClusterAgentError>;
40
41/// Errors returned by the cluster-aware agent entrypoint.
42#[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/// Transport used for node-to-node DCP sessions.
55#[derive(Clone)]
56pub enum NodeSessionTransport {
57    /// Plaintext TCP is accepted only for loopback development and tests.
58    TcpLoopback,
59    /// QUIC with mTLS. The same client config is used for every peer.
60    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/// Configuration for DCP node sessions maintained by a cluster-aware agent.
79#[derive(Clone, Debug)]
80pub struct NodeSessionConfig {
81    /// Role that marks a member as running a DCP agent endpoint.
82    pub agent_role: String,
83    /// Node-to-node DCP transport.
84    pub transport: NodeSessionTransport,
85    /// Initial reconnect backoff.
86    pub reconnect_min_backoff: Duration,
87    /// Maximum reconnect backoff.
88    pub reconnect_max_backoff: Duration,
89    /// Default per-peer fan-out timeout.
90    pub request_timeout: Duration,
91    /// Bounded per-session command queue capacity.
92    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/// Combined configuration for one cluster-aware Datum agent.
109#[derive(Clone, Default)]
110pub struct ClusterAgentConfig {
111    /// Local job registry configuration.
112    pub agent: AgentConfig,
113    /// Membership configuration.
114    pub cluster: ClusterConfig,
115    /// DCP listener configuration.
116    pub dcp: DcpServerConfig,
117    /// Node-session configuration.
118    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
133/// Entrypoint for starting membership, the local agent, DCP, and node sessions
134/// together.
135pub 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
189/// Running cluster-aware agent.
190pub 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}