datum_agent/dcp/
client.rs

1use std::{
2    collections::HashMap,
3    net::SocketAddr,
4    sync::{
5        Arc, Mutex,
6        atomic::{AtomicU64, Ordering},
7    },
8    time::Duration,
9};
10
11use datum_net::quic::quinn;
12use prost::Message as ProstMessage;
13use tokio::{
14    io::{AsyncRead, AsyncWrite, AsyncWriteExt},
15    net::TcpStream,
16    sync::{mpsc, oneshot},
17    task::JoinHandle,
18};
19
20use crate::dcp::{
21    DcpError, DcpResult,
22    frame::{read_frame, write_frame},
23    proto::{
24        ClientKind, ClusterJobList, ClusterNodeInfo, ClusterNodeList, CompleteShardingAsk,
25        ConfigValue, DcpFrame, Event, ForwardShardEnvelopes, GetConfig, GetShardAllocations, Hello,
26        JobStatus as WireJobStatus, JobStatusRequest, ListClusterJobs, ListJobs, MetricFrame,
27        MetricSample, OpenShardPipe, PlacementSpec, PutConfig, RememberClusterAssignment,
28        RememberShardAllocations, Request, Response, ResponseStatus, ShardAllocation,
29        ShardAllocationRequest, ShardAllocationTable, ShardEnvelopeBatchResult, ShardPipeFrame,
30        StartJob, StopJob, SubmitClusterJob, SubscribeEvents, SubscribeMetrics, UnsubscribeMetrics,
31        dcp_frame, request,
32    },
33    server::connect_quic_stream,
34};
35
36type PendingMap = HashMap<u64, oneshot::Sender<DcpResult<Response>>>;
37
38#[derive(Clone)]
39pub struct DcpClient {
40    inner: Arc<ClientInner>,
41    _reader_task: Arc<JoinHandle<()>>,
42    _writer_task: Arc<JoinHandle<()>>,
43    _quic_endpoint: Option<quinn::Endpoint>,
44    _quic_connection: Option<quinn::Connection>,
45}
46
47struct ClientInner {
48    outbound: mpsc::Sender<DcpFrame>,
49    pending: Mutex<PendingMap>,
50    event_subscriptions: Mutex<HashMap<u64, mpsc::Sender<Event>>>,
51    metric_subscriptions: Mutex<HashMap<u64, mpsc::Sender<MetricSample>>>,
52    next_request_id: AtomicU64,
53    disconnect_error: Arc<Mutex<Option<String>>>,
54}
55
56pub struct EventSubscription {
57    receiver: mpsc::Receiver<Event>,
58    disconnect_error: Arc<Mutex<Option<String>>>,
59}
60
61impl EventSubscription {
62    pub async fn recv(&mut self) -> Option<Event> {
63        self.receiver.recv().await
64    }
65
66    pub fn disconnect_reason(&self) -> Option<String> {
67        self.disconnect_error
68            .lock()
69            .unwrap_or_else(|p| p.into_inner())
70            .clone()
71    }
72}
73
74pub struct MetricSubscription {
75    receiver: mpsc::Receiver<MetricSample>,
76    disconnect_error: Arc<Mutex<Option<String>>>,
77    subscription_id: u64,
78    inner: Arc<ClientInner>,
79    cancelled: bool,
80}
81
82impl MetricSubscription {
83    pub async fn recv(&mut self) -> Option<MetricSample> {
84        self.receiver.recv().await
85    }
86
87    pub fn disconnect_reason(&self) -> Option<String> {
88        self.disconnect_error
89            .lock()
90            .unwrap_or_else(|p| p.into_inner())
91            .clone()
92    }
93
94    /// Stop this subscription and notify the server on the existing DCP
95    /// connection.
96    pub async fn cancel(mut self) -> DcpResult<()> {
97        self.remove_registration();
98        self.cancelled = true;
99        send_unsubscribe_metrics(&self.inner, self.subscription_id).await
100    }
101
102    fn remove_registration(&self) {
103        self.inner
104            .metric_subscriptions
105            .lock()
106            .unwrap_or_else(|p| p.into_inner())
107            .remove(&self.subscription_id);
108    }
109}
110
111impl Drop for MetricSubscription {
112    fn drop(&mut self) {
113        self.remove_registration();
114        if self.cancelled {
115            return;
116        }
117        let request_id = self.inner.next_request_id.fetch_add(1, Ordering::Relaxed);
118        let _ = self.inner.outbound.try_send(DcpFrame::request(Request {
119            request_id,
120            deadline_ms: 0,
121            command: Some(request::Command::UnsubscribeMetrics(UnsubscribeMetrics {
122                subscription_id: self.subscription_id,
123            })),
124        }));
125    }
126}
127
128#[derive(Clone)]
129pub struct ShardPipeClient {
130    outbound: mpsc::Sender<ShardPipeFrame>,
131    _reader_task: Arc<JoinHandle<()>>,
132    _writer_task: Arc<JoinHandle<()>>,
133    _quic_endpoint: Option<quinn::Endpoint>,
134    _quic_connection: Option<quinn::Connection>,
135}
136
137impl ShardPipeClient {
138    pub async fn connect_tcp(addr: SocketAddr, hello: Hello) -> DcpResult<Self> {
139        let mut stream = TcpStream::connect(addr).await?;
140        stream.set_nodelay(true)?;
141        {
142            let (mut reader, mut writer) = stream.split();
143            open_shard_pipe(&mut reader, &mut writer, hello).await?;
144        }
145        let (reader, writer) = stream.into_split();
146        Ok(Self::from_parts(reader, writer, None, None))
147    }
148
149    pub async fn connect_quic(
150        addr: SocketAddr,
151        server_name: &str,
152        client_config: quinn::ClientConfig,
153        hello: Hello,
154    ) -> DcpResult<Self> {
155        let (endpoint, connection, mut reader, mut writer) =
156            connect_quic_stream(addr, server_name, client_config).await?;
157        open_shard_pipe(&mut reader, &mut writer, hello).await?;
158        Ok(Self::from_parts(
159            reader,
160            writer,
161            Some(endpoint),
162            Some(connection),
163        ))
164    }
165
166    fn from_parts<R, W>(
167        reader: R,
168        writer: W,
169        quic_endpoint: Option<quinn::Endpoint>,
170        quic_connection: Option<quinn::Connection>,
171    ) -> Self
172    where
173        R: AsyncRead + Unpin + Send + 'static,
174        W: AsyncWrite + Unpin + Send + 'static,
175    {
176        let (outbound, outbound_receiver) = mpsc::channel(256);
177        let reader_task = tokio::spawn(async move {
178            read_shard_pipe_loop(reader).await;
179        });
180        let writer_task = tokio::spawn(async move {
181            let _ = write_shard_pipe_loop(writer, outbound_receiver).await;
182        });
183        Self {
184            outbound,
185            _reader_task: Arc::new(reader_task),
186            _writer_task: Arc::new(writer_task),
187            _quic_endpoint: quic_endpoint,
188            _quic_connection: quic_connection,
189        }
190    }
191
192    pub async fn send(&self, frame: ShardPipeFrame) -> DcpResult<()> {
193        self.outbound
194            .send(frame)
195            .await
196            .map_err(|_| DcpError::Closed)
197    }
198}
199
200const DEFAULT_LOCAL_REQUEST_TIMEOUT_MS: u64 = 30_000;
201
202impl DcpClient {
203    pub async fn connect_tcp(addr: SocketAddr, hello: Hello) -> DcpResult<Self> {
204        let stream = TcpStream::connect(addr).await?;
205        stream.set_nodelay(true)?;
206        let (reader, writer) = stream.into_split();
207        let client = Self::from_parts(reader, writer, None, None);
208        client.hello(hello).await?;
209        Ok(client)
210    }
211
212    pub async fn connect_quic(
213        addr: SocketAddr,
214        server_name: &str,
215        client_config: quinn::ClientConfig,
216        hello: Hello,
217    ) -> DcpResult<Self> {
218        let (endpoint, connection, reader, writer) =
219            connect_quic_stream(addr, server_name, client_config).await?;
220        let client = Self::from_parts(reader, writer, Some(endpoint), Some(connection));
221        client.hello(hello).await?;
222        Ok(client)
223    }
224
225    fn from_parts<R, W>(
226        reader: R,
227        writer: W,
228        quic_endpoint: Option<quinn::Endpoint>,
229        quic_connection: Option<quinn::Connection>,
230    ) -> Self
231    where
232        R: AsyncRead + Unpin + Send + 'static,
233        W: AsyncWrite + Unpin + Send + 'static,
234    {
235        let (outbound, outbound_receiver) = mpsc::channel(256);
236        let disconnect_error = Arc::new(Mutex::new(None));
237        let inner = Arc::new(ClientInner {
238            outbound,
239            pending: Mutex::new(HashMap::new()),
240            event_subscriptions: Mutex::new(HashMap::new()),
241            metric_subscriptions: Mutex::new(HashMap::new()),
242            next_request_id: AtomicU64::new(1),
243            disconnect_error: Arc::clone(&disconnect_error),
244        });
245        let reader_task = {
246            let inner = Arc::clone(&inner);
247            tokio::spawn(async move {
248                read_loop(reader, inner).await;
249            })
250        };
251        let writer_task = tokio::spawn(async move {
252            let _ = write_loop(writer, outbound_receiver).await;
253        });
254        Self {
255            inner,
256            _reader_task: Arc::new(reader_task),
257            _writer_task: Arc::new(writer_task),
258            _quic_endpoint: quic_endpoint,
259            _quic_connection: quic_connection,
260        }
261    }
262
263    pub async fn hello(&self, hello: Hello) -> DcpResult<()> {
264        let response = self
265            .send_frame_with_response(0, DcpFrame::hello(hello))
266            .await?;
267        ensure_ok(response).map(|_| ())
268    }
269
270    pub async fn request(&self, request: Request) -> DcpResult<Response> {
271        let response = self
272            .send_frame_with_response(request.request_id, DcpFrame::request(request))
273            .await?;
274        ensure_ok(response)
275    }
276
277    pub async fn list_jobs(&self) -> DcpResult<Vec<WireJobStatus>> {
278        let response = self
279            .send_command(
280                request::Command::ListJobs(ListJobs {}),
281                DEFAULT_LOCAL_REQUEST_TIMEOUT_MS,
282            )
283            .await?;
284        let list = crate::dcp::proto::JobList::decode(response.payload.as_slice())?;
285        Ok(list.jobs)
286    }
287
288    pub async fn list_cluster_jobs(&self, timeout_ms: u64) -> DcpResult<ClusterJobList> {
289        let request = self.send_command(
290            request::Command::ListClusterJobs(ListClusterJobs { timeout_ms }),
291            cluster_request_deadline(timeout_ms),
292        );
293        let response = timeout_cluster_request(timeout_ms, request).await?;
294        Ok(ClusterJobList::decode(response.payload.as_slice())?)
295    }
296
297    pub async fn cluster_node_info(&self, timeout_ms: u64) -> DcpResult<ClusterNodeList> {
298        let request = self.send_command(
299            request::Command::ClusterNodeInfo(ClusterNodeInfo { timeout_ms }),
300            cluster_request_deadline(timeout_ms),
301        );
302        let response = timeout_cluster_request(timeout_ms, request).await?;
303        Ok(ClusterNodeList::decode(response.payload.as_slice())?)
304    }
305
306    pub async fn start_job(
307        &self,
308        factory_name: impl Into<String>,
309        instance_name: impl Into<String>,
310        params: HashMap<String, String>,
311    ) -> DcpResult<WireJobStatus> {
312        self.job_status_response(request::Command::StartJob(StartJob {
313            factory_name: factory_name.into(),
314            instance_name: instance_name.into(),
315            params,
316            cluster: None,
317        }))
318        .await
319    }
320
321    pub(crate) async fn start_cluster_job_on_node(
322        &self,
323        factory_name: impl Into<String>,
324        instance_name: impl Into<String>,
325        params: HashMap<String, String>,
326        cluster: crate::dcp::proto::ClusterJobStart,
327    ) -> DcpResult<WireJobStatus> {
328        self.job_status_response(request::Command::StartJob(StartJob {
329            factory_name: factory_name.into(),
330            instance_name: instance_name.into(),
331            params,
332            cluster: Some(cluster),
333        }))
334        .await
335    }
336
337    pub async fn submit_cluster_job(
338        &self,
339        factory_name: impl Into<String>,
340        instance_name: impl Into<String>,
341        params: HashMap<String, String>,
342        placement: PlacementSpec,
343        timeout_ms: u64,
344    ) -> DcpResult<WireJobStatus> {
345        let request = self.send_command(
346            request::Command::SubmitClusterJob(SubmitClusterJob {
347                factory_name: factory_name.into(),
348                instance_name: instance_name.into(),
349                params,
350                placement: Some(placement),
351                timeout_ms,
352            }),
353            cluster_request_deadline(timeout_ms),
354        );
355        let response = timeout_cluster_request(timeout_ms, request).await?;
356        Ok(WireJobStatus::decode(response.payload.as_slice())?)
357    }
358
359    pub async fn drain_job(&self, name: impl Into<String>) -> DcpResult<WireJobStatus> {
360        self.job_status_response(request::Command::DrainJob(crate::dcp::proto::DrainJob {
361            name: name.into(),
362            cluster: false,
363        }))
364        .await
365    }
366
367    pub async fn drain_cluster_job(
368        &self,
369        name: impl Into<String>,
370        timeout_ms: u64,
371    ) -> DcpResult<WireJobStatus> {
372        let request = self.send_command(
373            request::Command::DrainJob(crate::dcp::proto::DrainJob {
374                name: name.into(),
375                cluster: true,
376            }),
377            cluster_request_deadline(timeout_ms),
378        );
379        let response = timeout_cluster_request(timeout_ms, request).await?;
380        Ok(WireJobStatus::decode(response.payload.as_slice())?)
381    }
382
383    pub async fn stop_job(&self, name: impl Into<String>) -> DcpResult<WireJobStatus> {
384        self.job_status_response(request::Command::StopJob(StopJob {
385            name: name.into(),
386            cluster: false,
387        }))
388        .await
389    }
390
391    pub async fn stop_cluster_job(
392        &self,
393        name: impl Into<String>,
394        timeout_ms: u64,
395    ) -> DcpResult<WireJobStatus> {
396        let request = self.send_command(
397            request::Command::StopJob(StopJob {
398                name: name.into(),
399                cluster: true,
400            }),
401            cluster_request_deadline(timeout_ms),
402        );
403        let response = timeout_cluster_request(timeout_ms, request).await?;
404        Ok(WireJobStatus::decode(response.payload.as_slice())?)
405    }
406
407    pub async fn restart_job(&self, name: impl Into<String>) -> DcpResult<WireJobStatus> {
408        self.job_status_response(request::Command::RestartJob(
409            crate::dcp::proto::RestartJob {
410                name: name.into(),
411                cluster: false,
412            },
413        ))
414        .await
415    }
416
417    pub async fn job_status(&self, name: impl Into<String>) -> DcpResult<WireJobStatus> {
418        self.job_status_response(request::Command::JobStatus(JobStatusRequest {
419            name: name.into(),
420            cluster: false,
421        }))
422        .await
423    }
424
425    pub async fn cluster_job_status(
426        &self,
427        name: impl Into<String>,
428        timeout_ms: u64,
429    ) -> DcpResult<WireJobStatus> {
430        let request = self.send_command(
431            request::Command::JobStatus(JobStatusRequest {
432                name: name.into(),
433                cluster: true,
434            }),
435            cluster_request_deadline(timeout_ms),
436        );
437        let response = timeout_cluster_request(timeout_ms, request).await?;
438        Ok(WireJobStatus::decode(response.payload.as_slice())?)
439    }
440
441    pub(crate) async fn remember_cluster_assignment(
442        &self,
443        instance_name: impl Into<String>,
444        assignment: Option<crate::dcp::proto::ClusterJobStart>,
445    ) -> DcpResult<()> {
446        let response = self
447            .send_command(
448                request::Command::RememberClusterAssignment(RememberClusterAssignment {
449                    instance_name: instance_name.into(),
450                    assignment,
451                }),
452                0,
453            )
454            .await?;
455        let _ = response;
456        Ok(())
457    }
458
459    pub async fn allocate_shard(
460        &self,
461        type_name: impl Into<String>,
462        shard_id: impl Into<String>,
463        timeout_ms: u64,
464    ) -> DcpResult<ShardAllocation> {
465        let request = self.send_command(
466            request::Command::AllocateShard(ShardAllocationRequest {
467                type_name: type_name.into(),
468                shard_id: shard_id.into(),
469                timeout_ms,
470            }),
471            cluster_request_deadline(timeout_ms),
472        );
473        let response = timeout_cluster_request(timeout_ms, request).await?;
474        Ok(ShardAllocation::decode(response.payload.as_slice())?)
475    }
476
477    pub async fn remember_shard_allocations(&self, table: ShardAllocationTable) -> DcpResult<()> {
478        let response = self
479            .send_command(
480                request::Command::RememberShardAllocations(RememberShardAllocations {
481                    table: Some(table),
482                }),
483                0,
484            )
485            .await?;
486        let _ = response;
487        Ok(())
488    }
489
490    pub async fn get_shard_allocations(
491        &self,
492        type_name: impl Into<String>,
493        timeout_ms: u64,
494    ) -> DcpResult<ShardAllocationTable> {
495        let request = self.send_command(
496            request::Command::GetShardAllocations(GetShardAllocations {
497                type_name: type_name.into(),
498            }),
499            cluster_request_deadline(timeout_ms),
500        );
501        let response = timeout_cluster_request(timeout_ms, request).await?;
502        Ok(ShardAllocationTable::decode(response.payload.as_slice())?)
503    }
504
505    pub async fn forward_shard_envelopes(
506        &self,
507        batch: ForwardShardEnvelopes,
508        timeout_ms: u64,
509    ) -> DcpResult<ShardEnvelopeBatchResult> {
510        let request = self.send_command(
511            request::Command::ForwardShardEnvelopes(batch),
512            cluster_request_deadline(timeout_ms),
513        );
514        let response = timeout_cluster_request(timeout_ms, request).await?;
515        Ok(ShardEnvelopeBatchResult::decode(
516            response.payload.as_slice(),
517        )?)
518    }
519
520    pub async fn complete_sharding_ask(&self, reply: CompleteShardingAsk) -> DcpResult<()> {
521        let response = self
522            .send_command(request::Command::CompleteShardingAsk(reply), 0)
523            .await?;
524        let _ = response;
525        Ok(())
526    }
527
528    pub async fn subscribe_events(&self) -> DcpResult<EventSubscription> {
529        let request_id = self.next_request_id();
530        let (sender, receiver) = mpsc::channel(256);
531        self.inner
532            .event_subscriptions
533            .lock()
534            .expect("DCP event subscriptions poisoned")
535            .insert(request_id, sender);
536        let request = Request {
537            request_id,
538            deadline_ms: DEFAULT_LOCAL_REQUEST_TIMEOUT_MS,
539            command: Some(request::Command::SubscribeEvents(SubscribeEvents {
540                buffer: 0,
541            })),
542        };
543        if let Err(error) = self.request(request).await {
544            self.inner
545                .event_subscriptions
546                .lock()
547                .expect("DCP event subscriptions poisoned")
548                .remove(&request_id);
549            return Err(error);
550        }
551        Ok(EventSubscription {
552            receiver,
553            disconnect_error: Arc::clone(&self.inner.disconnect_error),
554        })
555    }
556
557    pub async fn subscribe_metrics(&self, interval_ms: u64) -> DcpResult<MetricSubscription> {
558        self.subscribe_metrics_inner(interval_ms, Vec::new(), false)
559            .await
560    }
561
562    pub(crate) async fn subscribe_local_metrics(
563        &self,
564        interval_ms: u64,
565        job_names: Vec<String>,
566    ) -> DcpResult<MetricSubscription> {
567        self.subscribe_metrics_inner(interval_ms, job_names, true)
568            .await
569    }
570
571    async fn subscribe_metrics_inner(
572        &self,
573        interval_ms: u64,
574        job_names: Vec<String>,
575        local_only: bool,
576    ) -> DcpResult<MetricSubscription> {
577        let request_id = self.next_request_id();
578        let (sender, receiver) = mpsc::channel(16);
579        self.inner
580            .metric_subscriptions
581            .lock()
582            .expect("DCP metric subscriptions poisoned")
583            .insert(request_id, sender);
584        let request = Request {
585            request_id,
586            deadline_ms: DEFAULT_LOCAL_REQUEST_TIMEOUT_MS,
587            command: Some(request::Command::SubscribeMetrics(SubscribeMetrics {
588                interval_ms,
589                job_names,
590                local_only,
591            })),
592        };
593        if let Err(error) = self.request(request).await {
594            self.inner
595                .metric_subscriptions
596                .lock()
597                .expect("DCP metric subscriptions poisoned")
598                .remove(&request_id);
599            return Err(error);
600        }
601        Ok(MetricSubscription {
602            receiver,
603            disconnect_error: Arc::clone(&self.inner.disconnect_error),
604            subscription_id: request_id,
605            inner: Arc::clone(&self.inner),
606            cancelled: false,
607        })
608    }
609
610    pub async fn get_config(&self, key: impl Into<String>) -> DcpResult<ConfigValue> {
611        let response = self
612            .send_command(
613                request::Command::GetConfig(GetConfig { key: key.into() }),
614                DEFAULT_LOCAL_REQUEST_TIMEOUT_MS,
615            )
616            .await?;
617        Ok(ConfigValue::decode(response.payload.as_slice())?)
618    }
619
620    pub async fn put_config(
621        &self,
622        key: impl Into<String>,
623        value: impl Into<String>,
624    ) -> DcpResult<ConfigValue> {
625        let response = self
626            .send_command(
627                request::Command::PutConfig(PutConfig {
628                    key: key.into(),
629                    value: value.into(),
630                }),
631                DEFAULT_LOCAL_REQUEST_TIMEOUT_MS,
632            )
633            .await?;
634        Ok(ConfigValue::decode(response.payload.as_slice())?)
635    }
636
637    async fn job_status_response(&self, command: request::Command) -> DcpResult<WireJobStatus> {
638        let response = self
639            .send_command(command, DEFAULT_LOCAL_REQUEST_TIMEOUT_MS)
640            .await?;
641        Ok(WireJobStatus::decode(response.payload.as_slice())?)
642    }
643
644    async fn send_command(
645        &self,
646        command: request::Command,
647        deadline_ms: u64,
648    ) -> DcpResult<Response> {
649        let request_id = self.next_request_id();
650        let request = Request {
651            request_id,
652            deadline_ms,
653            command: Some(command),
654        };
655        self.request(request).await
656    }
657
658    async fn send_frame_with_response(
659        &self,
660        request_id: u64,
661        frame: DcpFrame,
662    ) -> DcpResult<Response> {
663        let (sender, receiver) = oneshot::channel();
664        self.inner
665            .pending
666            .lock()
667            .expect("DCP pending map poisoned")
668            .insert(request_id, sender);
669        if self.inner.outbound.send(frame).await.is_err() {
670            self.inner
671                .pending
672                .lock()
673                .expect("DCP pending map poisoned")
674                .remove(&request_id);
675            return Err(DcpError::Closed);
676        }
677        receiver.await.map_err(|_| DcpError::Closed)?
678    }
679
680    fn next_request_id(&self) -> u64 {
681        self.inner.next_request_id.fetch_add(1, Ordering::Relaxed)
682    }
683}
684
685async fn send_unsubscribe_metrics(inner: &Arc<ClientInner>, subscription_id: u64) -> DcpResult<()> {
686    let request_id = inner.next_request_id.fetch_add(1, Ordering::Relaxed);
687    inner
688        .outbound
689        .send(DcpFrame::request(Request {
690            request_id,
691            deadline_ms: 0,
692            command: Some(request::Command::UnsubscribeMetrics(UnsubscribeMetrics {
693                subscription_id,
694            })),
695        }))
696        .await
697        .map_err(|_| DcpError::Closed)
698}
699
700async fn read_loop<R>(mut reader: R, inner: Arc<ClientInner>)
701where
702    R: AsyncRead + Unpin,
703{
704    loop {
705        match read_frame(&mut reader).await {
706            Ok(Some(frame)) => handle_inbound_frame(frame, &inner).await,
707            Ok(None) => {
708                set_disconnect_error(&inner, DcpError::Closed);
709                fail_pending(&inner, DcpError::Closed);
710                clear_subscriptions(&inner);
711                break;
712            }
713            Err(error) => {
714                set_disconnect_error(
715                    &inner,
716                    DcpError::Protocol(format!("DCP read error: {}", error)),
717                );
718                fail_pending(&inner, error);
719                clear_subscriptions(&inner);
720                break;
721            }
722        }
723    }
724}
725
726fn clear_subscriptions(inner: &ClientInner) {
727    inner
728        .event_subscriptions
729        .lock()
730        .unwrap_or_else(|p| p.into_inner())
731        .clear();
732    inner
733        .metric_subscriptions
734        .lock()
735        .unwrap_or_else(|p| p.into_inner())
736        .clear();
737}
738
739fn set_disconnect_error(inner: &ClientInner, error: DcpError) {
740    *inner
741        .disconnect_error
742        .lock()
743        .unwrap_or_else(|p| p.into_inner()) = Some(error.to_string());
744}
745
746async fn handle_inbound_frame(frame: DcpFrame, inner: &Arc<ClientInner>) {
747    match frame.frame {
748        Some(dcp_frame::Frame::Response(response)) => {
749            if let Some(sender) = inner
750                .pending
751                .lock()
752                .expect("DCP pending map poisoned")
753                .remove(&response.request_id)
754            {
755                let _ = sender.send(Ok(response));
756            }
757        }
758        Some(dcp_frame::Frame::Event(event_frame)) => {
759            if let Some(event) = event_frame.event {
760                let sender = inner
761                    .event_subscriptions
762                    .lock()
763                    .expect("DCP event subscriptions poisoned")
764                    .get(&event_frame.subscription_id)
765                    .cloned();
766                if let Some(sender) = sender {
767                    let _ = sender.send(event).await;
768                }
769            }
770        }
771        Some(dcp_frame::Frame::Metric(MetricFrame {
772            subscription_id,
773            sample: Some(sample),
774        })) => {
775            let sender = inner
776                .metric_subscriptions
777                .lock()
778                .expect("DCP metric subscriptions poisoned")
779                .get(&subscription_id)
780                .cloned();
781            if let Some(sender) = sender {
782                let _ = sender.send(sample).await;
783            }
784        }
785        _ => {}
786    }
787}
788
789fn fail_pending(inner: &ClientInner, error: DcpError) {
790    let pending = std::mem::take(&mut *inner.pending.lock().expect("DCP pending map poisoned"));
791    let message = error.to_string();
792    for (_, sender) in pending {
793        let _ = sender.send(Err(DcpError::Protocol(message.clone())));
794    }
795}
796
797async fn write_loop<W>(mut writer: W, mut outbound: mpsc::Receiver<DcpFrame>) -> DcpResult<()>
798where
799    W: AsyncWrite + Unpin,
800{
801    while let Some(frame) = outbound.recv().await {
802        write_frame(&mut writer, &frame).await?;
803    }
804    let _ = writer.shutdown().await;
805    Ok(())
806}
807
808async fn open_shard_pipe<R, W>(reader: &mut R, writer: &mut W, hello: Hello) -> DcpResult<()>
809where
810    R: AsyncRead + Unpin,
811    W: AsyncWrite + Unpin,
812{
813    write_frame(writer, &DcpFrame::hello(hello)).await?;
814    let Some(response) = read_frame(reader).await? else {
815        return Err(DcpError::Closed);
816    };
817    match response.frame {
818        Some(dcp_frame::Frame::Response(response)) => {
819            ensure_ok(response)?;
820        }
821        _ => {
822            return Err(DcpError::Protocol(
823                "DCP shard pipe hello did not receive a response".to_owned(),
824            ));
825        }
826    }
827
828    let request = Request {
829        request_id: 1,
830        deadline_ms: 0,
831        command: Some(request::Command::OpenShardPipe(OpenShardPipe {})),
832    };
833    write_frame(writer, &DcpFrame::request(request)).await?;
834    let Some(response) = read_frame(reader).await? else {
835        return Err(DcpError::Closed);
836    };
837    match response.frame {
838        Some(dcp_frame::Frame::Response(response)) => ensure_ok(response).map(|_| ()),
839        _ => Err(DcpError::Protocol(
840            "DCP shard pipe open did not receive a response".to_owned(),
841        )),
842    }
843}
844
845async fn read_shard_pipe_loop<R>(mut reader: R)
846where
847    R: AsyncRead + Unpin,
848{
849    while let Ok(Some(_frame)) = read_frame(&mut reader).await {}
850}
851
852async fn write_shard_pipe_loop<W>(
853    mut writer: W,
854    mut outbound: mpsc::Receiver<ShardPipeFrame>,
855) -> DcpResult<()>
856where
857    W: AsyncWrite + Unpin,
858{
859    while let Some(frame) = outbound.recv().await {
860        write_frame(&mut writer, &DcpFrame::shard_pipe(frame)).await?;
861    }
862    let _ = writer.shutdown().await;
863    Ok(())
864}
865
866fn ensure_ok(response: Response) -> DcpResult<Response> {
867    let status = response.response_status();
868    if status == ResponseStatus::Ok {
869        Ok(response)
870    } else {
871        Err(DcpError::response(status, response.message))
872    }
873}
874
875async fn timeout_cluster_request<F>(timeout_ms: u64, request: F) -> DcpResult<Response>
876where
877    F: std::future::Future<Output = DcpResult<Response>>,
878{
879    if timeout_ms == 0 {
880        return request.await;
881    }
882    tokio::time::timeout(
883        Duration::from_millis(timeout_ms).saturating_add(Duration::from_millis(250)),
884        request,
885    )
886    .await
887    .map_err(|_| {
888        DcpError::response(
889            ResponseStatus::DeadlineExceeded,
890            "cluster request timed out",
891        )
892    })?
893}
894
895fn cluster_request_deadline(timeout_ms: u64) -> u64 {
896    if timeout_ms == 0 {
897        0
898    } else {
899        timeout_ms.saturating_add(250)
900    }
901}
902
903#[must_use]
904pub fn default_hello(node_id: impl Into<String>, client_kind: ClientKind) -> Hello {
905    Hello::new(node_id, client_kind)
906}
datum_agent/dcp/client.rs

datum_agent/dcp/
client.rs
