1use std::{
2 collections::HashMap,
3 net::SocketAddr,
4 sync::{
5 Arc, Mutex,
6 atomic::{AtomicU64, Ordering},
7 },
8};
9
10use datum_net::quic::quinn;
11use prost::Message as ProstMessage;
12use tokio::{
13 io::{AsyncRead, AsyncWrite, AsyncWriteExt},
14 net::TcpStream,
15 sync::{mpsc, oneshot},
16 task::JoinHandle,
17};
18
19use crate::dcp::{
20 DcpError, DcpResult,
21 frame::{read_frame, write_frame},
22 proto::{
23 ClientKind, ConfigValue, DcpFrame, Event, GetConfig, Hello, JobStatus as WireJobStatus,
24 JobStatusRequest, ListJobs, MetricFrame, MetricSample, PutConfig, Request, Response,
25 ResponseStatus, StartJob, StopJob, SubscribeEvents, SubscribeMetrics, dcp_frame, request,
26 },
27 server::connect_quic_stream,
28};
29
30type PendingMap = HashMap<u64, oneshot::Sender<DcpResult<Response>>>;
31
32#[derive(Clone)]
33pub struct DcpClient {
34 inner: Arc<ClientInner>,
35 _reader_task: Arc<JoinHandle<()>>,
36 _writer_task: Arc<JoinHandle<()>>,
37 _quic_endpoint: Option<quinn::Endpoint>,
38 _quic_connection: Option<quinn::Connection>,
39}
40
41struct ClientInner {
42 outbound: mpsc::Sender<DcpFrame>,
43 pending: Mutex<PendingMap>,
44 event_subscriptions: Mutex<HashMap<u64, mpsc::Sender<Event>>>,
45 metric_subscriptions: Mutex<HashMap<u64, mpsc::Sender<MetricSample>>>,
46 next_request_id: AtomicU64,
47}
48
49pub struct EventSubscription {
50 receiver: mpsc::Receiver<Event>,
51}
52
53impl EventSubscription {
54 pub async fn recv(&mut self) -> Option<Event> {
55 self.receiver.recv().await
56 }
57}
58
59pub struct MetricSubscription {
60 receiver: mpsc::Receiver<MetricSample>,
61}
62
63impl MetricSubscription {
64 pub async fn recv(&mut self) -> Option<MetricSample> {
65 self.receiver.recv().await
66 }
67}
68
69impl DcpClient {
70 pub async fn connect_tcp(addr: SocketAddr, hello: Hello) -> DcpResult<Self> {
71 let stream = TcpStream::connect(addr).await?;
72 stream.set_nodelay(true)?;
73 let (reader, writer) = stream.into_split();
74 let client = Self::from_parts(reader, writer, None, None);
75 client.hello(hello).await?;
76 Ok(client)
77 }
78
79 pub async fn connect_quic(
80 addr: SocketAddr,
81 server_name: &str,
82 client_config: quinn::ClientConfig,
83 hello: Hello,
84 ) -> DcpResult<Self> {
85 let (endpoint, connection, reader, writer) =
86 connect_quic_stream(addr, server_name, client_config).await?;
87 let client = Self::from_parts(reader, writer, Some(endpoint), Some(connection));
88 client.hello(hello).await?;
89 Ok(client)
90 }
91
92 fn from_parts<R, W>(
93 reader: R,
94 writer: W,
95 quic_endpoint: Option<quinn::Endpoint>,
96 quic_connection: Option<quinn::Connection>,
97 ) -> Self
98 where
99 R: AsyncRead + Unpin + Send + 'static,
100 W: AsyncWrite + Unpin + Send + 'static,
101 {
102 let (outbound, outbound_receiver) = mpsc::channel(256);
103 let inner = Arc::new(ClientInner {
104 outbound,
105 pending: Mutex::new(HashMap::new()),
106 event_subscriptions: Mutex::new(HashMap::new()),
107 metric_subscriptions: Mutex::new(HashMap::new()),
108 next_request_id: AtomicU64::new(1),
109 });
110 let reader_task = {
111 let inner = Arc::clone(&inner);
112 tokio::spawn(async move {
113 read_loop(reader, inner).await;
114 })
115 };
116 let writer_task = tokio::spawn(async move {
117 let _ = write_loop(writer, outbound_receiver).await;
118 });
119 Self {
120 inner,
121 _reader_task: Arc::new(reader_task),
122 _writer_task: Arc::new(writer_task),
123 _quic_endpoint: quic_endpoint,
124 _quic_connection: quic_connection,
125 }
126 }
127
128 pub async fn hello(&self, hello: Hello) -> DcpResult<()> {
129 let response = self
130 .send_frame_with_response(0, DcpFrame::hello(hello))
131 .await?;
132 ensure_ok(response).map(|_| ())
133 }
134
135 pub async fn request(&self, request: Request) -> DcpResult<Response> {
136 let response = self
137 .send_frame_with_response(request.request_id, DcpFrame::request(request))
138 .await?;
139 ensure_ok(response)
140 }
141
142 pub async fn list_jobs(&self) -> DcpResult<Vec<WireJobStatus>> {
143 let response = self
144 .send_command(request::Command::ListJobs(ListJobs {}), 0)
145 .await?;
146 let list = crate::dcp::proto::JobList::decode(response.payload.as_slice())?;
147 Ok(list.jobs)
148 }
149
150 pub async fn start_job(
151 &self,
152 factory_name: impl Into<String>,
153 instance_name: impl Into<String>,
154 params: HashMap<String, String>,
155 ) -> DcpResult<WireJobStatus> {
156 self.job_status_response(request::Command::StartJob(StartJob {
157 factory_name: factory_name.into(),
158 instance_name: instance_name.into(),
159 params,
160 }))
161 .await
162 }
163
164 pub async fn drain_job(&self, name: impl Into<String>) -> DcpResult<WireJobStatus> {
165 self.job_status_response(request::Command::DrainJob(crate::dcp::proto::DrainJob {
166 name: name.into(),
167 }))
168 .await
169 }
170
171 pub async fn stop_job(&self, name: impl Into<String>) -> DcpResult<WireJobStatus> {
172 self.job_status_response(request::Command::StopJob(StopJob { name: name.into() }))
173 .await
174 }
175
176 pub async fn restart_job(&self, name: impl Into<String>) -> DcpResult<WireJobStatus> {
177 self.job_status_response(request::Command::RestartJob(
178 crate::dcp::proto::RestartJob { name: name.into() },
179 ))
180 .await
181 }
182
183 pub async fn job_status(&self, name: impl Into<String>) -> DcpResult<WireJobStatus> {
184 self.job_status_response(request::Command::JobStatus(JobStatusRequest {
185 name: name.into(),
186 }))
187 .await
188 }
189
190 pub async fn subscribe_events(&self) -> DcpResult<EventSubscription> {
191 let request_id = self.next_request_id();
192 let (sender, receiver) = mpsc::channel(256);
193 self.inner
194 .event_subscriptions
195 .lock()
196 .expect("DCP event subscriptions poisoned")
197 .insert(request_id, sender);
198 let request = Request {
199 request_id,
200 deadline_ms: 0,
201 command: Some(request::Command::SubscribeEvents(SubscribeEvents {
202 buffer: 0,
203 })),
204 };
205 if let Err(error) = self.request(request).await {
206 self.inner
207 .event_subscriptions
208 .lock()
209 .expect("DCP event subscriptions poisoned")
210 .remove(&request_id);
211 return Err(error);
212 }
213 Ok(EventSubscription { receiver })
214 }
215
216 pub async fn subscribe_metrics(&self, interval_ms: u64) -> DcpResult<MetricSubscription> {
217 let request_id = self.next_request_id();
218 let (sender, receiver) = mpsc::channel(16);
219 self.inner
220 .metric_subscriptions
221 .lock()
222 .expect("DCP metric subscriptions poisoned")
223 .insert(request_id, sender);
224 let request = Request {
225 request_id,
226 deadline_ms: 0,
227 command: Some(request::Command::SubscribeMetrics(SubscribeMetrics {
228 interval_ms,
229 })),
230 };
231 if let Err(error) = self.request(request).await {
232 self.inner
233 .metric_subscriptions
234 .lock()
235 .expect("DCP metric subscriptions poisoned")
236 .remove(&request_id);
237 return Err(error);
238 }
239 Ok(MetricSubscription { receiver })
240 }
241
242 pub async fn get_config(&self, key: impl Into<String>) -> DcpResult<ConfigValue> {
243 let response = self
244 .send_command(
245 request::Command::GetConfig(GetConfig { key: key.into() }),
246 0,
247 )
248 .await?;
249 Ok(ConfigValue::decode(response.payload.as_slice())?)
250 }
251
252 pub async fn put_config(
253 &self,
254 key: impl Into<String>,
255 value: impl Into<String>,
256 ) -> DcpResult<ConfigValue> {
257 let response = self
258 .send_command(
259 request::Command::PutConfig(PutConfig {
260 key: key.into(),
261 value: value.into(),
262 }),
263 0,
264 )
265 .await?;
266 Ok(ConfigValue::decode(response.payload.as_slice())?)
267 }
268
269 async fn job_status_response(&self, command: request::Command) -> DcpResult<WireJobStatus> {
270 let response = self.send_command(command, 0).await?;
271 Ok(WireJobStatus::decode(response.payload.as_slice())?)
272 }
273
274 async fn send_command(
275 &self,
276 command: request::Command,
277 deadline_ms: u64,
278 ) -> DcpResult<Response> {
279 let request_id = self.next_request_id();
280 let request = Request {
281 request_id,
282 deadline_ms,
283 command: Some(command),
284 };
285 self.request(request).await
286 }
287
288 async fn send_frame_with_response(
289 &self,
290 request_id: u64,
291 frame: DcpFrame,
292 ) -> DcpResult<Response> {
293 let (sender, receiver) = oneshot::channel();
294 self.inner
295 .pending
296 .lock()
297 .expect("DCP pending map poisoned")
298 .insert(request_id, sender);
299 if self.inner.outbound.send(frame).await.is_err() {
300 self.inner
301 .pending
302 .lock()
303 .expect("DCP pending map poisoned")
304 .remove(&request_id);
305 return Err(DcpError::Closed);
306 }
307 receiver.await.map_err(|_| DcpError::Closed)?
308 }
309
310 fn next_request_id(&self) -> u64 {
311 self.inner.next_request_id.fetch_add(1, Ordering::Relaxed)
312 }
313}
314
315async fn read_loop<R>(mut reader: R, inner: Arc<ClientInner>)
316where
317 R: AsyncRead + Unpin,
318{
319 loop {
320 match read_frame(&mut reader).await {
321 Ok(Some(frame)) => handle_inbound_frame(frame, &inner).await,
322 Ok(None) => {
323 fail_pending(&inner, DcpError::Closed);
324 break;
325 }
326 Err(error) => {
327 fail_pending(&inner, error);
328 break;
329 }
330 }
331 }
332}
333
334async fn handle_inbound_frame(frame: DcpFrame, inner: &Arc<ClientInner>) {
335 match frame.frame {
336 Some(dcp_frame::Frame::Response(response)) => {
337 if let Some(sender) = inner
338 .pending
339 .lock()
340 .expect("DCP pending map poisoned")
341 .remove(&response.request_id)
342 {
343 let _ = sender.send(Ok(response));
344 }
345 }
346 Some(dcp_frame::Frame::Event(event_frame)) => {
347 if let Some(event) = event_frame.event {
348 let sender = inner
349 .event_subscriptions
350 .lock()
351 .expect("DCP event subscriptions poisoned")
352 .get(&event_frame.subscription_id)
353 .cloned();
354 if let Some(sender) = sender {
355 let _ = sender.send(event).await;
356 }
357 }
358 }
359 Some(dcp_frame::Frame::Metric(MetricFrame {
360 subscription_id,
361 sample: Some(sample),
362 })) => {
363 let sender = inner
364 .metric_subscriptions
365 .lock()
366 .expect("DCP metric subscriptions poisoned")
367 .get(&subscription_id)
368 .cloned();
369 if let Some(sender) = sender {
370 let _ = sender.send(sample).await;
371 }
372 }
373 _ => {}
374 }
375}
376
377fn fail_pending(inner: &ClientInner, error: DcpError) {
378 let pending = std::mem::take(&mut *inner.pending.lock().expect("DCP pending map poisoned"));
379 let message = error.to_string();
380 for (_, sender) in pending {
381 let _ = sender.send(Err(DcpError::Protocol(message.clone())));
382 }
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
396fn ensure_ok(response: Response) -> DcpResult<Response> {
397 let status = response.response_status();
398 if status == ResponseStatus::Ok {
399 Ok(response)
400 } else {
401 Err(DcpError::response(status, response.message))
402 }
403}
404
405impl Default for DcpClient {
406 fn default() -> Self {
407 panic!("DcpClient cannot be default-constructed; use connect_tcp or connect_quic")
408 }
409}
410
411#[must_use]
412pub fn default_hello(node_id: impl Into<String>, client_kind: ClientKind) -> Hello {
413 Hello::new(node_id, client_kind)
414}