use std::{
collections::HashMap,
net::SocketAddr,
sync::{
Arc, Mutex,
atomic::{AtomicU64, Ordering},
},
};
use datum_net::quic::quinn;
use prost::Message as ProstMessage;
use tokio::{
io::{AsyncRead, AsyncWrite, AsyncWriteExt},
net::TcpStream,
sync::{mpsc, oneshot},
task::JoinHandle,
};
use crate::dcp::{
DcpError, DcpResult,
frame::{read_frame, write_frame},
proto::{
ClientKind, ConfigValue, DcpFrame, Event, GetConfig, Hello, JobStatus as WireJobStatus,
JobStatusRequest, ListJobs, MetricFrame, MetricSample, PutConfig, Request, Response,
ResponseStatus, StartJob, StopJob, SubscribeEvents, SubscribeMetrics, dcp_frame, request,
},
server::connect_quic_stream,
};
type PendingMap = HashMap<u64, oneshot::Sender<DcpResult<Response>>>;
#[derive(Clone)]
pub struct DcpClient {
inner: Arc<ClientInner>,
_reader_task: Arc<JoinHandle<()>>,
_writer_task: Arc<JoinHandle<()>>,
_quic_endpoint: Option<quinn::Endpoint>,
_quic_connection: Option<quinn::Connection>,
}
struct ClientInner {
outbound: mpsc::Sender<DcpFrame>,
pending: Mutex<PendingMap>,
event_subscriptions: Mutex<HashMap<u64, mpsc::Sender<Event>>>,
metric_subscriptions: Mutex<HashMap<u64, mpsc::Sender<MetricSample>>>,
next_request_id: AtomicU64,
}
pub struct EventSubscription {
receiver: mpsc::Receiver<Event>,
}
impl EventSubscription {
pub async fn recv(&mut self) -> Option<Event> {
self.receiver.recv().await
}
}
pub struct MetricSubscription {
receiver: mpsc::Receiver<MetricSample>,
}
impl MetricSubscription {
pub async fn recv(&mut self) -> Option<MetricSample> {
self.receiver.recv().await
}
}
impl DcpClient {
pub async fn connect_tcp(addr: SocketAddr, hello: Hello) -> DcpResult<Self> {
let stream = TcpStream::connect(addr).await?;
stream.set_nodelay(true)?;
let (reader, writer) = stream.into_split();
let client = Self::from_parts(reader, writer, None, None);
client.hello(hello).await?;
Ok(client)
}
pub async fn connect_quic(
addr: SocketAddr,
server_name: &str,
client_config: quinn::ClientConfig,
hello: Hello,
) -> DcpResult<Self> {
let (endpoint, connection, reader, writer) =
connect_quic_stream(addr, server_name, client_config).await?;
let client = Self::from_parts(reader, writer, Some(endpoint), Some(connection));
client.hello(hello).await?;
Ok(client)
}
fn from_parts<R, W>(
reader: R,
writer: W,
quic_endpoint: Option<quinn::Endpoint>,
quic_connection: Option<quinn::Connection>,
) -> Self
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
let (outbound, outbound_receiver) = mpsc::channel(256);
let inner = Arc::new(ClientInner {
outbound,
pending: Mutex::new(HashMap::new()),
event_subscriptions: Mutex::new(HashMap::new()),
metric_subscriptions: Mutex::new(HashMap::new()),
next_request_id: AtomicU64::new(1),
});
let reader_task = {
let inner = Arc::clone(&inner);
tokio::spawn(async move {
read_loop(reader, inner).await;
})
};
let writer_task = tokio::spawn(async move {
let _ = write_loop(writer, outbound_receiver).await;
});
Self {
inner,
_reader_task: Arc::new(reader_task),
_writer_task: Arc::new(writer_task),
_quic_endpoint: quic_endpoint,
_quic_connection: quic_connection,
}
}
pub async fn hello(&self, hello: Hello) -> DcpResult<()> {
let response = self
.send_frame_with_response(0, DcpFrame::hello(hello))
.await?;
ensure_ok(response).map(|_| ())
}
pub async fn request(&self, request: Request) -> DcpResult<Response> {
let response = self
.send_frame_with_response(request.request_id, DcpFrame::request(request))
.await?;
ensure_ok(response)
}
pub async fn list_jobs(&self) -> DcpResult<Vec<WireJobStatus>> {
let response = self
.send_command(request::Command::ListJobs(ListJobs {}), 0)
.await?;
let list = crate::dcp::proto::JobList::decode(response.payload.as_slice())?;
Ok(list.jobs)
}
pub async fn start_job(
&self,
factory_name: impl Into<String>,
instance_name: impl Into<String>,
params: HashMap<String, String>,
) -> DcpResult<WireJobStatus> {
self.job_status_response(request::Command::StartJob(StartJob {
factory_name: factory_name.into(),
instance_name: instance_name.into(),
params,
}))
.await
}
pub async fn drain_job(&self, name: impl Into<String>) -> DcpResult<WireJobStatus> {
self.job_status_response(request::Command::DrainJob(crate::dcp::proto::DrainJob {
name: name.into(),
}))
.await
}
pub async fn stop_job(&self, name: impl Into<String>) -> DcpResult<WireJobStatus> {
self.job_status_response(request::Command::StopJob(StopJob { name: name.into() }))
.await
}
pub async fn restart_job(&self, name: impl Into<String>) -> DcpResult<WireJobStatus> {
self.job_status_response(request::Command::RestartJob(
crate::dcp::proto::RestartJob { name: name.into() },
))
.await
}
pub async fn job_status(&self, name: impl Into<String>) -> DcpResult<WireJobStatus> {
self.job_status_response(request::Command::JobStatus(JobStatusRequest {
name: name.into(),
}))
.await
}
pub async fn subscribe_events(&self) -> DcpResult<EventSubscription> {
let request_id = self.next_request_id();
let (sender, receiver) = mpsc::channel(256);
self.inner
.event_subscriptions
.lock()
.expect("DCP event subscriptions poisoned")
.insert(request_id, sender);
let request = Request {
request_id,
deadline_ms: 0,
command: Some(request::Command::SubscribeEvents(SubscribeEvents {
buffer: 0,
})),
};
if let Err(error) = self.request(request).await {
self.inner
.event_subscriptions
.lock()
.expect("DCP event subscriptions poisoned")
.remove(&request_id);
return Err(error);
}
Ok(EventSubscription { receiver })
}
pub async fn subscribe_metrics(&self, interval_ms: u64) -> DcpResult<MetricSubscription> {
let request_id = self.next_request_id();
let (sender, receiver) = mpsc::channel(16);
self.inner
.metric_subscriptions
.lock()
.expect("DCP metric subscriptions poisoned")
.insert(request_id, sender);
let request = Request {
request_id,
deadline_ms: 0,
command: Some(request::Command::SubscribeMetrics(SubscribeMetrics {
interval_ms,
})),
};
if let Err(error) = self.request(request).await {
self.inner
.metric_subscriptions
.lock()
.expect("DCP metric subscriptions poisoned")
.remove(&request_id);
return Err(error);
}
Ok(MetricSubscription { receiver })
}
pub async fn get_config(&self, key: impl Into<String>) -> DcpResult<ConfigValue> {
let response = self
.send_command(
request::Command::GetConfig(GetConfig { key: key.into() }),
0,
)
.await?;
Ok(ConfigValue::decode(response.payload.as_slice())?)
}
pub async fn put_config(
&self,
key: impl Into<String>,
value: impl Into<String>,
) -> DcpResult<ConfigValue> {
let response = self
.send_command(
request::Command::PutConfig(PutConfig {
key: key.into(),
value: value.into(),
}),
0,
)
.await?;
Ok(ConfigValue::decode(response.payload.as_slice())?)
}
async fn job_status_response(&self, command: request::Command) -> DcpResult<WireJobStatus> {
let response = self.send_command(command, 0).await?;
Ok(WireJobStatus::decode(response.payload.as_slice())?)
}
async fn send_command(
&self,
command: request::Command,
deadline_ms: u64,
) -> DcpResult<Response> {
let request_id = self.next_request_id();
let request = Request {
request_id,
deadline_ms,
command: Some(command),
};
self.request(request).await
}
async fn send_frame_with_response(
&self,
request_id: u64,
frame: DcpFrame,
) -> DcpResult<Response> {
let (sender, receiver) = oneshot::channel();
self.inner
.pending
.lock()
.expect("DCP pending map poisoned")
.insert(request_id, sender);
if self.inner.outbound.send(frame).await.is_err() {
self.inner
.pending
.lock()
.expect("DCP pending map poisoned")
.remove(&request_id);
return Err(DcpError::Closed);
}
receiver.await.map_err(|_| DcpError::Closed)?
}
fn next_request_id(&self) -> u64 {
self.inner.next_request_id.fetch_add(1, Ordering::Relaxed)
}
}
async fn read_loop<R>(mut reader: R, inner: Arc<ClientInner>)
where
R: AsyncRead + Unpin,
{
loop {
match read_frame(&mut reader).await {
Ok(Some(frame)) => handle_inbound_frame(frame, &inner).await,
Ok(None) => {
fail_pending(&inner, DcpError::Closed);
break;
}
Err(error) => {
fail_pending(&inner, error);
break;
}
}
}
}
async fn handle_inbound_frame(frame: DcpFrame, inner: &Arc<ClientInner>) {
match frame.frame {
Some(dcp_frame::Frame::Response(response)) => {
if let Some(sender) = inner
.pending
.lock()
.expect("DCP pending map poisoned")
.remove(&response.request_id)
{
let _ = sender.send(Ok(response));
}
}
Some(dcp_frame::Frame::Event(event_frame)) => {
if let Some(event) = event_frame.event {
let sender = inner
.event_subscriptions
.lock()
.expect("DCP event subscriptions poisoned")
.get(&event_frame.subscription_id)
.cloned();
if let Some(sender) = sender {
let _ = sender.send(event).await;
}
}
}
Some(dcp_frame::Frame::Metric(MetricFrame {
subscription_id,
sample: Some(sample),
})) => {
let sender = inner
.metric_subscriptions
.lock()
.expect("DCP metric subscriptions poisoned")
.get(&subscription_id)
.cloned();
if let Some(sender) = sender {
let _ = sender.send(sample).await;
}
}
_ => {}
}
}
fn fail_pending(inner: &ClientInner, error: DcpError) {
let pending = std::mem::take(&mut *inner.pending.lock().expect("DCP pending map poisoned"));
let message = error.to_string();
for (_, sender) in pending {
let _ = sender.send(Err(DcpError::Protocol(message.clone())));
}
}
async fn write_loop<W>(mut writer: W, mut outbound: mpsc::Receiver<DcpFrame>) -> DcpResult<()>
where
W: AsyncWrite + Unpin,
{
while let Some(frame) = outbound.recv().await {
write_frame(&mut writer, &frame).await?;
}
let _ = writer.shutdown().await;
Ok(())
}
fn ensure_ok(response: Response) -> DcpResult<Response> {
let status = response.response_status();
if status == ResponseStatus::Ok {
Ok(response)
} else {
Err(DcpError::response(status, response.message))
}
}
impl Default for DcpClient {
fn default() -> Self {
panic!("DcpClient cannot be default-constructed; use connect_tcp or connect_quic")
}
}
#[must_use]
pub fn default_hello(node_id: impl Into<String>, client_kind: ClientKind) -> Hello {
Hello::new(node_id, client_kind)
}