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

datum_agent/dcp/
client.rs
