Skip to main content

datum_agent/dcp/
server.rs

1use std::{
2    collections::{BTreeMap, HashMap},
3    future::Future,
4    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
5    pin::Pin,
6    sync::{Arc, Mutex},
7    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
8};
9
10use datum::{
11    Sink, StreamInstrumentationRegistry, StreamInstrumentationSnapshot, StreamInstrumentationState,
12};
13use datum_net::quic::quinn;
14use prost::Message as ProstMessage;
15use tokio::{
16    io::{AsyncRead, AsyncWrite, AsyncWriteExt},
17    net::{TcpListener, TcpStream},
18    sync::{mpsc, watch},
19    task::JoinHandle,
20};
21
22use crate::{
23    AgentError, AgentHandle, AgentResult, ClusterJobMetadata, ClusterPlacementHistory, JobEvent,
24    JobEventKind, JobExitReason, JobRegistryHandle, JobSpec, JobStatus as RegistryJobStatus,
25    PlacementSpec as RegistryPlacementSpec, PlacementStrategy as RegistryPlacementStrategy,
26    dcp::{
27        DcpError, DcpResult,
28        frame::{read_frame, write_frame},
29        proto::{
30            ClusterEvent, ClusterJobList, ClusterJobStart, ClusterNodeList, CompleteShardingAsk,
31            ConfigValue, DCP_PROTOCOL_MAJOR, DcpFrame, Event, ForwardShardEnvelopes, Hello,
32            JobList, JobStatus as WireJobStatus, MetricSample, NodeMetric, PlacementSpec,
33            PlacementStrategy, RememberClusterAssignment, RememberShardAllocations, Request,
34            Response, ResponseStatus, ShardAllocation, ShardAllocationRequest,
35            ShardAllocationTable, ShardEnvelopeBatchResult, ShardPipeFrame, StreamMetric,
36            SubmitClusterJob, dcp_frame, request,
37        },
38    },
39};
40
41type DcpJobFactory =
42    dyn Fn(String, HashMap<String, String>) -> AgentResult<JobSpec> + Send + Sync + 'static;
43pub type ClusterViewFuture<'a, T> = Pin<Box<dyn Future<Output = DcpResult<T>> + Send + 'a>>;
44const DEFAULT_CLUSTER_REQUEST_TIMEOUT: Duration = Duration::from_millis(750);
45const TERMINAL_METRIC_GRACE_PERIOD: Duration = Duration::from_secs(3);
46
47/// Provider installed by cluster-aware agents to serve read-only cluster DCP
48/// requests.
49pub trait ClusterViewProvider: Send + Sync + 'static {
50    /// Subscribe to membership, coordinator, and placement events observed by
51    /// this cluster-aware agent.
52    fn subscribe_cluster_events(&self) -> DcpResult<mpsc::Receiver<ClusterEvent>> {
53        Err(DcpError::response(
54            ResponseStatus::Failed,
55            "cluster event subscriptions are unavailable",
56        ))
57    }
58
59    /// Proxy metrics from remote cluster-job placement nodes over the existing
60    /// authenticated peer sessions. Implementations return `None` when remote
61    /// metrics routing is unavailable.
62    fn subscribe_cluster_metrics(
63        &self,
64        _interval: Duration,
65    ) -> Option<mpsc::Receiver<MetricSample>> {
66        None
67    }
68
69    /// Submit a cluster-scoped job to the deterministic placement coordinator.
70    fn submit_cluster_job(
71        &self,
72        request: SubmitClusterJob,
73        timeout: Duration,
74    ) -> ClusterViewFuture<'_, WireJobStatus>;
75
76    /// Aggregate job snapshots from local and peer nodes.
77    fn list_cluster_jobs(&self, timeout: Duration) -> ClusterViewFuture<'_, ClusterJobList>;
78
79    /// Return the local cluster membership and DCP node-session view.
80    fn cluster_node_info(&self, timeout: Duration) -> ClusterViewFuture<'_, ClusterNodeList>;
81
82    /// Return status for a cluster job, routed through the coordinator table.
83    fn cluster_job_status(
84        &self,
85        name: String,
86        timeout: Duration,
87    ) -> ClusterViewFuture<'_, WireJobStatus>;
88
89    /// Drain a cluster job on its current placement node.
90    fn drain_cluster_job(
91        &self,
92        name: String,
93        timeout: Duration,
94    ) -> ClusterViewFuture<'_, WireJobStatus>;
95
96    /// Stop a cluster job on its current placement node.
97    fn stop_cluster_job(
98        &self,
99        name: String,
100        timeout: Duration,
101    ) -> ClusterViewFuture<'_, WireJobStatus>;
102
103    /// Remember an assignment replicated by the active coordinator.
104    fn remember_cluster_assignment(
105        &self,
106        request: RememberClusterAssignment,
107    ) -> ClusterViewFuture<'_, ()>;
108
109    /// Re-register intent when an owning node locally restarts a cluster job.
110    fn register_restarted_cluster_assignment(
111        &self,
112        instance_name: String,
113        assignment: ClusterJobStart,
114        timeout: Duration,
115    ) -> ClusterViewFuture<'_, ()>;
116}
117
118/// Provider installed by cluster sharding to serve shard allocation, region
119/// forwarding, and remote ask replies over the existing node-to-node DCP
120/// sessions.
121pub trait ShardingViewProvider: Send + Sync + 'static {
122    /// Allocate or return the cached owner for one shard. Only the deterministic
123    /// coordinator should accept this request.
124    fn allocate_shard(
125        &self,
126        request: ShardAllocationRequest,
127        timeout: Duration,
128    ) -> ClusterViewFuture<'_, ShardAllocation>;
129
130    /// Remember allocation-table rows replicated by the active coordinator.
131    fn remember_shard_allocations(
132        &self,
133        request: RememberShardAllocations,
134    ) -> ClusterViewFuture<'_, ()>;
135
136    /// Return this region's cached allocation table for coordinator takeover.
137    fn get_shard_allocations(
138        &self,
139        type_name: String,
140        timeout: Duration,
141    ) -> ClusterViewFuture<'_, ShardAllocationTable>;
142
143    /// Enqueue a batch of serialized envelopes at this region.
144    fn forward_shard_envelopes(
145        &self,
146        request: ForwardShardEnvelopes,
147        timeout: Duration,
148    ) -> ClusterViewFuture<'_, ShardEnvelopeBatchResult>;
149
150    /// Complete an ask whose reply port originated on this region.
151    fn complete_sharding_ask(&self, request: CompleteShardingAsk) -> ClusterViewFuture<'_, ()>;
152}
153
154#[derive(Clone, Default)]
155pub struct DcpJobFactories {
156    factories: Arc<Mutex<HashMap<String, Arc<DcpJobFactory>>>>,
157}
158
159impl DcpJobFactories {
160    #[must_use]
161    pub fn new() -> Self {
162        Self::default()
163    }
164
165    pub fn register<F>(&self, name: impl Into<String>, factory: F) -> AgentResult<()>
166    where
167        F: Fn(String, HashMap<String, String>) -> AgentResult<JobSpec> + Send + Sync + 'static,
168    {
169        let name = name.into();
170        if name.trim().is_empty() {
171            return Err(AgentError::InvalidJobName);
172        }
173        self.factories
174            .lock()
175            .expect("DCP job factories poisoned")
176            .insert(name, Arc::new(factory));
177        Ok(())
178    }
179
180    pub(crate) fn build(
181        &self,
182        factory_name: &str,
183        instance_name: String,
184        params: HashMap<String, String>,
185    ) -> DcpResult<JobSpec> {
186        let factory = self
187            .factories
188            .lock()
189            .expect("DCP job factories poisoned")
190            .get(factory_name)
191            .cloned()
192            .ok_or_else(|| {
193                DcpError::response(
194                    ResponseStatus::NotFound,
195                    format!("job factory not found: {factory_name}"),
196                )
197            })?;
198        Ok(factory(instance_name, params)?)
199    }
200}
201
202#[derive(Clone)]
203pub struct DcpTcpServerConfig {
204    pub addr: SocketAddr,
205}
206
207#[derive(Clone)]
208pub struct DcpQuicServerConfig {
209    pub addr: SocketAddr,
210    pub server_config: quinn::ServerConfig,
211}
212
213#[derive(Clone)]
214pub struct DcpServerConfig {
215    pub node_id: String,
216    pub tcp: Option<DcpTcpServerConfig>,
217    pub quic: Option<DcpQuicServerConfig>,
218    pub auth_token: Option<String>,
219    pub metrics_interval: Duration,
220    pub frame_buffer: usize,
221}
222
223impl Default for DcpServerConfig {
224    fn default() -> Self {
225        Self {
226            node_id: format!("datum-agent-{}", std::process::id()),
227            tcp: Some(DcpTcpServerConfig {
228                addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
229            }),
230            quic: None,
231            auth_token: None,
232            metrics_interval: Duration::from_secs(1),
233            frame_buffer: 256,
234        }
235    }
236}
237
238#[derive(Clone)]
239pub struct DcpServer {
240    state: Arc<DcpServerState>,
241}
242
243struct DcpServerState {
244    registry: JobRegistryHandle,
245    instrumentation: StreamInstrumentationRegistry,
246    factories: DcpJobFactories,
247    config: DcpServerConfig,
248    config_store: Mutex<HashMap<String, String>>,
249    cluster_view: Mutex<Option<Arc<dyn ClusterViewProvider>>>,
250    sharding_view: Mutex<Option<Arc<dyn ShardingViewProvider>>>,
251}
252
253impl DcpServer {
254    #[must_use]
255    pub fn new(
256        registry: JobRegistryHandle,
257        instrumentation: StreamInstrumentationRegistry,
258        factories: DcpJobFactories,
259        config: DcpServerConfig,
260    ) -> Self {
261        Self {
262            state: Arc::new(DcpServerState {
263                registry,
264                instrumentation,
265                factories,
266                config,
267                config_store: Mutex::new(HashMap::new()),
268                cluster_view: Mutex::new(None),
269                sharding_view: Mutex::new(None),
270            }),
271        }
272    }
273
274    #[must_use]
275    pub fn from_agent(
276        agent: &AgentHandle,
277        factories: DcpJobFactories,
278        config: DcpServerConfig,
279    ) -> Self {
280        Self::new(
281            agent.registry().clone(),
282            agent.instrumentation_registry().clone(),
283            factories,
284            config,
285        )
286    }
287
288    /// Install or replace the provider used for cluster-wide read-only DCP
289    /// requests.
290    pub fn set_cluster_view(&self, provider: Arc<dyn ClusterViewProvider>) {
291        *self
292            .state
293            .cluster_view
294            .lock()
295            .expect("DCP cluster view provider poisoned") = Some(provider);
296    }
297
298    /// Remove the cluster view provider. Cluster DCP requests will fail until a
299    /// provider is installed again.
300    pub fn clear_cluster_view(&self) {
301        *self
302            .state
303            .cluster_view
304            .lock()
305            .expect("DCP cluster view provider poisoned") = None;
306    }
307
308    /// Install or replace the provider used for cluster-sharding DCP requests.
309    pub fn set_sharding_view(&self, provider: Arc<dyn ShardingViewProvider>) {
310        *self
311            .state
312            .sharding_view
313            .lock()
314            .expect("DCP sharding view provider poisoned") = Some(provider);
315    }
316
317    /// Remove the sharding provider. Sharding DCP requests will fail until a
318    /// provider is installed again.
319    pub fn clear_sharding_view(&self) {
320        *self
321            .state
322            .sharding_view
323            .lock()
324            .expect("DCP sharding view provider poisoned") = None;
325    }
326
327    pub async fn start(&self) -> DcpResult<DcpServerHandle> {
328        let (shutdown_sender, shutdown_receiver) = watch::channel(false);
329        let mut tasks = Vec::new();
330        let mut tcp_addr = None;
331        let mut quic_addr = None;
332
333        if let Some(tcp) = &self.state.config.tcp {
334            ensure_loopback(tcp.addr)?;
335            let listener = TcpListener::bind(tcp.addr).await?;
336            tcp_addr = Some(listener.local_addr()?);
337            let state = Arc::clone(&self.state);
338            let shutdown = shutdown_receiver.clone();
339            tasks.push(tokio::spawn(async move {
340                run_tcp_listener(listener, state, shutdown).await;
341            }));
342        }
343
344        if let Some(quic) = &self.state.config.quic {
345            let endpoint = quinn::Endpoint::server(quic.server_config.clone(), quic.addr)?;
346            quic_addr = Some(endpoint.local_addr()?);
347            let state = Arc::clone(&self.state);
348            let shutdown = shutdown_receiver.clone();
349            tasks.push(tokio::spawn(async move {
350                run_quic_listener(endpoint, state, shutdown).await;
351            }));
352        }
353
354        if tcp_addr.is_none() && quic_addr.is_none() {
355            return Err(DcpError::Protocol(
356                "DCP server has no configured listeners".to_owned(),
357            ));
358        }
359
360        Ok(DcpServerHandle {
361            tcp_addr,
362            quic_addr,
363            shutdown: shutdown_sender,
364            tasks,
365        })
366    }
367}
368
369#[must_use = "dropping DcpServerHandle immediately aborts the DCP listeners; hold it for the server's lifetime"]
370pub struct DcpServerHandle {
371    tcp_addr: Option<SocketAddr>,
372    quic_addr: Option<SocketAddr>,
373    shutdown: watch::Sender<bool>,
374    tasks: Vec<JoinHandle<()>>,
375}
376
377impl DcpServerHandle {
378    #[must_use]
379    pub fn tcp_addr(&self) -> Option<SocketAddr> {
380        self.tcp_addr
381    }
382
383    #[must_use]
384    pub fn quic_addr(&self) -> Option<SocketAddr> {
385        self.quic_addr
386    }
387
388    pub async fn shutdown(mut self) {
389        let _ = self.shutdown.send(true);
390        for task in self.tasks.drain(..) {
391            task.abort();
392            let _ = task.await;
393        }
394    }
395}
396
397impl Drop for DcpServerHandle {
398    fn drop(&mut self) {
399        let _ = self.shutdown.send(true);
400        for task in &self.tasks {
401            task.abort();
402        }
403    }
404}
405
406async fn run_tcp_listener(
407    listener: TcpListener,
408    state: Arc<DcpServerState>,
409    mut shutdown: watch::Receiver<bool>,
410) {
411    loop {
412        tokio::select! {
413            changed = shutdown.changed() => {
414                if changed.is_err() || *shutdown.borrow() {
415                    break;
416                }
417            }
418            accepted = listener.accept() => {
419                let Ok((stream, _peer)) = accepted else {
420                    break;
421                };
422                let state = Arc::clone(&state);
423                tokio::spawn(async move {
424                    let _ = run_tcp_connection(stream, state).await;
425                });
426            }
427        }
428    }
429}
430
431async fn run_tcp_connection(stream: TcpStream, state: Arc<DcpServerState>) -> DcpResult<()> {
432    stream.set_nodelay(true)?;
433    let (reader, writer) = stream.into_split();
434    run_connection(reader, writer, state).await
435}
436
437async fn run_quic_listener(
438    endpoint: quinn::Endpoint,
439    state: Arc<DcpServerState>,
440    mut shutdown: watch::Receiver<bool>,
441) {
442    loop {
443        tokio::select! {
444            changed = shutdown.changed() => {
445                if changed.is_err() || *shutdown.borrow() {
446                    endpoint.close(quinn::VarInt::from_u32(0), b"DCP shutdown");
447                    break;
448                }
449            }
450            incoming = endpoint.accept() => {
451                let Some(incoming) = incoming else {
452                    break;
453                };
454                let state = Arc::clone(&state);
455                let shutdown = shutdown.clone();
456                tokio::spawn(async move {
457                    if let Ok(connection) = incoming.await {
458                        run_quic_connection(connection, state, shutdown).await;
459                    }
460                });
461            }
462        }
463    }
464}
465
466async fn run_quic_connection(
467    connection: quinn::Connection,
468    state: Arc<DcpServerState>,
469    mut shutdown: watch::Receiver<bool>,
470) {
471    loop {
472        tokio::select! {
473            changed = shutdown.changed() => {
474                if changed.is_err() || *shutdown.borrow() {
475                    connection.close(quinn::VarInt::from_u32(0), b"DCP shutdown");
476                    break;
477                }
478            }
479            accepted = connection.accept_bi() => {
480                let Ok((send, recv)) = accepted else {
481                    break;
482                };
483                let state = Arc::clone(&state);
484                tokio::spawn(async move {
485                    let _ = run_connection(recv, send, state).await;
486                });
487            }
488        }
489    }
490}
491
492async fn run_connection<R, W>(mut reader: R, writer: W, state: Arc<DcpServerState>) -> DcpResult<()>
493where
494    R: AsyncRead + Unpin + Send + 'static,
495    W: AsyncWrite + Unpin + Send + 'static,
496{
497    let (outbound, outbound_receiver) = mpsc::channel(state.config.frame_buffer.max(1));
498    let writer_task = tokio::spawn(write_loop(writer, outbound_receiver));
499    let mut subscriptions = SubscriptionTasks::default();
500
501    let Some(first) = read_frame(&mut reader).await? else {
502        return Err(DcpError::Closed);
503    };
504    let hello = match first.frame {
505        Some(dcp_frame::Frame::Hello(hello)) => hello,
506        _ => {
507            let response =
508                Response::error(0, ResponseStatus::BadRequest, "first frame must be Hello");
509            send_frame(&outbound, DcpFrame::response(response)).await?;
510            return Err(DcpError::Protocol("first frame must be Hello".to_owned()));
511        }
512    };
513
514    let hello_response = negotiate_hello(&state, &hello);
515    let accepted = hello_response.response_status() == ResponseStatus::Ok;
516    send_frame(&outbound, DcpFrame::response(hello_response)).await?;
517    if !accepted {
518        return Ok(());
519    }
520
521    while let Some(frame) = read_frame(&mut reader).await? {
522        let request = match frame.frame {
523            Some(dcp_frame::Frame::Request(request)) => request,
524            _ => {
525                return Err(DcpError::Protocol(
526                    "client sent non-request frame after hello".to_owned(),
527                ));
528            }
529        };
530        if matches!(&request.command, Some(request::Command::OpenShardPipe(_))) {
531            send_frame(
532                &outbound,
533                DcpFrame::response(Response::ok(request.request_id, Vec::new())),
534            )
535            .await?;
536            run_shard_pipe_connection(reader, outbound.clone(), Arc::clone(&state)).await?;
537            break;
538        }
539        let (response, subscription_action) =
540            dispatch_request(Arc::clone(&state), request, outbound.clone()).await;
541        send_frame(&outbound, DcpFrame::response(response)).await?;
542        if let Some(action) = subscription_action {
543            match action {
544                SubscriptionAction::Add {
545                    subscription_id,
546                    task,
547                } => subscriptions.insert(subscription_id, task),
548                SubscriptionAction::Remove { subscription_id } => {
549                    subscriptions.remove(subscription_id);
550                }
551            }
552        }
553    }
554
555    drop(subscriptions);
556    drop(outbound);
557    writer_task.await??;
558    Ok(())
559}
560
561async fn run_shard_pipe_connection<R>(
562    mut reader: R,
563    outbound: mpsc::Sender<DcpFrame>,
564    state: Arc<DcpServerState>,
565) -> DcpResult<()>
566where
567    R: AsyncRead + Unpin,
568{
569    while let Some(frame) = read_frame(&mut reader).await? {
570        let pipe = match frame.frame {
571            Some(dcp_frame::Frame::ShardPipe(pipe)) => pipe,
572            _ => {
573                return Err(DcpError::Protocol(
574                    "shard pipe connection received a non-pipe frame".to_owned(),
575                ));
576            }
577        };
578        dispatch_shard_pipe(Arc::clone(&state), pipe).await?;
579    }
580    drop(outbound);
581    Ok(())
582}
583
584async fn dispatch_shard_pipe(state: Arc<DcpServerState>, pipe: ShardPipeFrame) -> DcpResult<()> {
585    let provider = sharding_view_provider(&state)?;
586    for batch in pipe.forwards {
587        provider
588            .forward_shard_envelopes(batch, DEFAULT_CLUSTER_REQUEST_TIMEOUT)
589            .await?;
590    }
591    for reply in pipe.replies {
592        provider.complete_sharding_ask(reply).await?;
593    }
594    Ok(())
595}
596
597async fn write_loop<W>(mut writer: W, mut outbound: mpsc::Receiver<DcpFrame>) -> DcpResult<()>
598where
599    W: AsyncWrite + Unpin,
600{
601    while let Some(frame) = outbound.recv().await {
602        write_frame(&mut writer, &frame).await?;
603    }
604    let _ = writer.shutdown().await;
605    Ok(())
606}
607
608async fn send_frame(sender: &mpsc::Sender<DcpFrame>, frame: DcpFrame) -> DcpResult<()> {
609    sender.send(frame).await.map_err(|_| DcpError::Closed)
610}
611
612enum SubscriptionAction {
613    Add {
614        subscription_id: u64,
615        task: JoinHandle<()>,
616    },
617    Remove {
618        subscription_id: u64,
619    },
620}
621
622#[derive(Default)]
623struct SubscriptionTasks {
624    tasks: HashMap<u64, JoinHandle<()>>,
625}
626
627impl SubscriptionTasks {
628    fn insert(&mut self, subscription_id: u64, task: JoinHandle<()>) {
629        if let Some(previous) = self.tasks.insert(subscription_id, task) {
630            previous.abort();
631        }
632    }
633
634    fn remove(&mut self, subscription_id: u64) {
635        if let Some(task) = self.tasks.remove(&subscription_id) {
636            task.abort();
637        }
638    }
639}
640
641impl Drop for SubscriptionTasks {
642    fn drop(&mut self) {
643        for (_, task) in self.tasks.drain() {
644            task.abort();
645        }
646    }
647}
648
649fn negotiate_hello(state: &DcpServerState, hello: &Hello) -> Response {
650    match protocol_major(&hello.protocol_version) {
651        Some(DCP_PROTOCOL_MAJOR) => {}
652        Some(other) => {
653            return Response::error(
654                0,
655                ResponseStatus::ProtocolMismatch,
656                format!("unsupported DCP major version: {other}"),
657            );
658        }
659        None => {
660            return Response::error(
661                0,
662                ResponseStatus::ProtocolMismatch,
663                format!("invalid DCP protocol version: {}", hello.protocol_version),
664            );
665        }
666    }
667
668    if let Some(expected) = &state.config.auth_token {
669        let actual = hello
670            .auth
671            .as_ref()
672            .map(|auth| auth.bearer_token.as_str())
673            .unwrap_or_default();
674        if actual != expected {
675            return Response::error(0, ResponseStatus::Unauthorized, "invalid DCP token");
676        }
677    }
678
679    Response::ok(0, state.config.node_id.as_bytes().to_vec())
680}
681
682fn protocol_major(version: &str) -> Option<u32> {
683    version.split('.').next()?.parse().ok()
684}
685
686async fn dispatch_request(
687    state: Arc<DcpServerState>,
688    request: Request,
689    outbound: mpsc::Sender<DcpFrame>,
690) -> (Response, Option<SubscriptionAction>) {
691    let request_id = request.request_id;
692    let deadline = request.deadline_ms;
693    let dispatch = async {
694        let command = request.command.ok_or_else(|| {
695            DcpError::response(ResponseStatus::BadRequest, "request missing command")
696        })?;
697        dispatch_command(state, request_id, command, outbound).await
698    };
699
700    let result = if deadline == 0 {
701        dispatch.await
702    } else {
703        match tokio::time::timeout(Duration::from_millis(deadline), dispatch).await {
704            Ok(result) => result,
705            Err(_) => Err(DcpError::response(
706                ResponseStatus::DeadlineExceeded,
707                "request deadline exceeded",
708            )),
709        }
710    };
711
712    match result {
713        Ok(result) => result,
714        Err(error) => (response_for_error(request_id, error), None),
715    }
716}
717
718async fn dispatch_command(
719    state: Arc<DcpServerState>,
720    request_id: u64,
721    command: request::Command,
722    outbound: mpsc::Sender<DcpFrame>,
723) -> DcpResult<(Response, Option<SubscriptionAction>)> {
724    match command {
725        request::Command::ListJobs(_) => {
726            let registry = state.registry.clone();
727            let jobs = registry_call(move || registry.list()).await?;
728            let payload = JobList {
729                jobs: jobs.iter().map(wire_job_status).collect(),
730            }
731            .encode_to_vec();
732            Ok((Response::ok(request_id, payload), None))
733        }
734        request::Command::StartJob(start) => {
735            if start.factory_name.trim().is_empty() || start.instance_name.trim().is_empty() {
736                return Err(DcpError::response(
737                    ResponseStatus::BadRequest,
738                    "StartJob requires factory_name and instance_name",
739                ));
740            }
741            let mut spec = state.factories.build(
742                &start.factory_name,
743                start.instance_name.clone(),
744                start.params,
745            )?;
746            if let Some(cluster) = start.cluster {
747                spec = spec.with_cluster_metadata(cluster_metadata_from_wire(cluster)?);
748            }
749            let registry = state.registry.clone();
750            let name = spec.name.clone();
751            let status = registry_call(move || {
752                registry.submit(spec)?;
753                registry.start(name)
754            })
755            .await?;
756            Ok((status_response(request_id, &status), None))
757        }
758        request::Command::DrainJob(drain) => {
759            if drain.cluster {
760                let provider = cluster_view_provider(&state)?;
761                let timeout = request_timeout(0);
762                let payload = provider
763                    .drain_cluster_job(drain.name, timeout)
764                    .await?
765                    .encode_to_vec();
766                return Ok((Response::ok(request_id, payload), None));
767            }
768            let registry = state.registry.clone();
769            let status = registry_call(move || registry.drain(drain.name)).await?;
770            Ok((status_response(request_id, &status), None))
771        }
772        request::Command::StopJob(stop) => {
773            if stop.cluster {
774                let provider = cluster_view_provider(&state)?;
775                let timeout = request_timeout(0);
776                let payload = provider
777                    .stop_cluster_job(stop.name, timeout)
778                    .await?
779                    .encode_to_vec();
780                return Ok((Response::ok(request_id, payload), None));
781            }
782            let registry = state.registry.clone();
783            let status = registry_call(move || registry.stop(stop.name)).await?;
784            Ok((status_response(request_id, &status), None))
785        }
786        request::Command::RestartJob(restart) => {
787            if restart.cluster {
788                return Err(DcpError::response(
789                    ResponseStatus::BadRequest,
790                    "cluster restart is not implemented in v0.10; drain/stop and submit again",
791                ));
792            }
793            let registry = state.registry.clone();
794            let status = registry_call(move || registry.restart(restart.name)).await?;
795            if let Some(metadata) = status.cluster.as_ref() {
796                let provider = cluster_view_provider(&state)?;
797                provider
798                    .register_restarted_cluster_assignment(
799                        status.name.clone(),
800                        wire_cluster_job_start(metadata),
801                        request_timeout(0),
802                    )
803                    .await?;
804            }
805            Ok((status_response(request_id, &status), None))
806        }
807        request::Command::JobStatus(status) => {
808            if status.cluster {
809                let provider = cluster_view_provider(&state)?;
810                let timeout = request_timeout(0);
811                let payload = provider
812                    .cluster_job_status(status.name, timeout)
813                    .await?
814                    .encode_to_vec();
815                return Ok((Response::ok(request_id, payload), None));
816            }
817            let registry = state.registry.clone();
818            let status = registry_call(move || registry.status(status.name)).await?;
819            Ok((status_response(request_id, &status), None))
820        }
821        request::Command::SubscribeEvents(_) => {
822            let registry = state.registry.clone();
823            let subscription = spawn_event_subscription(request_id, registry, outbound).await?;
824            Ok((
825                Response::ok(request_id, Vec::new()),
826                Some(SubscriptionAction::Add {
827                    subscription_id: request_id,
828                    task: subscription,
829                }),
830            ))
831        }
832        request::Command::SubscribeClusterEvents(_) => {
833            let provider = cluster_view_provider(&state)?;
834            let events = provider.subscribe_cluster_events()?;
835            let subscription =
836                spawn_cluster_event_subscription(request_id, events, outbound).await?;
837            Ok((
838                Response::ok(request_id, Vec::new()),
839                Some(SubscriptionAction::Add {
840                    subscription_id: request_id,
841                    task: subscription,
842                }),
843            ))
844        }
845        request::Command::SubscribeMetrics(metrics) => {
846            let interval = if metrics.interval_ms == 0 {
847                state.config.metrics_interval
848            } else {
849                Duration::from_millis(metrics.interval_ms)
850            }
851            .max(Duration::from_millis(1));
852            let cluster_metrics = if metrics.local_only {
853                None
854            } else {
855                state
856                    .cluster_view
857                    .lock()
858                    .expect("DCP cluster view provider poisoned")
859                    .clone()
860                    .and_then(|provider| provider.subscribe_cluster_metrics(interval))
861            };
862            let subscription = spawn_metrics_subscription(
863                request_id,
864                state.config.node_id.clone(),
865                state.instrumentation.clone(),
866                interval,
867                metrics.job_names,
868                cluster_metrics,
869                outbound,
870            );
871            Ok((
872                Response::ok(request_id, Vec::new()),
873                Some(SubscriptionAction::Add {
874                    subscription_id: request_id,
875                    task: subscription,
876                }),
877            ))
878        }
879        request::Command::UnsubscribeMetrics(unsubscribe) => Ok((
880            Response::ok(request_id, Vec::new()),
881            Some(SubscriptionAction::Remove {
882                subscription_id: unsubscribe.subscription_id,
883            }),
884        )),
885        request::Command::GetConfig(get) => {
886            let value = state
887                .config_store
888                .lock()
889                .expect("DCP config store poisoned")
890                .get(&get.key)
891                .cloned();
892            let payload = ConfigValue {
893                key: get.key,
894                value: value.clone().unwrap_or_default(),
895                existed: value.is_some(),
896            }
897            .encode_to_vec();
898            Ok((Response::ok(request_id, payload), None))
899        }
900        request::Command::PutConfig(put) => {
901            let existed = state
902                .config_store
903                .lock()
904                .expect("DCP config store poisoned")
905                .insert(put.key.clone(), put.value.clone())
906                .is_some();
907            let payload = ConfigValue {
908                key: put.key,
909                value: put.value,
910                existed,
911            }
912            .encode_to_vec();
913            Ok((Response::ok(request_id, payload), None))
914        }
915        request::Command::ListClusterJobs(list) => {
916            let provider = cluster_view_provider(&state)?;
917            let timeout = request_timeout(list.timeout_ms);
918            let payload = provider.list_cluster_jobs(timeout).await?.encode_to_vec();
919            Ok((Response::ok(request_id, payload), None))
920        }
921        request::Command::ClusterNodeInfo(info) => {
922            let provider = cluster_view_provider(&state)?;
923            let timeout = request_timeout(info.timeout_ms);
924            let payload = provider.cluster_node_info(timeout).await?.encode_to_vec();
925            Ok((Response::ok(request_id, payload), None))
926        }
927        request::Command::SubmitClusterJob(submit) => {
928            let provider = cluster_view_provider(&state)?;
929            let timeout = request_timeout(submit.timeout_ms);
930            let payload = provider
931                .submit_cluster_job(submit, timeout)
932                .await?
933                .encode_to_vec();
934            Ok((Response::ok(request_id, payload), None))
935        }
936        request::Command::RememberClusterAssignment(remember) => {
937            let provider = cluster_view_provider(&state)?;
938            provider.remember_cluster_assignment(remember).await?;
939            Ok((Response::ok(request_id, Vec::new()), None))
940        }
941        request::Command::AllocateShard(allocate) => {
942            let provider = sharding_view_provider(&state)?;
943            let timeout = request_timeout(allocate.timeout_ms);
944            let payload = provider
945                .allocate_shard(allocate, timeout)
946                .await?
947                .encode_to_vec();
948            Ok((Response::ok(request_id, payload), None))
949        }
950        request::Command::RememberShardAllocations(remember) => {
951            let provider = sharding_view_provider(&state)?;
952            provider.remember_shard_allocations(remember).await?;
953            Ok((Response::ok(request_id, Vec::new()), None))
954        }
955        request::Command::GetShardAllocations(get) => {
956            let provider = sharding_view_provider(&state)?;
957            let timeout = request_timeout(0);
958            let payload = provider
959                .get_shard_allocations(get.type_name, timeout)
960                .await?
961                .encode_to_vec();
962            Ok((Response::ok(request_id, payload), None))
963        }
964        request::Command::ForwardShardEnvelopes(batch) => {
965            let provider = sharding_view_provider(&state)?;
966            let timeout = request_timeout(0);
967            let payload = provider
968                .forward_shard_envelopes(batch, timeout)
969                .await?
970                .encode_to_vec();
971            Ok((Response::ok(request_id, payload), None))
972        }
973        request::Command::CompleteShardingAsk(reply) => {
974            let provider = sharding_view_provider(&state)?;
975            provider.complete_sharding_ask(reply).await?;
976            Ok((Response::ok(request_id, Vec::new()), None))
977        }
978        request::Command::OpenShardPipe(_) => Err(DcpError::response(
979            ResponseStatus::BadRequest,
980            "OpenShardPipe must be handled by the DCP connection loop",
981        )),
982    }
983}
984
985fn cluster_view_provider(state: &DcpServerState) -> DcpResult<Arc<dyn ClusterViewProvider>> {
986    state
987        .cluster_view
988        .lock()
989        .expect("DCP cluster view provider poisoned")
990        .clone()
991        .ok_or_else(|| {
992            DcpError::response(
993                ResponseStatus::Failed,
994                "cluster view is not configured on this datum-agent",
995            )
996        })
997}
998
999fn sharding_view_provider(state: &DcpServerState) -> DcpResult<Arc<dyn ShardingViewProvider>> {
1000    state
1001        .sharding_view
1002        .lock()
1003        .expect("DCP sharding view provider poisoned")
1004        .clone()
1005        .ok_or_else(|| {
1006            DcpError::response(
1007                ResponseStatus::Failed,
1008                "cluster sharding is not configured on this datum-agent",
1009            )
1010        })
1011}
1012
1013fn request_timeout(timeout_ms: u64) -> Duration {
1014    if timeout_ms == 0 {
1015        DEFAULT_CLUSTER_REQUEST_TIMEOUT
1016    } else {
1017        Duration::from_millis(timeout_ms)
1018    }
1019}
1020
1021async fn registry_call<T, F>(call: F) -> DcpResult<T>
1022where
1023    T: Send + 'static,
1024    F: FnOnce() -> AgentResult<T> + Send + 'static,
1025{
1026    Ok(tokio::task::spawn_blocking(call).await??)
1027}
1028
1029fn status_response(request_id: u64, status: &RegistryJobStatus) -> Response {
1030    Response::ok(request_id, wire_job_status(status).encode_to_vec())
1031}
1032
1033fn response_for_error(request_id: u64, error: DcpError) -> Response {
1034    match error {
1035        DcpError::Response { status, message } => Response::error(request_id, status, message),
1036        DcpError::Agent(error) => {
1037            let status = match &error {
1038                AgentError::InvalidJobName | AgentError::InvalidFactoryParameter { .. } => {
1039                    ResponseStatus::BadRequest
1040                }
1041                AgentError::JobNotFound(_) => ResponseStatus::NotFound,
1042                AgentError::JobAlreadyExists(_) | AgentError::JobAlreadyRunning(_) => {
1043                    ResponseStatus::Conflict
1044                }
1045                AgentError::DrainUnsupported(_)
1046                | AgentError::JobNotRunning(_)
1047                | AgentError::RestartLimitExceeded(_) => ResponseStatus::Failed,
1048                AgentError::RegistryClosed | AgentError::Actor(_) | AgentError::Stream(_) => {
1049                    ResponseStatus::Failed
1050                }
1051            };
1052            Response::error(request_id, status, error.to_string())
1053        }
1054        other => Response::error(request_id, ResponseStatus::Failed, other.to_string()),
1055    }
1056}
1057
1058async fn spawn_event_subscription(
1059    subscription_id: u64,
1060    registry: JobRegistryHandle,
1061    outbound: mpsc::Sender<DcpFrame>,
1062) -> DcpResult<JoinHandle<()>> {
1063    let queue =
1064        tokio::task::spawn_blocking(move || registry.events().run_with(Sink::queue())).await??;
1065    Ok(tokio::task::spawn_blocking(move || {
1066        while let Ok(Some(event)) = queue.pull() {
1067            let frame = DcpFrame::event(subscription_id, wire_event(&event));
1068            if outbound.blocking_send(frame).is_err() {
1069                break;
1070            }
1071        }
1072    }))
1073}
1074
1075async fn spawn_cluster_event_subscription(
1076    subscription_id: u64,
1077    mut events: mpsc::Receiver<ClusterEvent>,
1078    outbound: mpsc::Sender<DcpFrame>,
1079) -> DcpResult<JoinHandle<()>> {
1080    Ok(tokio::spawn(async move {
1081        while let Some(event) = events.recv().await {
1082            if outbound
1083                .send(DcpFrame::cluster_event(subscription_id, event))
1084                .await
1085                .is_err()
1086            {
1087                break;
1088            }
1089        }
1090    }))
1091}
1092
1093fn spawn_metrics_subscription(
1094    subscription_id: u64,
1095    node_id: String,
1096    instrumentation: StreamInstrumentationRegistry,
1097    interval: Duration,
1098    job_names: Vec<String>,
1099    mut cluster_metrics: Option<mpsc::Receiver<MetricSample>>,
1100    outbound: mpsc::Sender<DcpFrame>,
1101) -> JoinHandle<()> {
1102    tokio::spawn(async move {
1103        let mut ticker = tokio::time::interval(interval);
1104        let mut sample_state = MetricSampleState::default();
1105        loop {
1106            if let Some(remote) = cluster_metrics.as_mut() {
1107                tokio::select! {
1108                    _ = ticker.tick() => {
1109                        if !send_metric_sample(
1110                            subscription_id,
1111                            &node_id,
1112                            &instrumentation,
1113                            &job_names,
1114                            &mut sample_state,
1115                            &outbound,
1116                        ) {
1117                            break;
1118                        }
1119                    }
1120                    sample = remote.recv() => {
1121                        match sample {
1122                            Some(sample) => {
1123                                sample_state.remote_streams = sample.streams;
1124                                sample_state.remote_nodes = sample.nodes;
1125                            }
1126                            None => {
1127                                cluster_metrics = None;
1128                                sample_state.remote_streams.clear();
1129                                sample_state.remote_nodes.clear();
1130                            }
1131                        }
1132                    }
1133                }
1134            } else {
1135                ticker.tick().await;
1136                if !send_metric_sample(
1137                    subscription_id,
1138                    &node_id,
1139                    &instrumentation,
1140                    &job_names,
1141                    &mut sample_state,
1142                    &outbound,
1143                ) {
1144                    break;
1145                }
1146            }
1147        }
1148    })
1149}
1150
1151#[derive(Default)]
1152struct MetricSampleState {
1153    remote_streams: Vec<StreamMetric>,
1154    remote_nodes: Vec<NodeMetric>,
1155    previous_cpu: Option<NodeCpuBaseline>,
1156}
1157
1158fn send_metric_sample(
1159    subscription_id: u64,
1160    node_id: &str,
1161    instrumentation: &StreamInstrumentationRegistry,
1162    job_names: &[String],
1163    state: &mut MetricSampleState,
1164    outbound: &mpsc::Sender<DcpFrame>,
1165) -> bool {
1166    let mut streams = sample_local_metrics(instrumentation, job_names);
1167    if job_names.is_empty() {
1168        streams.extend(state.remote_streams.iter().cloned());
1169    } else {
1170        streams.extend(
1171            state
1172                .remote_streams
1173                .iter()
1174                .filter(|metric| metric_matches_jobs(&metric.name, job_names))
1175                .cloned(),
1176        );
1177    }
1178    let local_node = sample_node_metrics(node_id, &mut state.previous_cpu);
1179    let mut nodes = Vec::with_capacity(state.remote_nodes.len().saturating_add(1));
1180    nodes.push(local_node);
1181    nodes.extend_from_slice(&state.remote_nodes);
1182    let sample = MetricSample {
1183        timestamp_ms: system_time_ms(SystemTime::now()),
1184        streams,
1185        nodes,
1186    };
1187    match outbound.try_send(DcpFrame::metric(subscription_id, sample)) {
1188        Ok(()) | Err(mpsc::error::TrySendError::Full(_)) => true,
1189        Err(mpsc::error::TrySendError::Closed(_)) => false,
1190    }
1191}
1192
1193const USER_HZ: u64 = 100;
1194
1195#[derive(Clone, Copy)]
1196struct NodeCpuBaseline {
1197    observed_at: Instant,
1198    cpu_ms: u64,
1199}
1200
1201#[derive(Default)]
1202struct NodeProcessSample {
1203    cpu_ms: u64,
1204    cpu_valid: bool,
1205    rss_bytes: u64,
1206    mem_total_bytes: Option<u64>,
1207    threads: u32,
1208}
1209
1210fn sample_node_metrics(node_id: &str, previous_cpu: &mut Option<NodeCpuBaseline>) -> NodeMetric {
1211    let observed_at = Instant::now();
1212    let process = sample_process_metrics();
1213    let cpu_percent = if process.cpu_valid {
1214        previous_cpu.map_or(0.0, |previous| {
1215            let elapsed_ms = observed_at
1216                .duration_since(previous.observed_at)
1217                .as_secs_f64()
1218                * 1_000.0;
1219            if elapsed_ms <= f64::EPSILON {
1220                0.0
1221            } else {
1222                process.cpu_ms.saturating_sub(previous.cpu_ms) as f64 * 100.0 / elapsed_ms
1223            }
1224        })
1225    } else {
1226        0.0
1227    };
1228    *previous_cpu = process.cpu_valid.then_some(NodeCpuBaseline {
1229        observed_at,
1230        cpu_ms: process.cpu_ms,
1231    });
1232
1233    NodeMetric {
1234        node_id: node_id.to_owned(),
1235        cpu_percent,
1236        cpu_ms: process.cpu_ms,
1237        rss_bytes: process.rss_bytes,
1238        mem_total_bytes: process.mem_total_bytes,
1239        threads: process.threads,
1240        sample_ts_ms: system_time_ms(SystemTime::now()),
1241    }
1242}
1243
1244#[cfg(target_os = "linux")]
1245fn sample_process_metrics() -> NodeProcessSample {
1246    let cpu_ticks = std::fs::read_to_string("/proc/self/stat")
1247        .ok()
1248        .and_then(|stat| {
1249            let after_comm = stat.rsplit_once(") ")?.1;
1250            let fields = after_comm.split_whitespace().collect::<Vec<_>>();
1251            let utime = fields.get(11)?.parse::<u64>().ok()?;
1252            let stime = fields.get(12)?.parse::<u64>().ok()?;
1253            Some(utime.saturating_add(stime))
1254        });
1255    let status = std::fs::read_to_string("/proc/self/status").unwrap_or_default();
1256    let meminfo = std::fs::read_to_string("/proc/meminfo").unwrap_or_default();
1257    NodeProcessSample {
1258        cpu_ms: cpu_ticks
1259            .unwrap_or(0)
1260            .saturating_mul(1_000)
1261            .checked_div(USER_HZ)
1262            .unwrap_or(0),
1263        cpu_valid: cpu_ticks.is_some(),
1264        rss_bytes: proc_status_value(&status, "VmRSS")
1265            .unwrap_or(0)
1266            .saturating_mul(1_024),
1267        mem_total_bytes: proc_status_value(&meminfo, "MemTotal")
1268            .map(|kib| kib.saturating_mul(1_024)),
1269        threads: proc_status_value(&status, "Threads")
1270            .unwrap_or(0)
1271            .min(u64::from(u32::MAX)) as u32,
1272    }
1273}
1274
1275#[cfg(not(target_os = "linux"))]
1276fn sample_process_metrics() -> NodeProcessSample {
1277    NodeProcessSample::default()
1278}
1279
1280#[cfg(target_os = "linux")]
1281fn proc_status_value(contents: &str, key: &str) -> Option<u64> {
1282    contents.lines().find_map(|line| {
1283        let rest = line.strip_prefix(key)?.trim_start_matches(':').trim();
1284        rest.split_whitespace().next()?.parse().ok()
1285    })
1286}
1287
1288fn sample_local_metrics(
1289    instrumentation: &StreamInstrumentationRegistry,
1290    job_names: &[String],
1291) -> Vec<StreamMetric> {
1292    let now = SystemTime::now();
1293    instrumentation
1294        .snapshots()
1295        .into_iter()
1296        .filter_map(|snapshot| {
1297            let terminal_is_stale = snapshot.state != StreamInstrumentationState::Running
1298                && now
1299                    .duration_since(snapshot.state_changed_at)
1300                    .is_ok_and(|age| age >= TERMINAL_METRIC_GRACE_PERIOD);
1301            if terminal_is_stale {
1302                instrumentation.remove(snapshot.id);
1303                return None;
1304            }
1305            metric_matches_jobs(&snapshot.name, job_names).then(|| wire_stream_metric(&snapshot))
1306        })
1307        .collect()
1308}
1309
1310fn metric_matches_jobs(metric_name: &str, job_names: &[String]) -> bool {
1311    job_names.is_empty()
1312        || job_names.iter().any(|job_name| {
1313            metric_name == job_name
1314                || metric_name
1315                    .strip_prefix(job_name)
1316                    .is_some_and(|suffix| suffix.starts_with(':'))
1317        })
1318}
1319
1320pub(crate) fn wire_job_status(status: &RegistryJobStatus) -> WireJobStatus {
1321    let now = Instant::now();
1322    let cluster = status.cluster.as_ref();
1323    WireJobStatus {
1324        name: status.name.clone(),
1325        job_id: status.job_id.0,
1326        state: format!("{:?}", status.state),
1327        desired_state: format!("{:?}", status.desired_state),
1328        generation: status.generation,
1329        starts_total: status.starts_total,
1330        restarts_total: status.restarts_total,
1331        last_start_at_ms: status.last_start_at.map(system_time_ms),
1332        last_exit_at_ms: status.last_exit_at.map(system_time_ms),
1333        last_exit_reason: status
1334            .last_exit_reason
1335            .as_ref()
1336            .map(exit_reason_text)
1337            .unwrap_or_default(),
1338        backoff_remaining_ms: status
1339            .backoff_until
1340            .map(|deadline| duration_ms(deadline.saturating_duration_since(now))),
1341        drain_remaining_ms: status
1342            .drain_deadline
1343            .map(|deadline| duration_ms(deadline.saturating_duration_since(now))),
1344        drain_supported: status.drain_supported,
1345        active_streams: status.active_streams.map(|streams| streams as u64),
1346        cluster_job: cluster.is_some(),
1347        factory_name: cluster
1348            .map(|metadata| metadata.factory_name.clone())
1349            .unwrap_or_default(),
1350        placement: cluster.map(|metadata| wire_placement_spec(&metadata.placement)),
1351        coordinator_node_id: cluster
1352            .map(|metadata| metadata.coordinator_node.clone())
1353            .unwrap_or_default(),
1354        placement_node_id: cluster
1355            .map(|metadata| metadata.assigned_node.clone())
1356            .unwrap_or_default(),
1357        placement_generation: cluster
1358            .map(|metadata| metadata.placement_generation)
1359            .unwrap_or(0),
1360        placement_history: cluster
1361            .map(|metadata| {
1362                metadata
1363                    .history
1364                    .iter()
1365                    .map(wire_placement_history)
1366                    .collect()
1367            })
1368            .unwrap_or_default(),
1369        params: cluster
1370            .map(|metadata| metadata.params.clone().into_iter().collect())
1371            .unwrap_or_default(),
1372    }
1373}
1374
1375pub(crate) fn wire_placement_spec(spec: &RegistryPlacementSpec) -> PlacementSpec {
1376    let (strategy, pinned_node_id) = match &spec.strategy {
1377        RegistryPlacementStrategy::LeastJobs => (PlacementStrategy::LeastJobs, String::new()),
1378        RegistryPlacementStrategy::Pinned { node_id } => {
1379            (PlacementStrategy::Pinned, node_id.clone())
1380        }
1381    };
1382    PlacementSpec {
1383        role_constraint: spec.role_constraint.clone().unwrap_or_default(),
1384        strategy: strategy as i32,
1385        pinned_node_id,
1386    }
1387}
1388
1389pub(crate) fn placement_spec_from_wire(
1390    spec: Option<PlacementSpec>,
1391) -> DcpResult<RegistryPlacementSpec> {
1392    let spec = spec.unwrap_or(PlacementSpec {
1393        role_constraint: String::new(),
1394        strategy: PlacementStrategy::LeastJobs as i32,
1395        pinned_node_id: String::new(),
1396    });
1397    let role_constraint = normalize_optional_string(spec.role_constraint);
1398    let strategy =
1399        match PlacementStrategy::try_from(spec.strategy).unwrap_or(PlacementStrategy::LeastJobs) {
1400            PlacementStrategy::LeastJobs => RegistryPlacementStrategy::LeastJobs,
1401            PlacementStrategy::Pinned => {
1402                let node_id = spec.pinned_node_id.trim().to_owned();
1403                if node_id.is_empty() {
1404                    return Err(DcpError::response(
1405                        ResponseStatus::BadRequest,
1406                        "Pinned placement requires pinned_node_id",
1407                    ));
1408                }
1409                RegistryPlacementStrategy::Pinned { node_id }
1410            }
1411        };
1412    Ok(RegistryPlacementSpec {
1413        role_constraint,
1414        strategy,
1415    })
1416}
1417
1418fn wire_placement_history(
1419    history: &ClusterPlacementHistory,
1420) -> crate::dcp::proto::ClusterPlacementHistory {
1421    crate::dcp::proto::ClusterPlacementHistory {
1422        generation: history.generation,
1423        from_node_id: history.from_node.clone().unwrap_or_default(),
1424        to_node_id: history.to_node.clone(),
1425        reason: history.reason.clone(),
1426        timestamp_ms: system_time_ms(history.timestamp),
1427    }
1428}
1429
1430fn placement_history_from_wire(
1431    history: crate::dcp::proto::ClusterPlacementHistory,
1432) -> ClusterPlacementHistory {
1433    ClusterPlacementHistory {
1434        generation: history.generation,
1435        from_node: normalize_optional_string(history.from_node_id),
1436        to_node: history.to_node_id,
1437        reason: history.reason,
1438        timestamp: UNIX_EPOCH + Duration::from_millis(history.timestamp_ms),
1439    }
1440}
1441
1442pub(crate) fn cluster_metadata_from_wire(
1443    cluster: ClusterJobStart,
1444) -> DcpResult<ClusterJobMetadata> {
1445    Ok(ClusterJobMetadata {
1446        factory_name: cluster.factory_name,
1447        params: cluster.params.into_iter().collect::<BTreeMap<_, _>>(),
1448        placement: placement_spec_from_wire(cluster.placement)?,
1449        coordinator_node: cluster.coordinator_node_id,
1450        assigned_node: cluster.assigned_node_id,
1451        placement_generation: cluster.placement_generation,
1452        history: cluster
1453            .history
1454            .into_iter()
1455            .map(placement_history_from_wire)
1456            .collect(),
1457    })
1458}
1459
1460pub(crate) fn wire_cluster_job_start(metadata: &ClusterJobMetadata) -> ClusterJobStart {
1461    ClusterJobStart {
1462        factory_name: metadata.factory_name.clone(),
1463        params: metadata.params.clone().into_iter().collect(),
1464        placement: Some(wire_placement_spec(&metadata.placement)),
1465        coordinator_node_id: metadata.coordinator_node.clone(),
1466        assigned_node_id: metadata.assigned_node.clone(),
1467        placement_generation: metadata.placement_generation,
1468        history: metadata
1469            .history
1470            .iter()
1471            .map(wire_placement_history)
1472            .collect(),
1473    }
1474}
1475
1476fn wire_event(event: &JobEvent) -> Event {
1477    let (kind, detail) = event_kind_text(&event.kind);
1478    Event {
1479        sequence: event.sequence,
1480        timestamp_ms: system_time_ms(event.timestamp),
1481        name: event.name.clone(),
1482        job_id: event.job_id.0,
1483        generation: event.generation,
1484        kind,
1485        detail,
1486    }
1487}
1488
1489fn wire_stream_metric(snapshot: &StreamInstrumentationSnapshot) -> StreamMetric {
1490    StreamMetric {
1491        id: snapshot.id.get(),
1492        name: snapshot.name.clone(),
1493        elements_through: snapshot.elements_through,
1494        restarts: snapshot.restarts,
1495        state: instrumentation_state_text(snapshot.state).to_owned(),
1496        started_at_ms: system_time_ms(snapshot.started_at),
1497        state_changed_at_ms: system_time_ms(snapshot.state_changed_at),
1498        finished_at_ms: snapshot.finished_at.map(system_time_ms),
1499        uptime_ms: duration_ms(snapshot.uptime),
1500    }
1501}
1502
1503fn instrumentation_state_text(state: StreamInstrumentationState) -> &'static str {
1504    match state {
1505        StreamInstrumentationState::Running => "Running",
1506        StreamInstrumentationState::Draining => "Draining",
1507        StreamInstrumentationState::Completed => "Completed",
1508        StreamInstrumentationState::Failed => "Failed",
1509    }
1510}
1511
1512fn event_kind_text(kind: &JobEventKind) -> (String, String) {
1513    match kind {
1514        JobEventKind::Submitted => ("Submitted".to_owned(), String::new()),
1515        JobEventKind::Started => ("Started".to_owned(), String::new()),
1516        JobEventKind::Failed { reason } => ("Failed".to_owned(), exit_reason_text(reason)),
1517        JobEventKind::RestartScheduled { delay } => (
1518            "RestartScheduled".to_owned(),
1519            format!("delay_ms={}", duration_ms(*delay)),
1520        ),
1521        JobEventKind::Restarted {
1522            previous_generation,
1523        } => (
1524            "Restarted".to_owned(),
1525            format!("previous_generation={previous_generation}"),
1526        ),
1527        JobEventKind::Draining => ("Draining".to_owned(), String::new()),
1528        JobEventKind::Drained => ("Drained".to_owned(), String::new()),
1529        JobEventKind::Stopped { reason } => ("Stopped".to_owned(), exit_reason_text(reason)),
1530        JobEventKind::Completed => ("Completed".to_owned(), String::new()),
1531    }
1532}
1533
1534fn exit_reason_text(reason: &JobExitReason) -> String {
1535    match reason {
1536        JobExitReason::Completed => "Completed".to_owned(),
1537        JobExitReason::Failed(error) => format!("Failed({error})"),
1538        JobExitReason::Drained => "Drained".to_owned(),
1539        JobExitReason::Stopped => "Stopped".to_owned(),
1540        JobExitReason::DrainTimedOut => "DrainTimedOut".to_owned(),
1541    }
1542}
1543
1544fn system_time_ms(time: SystemTime) -> u64 {
1545    time.duration_since(UNIX_EPOCH)
1546        .map(duration_ms)
1547        .unwrap_or(0)
1548}
1549
1550fn duration_ms(duration: Duration) -> u64 {
1551    duration.as_millis().min(u128::from(u64::MAX)) as u64
1552}
1553
1554fn normalize_optional_string(value: String) -> Option<String> {
1555    let value = value.trim().to_owned();
1556    if value.is_empty() { None } else { Some(value) }
1557}
1558
1559fn ensure_loopback(addr: SocketAddr) -> DcpResult<()> {
1560    if addr.ip().is_loopback() {
1561        return Ok(());
1562    }
1563    Err(DcpError::Protocol(format!(
1564        "plaintext DCP TCP listener must bind loopback, got {addr}"
1565    )))
1566}
1567
1568fn client_bind_addr(remote_addr: SocketAddr) -> SocketAddr {
1569    if remote_addr.is_ipv6() {
1570        SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 0)
1571    } else {
1572        SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0)
1573    }
1574}
1575
1576pub(crate) async fn connect_quic_stream(
1577    addr: SocketAddr,
1578    server_name: &str,
1579    client_config: quinn::ClientConfig,
1580) -> DcpResult<(
1581    quinn::Endpoint,
1582    quinn::Connection,
1583    quinn::RecvStream,
1584    quinn::SendStream,
1585)> {
1586    let mut endpoint = quinn::Endpoint::client(client_bind_addr(addr))?;
1587    endpoint.set_default_client_config(client_config);
1588    let connection = endpoint
1589        .connect(addr, server_name)
1590        .map_err(|error| DcpError::Protocol(error.to_string()))?
1591        .await
1592        .map_err(|error| DcpError::Protocol(error.to_string()))?;
1593    let (send, recv) = connection
1594        .open_bi()
1595        .await
1596        .map_err(|error| DcpError::Protocol(error.to_string()))?;
1597    Ok((endpoint, connection, recv, send))
1598}