1use std::{
2 collections::HashMap,
3 net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
4 sync::{Arc, Mutex},
5 time::{Duration, Instant, SystemTime, UNIX_EPOCH},
6};
7
8use datum::{
9 Sink, StreamInstrumentationRegistry, StreamInstrumentationSnapshot, StreamInstrumentationState,
10};
11use datum_net::quic::quinn;
12use prost::Message as ProstMessage;
13use tokio::{
14 io::{AsyncRead, AsyncWrite, AsyncWriteExt},
15 net::{TcpListener, TcpStream},
16 sync::{mpsc, watch},
17 task::JoinHandle,
18};
19
20use crate::{
21 AgentError, AgentHandle, AgentResult, JobEvent, JobEventKind, JobExitReason, JobRegistryHandle,
22 JobSpec, JobStatus as RegistryJobStatus,
23 dcp::{
24 DcpError, DcpResult,
25 frame::{read_frame, write_frame},
26 proto::{
27 ConfigValue, DCP_PROTOCOL_MAJOR, DcpFrame, Event, Hello, JobList,
28 JobStatus as WireJobStatus, MetricSample, Request, Response, ResponseStatus,
29 StreamMetric, dcp_frame, request,
30 },
31 },
32};
33
34type DcpJobFactory =
35 dyn Fn(String, HashMap<String, String>) -> AgentResult<JobSpec> + Send + Sync + 'static;
36
37#[derive(Clone, Default)]
38pub struct DcpJobFactories {
39 factories: Arc<Mutex<HashMap<String, Arc<DcpJobFactory>>>>,
40}
41
42impl DcpJobFactories {
43 #[must_use]
44 pub fn new() -> Self {
45 Self::default()
46 }
47
48 pub fn register<F>(&self, name: impl Into<String>, factory: F) -> AgentResult<()>
49 where
50 F: Fn(String, HashMap<String, String>) -> AgentResult<JobSpec> + Send + Sync + 'static,
51 {
52 let name = name.into();
53 if name.trim().is_empty() {
54 return Err(AgentError::InvalidJobName);
55 }
56 self.factories
57 .lock()
58 .expect("DCP job factories poisoned")
59 .insert(name, Arc::new(factory));
60 Ok(())
61 }
62
63 fn build(
64 &self,
65 factory_name: &str,
66 instance_name: String,
67 params: HashMap<String, String>,
68 ) -> DcpResult<JobSpec> {
69 let factory = self
70 .factories
71 .lock()
72 .expect("DCP job factories poisoned")
73 .get(factory_name)
74 .cloned()
75 .ok_or_else(|| {
76 DcpError::response(
77 ResponseStatus::NotFound,
78 format!("job factory not found: {factory_name}"),
79 )
80 })?;
81 Ok(factory(instance_name, params)?)
82 }
83}
84
85#[derive(Clone)]
86pub struct DcpTcpServerConfig {
87 pub addr: SocketAddr,
88}
89
90#[derive(Clone)]
91pub struct DcpQuicServerConfig {
92 pub addr: SocketAddr,
93 pub server_config: quinn::ServerConfig,
94}
95
96#[derive(Clone)]
97pub struct DcpServerConfig {
98 pub node_id: String,
99 pub tcp: Option<DcpTcpServerConfig>,
100 pub quic: Option<DcpQuicServerConfig>,
101 pub auth_token: Option<String>,
102 pub metrics_interval: Duration,
103 pub frame_buffer: usize,
104}
105
106impl Default for DcpServerConfig {
107 fn default() -> Self {
108 Self {
109 node_id: format!("datum-agent-{}", std::process::id()),
110 tcp: Some(DcpTcpServerConfig {
111 addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0),
112 }),
113 quic: None,
114 auth_token: None,
115 metrics_interval: Duration::from_secs(1),
116 frame_buffer: 256,
117 }
118 }
119}
120
121#[derive(Clone)]
122pub struct DcpServer {
123 state: Arc<DcpServerState>,
124}
125
126struct DcpServerState {
127 registry: JobRegistryHandle,
128 instrumentation: StreamInstrumentationRegistry,
129 factories: DcpJobFactories,
130 config: DcpServerConfig,
131 config_store: Mutex<HashMap<String, String>>,
132}
133
134impl DcpServer {
135 #[must_use]
136 pub fn new(
137 registry: JobRegistryHandle,
138 instrumentation: StreamInstrumentationRegistry,
139 factories: DcpJobFactories,
140 config: DcpServerConfig,
141 ) -> Self {
142 Self {
143 state: Arc::new(DcpServerState {
144 registry,
145 instrumentation,
146 factories,
147 config,
148 config_store: Mutex::new(HashMap::new()),
149 }),
150 }
151 }
152
153 #[must_use]
154 pub fn from_agent(
155 agent: &AgentHandle,
156 factories: DcpJobFactories,
157 config: DcpServerConfig,
158 ) -> Self {
159 Self::new(
160 agent.registry().clone(),
161 agent.instrumentation_registry().clone(),
162 factories,
163 config,
164 )
165 }
166
167 pub async fn start(&self) -> DcpResult<DcpServerHandle> {
168 let (shutdown_sender, shutdown_receiver) = watch::channel(false);
169 let mut tasks = Vec::new();
170 let mut tcp_addr = None;
171 let mut quic_addr = None;
172
173 if let Some(tcp) = &self.state.config.tcp {
174 ensure_loopback(tcp.addr)?;
175 let listener = TcpListener::bind(tcp.addr).await?;
176 tcp_addr = Some(listener.local_addr()?);
177 let state = Arc::clone(&self.state);
178 let shutdown = shutdown_receiver.clone();
179 tasks.push(tokio::spawn(async move {
180 run_tcp_listener(listener, state, shutdown).await;
181 }));
182 }
183
184 if let Some(quic) = &self.state.config.quic {
185 let endpoint = quinn::Endpoint::server(quic.server_config.clone(), quic.addr)?;
186 quic_addr = Some(endpoint.local_addr()?);
187 let state = Arc::clone(&self.state);
188 let shutdown = shutdown_receiver.clone();
189 tasks.push(tokio::spawn(async move {
190 run_quic_listener(endpoint, state, shutdown).await;
191 }));
192 }
193
194 if tcp_addr.is_none() && quic_addr.is_none() {
195 return Err(DcpError::Protocol(
196 "DCP server has no configured listeners".to_owned(),
197 ));
198 }
199
200 Ok(DcpServerHandle {
201 tcp_addr,
202 quic_addr,
203 shutdown: shutdown_sender,
204 tasks,
205 })
206 }
207}
208
209pub struct DcpServerHandle {
210 tcp_addr: Option<SocketAddr>,
211 quic_addr: Option<SocketAddr>,
212 shutdown: watch::Sender<bool>,
213 tasks: Vec<JoinHandle<()>>,
214}
215
216impl DcpServerHandle {
217 #[must_use]
218 pub fn tcp_addr(&self) -> Option<SocketAddr> {
219 self.tcp_addr
220 }
221
222 #[must_use]
223 pub fn quic_addr(&self) -> Option<SocketAddr> {
224 self.quic_addr
225 }
226
227 pub async fn shutdown(mut self) {
228 let _ = self.shutdown.send(true);
229 for task in self.tasks.drain(..) {
230 task.abort();
231 let _ = task.await;
232 }
233 }
234}
235
236impl Drop for DcpServerHandle {
237 fn drop(&mut self) {
238 let _ = self.shutdown.send(true);
239 for task in &self.tasks {
240 task.abort();
241 }
242 }
243}
244
245async fn run_tcp_listener(
246 listener: TcpListener,
247 state: Arc<DcpServerState>,
248 mut shutdown: watch::Receiver<bool>,
249) {
250 loop {
251 tokio::select! {
252 changed = shutdown.changed() => {
253 if changed.is_err() || *shutdown.borrow() {
254 break;
255 }
256 }
257 accepted = listener.accept() => {
258 let Ok((stream, _peer)) = accepted else {
259 break;
260 };
261 let state = Arc::clone(&state);
262 tokio::spawn(async move {
263 let _ = run_tcp_connection(stream, state).await;
264 });
265 }
266 }
267 }
268}
269
270async fn run_tcp_connection(stream: TcpStream, state: Arc<DcpServerState>) -> DcpResult<()> {
271 stream.set_nodelay(true)?;
272 let (reader, writer) = stream.into_split();
273 run_connection(reader, writer, state).await
274}
275
276async fn run_quic_listener(
277 endpoint: quinn::Endpoint,
278 state: Arc<DcpServerState>,
279 mut shutdown: watch::Receiver<bool>,
280) {
281 loop {
282 tokio::select! {
283 changed = shutdown.changed() => {
284 if changed.is_err() || *shutdown.borrow() {
285 endpoint.close(quinn::VarInt::from_u32(0), b"DCP shutdown");
286 break;
287 }
288 }
289 incoming = endpoint.accept() => {
290 let Some(incoming) = incoming else {
291 break;
292 };
293 let state = Arc::clone(&state);
294 let shutdown = shutdown.clone();
295 tokio::spawn(async move {
296 if let Ok(connection) = incoming.await {
297 run_quic_connection(connection, state, shutdown).await;
298 }
299 });
300 }
301 }
302 }
303}
304
305async fn run_quic_connection(
306 connection: quinn::Connection,
307 state: Arc<DcpServerState>,
308 mut shutdown: watch::Receiver<bool>,
309) {
310 loop {
311 tokio::select! {
312 changed = shutdown.changed() => {
313 if changed.is_err() || *shutdown.borrow() {
314 connection.close(quinn::VarInt::from_u32(0), b"DCP shutdown");
315 break;
316 }
317 }
318 accepted = connection.accept_bi() => {
319 let Ok((send, recv)) = accepted else {
320 break;
321 };
322 let state = Arc::clone(&state);
323 tokio::spawn(async move {
324 let _ = run_connection(recv, send, state).await;
325 });
326 }
327 }
328 }
329}
330
331async fn run_connection<R, W>(mut reader: R, writer: W, state: Arc<DcpServerState>) -> DcpResult<()>
332where
333 R: AsyncRead + Unpin + Send + 'static,
334 W: AsyncWrite + Unpin + Send + 'static,
335{
336 let (outbound, outbound_receiver) = mpsc::channel(state.config.frame_buffer.max(1));
337 let writer_task = tokio::spawn(write_loop(writer, outbound_receiver));
338 let mut subscriptions = Vec::new();
339
340 let Some(first) = read_frame(&mut reader).await? else {
341 return Err(DcpError::Closed);
342 };
343 let hello = match first.frame {
344 Some(dcp_frame::Frame::Hello(hello)) => hello,
345 _ => {
346 let response =
347 Response::error(0, ResponseStatus::BadRequest, "first frame must be Hello");
348 send_frame(&outbound, DcpFrame::response(response)).await?;
349 return Err(DcpError::Protocol("first frame must be Hello".to_owned()));
350 }
351 };
352
353 let hello_response = negotiate_hello(&state, &hello);
354 let accepted = hello_response.response_status() == ResponseStatus::Ok;
355 send_frame(&outbound, DcpFrame::response(hello_response)).await?;
356 if !accepted {
357 return Ok(());
358 }
359
360 while let Some(frame) = read_frame(&mut reader).await? {
361 let request = match frame.frame {
362 Some(dcp_frame::Frame::Request(request)) => request,
363 _ => {
364 return Err(DcpError::Protocol(
365 "client sent non-request frame after hello".to_owned(),
366 ));
367 }
368 };
369 let (response, subscription) =
370 dispatch_request(Arc::clone(&state), request, outbound.clone()).await;
371 send_frame(&outbound, DcpFrame::response(response)).await?;
372 if let Some(subscription) = subscription {
373 subscriptions.push(subscription);
374 }
375 }
376
377 for subscription in subscriptions {
378 subscription.abort();
379 }
380 drop(outbound);
381 writer_task.await??;
382 Ok(())
383}
384
385async fn write_loop<W>(mut writer: W, mut outbound: mpsc::Receiver<DcpFrame>) -> DcpResult<()>
386where
387 W: AsyncWrite + Unpin,
388{
389 while let Some(frame) = outbound.recv().await {
390 write_frame(&mut writer, &frame).await?;
391 }
392 let _ = writer.shutdown().await;
393 Ok(())
394}
395
396async fn send_frame(sender: &mpsc::Sender<DcpFrame>, frame: DcpFrame) -> DcpResult<()> {
397 sender.send(frame).await.map_err(|_| DcpError::Closed)
398}
399
400fn negotiate_hello(state: &DcpServerState, hello: &Hello) -> Response {
401 match protocol_major(&hello.protocol_version) {
402 Some(DCP_PROTOCOL_MAJOR) => {}
403 Some(other) => {
404 return Response::error(
405 0,
406 ResponseStatus::ProtocolMismatch,
407 format!("unsupported DCP major version: {other}"),
408 );
409 }
410 None => {
411 return Response::error(
412 0,
413 ResponseStatus::ProtocolMismatch,
414 format!("invalid DCP protocol version: {}", hello.protocol_version),
415 );
416 }
417 }
418
419 if let Some(expected) = &state.config.auth_token {
420 let actual = hello
421 .auth
422 .as_ref()
423 .map(|auth| auth.bearer_token.as_str())
424 .unwrap_or_default();
425 if actual != expected {
426 return Response::error(0, ResponseStatus::Unauthorized, "invalid DCP token");
427 }
428 }
429
430 Response::ok(0, state.config.node_id.as_bytes().to_vec())
431}
432
433fn protocol_major(version: &str) -> Option<u32> {
434 version.split('.').next()?.parse().ok()
435}
436
437async fn dispatch_request(
438 state: Arc<DcpServerState>,
439 request: Request,
440 outbound: mpsc::Sender<DcpFrame>,
441) -> (Response, Option<JoinHandle<()>>) {
442 let request_id = request.request_id;
443 let deadline = request.deadline_ms;
444 let dispatch = async {
445 let command = request.command.ok_or_else(|| {
446 DcpError::response(ResponseStatus::BadRequest, "request missing command")
447 })?;
448 dispatch_command(state, request_id, command, outbound).await
449 };
450
451 let result = if deadline == 0 {
452 dispatch.await
453 } else {
454 match tokio::time::timeout(Duration::from_millis(deadline), dispatch).await {
455 Ok(result) => result,
456 Err(_) => Err(DcpError::response(
457 ResponseStatus::DeadlineExceeded,
458 "request deadline exceeded",
459 )),
460 }
461 };
462
463 match result {
464 Ok(result) => result,
465 Err(error) => (response_for_error(request_id, error), None),
466 }
467}
468
469async fn dispatch_command(
470 state: Arc<DcpServerState>,
471 request_id: u64,
472 command: request::Command,
473 outbound: mpsc::Sender<DcpFrame>,
474) -> DcpResult<(Response, Option<JoinHandle<()>>)> {
475 match command {
476 request::Command::ListJobs(_) => {
477 let registry = state.registry.clone();
478 let jobs = registry_call(move || registry.list()).await?;
479 let payload = JobList {
480 jobs: jobs.iter().map(wire_job_status).collect(),
481 }
482 .encode_to_vec();
483 Ok((Response::ok(request_id, payload), None))
484 }
485 request::Command::StartJob(start) => {
486 if start.factory_name.trim().is_empty() || start.instance_name.trim().is_empty() {
487 return Err(DcpError::response(
488 ResponseStatus::BadRequest,
489 "StartJob requires factory_name and instance_name",
490 ));
491 }
492 let spec = state.factories.build(
493 &start.factory_name,
494 start.instance_name.clone(),
495 start.params,
496 )?;
497 let registry = state.registry.clone();
498 let name = spec.name.clone();
499 let status = registry_call(move || {
500 registry.submit(spec)?;
501 registry.start(name)
502 })
503 .await?;
504 Ok((status_response(request_id, &status), None))
505 }
506 request::Command::DrainJob(drain) => {
507 let registry = state.registry.clone();
508 let status = registry_call(move || registry.drain(drain.name)).await?;
509 Ok((status_response(request_id, &status), None))
510 }
511 request::Command::StopJob(stop) => {
512 let registry = state.registry.clone();
513 let status = registry_call(move || registry.stop(stop.name)).await?;
514 Ok((status_response(request_id, &status), None))
515 }
516 request::Command::RestartJob(restart) => {
517 let registry = state.registry.clone();
518 let status = registry_call(move || registry.restart(restart.name)).await?;
519 Ok((status_response(request_id, &status), None))
520 }
521 request::Command::JobStatus(status) => {
522 let registry = state.registry.clone();
523 let status = registry_call(move || registry.status(status.name)).await?;
524 Ok((status_response(request_id, &status), None))
525 }
526 request::Command::SubscribeEvents(_) => {
527 let registry = state.registry.clone();
528 let subscription = spawn_event_subscription(request_id, registry, outbound);
529 Ok((Response::ok(request_id, Vec::new()), Some(subscription)))
530 }
531 request::Command::SubscribeMetrics(metrics) => {
532 let interval = if metrics.interval_ms == 0 {
533 state.config.metrics_interval
534 } else {
535 Duration::from_millis(metrics.interval_ms)
536 };
537 let subscription = spawn_metrics_subscription(
538 request_id,
539 state.instrumentation.clone(),
540 interval.max(Duration::from_millis(1)),
541 outbound,
542 );
543 Ok((Response::ok(request_id, Vec::new()), Some(subscription)))
544 }
545 request::Command::GetConfig(get) => {
546 let value = state
547 .config_store
548 .lock()
549 .expect("DCP config store poisoned")
550 .get(&get.key)
551 .cloned();
552 let payload = ConfigValue {
553 key: get.key,
554 value: value.clone().unwrap_or_default(),
555 existed: value.is_some(),
556 }
557 .encode_to_vec();
558 Ok((Response::ok(request_id, payload), None))
559 }
560 request::Command::PutConfig(put) => {
561 let existed = state
562 .config_store
563 .lock()
564 .expect("DCP config store poisoned")
565 .insert(put.key.clone(), put.value.clone())
566 .is_some();
567 let payload = ConfigValue {
568 key: put.key,
569 value: put.value,
570 existed,
571 }
572 .encode_to_vec();
573 Ok((Response::ok(request_id, payload), None))
574 }
575 }
576}
577
578async fn registry_call<T, F>(call: F) -> DcpResult<T>
579where
580 T: Send + 'static,
581 F: FnOnce() -> AgentResult<T> + Send + 'static,
582{
583 Ok(tokio::task::spawn_blocking(call).await??)
584}
585
586fn status_response(request_id: u64, status: &RegistryJobStatus) -> Response {
587 Response::ok(request_id, wire_job_status(status).encode_to_vec())
588}
589
590fn response_for_error(request_id: u64, error: DcpError) -> Response {
591 match error {
592 DcpError::Response { status, message } => Response::error(request_id, status, message),
593 DcpError::Agent(error) => {
594 let status = match &error {
595 AgentError::InvalidJobName => ResponseStatus::BadRequest,
596 AgentError::JobNotFound(_) => ResponseStatus::NotFound,
597 AgentError::JobAlreadyExists(_) | AgentError::JobAlreadyRunning(_) => {
598 ResponseStatus::Conflict
599 }
600 AgentError::DrainUnsupported(_)
601 | AgentError::JobNotRunning(_)
602 | AgentError::RestartLimitExceeded(_) => ResponseStatus::Failed,
603 AgentError::RegistryClosed | AgentError::Actor(_) | AgentError::Stream(_) => {
604 ResponseStatus::Failed
605 }
606 };
607 Response::error(request_id, status, error.to_string())
608 }
609 other => Response::error(request_id, ResponseStatus::Failed, other.to_string()),
610 }
611}
612
613fn spawn_event_subscription(
614 subscription_id: u64,
615 registry: JobRegistryHandle,
616 outbound: mpsc::Sender<DcpFrame>,
617) -> JoinHandle<()> {
618 tokio::task::spawn_blocking(move || {
619 let Ok(queue) = registry.events().run_with(Sink::queue()) else {
620 return;
621 };
622 while let Ok(Some(event)) = queue.pull() {
623 let frame = DcpFrame::event(subscription_id, wire_event(&event));
624 if outbound.blocking_send(frame).is_err() {
625 break;
626 }
627 }
628 })
629}
630
631fn spawn_metrics_subscription(
632 subscription_id: u64,
633 instrumentation: StreamInstrumentationRegistry,
634 interval: Duration,
635 outbound: mpsc::Sender<DcpFrame>,
636) -> JoinHandle<()> {
637 tokio::spawn(async move {
638 let mut ticker = tokio::time::interval(interval);
639 loop {
640 ticker.tick().await;
641 let sample = MetricSample {
642 timestamp_ms: system_time_ms(SystemTime::now()),
643 streams: instrumentation
644 .snapshots()
645 .iter()
646 .map(wire_stream_metric)
647 .collect(),
648 };
649 let frame = DcpFrame::metric(subscription_id, sample);
650 match outbound.try_send(frame) {
651 Ok(()) => {}
652 Err(mpsc::error::TrySendError::Full(_)) => {}
653 Err(mpsc::error::TrySendError::Closed(_)) => break,
654 }
655 }
656 })
657}
658
659fn wire_job_status(status: &RegistryJobStatus) -> WireJobStatus {
660 let now = Instant::now();
661 WireJobStatus {
662 name: status.name.clone(),
663 job_id: status.job_id.0,
664 state: format!("{:?}", status.state),
665 desired_state: format!("{:?}", status.desired_state),
666 generation: status.generation,
667 starts_total: status.starts_total,
668 restarts_total: status.restarts_total,
669 last_start_at_ms: status.last_start_at.map(system_time_ms),
670 last_exit_at_ms: status.last_exit_at.map(system_time_ms),
671 last_exit_reason: status
672 .last_exit_reason
673 .as_ref()
674 .map(exit_reason_text)
675 .unwrap_or_default(),
676 backoff_remaining_ms: status
677 .backoff_until
678 .map(|deadline| duration_ms(deadline.saturating_duration_since(now))),
679 drain_remaining_ms: status
680 .drain_deadline
681 .map(|deadline| duration_ms(deadline.saturating_duration_since(now))),
682 drain_supported: status.drain_supported,
683 active_streams: status.active_streams.map(|streams| streams as u64),
684 }
685}
686
687fn wire_event(event: &JobEvent) -> Event {
688 let (kind, detail) = event_kind_text(&event.kind);
689 Event {
690 sequence: event.sequence,
691 timestamp_ms: system_time_ms(event.timestamp),
692 name: event.name.clone(),
693 job_id: event.job_id.0,
694 generation: event.generation,
695 kind,
696 detail,
697 }
698}
699
700fn wire_stream_metric(snapshot: &StreamInstrumentationSnapshot) -> StreamMetric {
701 StreamMetric {
702 id: snapshot.id.get(),
703 name: snapshot.name.clone(),
704 elements_through: snapshot.elements_through,
705 restarts: snapshot.restarts,
706 state: instrumentation_state_text(snapshot.state).to_owned(),
707 started_at_ms: system_time_ms(snapshot.started_at),
708 state_changed_at_ms: system_time_ms(snapshot.state_changed_at),
709 finished_at_ms: snapshot.finished_at.map(system_time_ms),
710 uptime_ms: duration_ms(snapshot.uptime),
711 }
712}
713
714fn instrumentation_state_text(state: StreamInstrumentationState) -> &'static str {
715 match state {
716 StreamInstrumentationState::Running => "Running",
717 StreamInstrumentationState::Draining => "Draining",
718 StreamInstrumentationState::Completed => "Completed",
719 StreamInstrumentationState::Failed => "Failed",
720 }
721}
722
723fn event_kind_text(kind: &JobEventKind) -> (String, String) {
724 match kind {
725 JobEventKind::Submitted => ("Submitted".to_owned(), String::new()),
726 JobEventKind::Started => ("Started".to_owned(), String::new()),
727 JobEventKind::Failed { reason } => ("Failed".to_owned(), exit_reason_text(reason)),
728 JobEventKind::RestartScheduled { delay } => (
729 "RestartScheduled".to_owned(),
730 format!("delay_ms={}", duration_ms(*delay)),
731 ),
732 JobEventKind::Restarted {
733 previous_generation,
734 } => (
735 "Restarted".to_owned(),
736 format!("previous_generation={previous_generation}"),
737 ),
738 JobEventKind::Draining => ("Draining".to_owned(), String::new()),
739 JobEventKind::Drained => ("Drained".to_owned(), String::new()),
740 JobEventKind::Stopped { reason } => ("Stopped".to_owned(), exit_reason_text(reason)),
741 JobEventKind::Completed => ("Completed".to_owned(), String::new()),
742 }
743}
744
745fn exit_reason_text(reason: &JobExitReason) -> String {
746 match reason {
747 JobExitReason::Completed => "Completed".to_owned(),
748 JobExitReason::Failed(error) => format!("Failed({error})"),
749 JobExitReason::Drained => "Drained".to_owned(),
750 JobExitReason::Stopped => "Stopped".to_owned(),
751 JobExitReason::DrainTimedOut => "DrainTimedOut".to_owned(),
752 }
753}
754
755fn system_time_ms(time: SystemTime) -> u64 {
756 time.duration_since(UNIX_EPOCH)
757 .map(duration_ms)
758 .unwrap_or(0)
759}
760
761fn duration_ms(duration: Duration) -> u64 {
762 duration.as_millis().min(u128::from(u64::MAX)) as u64
763}
764
765fn ensure_loopback(addr: SocketAddr) -> DcpResult<()> {
766 if addr.ip().is_loopback() {
767 return Ok(());
768 }
769 Err(DcpError::Protocol(format!(
770 "plaintext DCP TCP listener must bind loopback, got {addr}"
771 )))
772}
773
774fn client_bind_addr(remote_addr: SocketAddr) -> SocketAddr {
775 if remote_addr.is_ipv6() {
776 SocketAddr::new(IpAddr::V6(Ipv6Addr::UNSPECIFIED), 0)
777 } else {
778 SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0)
779 }
780}
781
782pub(crate) async fn connect_quic_stream(
783 addr: SocketAddr,
784 server_name: &str,
785 client_config: quinn::ClientConfig,
786) -> DcpResult<(
787 quinn::Endpoint,
788 quinn::Connection,
789 quinn::RecvStream,
790 quinn::SendStream,
791)> {
792 let mut endpoint = quinn::Endpoint::client(client_bind_addr(addr))?;
793 endpoint.set_default_client_config(client_config);
794 let connection = endpoint
795 .connect(addr, server_name)
796 .map_err(|error| DcpError::Protocol(error.to_string()))?
797 .await
798 .map_err(|error| DcpError::Protocol(error.to_string()))?;
799 let (send, recv) = connection
800 .open_bi()
801 .await
802 .map_err(|error| DcpError::Protocol(error.to_string()))?;
803 Ok((endpoint, connection, recv, send))
804}