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