use std::{
collections::HashMap,
net::SocketAddr,
sync::{
Arc, Mutex,
atomic::{AtomicU64, Ordering},
},
time::Duration,
};
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, ClusterJobList, ClusterNodeInfo, ClusterNodeList, CompleteShardingAsk,
ConfigValue, DcpFrame, Event, ForwardShardEnvelopes, GetConfig, GetShardAllocations, Hello,
JobStatus as WireJobStatus, JobStatusRequest, ListClusterJobs, ListJobs, MetricFrame,
MetricSample, OpenShardPipe, PlacementSpec, PutConfig, RememberClusterAssignment,
RememberShardAllocations, Request, Response, ResponseStatus, ShardAllocation,
ShardAllocationRequest, ShardAllocationTable, ShardEnvelopeBatchResult, ShardPipeFrame,
StartJob, StopJob, SubmitClusterJob, 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,
disconnect_error: Arc<Mutex<Option<String>>>,
}
pub struct EventSubscription {
receiver: mpsc::Receiver<Event>,
disconnect_error: Arc<Mutex<Option<String>>>,
}
impl EventSubscription {
pub async fn recv(&mut self) -> Option<Event> {
self.receiver.recv().await
}
pub fn disconnect_reason(&self) -> Option<String> {
self.disconnect_error
.lock()
.unwrap_or_else(|p| p.into_inner())
.clone()
}
}
pub struct MetricSubscription {
receiver: mpsc::Receiver<MetricSample>,
disconnect_error: Arc<Mutex<Option<String>>>,
}
impl MetricSubscription {
pub async fn recv(&mut self) -> Option<MetricSample> {
self.receiver.recv().await
}
pub fn disconnect_reason(&self) -> Option<String> {
self.disconnect_error
.lock()
.unwrap_or_else(|p| p.into_inner())
.clone()
}
}
#[derive(Clone)]
pub struct ShardPipeClient {
outbound: mpsc::Sender<ShardPipeFrame>,
_reader_task: Arc<JoinHandle<()>>,
_writer_task: Arc<JoinHandle<()>>,
_quic_endpoint: Option<quinn::Endpoint>,
_quic_connection: Option<quinn::Connection>,
}
impl ShardPipeClient {
pub async fn connect_tcp(addr: SocketAddr, hello: Hello) -> DcpResult<Self> {
let mut stream = TcpStream::connect(addr).await?;
stream.set_nodelay(true)?;
{
let (mut reader, mut writer) = stream.split();
open_shard_pipe(&mut reader, &mut writer, hello).await?;
}
let (reader, writer) = stream.into_split();
Ok(Self::from_parts(reader, writer, None, None))
}
pub async fn connect_quic(
addr: SocketAddr,
server_name: &str,
client_config: quinn::ClientConfig,
hello: Hello,
) -> DcpResult<Self> {
let (endpoint, connection, mut reader, mut writer) =
connect_quic_stream(addr, server_name, client_config).await?;
open_shard_pipe(&mut reader, &mut writer, hello).await?;
Ok(Self::from_parts(
reader,
writer,
Some(endpoint),
Some(connection),
))
}
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 reader_task = tokio::spawn(async move {
read_shard_pipe_loop(reader).await;
});
let writer_task = tokio::spawn(async move {
let _ = write_shard_pipe_loop(writer, outbound_receiver).await;
});
Self {
outbound,
_reader_task: Arc::new(reader_task),
_writer_task: Arc::new(writer_task),
_quic_endpoint: quic_endpoint,
_quic_connection: quic_connection,
}
}
pub async fn send(&self, frame: ShardPipeFrame) -> DcpResult<()> {
self.outbound
.send(frame)
.await
.map_err(|_| DcpError::Closed)
}
}
const DEFAULT_LOCAL_REQUEST_TIMEOUT_MS: u64 = 30_000;
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 disconnect_error = Arc::new(Mutex::new(None));
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),
disconnect_error: Arc::clone(&disconnect_error),
});
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 {}),
DEFAULT_LOCAL_REQUEST_TIMEOUT_MS,
)
.await?;
let list = crate::dcp::proto::JobList::decode(response.payload.as_slice())?;
Ok(list.jobs)
}
pub async fn list_cluster_jobs(&self, timeout_ms: u64) -> DcpResult<ClusterJobList> {
let request = self.send_command(
request::Command::ListClusterJobs(ListClusterJobs { timeout_ms }),
cluster_request_deadline(timeout_ms),
);
let response = timeout_cluster_request(timeout_ms, request).await?;
Ok(ClusterJobList::decode(response.payload.as_slice())?)
}
pub async fn cluster_node_info(&self, timeout_ms: u64) -> DcpResult<ClusterNodeList> {
let request = self.send_command(
request::Command::ClusterNodeInfo(ClusterNodeInfo { timeout_ms }),
cluster_request_deadline(timeout_ms),
);
let response = timeout_cluster_request(timeout_ms, request).await?;
Ok(ClusterNodeList::decode(response.payload.as_slice())?)
}
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,
cluster: None,
}))
.await
}
pub(crate) async fn start_cluster_job_on_node(
&self,
factory_name: impl Into<String>,
instance_name: impl Into<String>,
params: HashMap<String, String>,
cluster: crate::dcp::proto::ClusterJobStart,
) -> DcpResult<WireJobStatus> {
self.job_status_response(request::Command::StartJob(StartJob {
factory_name: factory_name.into(),
instance_name: instance_name.into(),
params,
cluster: Some(cluster),
}))
.await
}
pub async fn submit_cluster_job(
&self,
factory_name: impl Into<String>,
instance_name: impl Into<String>,
params: HashMap<String, String>,
placement: PlacementSpec,
timeout_ms: u64,
) -> DcpResult<WireJobStatus> {
let request = self.send_command(
request::Command::SubmitClusterJob(SubmitClusterJob {
factory_name: factory_name.into(),
instance_name: instance_name.into(),
params,
placement: Some(placement),
timeout_ms,
}),
cluster_request_deadline(timeout_ms),
);
let response = timeout_cluster_request(timeout_ms, request).await?;
Ok(WireJobStatus::decode(response.payload.as_slice())?)
}
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(),
cluster: false,
}))
.await
}
pub async fn drain_cluster_job(
&self,
name: impl Into<String>,
timeout_ms: u64,
) -> DcpResult<WireJobStatus> {
let request = self.send_command(
request::Command::DrainJob(crate::dcp::proto::DrainJob {
name: name.into(),
cluster: true,
}),
cluster_request_deadline(timeout_ms),
);
let response = timeout_cluster_request(timeout_ms, request).await?;
Ok(WireJobStatus::decode(response.payload.as_slice())?)
}
pub async fn stop_job(&self, name: impl Into<String>) -> DcpResult<WireJobStatus> {
self.job_status_response(request::Command::StopJob(StopJob {
name: name.into(),
cluster: false,
}))
.await
}
pub async fn stop_cluster_job(
&self,
name: impl Into<String>,
timeout_ms: u64,
) -> DcpResult<WireJobStatus> {
let request = self.send_command(
request::Command::StopJob(StopJob {
name: name.into(),
cluster: true,
}),
cluster_request_deadline(timeout_ms),
);
let response = timeout_cluster_request(timeout_ms, request).await?;
Ok(WireJobStatus::decode(response.payload.as_slice())?)
}
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(),
cluster: false,
},
))
.await
}
pub async fn job_status(&self, name: impl Into<String>) -> DcpResult<WireJobStatus> {
self.job_status_response(request::Command::JobStatus(JobStatusRequest {
name: name.into(),
cluster: false,
}))
.await
}
pub async fn cluster_job_status(
&self,
name: impl Into<String>,
timeout_ms: u64,
) -> DcpResult<WireJobStatus> {
let request = self.send_command(
request::Command::JobStatus(JobStatusRequest {
name: name.into(),
cluster: true,
}),
cluster_request_deadline(timeout_ms),
);
let response = timeout_cluster_request(timeout_ms, request).await?;
Ok(WireJobStatus::decode(response.payload.as_slice())?)
}
pub(crate) async fn remember_cluster_assignment(
&self,
instance_name: impl Into<String>,
assignment: crate::dcp::proto::ClusterJobStart,
) -> DcpResult<()> {
let response = self
.send_command(
request::Command::RememberClusterAssignment(RememberClusterAssignment {
instance_name: instance_name.into(),
assignment: Some(assignment),
}),
0,
)
.await?;
let _ = response;
Ok(())
}
pub async fn allocate_shard(
&self,
type_name: impl Into<String>,
shard_id: impl Into<String>,
timeout_ms: u64,
) -> DcpResult<ShardAllocation> {
let request = self.send_command(
request::Command::AllocateShard(ShardAllocationRequest {
type_name: type_name.into(),
shard_id: shard_id.into(),
timeout_ms,
}),
cluster_request_deadline(timeout_ms),
);
let response = timeout_cluster_request(timeout_ms, request).await?;
Ok(ShardAllocation::decode(response.payload.as_slice())?)
}
pub async fn remember_shard_allocations(&self, table: ShardAllocationTable) -> DcpResult<()> {
let response = self
.send_command(
request::Command::RememberShardAllocations(RememberShardAllocations {
table: Some(table),
}),
0,
)
.await?;
let _ = response;
Ok(())
}
pub async fn get_shard_allocations(
&self,
type_name: impl Into<String>,
timeout_ms: u64,
) -> DcpResult<ShardAllocationTable> {
let request = self.send_command(
request::Command::GetShardAllocations(GetShardAllocations {
type_name: type_name.into(),
}),
cluster_request_deadline(timeout_ms),
);
let response = timeout_cluster_request(timeout_ms, request).await?;
Ok(ShardAllocationTable::decode(response.payload.as_slice())?)
}
pub async fn forward_shard_envelopes(
&self,
batch: ForwardShardEnvelopes,
timeout_ms: u64,
) -> DcpResult<ShardEnvelopeBatchResult> {
let request = self.send_command(
request::Command::ForwardShardEnvelopes(batch),
cluster_request_deadline(timeout_ms),
);
let response = timeout_cluster_request(timeout_ms, request).await?;
Ok(ShardEnvelopeBatchResult::decode(
response.payload.as_slice(),
)?)
}
pub async fn complete_sharding_ask(&self, reply: CompleteShardingAsk) -> DcpResult<()> {
let response = self
.send_command(request::Command::CompleteShardingAsk(reply), 0)
.await?;
let _ = response;
Ok(())
}
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: DEFAULT_LOCAL_REQUEST_TIMEOUT_MS,
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,
disconnect_error: Arc::clone(&self.inner.disconnect_error),
})
}
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: DEFAULT_LOCAL_REQUEST_TIMEOUT_MS,
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,
disconnect_error: Arc::clone(&self.inner.disconnect_error),
})
}
pub async fn get_config(&self, key: impl Into<String>) -> DcpResult<ConfigValue> {
let response = self
.send_command(
request::Command::GetConfig(GetConfig { key: key.into() }),
DEFAULT_LOCAL_REQUEST_TIMEOUT_MS,
)
.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(),
}),
DEFAULT_LOCAL_REQUEST_TIMEOUT_MS,
)
.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, DEFAULT_LOCAL_REQUEST_TIMEOUT_MS)
.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) => {
set_disconnect_error(&inner, DcpError::Closed);
fail_pending(&inner, DcpError::Closed);
clear_subscriptions(&inner);
break;
}
Err(error) => {
set_disconnect_error(
&inner,
DcpError::Protocol(format!("DCP read error: {}", error)),
);
fail_pending(&inner, error);
clear_subscriptions(&inner);
break;
}
}
}
}
fn clear_subscriptions(inner: &ClientInner) {
inner
.event_subscriptions
.lock()
.unwrap_or_else(|p| p.into_inner())
.clear();
inner
.metric_subscriptions
.lock()
.unwrap_or_else(|p| p.into_inner())
.clear();
}
fn set_disconnect_error(inner: &ClientInner, error: DcpError) {
*inner
.disconnect_error
.lock()
.unwrap_or_else(|p| p.into_inner()) = Some(error.to_string());
}
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(())
}
async fn open_shard_pipe<R, W>(reader: &mut R, writer: &mut W, hello: Hello) -> DcpResult<()>
where
R: AsyncRead + Unpin,
W: AsyncWrite + Unpin,
{
write_frame(writer, &DcpFrame::hello(hello)).await?;
let Some(response) = read_frame(reader).await? else {
return Err(DcpError::Closed);
};
match response.frame {
Some(dcp_frame::Frame::Response(response)) => {
ensure_ok(response)?;
}
_ => {
return Err(DcpError::Protocol(
"DCP shard pipe hello did not receive a response".to_owned(),
));
}
}
let request = Request {
request_id: 1,
deadline_ms: 0,
command: Some(request::Command::OpenShardPipe(OpenShardPipe {})),
};
write_frame(writer, &DcpFrame::request(request)).await?;
let Some(response) = read_frame(reader).await? else {
return Err(DcpError::Closed);
};
match response.frame {
Some(dcp_frame::Frame::Response(response)) => ensure_ok(response).map(|_| ()),
_ => Err(DcpError::Protocol(
"DCP shard pipe open did not receive a response".to_owned(),
)),
}
}
async fn read_shard_pipe_loop<R>(mut reader: R)
where
R: AsyncRead + Unpin,
{
while let Ok(Some(_frame)) = read_frame(&mut reader).await {}
}
async fn write_shard_pipe_loop<W>(
mut writer: W,
mut outbound: mpsc::Receiver<ShardPipeFrame>,
) -> DcpResult<()>
where
W: AsyncWrite + Unpin,
{
while let Some(frame) = outbound.recv().await {
write_frame(&mut writer, &DcpFrame::shard_pipe(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))
}
}
async fn timeout_cluster_request<F>(timeout_ms: u64, request: F) -> DcpResult<Response>
where
F: std::future::Future<Output = DcpResult<Response>>,
{
if timeout_ms == 0 {
return request.await;
}
tokio::time::timeout(
Duration::from_millis(timeout_ms).saturating_add(Duration::from_millis(250)),
request,
)
.await
.map_err(|_| {
DcpError::response(
ResponseStatus::DeadlineExceeded,
"cluster request timed out",
)
})?
}
fn cluster_request_deadline(timeout_ms: u64) -> u64 {
if timeout_ms == 0 {
0
} else {
timeout_ms.saturating_add(250)
}
}
#[must_use]
pub fn default_hello(node_id: impl Into<String>, client_kind: ClientKind) -> Hello {
Hello::new(node_id, client_kind)
}