pub mod pool;
pub use pool::{Pool, PoolConfig};
use std::cell::Cell;
use std::io;
use std::time::Instant;
use ping_proto::{Request as PingRequest, Response as PingResponse};
use ringline::{ConnCtx, ParseResult};
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum Error {
#[error("connection closed")]
ConnectionClosed,
#[error("unexpected response")]
UnexpectedResponse,
#[error("protocol error: {0}")]
Protocol(#[from] ping_proto::ParseError),
#[error("io error: {0}")]
Io(#[from] io::Error),
#[error("all connections failed")]
AllConnectionsFailed,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum CommandType {
Ping,
}
#[derive(Debug, Clone)]
pub struct CommandResult {
pub command: CommandType,
pub latency_ns: u64,
pub success: bool,
pub ttfb_ns: Option<u64>,
pub tx_bytes: u32,
pub rx_bytes: u32,
}
#[cfg(feature = "metrics")]
pub struct ClientMetrics {
pub latency: histogram::Histogram,
pub requests: u64,
pub errors: u64,
}
#[cfg(feature = "metrics")]
impl ClientMetrics {
fn new() -> Self {
Self {
latency: histogram::Histogram::new(7, 64).unwrap(),
requests: 0,
errors: 0,
}
}
fn record(&mut self, result: &CommandResult) {
self.requests += 1;
let _ = self.latency.increment(result.latency_ns);
if !result.success {
self.errors += 1;
}
}
}
type ResultCallback = Box<dyn Fn(&CommandResult)>;
pub struct ClientBuilder {
conn: ConnCtx,
on_result: Option<ResultCallback>,
#[cfg(feature = "timestamps")]
use_kernel_ts: bool,
#[cfg(feature = "metrics")]
with_metrics: bool,
}
impl ClientBuilder {
pub(crate) fn new(conn: ConnCtx) -> Self {
Self {
conn,
on_result: None,
#[cfg(feature = "timestamps")]
use_kernel_ts: false,
#[cfg(feature = "metrics")]
with_metrics: false,
}
}
pub fn on_result<F: Fn(&CommandResult) + 'static>(mut self, f: F) -> Self {
self.on_result = Some(Box::new(f));
self
}
#[cfg(feature = "timestamps")]
pub fn kernel_timestamps(mut self, enabled: bool) -> Self {
self.use_kernel_ts = enabled;
self
}
#[cfg(feature = "metrics")]
pub fn with_metrics(mut self) -> Self {
self.with_metrics = true;
self
}
pub fn build(self) -> Client {
Client {
conn: self.conn,
on_result: self.on_result,
last_rx_bytes: Cell::new(0),
#[cfg(feature = "timestamps")]
use_kernel_ts: self.use_kernel_ts,
#[cfg(feature = "metrics")]
metrics: if self.with_metrics {
Some(ClientMetrics::new())
} else {
None
},
}
}
}
pub struct Client {
conn: ConnCtx,
on_result: Option<ResultCallback>,
last_rx_bytes: Cell<u32>,
#[cfg(feature = "timestamps")]
use_kernel_ts: bool,
#[cfg(feature = "metrics")]
metrics: Option<ClientMetrics>,
}
impl Client {
pub fn new(conn: ConnCtx) -> Self {
Self {
conn,
on_result: None,
last_rx_bytes: Cell::new(0),
#[cfg(feature = "timestamps")]
use_kernel_ts: false,
#[cfg(feature = "metrics")]
metrics: None,
}
}
pub fn builder(conn: ConnCtx) -> ClientBuilder {
ClientBuilder::new(conn)
}
pub fn conn(&self) -> ConnCtx {
self.conn
}
#[cfg(feature = "metrics")]
pub fn metrics(&self) -> Option<&ClientMetrics> {
self.metrics.as_ref()
}
#[cfg(feature = "metrics")]
pub fn metrics_mut(&mut self) -> Option<&mut ClientMetrics> {
self.metrics.as_mut()
}
#[inline]
fn is_instrumented(&self) -> bool {
if self.on_result.is_some() {
return true;
}
#[cfg(feature = "metrics")]
if self.metrics.is_some() {
return true;
}
false
}
#[cfg(feature = "timestamps")]
#[inline]
fn send_timestamp(&self) -> u64 {
if self.use_kernel_ts {
now_realtime_ns()
} else {
0
}
}
#[cfg(not(feature = "timestamps"))]
#[inline]
fn send_timestamp(&self) -> u64 {
0
}
#[cfg(feature = "timestamps")]
#[inline]
fn finish_timing(&self, send_ts: u64, start: Instant) -> u64 {
if self.use_kernel_ts {
let recv_ts = self.conn.recv_timestamp();
if recv_ts > 0 && recv_ts > send_ts {
return recv_ts - send_ts;
}
}
start.elapsed().as_nanos() as u64
}
#[cfg(not(feature = "timestamps"))]
#[inline]
fn finish_timing(&self, _send_ts: u64, start: Instant) -> u64 {
start.elapsed().as_nanos() as u64
}
fn record(&mut self, result: &CommandResult) {
if let Some(ref cb) = self.on_result {
cb(result);
}
#[cfg(feature = "metrics")]
if let Some(ref mut m) = self.metrics {
m.record(result);
}
}
pub(crate) async fn read_response(&self) -> Result<PingResponse, Error> {
let mut result: Option<Result<PingResponse, Error>> = None;
let n = self
.conn
.with_data(|data| match PingResponse::parse(data) {
Ok((response, consumed)) => {
result = Some(Ok(response));
ParseResult::Consumed(consumed)
}
Err(ping_proto::ParseError::Incomplete) => ParseResult::Consumed(0),
Err(e) => {
result = Some(Err(Error::Protocol(e)));
ParseResult::Consumed(data.len())
}
})
.await;
self.last_rx_bytes.set(n as u32);
if n == 0 {
return result.unwrap_or(Err(Error::ConnectionClosed));
}
let r = result.unwrap();
if matches!(r, Err(Error::Protocol(_))) {
self.conn.close();
}
r
}
async fn execute(&self, encoded: &[u8]) -> Result<PingResponse, Error> {
self.conn.send(encoded)?;
self.read_response().await
}
pub async fn ping(&mut self) -> Result<(), Error> {
let mut buf = [0u8; 6];
let len = PingRequest::Ping.encode(&mut buf);
if !self.is_instrumented() {
let response = self.execute(&buf[..len]).await?;
return match response {
PingResponse::Pong => Ok(()),
#[allow(unreachable_patterns)]
_ => Err(Error::UnexpectedResponse),
};
}
let tx_bytes = len as u32;
let send_ts = self.send_timestamp();
let start = Instant::now();
let response = self.execute(&buf[..len]).await;
let latency_ns = self.finish_timing(send_ts, start);
let rx_bytes = self.last_rx_bytes.get();
let result = match response {
Ok(PingResponse::Pong) => Ok(()),
Ok(_) => Err(Error::UnexpectedResponse),
Err(e) => Err(e),
};
self.record(&CommandResult {
command: CommandType::Ping,
latency_ns,
success: result.is_ok(),
ttfb_ns: None,
tx_bytes,
rx_bytes,
});
result
}
}
#[cfg(feature = "timestamps")]
fn now_realtime_ns() -> u64 {
let mut ts = libc::timespec {
tv_sec: 0,
tv_nsec: 0,
};
unsafe {
libc::clock_gettime(libc::CLOCK_REALTIME, &mut ts);
}
ts.tv_sec as u64 * 1_000_000_000 + ts.tv_nsec as u64
}