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
47pub trait ClusterViewProvider: Send + Sync + 'static {
50 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 fn subscribe_cluster_metrics(
63 &self,
64 _interval: Duration,
65 ) -> Option<mpsc::Receiver<MetricSample>> {
66 None
67 }
68
69 fn submit_cluster_job(
71 &self,
72 request: SubmitClusterJob,
73 timeout: Duration,
74 ) -> ClusterViewFuture<'_, WireJobStatus>;
75
76 fn list_cluster_jobs(&self, timeout: Duration) -> ClusterViewFuture<'_, ClusterJobList>;
78
79 fn cluster_node_info(&self, timeout: Duration) -> ClusterViewFuture<'_, ClusterNodeList>;
81
82 fn cluster_job_status(
84 &self,
85 name: String,
86 timeout: Duration,
87 ) -> ClusterViewFuture<'_, WireJobStatus>;
88
89 fn drain_cluster_job(
91 &self,
92 name: String,
93 timeout: Duration,
94 ) -> ClusterViewFuture<'_, WireJobStatus>;
95
96 fn stop_cluster_job(
98 &self,
99 name: String,
100 timeout: Duration,
101 ) -> ClusterViewFuture<'_, WireJobStatus>;
102
103 fn remember_cluster_assignment(
105 &self,
106 request: RememberClusterAssignment,
107 ) -> ClusterViewFuture<'_, ()>;
108
109 fn register_restarted_cluster_assignment(
111 &self,
112 instance_name: String,
113 assignment: ClusterJobStart,
114 timeout: Duration,
115 ) -> ClusterViewFuture<'_, ()>;
116}
117
118pub trait ShardingViewProvider: Send + Sync + 'static {
122 fn allocate_shard(
125 &self,
126 request: ShardAllocationRequest,
127 timeout: Duration,
128 ) -> ClusterViewFuture<'_, ShardAllocation>;
129
130 fn remember_shard_allocations(
132 &self,
133 request: RememberShardAllocations,
134 ) -> ClusterViewFuture<'_, ()>;
135
136 fn get_shard_allocations(
138 &self,
139 type_name: String,
140 timeout: Duration,
141 ) -> ClusterViewFuture<'_, ShardAllocationTable>;
142
143 fn forward_shard_envelopes(
145 &self,
146 request: ForwardShardEnvelopes,
147 timeout: Duration,
148 ) -> ClusterViewFuture<'_, ShardEnvelopeBatchResult>;
149
150 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 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 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 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 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}