use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
use std::sync::{Arc, Mutex as StdMutex, MutexGuard, PoisonError};
use std::time::Duration;
use tokio::io::{AsyncWriteExt, BufReader, BufWriter};
use tokio::net::tcp::{OwnedReadHalf, OwnedWriteHalf};
use tokio::net::TcpStream;
use tokio::sync::{mpsc, oneshot, Mutex as TokioMutex, Semaphore};
use tokio::task::JoinHandle;
use crate::wire::config::{Handshake, HelloStyle, PushPolicy};
use crate::wire::{
encode_frame, read_response_with_limit, Config, Request, Response, Value, PUSH_ID,
};
use crate::client::endpoint::{parse_endpoint, Endpoint};
use crate::client::error::ClientError;
const BACKOFF_BASE: Duration = Duration::from_millis(50);
const BACKOFF_CAP: Duration = Duration::from_millis(500);
const RECONNECT_ATTEMPTS: u32 = 2;
#[derive(Debug, Clone)]
pub enum Credentials {
Token(String),
ApiKey(String),
UserPass {
user: String,
pass: String,
},
}
#[derive(Debug, Clone)]
pub struct ClientConfig {
pub connect_timeout: Duration,
pub call_timeout: Duration,
pub credentials: Option<Credentials>,
pub client_name: Option<String>,
}
impl Default for ClientConfig {
fn default() -> Self {
Self {
connect_timeout: Duration::from_secs(10),
call_timeout: Duration::from_secs(30),
credentials: None,
client_name: None,
}
}
}
impl ClientConfig {
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn connect_timeout(mut self, timeout: Duration) -> Self {
self.connect_timeout = timeout;
self
}
#[must_use]
pub fn call_timeout(mut self, timeout: Duration) -> Self {
self.call_timeout = timeout;
self
}
#[must_use]
pub fn token(mut self, token: impl Into<String>) -> Self {
self.credentials = Some(Credentials::Token(token.into()));
self
}
#[must_use]
pub fn api_key(mut self, api_key: impl Into<String>) -> Self {
self.credentials = Some(Credentials::ApiKey(api_key.into()));
self
}
#[must_use]
pub fn user_pass(mut self, user: impl Into<String>, pass: impl Into<String>) -> Self {
self.credentials = Some(Credentials::UserPass {
user: user.into(),
pass: pass.into(),
});
self
}
#[must_use]
pub fn client_name(mut self, name: impl Into<String>) -> Self {
self.client_name = Some(name.into());
self
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct HandshakeInfo {
pub authenticated: bool,
pub capabilities: Vec<String>,
}
type PushHandler = Arc<dyn Fn(Value) + Send + Sync>;
type PendingTx = oneshot::Sender<Result<Response, ClientError>>;
struct ConnShared {
pending: StdMutex<HashMap<u32, PendingTx>>,
alive: AtomicBool,
}
impl ConnShared {
fn poison(&self, err: &ClientError) {
self.alive.store(false, Ordering::SeqCst);
let drained: Vec<PendingTx> = {
let mut pending = lock(&self.pending);
pending.drain().map(|(_, tx)| tx).collect()
};
for tx in drained {
let _ = tx.send(Err(err.clone()));
}
}
}
struct Conn {
shared: Arc<ConnShared>,
write_tx: mpsc::Sender<Vec<u8>>,
reader_task: JoinHandle<()>,
writer_task: JoinHandle<()>,
}
impl Conn {
fn is_alive(&self) -> bool {
self.shared.alive.load(Ordering::SeqCst)
}
fn kill(&self, err: &ClientError) {
self.reader_task.abort();
self.writer_task.abort();
self.shared.poison(err);
}
}
async fn writer_loop(
write_half: OwnedWriteHalf,
mut rx: mpsc::Receiver<Vec<u8>>,
shared: Arc<ConnShared>,
) {
let mut writer = BufWriter::new(write_half);
while let Some(frame) = rx.recv().await {
if writer.write_all(&frame).await.is_err() {
break;
}
while let Ok(next) = rx.try_recv() {
if writer.write_all(&next).await.is_err() {
shared.poison(&ClientError::Connection {
message: "write failed".to_owned(),
});
return;
}
}
if writer.flush().await.is_err() {
break;
}
}
let _ = writer.flush().await;
let _ = writer.shutdown().await;
}
impl Drop for Conn {
fn drop(&mut self) {
self.kill(&ClientError::Connection {
message: "connection dropped".to_owned(),
});
}
}
enum DispatchError {
WriteFailed(ClientError),
Fatal(ClientError),
}
impl DispatchError {
fn into_error(self) -> ClientError {
match self {
Self::WriteFailed(e) | Self::Fatal(e) => e,
}
}
}
pub struct Client {
config: Config,
client_config: ClientConfig,
endpoint: Endpoint,
next_id: AtomicU32,
in_flight: Semaphore,
conn: StdMutex<Option<Arc<Conn>>>,
reconnect: TokioMutex<()>,
closed: AtomicBool,
push_handler: Arc<StdMutex<Option<PushHandler>>>,
unknown_drops: Arc<AtomicU64>,
handshake_info: StdMutex<HandshakeInfo>,
}
impl Client {
pub async fn connect(endpoint: &str, config: Config) -> Result<Self, ClientError> {
Self::connect_with(endpoint, config, ClientConfig::default()).await
}
pub async fn connect_with(
endpoint: &str,
config: Config,
client_config: ClientConfig,
) -> Result<Self, ClientError> {
let endpoint = parse_endpoint(endpoint, &config)?;
let client = Self {
next_id: AtomicU32::new(1),
in_flight: Semaphore::new(config.max_in_flight),
conn: StdMutex::new(None),
reconnect: TokioMutex::new(()),
closed: AtomicBool::new(false),
push_handler: Arc::new(StdMutex::new(None)),
unknown_drops: Arc::new(AtomicU64::new(0)),
handshake_info: StdMutex::new(HandshakeInfo::default()),
endpoint,
config,
client_config,
};
let conn = client.establish().await?;
*lock(&client.conn) = Some(conn);
Ok(client)
}
pub async fn call(
&self,
command: impl Into<String>,
args: Vec<Value>,
) -> Result<Value, ClientError> {
let command = command.into();
self.call_with_timeout(&command, args, self.client_config.call_timeout)
.await
}
pub async fn call_with_timeout(
&self,
command: &str,
args: Vec<Value>,
timeout: Duration,
) -> Result<Value, ClientError> {
let _permit = self
.in_flight
.acquire()
.await
.map_err(|_| Self::closed_error())?;
let mut redials_left = RECONNECT_ATTEMPTS;
loop {
let conn = self.live_conn(&mut redials_left).await?;
match self.dispatch(&conn, command, args.clone(), timeout).await {
Ok(value) => return Ok(value),
Err(DispatchError::Fatal(err)) => return Err(err),
Err(DispatchError::WriteFailed(err)) => {
if redials_left == 0 {
return Err(err);
}
}
}
}
}
pub fn on_push<F>(&self, handler: F)
where
F: Fn(Value) + Send + Sync + 'static,
{
*lock(&self.push_handler) = Some(Arc::new(handler));
}
pub async fn close(&self) {
self.closed.store(true, Ordering::SeqCst);
self.in_flight.close();
let conn = lock(&self.conn).take();
if let Some(conn) = conn {
conn.kill(&Self::closed_error());
}
}
pub fn is_authenticated(&self) -> bool {
lock(&self.handshake_info).authenticated
}
pub fn capabilities(&self) -> Vec<String> {
lock(&self.handshake_info).capabilities.clone()
}
pub fn handshake_info(&self) -> HandshakeInfo {
lock(&self.handshake_info).clone()
}
pub fn unknown_response_drops(&self) -> u64 {
self.unknown_drops.load(Ordering::Relaxed)
}
pub fn config(&self) -> &Config {
&self.config
}
fn closed_error() -> ClientError {
ClientError::Connection {
message: "client is closed".to_owned(),
}
}
fn alloc_id(&self) -> u32 {
loop {
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
if id != PUSH_ID {
return id;
}
}
}
async fn live_conn(&self, redials_left: &mut u32) -> Result<Arc<Conn>, ClientError> {
if self.closed.load(Ordering::SeqCst) {
return Err(Self::closed_error());
}
let current = { lock(&self.conn).clone() };
if let Some(conn) = current {
if conn.is_alive() {
return Ok(conn);
}
}
let _guard = self.reconnect.lock().await;
if self.closed.load(Ordering::SeqCst) {
return Err(Self::closed_error());
}
let current = { lock(&self.conn).clone() };
if let Some(conn) = current {
if conn.is_alive() {
return Ok(conn);
}
}
let mut last_err = ClientError::Connection {
message: "connection is dead".to_owned(),
};
let mut backoff = BACKOFF_BASE;
while *redials_left > 0 {
*redials_left -= 1;
match self.establish().await {
Ok(conn) => {
*lock(&self.conn) = Some(Arc::clone(&conn));
return Ok(conn);
}
Err(err @ ClientError::Auth { .. }) => return Err(err),
Err(err) => {
last_err = err;
if *redials_left > 0 {
tokio::time::sleep(backoff).await;
backoff = (backoff * 2).min(BACKOFF_CAP);
}
}
}
}
Err(last_err)
}
async fn establish(&self) -> Result<Arc<Conn>, ClientError> {
let addr = (self.endpoint.host.as_str(), self.endpoint.port);
let stream =
tokio::time::timeout(self.client_config.connect_timeout, TcpStream::connect(addr))
.await
.map_err(|_| ClientError::Timeout)?
.map_err(|e| ClientError::Connection {
message: format!(
"connect to {}:{} failed: {e}",
self.endpoint.host, self.endpoint.port
),
})?;
stream
.set_nodelay(true)
.map_err(|e| ClientError::Connection {
message: format!("TCP_NODELAY failed: {e}"),
})?;
let (read_half, write_half) = stream.into_split();
let shared = Arc::new(ConnShared {
pending: StdMutex::new(HashMap::new()),
alive: AtomicBool::new(true),
});
let reader_task = tokio::spawn(reader_loop(
BufReader::new(read_half),
Arc::clone(&shared),
self.config.max_frame_bytes,
self.config.push,
Arc::clone(&self.push_handler),
Arc::clone(&self.unknown_drops),
));
let (write_tx, write_rx) = mpsc::channel::<Vec<u8>>(1024);
let writer_task = tokio::spawn(writer_loop(write_half, write_rx, Arc::clone(&shared)));
let conn = Arc::new(Conn {
shared,
write_tx,
reader_task,
writer_task,
});
let info = self.handshake(&conn).await?;
*lock(&self.handshake_info) = info;
Ok(conn)
}
async fn handshake(&self, conn: &Arc<Conn>) -> Result<HandshakeInfo, ClientError> {
match self.config.handshake {
Handshake::None => Ok(HandshakeInfo::default()),
Handshake::AuthCommand => {
let Some(credentials) = self.client_config.credentials.clone() else {
return Ok(HandshakeInfo::default());
};
if self.config.hello_style == HelloStyle::ArgLess {
self.handshake_call(conn, "HELLO", Vec::new()).await?;
}
let args = match credentials {
Credentials::Token(token) => vec![Value::Str(token)],
Credentials::ApiKey(api_key) => vec![Value::Str(api_key)],
Credentials::UserPass { user, pass } => {
vec![Value::Str(user), Value::Str(pass)]
}
};
self.handshake_call(conn, "AUTH", args).await?;
Ok(HandshakeInfo {
authenticated: true,
capabilities: Vec::new(),
})
}
Handshake::HelloMandatory => {
let mut pairs = vec![(Value::Str("version".to_owned()), Value::Int(1))];
match &self.client_config.credentials {
Some(Credentials::Token(token)) => {
pairs.push((Value::Str("token".to_owned()), Value::Str(token.clone())));
}
Some(Credentials::ApiKey(api_key)) => {
pairs.push((
Value::Str("api_key".to_owned()),
Value::Str(api_key.clone()),
));
}
Some(Credentials::UserPass { .. }) => {
return Err(ClientError::Auth {
message: "user/password credentials are not supported by \
HelloMandatory profiles — use a token or api_key (PRO-001)"
.to_owned(),
});
}
None => {}
}
let name = self
.client_config
.client_name
.clone()
.unwrap_or_else(|| "thunder-client".to_owned());
pairs.push((Value::Str("client_name".to_owned()), Value::Str(name)));
let reply = self
.handshake_call(conn, "HELLO", vec![Value::Map(pairs)])
.await?;
Ok(HandshakeInfo {
authenticated: reply
.map_get("authenticated")
.and_then(Value::as_bool)
.unwrap_or(false),
capabilities: reply
.map_get("capabilities")
.and_then(Value::as_array)
.map(|caps| {
caps.iter()
.filter_map(|v| v.as_str().map(str::to_owned))
.collect()
})
.unwrap_or_default(),
})
}
}
}
async fn handshake_call(
&self,
conn: &Arc<Conn>,
command: &str,
args: Vec<Value>,
) -> Result<Value, ClientError> {
self.dispatch(conn, command, args, self.client_config.call_timeout)
.await
.map_err(|e| match e.into_error() {
ClientError::Server { message, .. } | ClientError::Auth { message } => {
ClientError::Auth { message }
}
other => other,
})
}
async fn dispatch(
&self,
conn: &Arc<Conn>,
command: &str,
args: Vec<Value>,
timeout: Duration,
) -> Result<Value, DispatchError> {
let id = self.alloc_id();
let (tx, rx) = oneshot::channel();
{
let mut pending = lock(&conn.shared.pending);
if !conn.shared.alive.load(Ordering::SeqCst) {
return Err(DispatchError::WriteFailed(ClientError::Connection {
message: "connection is dead".to_owned(),
}));
}
pending.insert(id, tx);
}
let request = Request {
id,
command: command.to_owned(),
args,
};
let frame = match encode_frame(&request) {
Ok(frame) => frame,
Err(e) => {
lock(&conn.shared.pending).remove(&id);
let err = ClientError::Connection {
message: format!("encode failed: {e}"),
};
return Err(DispatchError::WriteFailed(err));
}
};
if conn.write_tx.send(frame).await.is_err() {
lock(&conn.shared.pending).remove(&id);
let err = ClientError::Connection {
message: "write failed: connection closed".to_owned(),
};
conn.kill(&err);
return Err(DispatchError::WriteFailed(err));
}
match tokio::time::timeout(timeout, rx).await {
Err(_elapsed) => {
lock(&conn.shared.pending).remove(&id);
Err(DispatchError::Fatal(ClientError::Timeout))
}
Ok(Err(_recv)) => Err(DispatchError::Fatal(ClientError::Connection {
message: "connection closed before response".to_owned(),
})),
Ok(Ok(Err(poison))) => Err(DispatchError::Fatal(poison)),
Ok(Ok(Ok(response))) => match response.result {
Ok(value) => Ok(value),
Err(message) => Err(DispatchError::Fatal(ClientError::from_server_message(
message,
self.config.error_codes,
))),
},
}
}
}
impl std::fmt::Debug for Client {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Client")
.field("scheme", &self.config.scheme)
.field("endpoint", &self.endpoint)
.field("closed", &self.closed.load(Ordering::Relaxed))
.finish_non_exhaustive()
}
}
impl Drop for Client {
fn drop(&mut self) {
if let Ok(mut guard) = self.conn.lock() {
if let Some(conn) = guard.take() {
conn.kill(&Self::closed_error());
}
}
}
}
async fn reader_loop(
mut reader: BufReader<OwnedReadHalf>,
shared: Arc<ConnShared>,
max_frame_bytes: usize,
push: PushPolicy,
push_handler: Arc<StdMutex<Option<PushHandler>>>,
unknown_drops: Arc<AtomicU64>,
) {
let err = loop {
match read_response_with_limit(&mut reader, max_frame_bytes).await {
Ok((response, _frame_bytes)) => {
if response.id == PUSH_ID {
match push {
PushPolicy::Enabled => {
let handler = { lock(&push_handler).clone() };
if let (Some(handler), Ok(value)) = (handler, response.result) {
handler(value);
}
}
PushPolicy::Reserved => {
break ClientError::Decode {
message: "server sent a push frame but the profile reserves \
PUSH_ID (CLT-060)"
.to_owned(),
};
}
}
continue;
}
let tx = lock(&shared.pending).remove(&response.id);
match tx {
Some(tx) => {
let _ = tx.send(Ok(response));
}
None => {
unknown_drops.fetch_add(1, Ordering::Relaxed);
}
}
}
Err(e) => break classify_read_error(&e),
}
};
shared.poison(&err);
}
fn classify_read_error(e: &std::io::Error) -> ClientError {
if e.kind() == std::io::ErrorKind::InvalidData {
let message = e.to_string();
if message.contains("exceeds limit") {
ClientError::FrameTooLarge { message }
} else {
ClientError::Decode { message }
}
} else {
ClientError::Connection {
message: format!("connection lost: {e}"),
}
}
}
fn lock<T>(mutex: &StdMutex<T>) -> MutexGuard<'_, T> {
mutex.lock().unwrap_or_else(PoisonError::into_inner)
}