1pub mod pool;
27pub use pool::{Pool, PoolConfig};
28
29use std::cell::Cell;
30use std::io;
31use std::time::Instant;
32
33use ping_proto::{Request as PingRequest, Response as PingResponse};
34use ringline::{ConnCtx, ParseResult};
35
36#[derive(Debug, thiserror::Error)]
44#[non_exhaustive]
45pub enum Error {
46 #[error("connection closed")]
48 ConnectionClosed,
49
50 #[error("unexpected response")]
52 UnexpectedResponse,
53
54 #[error("protocol error: {0}")]
56 Protocol(#[from] ping_proto::ParseError),
57
58 #[error("io error: {0}")]
60 Io(#[from] io::Error),
61
62 #[error("all connections failed")]
64 AllConnectionsFailed,
65}
66
67#[derive(Debug, Clone, Copy, PartialEq, Eq)]
71#[non_exhaustive]
72pub enum CommandType {
73 Ping,
74}
75
76#[derive(Debug, Clone)]
78pub struct CommandResult {
79 pub command: CommandType,
81 pub latency_ns: u64,
83 pub success: bool,
85 pub ttfb_ns: Option<u64>,
87 pub tx_bytes: u32,
89 pub rx_bytes: u32,
91}
92
93#[cfg(feature = "metrics")]
98pub struct ClientMetrics {
99 pub latency: histogram::Histogram,
101 pub requests: u64,
103 pub errors: u64,
105}
106
107#[cfg(feature = "metrics")]
108impl ClientMetrics {
109 fn new() -> Self {
110 Self {
111 latency: histogram::Histogram::new(7, 64).unwrap(),
112 requests: 0,
113 errors: 0,
114 }
115 }
116
117 fn record(&mut self, result: &CommandResult) {
118 self.requests += 1;
119 let _ = self.latency.increment(result.latency_ns);
120
121 if !result.success {
122 self.errors += 1;
123 }
124 }
125}
126
127type ResultCallback = Box<dyn Fn(&CommandResult)>;
130
131pub struct ClientBuilder {
133 conn: ConnCtx,
134 on_result: Option<ResultCallback>,
135 #[cfg(feature = "timestamps")]
136 use_kernel_ts: bool,
137 #[cfg(feature = "metrics")]
138 with_metrics: bool,
139}
140
141impl ClientBuilder {
142 pub(crate) fn new(conn: ConnCtx) -> Self {
143 Self {
144 conn,
145 on_result: None,
146 #[cfg(feature = "timestamps")]
147 use_kernel_ts: false,
148 #[cfg(feature = "metrics")]
149 with_metrics: false,
150 }
151 }
152
153 pub fn on_result<F: Fn(&CommandResult) + 'static>(mut self, f: F) -> Self {
155 self.on_result = Some(Box::new(f));
156 self
157 }
158
159 #[cfg(feature = "timestamps")]
161 pub fn kernel_timestamps(mut self, enabled: bool) -> Self {
162 self.use_kernel_ts = enabled;
163 self
164 }
165
166 #[cfg(feature = "metrics")]
168 pub fn with_metrics(mut self) -> Self {
169 self.with_metrics = true;
170 self
171 }
172
173 pub fn build(self) -> Client {
175 Client {
176 conn: self.conn,
177 on_result: self.on_result,
178 last_rx_bytes: Cell::new(0),
179 #[cfg(feature = "timestamps")]
180 use_kernel_ts: self.use_kernel_ts,
181 #[cfg(feature = "metrics")]
182 metrics: if self.with_metrics {
183 Some(ClientMetrics::new())
184 } else {
185 None
186 },
187 }
188 }
189}
190
191pub struct Client {
199 conn: ConnCtx,
200 on_result: Option<ResultCallback>,
201 last_rx_bytes: Cell<u32>,
202 #[cfg(feature = "timestamps")]
203 use_kernel_ts: bool,
204 #[cfg(feature = "metrics")]
205 metrics: Option<ClientMetrics>,
206}
207
208impl Client {
209 pub fn new(conn: ConnCtx) -> Self {
213 Self {
214 conn,
215 on_result: None,
216 last_rx_bytes: Cell::new(0),
217 #[cfg(feature = "timestamps")]
218 use_kernel_ts: false,
219 #[cfg(feature = "metrics")]
220 metrics: None,
221 }
222 }
223
224 pub fn builder(conn: ConnCtx) -> ClientBuilder {
226 ClientBuilder::new(conn)
227 }
228
229 pub fn conn(&self) -> ConnCtx {
231 self.conn
232 }
233
234 #[cfg(feature = "metrics")]
236 pub fn metrics(&self) -> Option<&ClientMetrics> {
237 self.metrics.as_ref()
238 }
239
240 #[cfg(feature = "metrics")]
242 pub fn metrics_mut(&mut self) -> Option<&mut ClientMetrics> {
243 self.metrics.as_mut()
244 }
245
246 #[inline]
249 fn is_instrumented(&self) -> bool {
250 if self.on_result.is_some() {
251 return true;
252 }
253 #[cfg(feature = "metrics")]
254 if self.metrics.is_some() {
255 return true;
256 }
257 false
258 }
259
260 #[cfg(feature = "timestamps")]
261 #[inline]
262 fn send_timestamp(&self) -> u64 {
263 if self.use_kernel_ts {
264 now_realtime_ns()
265 } else {
266 0
267 }
268 }
269
270 #[cfg(not(feature = "timestamps"))]
271 #[inline]
272 fn send_timestamp(&self) -> u64 {
273 0
274 }
275
276 #[cfg(feature = "timestamps")]
277 #[inline]
278 fn finish_timing(&self, send_ts: u64, start: Instant) -> u64 {
279 if self.use_kernel_ts {
280 let recv_ts = self.conn.recv_timestamp();
281 if recv_ts > 0 && recv_ts > send_ts {
282 return recv_ts - send_ts;
283 }
284 }
285 start.elapsed().as_nanos() as u64
286 }
287
288 #[cfg(not(feature = "timestamps"))]
289 #[inline]
290 fn finish_timing(&self, _send_ts: u64, start: Instant) -> u64 {
291 start.elapsed().as_nanos() as u64
292 }
293
294 fn record(&mut self, result: &CommandResult) {
295 if let Some(ref cb) = self.on_result {
296 cb(result);
297 }
298 #[cfg(feature = "metrics")]
299 if let Some(ref mut m) = self.metrics {
300 m.record(result);
301 }
302 }
303
304 pub(crate) async fn read_response(&self) -> Result<PingResponse, Error> {
315 let mut result: Option<Result<PingResponse, Error>> = None;
316 let n = self
317 .conn
318 .with_data(|data| match PingResponse::parse(data) {
319 Ok((response, consumed)) => {
320 result = Some(Ok(response));
321 ParseResult::Consumed(consumed)
322 }
323 Err(ping_proto::ParseError::Incomplete) => ParseResult::Consumed(0),
324 Err(e) => {
325 result = Some(Err(Error::Protocol(e)));
326 ParseResult::Consumed(data.len())
327 }
328 })
329 .await;
330 self.last_rx_bytes.set(n as u32);
331 if n == 0 {
332 return result.unwrap_or(Err(Error::ConnectionClosed));
333 }
334 let r = result.unwrap();
335 if matches!(r, Err(Error::Protocol(_))) {
336 self.conn.close();
337 }
338 r
339 }
340
341 async fn execute(&self, encoded: &[u8]) -> Result<PingResponse, Error> {
343 self.conn.send(encoded)?;
344 self.read_response().await
345 }
346
347 pub async fn ping(&mut self) -> Result<(), Error> {
351 let mut buf = [0u8; 6];
352 let len = PingRequest::Ping.encode(&mut buf);
353
354 if !self.is_instrumented() {
355 let response = self.execute(&buf[..len]).await?;
356 return match response {
357 PingResponse::Pong => Ok(()),
358 #[allow(unreachable_patterns)]
359 _ => Err(Error::UnexpectedResponse),
360 };
361 }
362
363 let tx_bytes = len as u32;
364 let send_ts = self.send_timestamp();
365 let start = Instant::now();
366 let response = self.execute(&buf[..len]).await;
367 let latency_ns = self.finish_timing(send_ts, start);
368 let rx_bytes = self.last_rx_bytes.get();
369
370 let result = match response {
371 Ok(PingResponse::Pong) => Ok(()),
372 Ok(_) => Err(Error::UnexpectedResponse),
373 Err(e) => Err(e),
374 };
375 self.record(&CommandResult {
376 command: CommandType::Ping,
377 latency_ns,
378 success: result.is_ok(),
379 ttfb_ns: None,
380 tx_bytes,
381 rx_bytes,
382 });
383 result
384 }
385}
386
387#[cfg(feature = "timestamps")]
390fn now_realtime_ns() -> u64 {
391 let mut ts = libc::timespec {
392 tv_sec: 0,
393 tv_nsec: 0,
394 };
395 unsafe {
396 libc::clock_gettime(libc::CLOCK_REALTIME, &mut ts);
397 }
398 ts.tv_sec as u64 * 1_000_000_000 + ts.tv_nsec as u64
399}