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)]
40pub enum Error {
41 #[error("connection closed")]
43 ConnectionClosed,
44
45 #[error("unexpected response")]
47 UnexpectedResponse,
48
49 #[error("protocol error: {0}")]
51 Protocol(#[from] ping_proto::ParseError),
52
53 #[error("io error: {0}")]
55 Io(#[from] io::Error),
56
57 #[error("all connections failed")]
59 AllConnectionsFailed,
60}
61
62#[derive(Debug, Clone, Copy, PartialEq, Eq)]
66pub enum CommandType {
67 Ping,
68}
69
70#[derive(Debug, Clone)]
72pub struct CommandResult {
73 pub command: CommandType,
75 pub latency_ns: u64,
77 pub success: bool,
79 pub ttfb_ns: Option<u64>,
81 pub tx_bytes: u32,
83 pub rx_bytes: u32,
85}
86
87#[cfg(feature = "metrics")]
92pub struct ClientMetrics {
93 pub latency: histogram::Histogram,
95 pub requests: u64,
97 pub errors: u64,
99}
100
101#[cfg(feature = "metrics")]
102impl ClientMetrics {
103 fn new() -> Self {
104 Self {
105 latency: histogram::Histogram::new(7, 64).unwrap(),
106 requests: 0,
107 errors: 0,
108 }
109 }
110
111 fn record(&mut self, result: &CommandResult) {
112 self.requests += 1;
113 let _ = self.latency.increment(result.latency_ns);
114
115 if !result.success {
116 self.errors += 1;
117 }
118 }
119}
120
121type ResultCallback = Box<dyn Fn(&CommandResult)>;
124
125pub struct ClientBuilder {
127 conn: ConnCtx,
128 on_result: Option<ResultCallback>,
129 #[cfg(feature = "timestamps")]
130 use_kernel_ts: bool,
131 #[cfg(feature = "metrics")]
132 with_metrics: bool,
133}
134
135impl ClientBuilder {
136 pub(crate) fn new(conn: ConnCtx) -> Self {
137 Self {
138 conn,
139 on_result: None,
140 #[cfg(feature = "timestamps")]
141 use_kernel_ts: false,
142 #[cfg(feature = "metrics")]
143 with_metrics: false,
144 }
145 }
146
147 pub fn on_result<F: Fn(&CommandResult) + 'static>(mut self, f: F) -> Self {
149 self.on_result = Some(Box::new(f));
150 self
151 }
152
153 #[cfg(feature = "timestamps")]
155 pub fn kernel_timestamps(mut self, enabled: bool) -> Self {
156 self.use_kernel_ts = enabled;
157 self
158 }
159
160 #[cfg(feature = "metrics")]
162 pub fn with_metrics(mut self) -> Self {
163 self.with_metrics = true;
164 self
165 }
166
167 pub fn build(self) -> Client {
169 Client {
170 conn: self.conn,
171 on_result: self.on_result,
172 last_rx_bytes: Cell::new(0),
173 #[cfg(feature = "timestamps")]
174 use_kernel_ts: self.use_kernel_ts,
175 #[cfg(feature = "metrics")]
176 metrics: if self.with_metrics {
177 Some(ClientMetrics::new())
178 } else {
179 None
180 },
181 }
182 }
183}
184
185pub struct Client {
193 conn: ConnCtx,
194 on_result: Option<ResultCallback>,
195 last_rx_bytes: Cell<u32>,
196 #[cfg(feature = "timestamps")]
197 use_kernel_ts: bool,
198 #[cfg(feature = "metrics")]
199 metrics: Option<ClientMetrics>,
200}
201
202impl Client {
203 pub fn new(conn: ConnCtx) -> Self {
207 Self {
208 conn,
209 on_result: None,
210 last_rx_bytes: Cell::new(0),
211 #[cfg(feature = "timestamps")]
212 use_kernel_ts: false,
213 #[cfg(feature = "metrics")]
214 metrics: None,
215 }
216 }
217
218 pub fn builder(conn: ConnCtx) -> ClientBuilder {
220 ClientBuilder::new(conn)
221 }
222
223 pub fn conn(&self) -> ConnCtx {
225 self.conn
226 }
227
228 #[cfg(feature = "metrics")]
230 pub fn metrics(&self) -> Option<&ClientMetrics> {
231 self.metrics.as_ref()
232 }
233
234 #[cfg(feature = "metrics")]
236 pub fn metrics_mut(&mut self) -> Option<&mut ClientMetrics> {
237 self.metrics.as_mut()
238 }
239
240 #[inline]
243 fn is_instrumented(&self) -> bool {
244 if self.on_result.is_some() {
245 return true;
246 }
247 #[cfg(feature = "metrics")]
248 if self.metrics.is_some() {
249 return true;
250 }
251 false
252 }
253
254 #[cfg(feature = "timestamps")]
255 #[inline]
256 fn send_timestamp(&self) -> u64 {
257 if self.use_kernel_ts {
258 now_realtime_ns()
259 } else {
260 0
261 }
262 }
263
264 #[cfg(not(feature = "timestamps"))]
265 #[inline]
266 fn send_timestamp(&self) -> u64 {
267 0
268 }
269
270 #[cfg(feature = "timestamps")]
271 #[inline]
272 fn finish_timing(&self, send_ts: u64, start: Instant) -> u64 {
273 if self.use_kernel_ts {
274 let recv_ts = self.conn.recv_timestamp();
275 if recv_ts > 0 && recv_ts > send_ts {
276 return recv_ts - send_ts;
277 }
278 }
279 start.elapsed().as_nanos() as u64
280 }
281
282 #[cfg(not(feature = "timestamps"))]
283 #[inline]
284 fn finish_timing(&self, _send_ts: u64, start: Instant) -> u64 {
285 start.elapsed().as_nanos() as u64
286 }
287
288 fn record(&mut self, result: &CommandResult) {
289 if let Some(ref cb) = self.on_result {
290 cb(result);
291 }
292 #[cfg(feature = "metrics")]
293 if let Some(ref mut m) = self.metrics {
294 m.record(result);
295 }
296 }
297
298 pub(crate) async fn read_response(&self) -> Result<PingResponse, Error> {
302 let mut result: Option<Result<PingResponse, Error>> = None;
303 let n = self
304 .conn
305 .with_data(|data| match PingResponse::parse(data) {
306 Ok((response, consumed)) => {
307 result = Some(Ok(response));
308 ParseResult::Consumed(consumed)
309 }
310 Err(ping_proto::ParseError::Incomplete) => ParseResult::Consumed(0),
311 Err(e) => {
312 result = Some(Err(Error::Protocol(e)));
313 ParseResult::Consumed(data.len())
314 }
315 })
316 .await;
317 self.last_rx_bytes.set(n as u32);
318 if n == 0 {
319 return result.unwrap_or(Err(Error::ConnectionClosed));
320 }
321 result.unwrap()
322 }
323
324 async fn execute(&self, encoded: &[u8]) -> Result<PingResponse, Error> {
326 self.conn.send(encoded)?;
327 self.read_response().await
328 }
329
330 pub async fn ping(&mut self) -> Result<(), Error> {
334 let mut buf = [0u8; 6];
335 let len = PingRequest::Ping.encode(&mut buf);
336
337 if !self.is_instrumented() {
338 let response = self.execute(&buf[..len]).await?;
339 return match response {
340 PingResponse::Pong => Ok(()),
341 #[allow(unreachable_patterns)]
342 _ => Err(Error::UnexpectedResponse),
343 };
344 }
345
346 let tx_bytes = len as u32;
347 let send_ts = self.send_timestamp();
348 let start = Instant::now();
349 let response = self.execute(&buf[..len]).await;
350 let latency_ns = self.finish_timing(send_ts, start);
351 let rx_bytes = self.last_rx_bytes.get();
352
353 let result = match response {
354 Ok(PingResponse::Pong) => Ok(()),
355 Ok(_) => Err(Error::UnexpectedResponse),
356 Err(e) => Err(e),
357 };
358 self.record(&CommandResult {
359 command: CommandType::Ping,
360 latency_ns,
361 success: result.is_ok(),
362 ttfb_ns: None,
363 tx_bytes,
364 rx_bytes,
365 });
366 result
367 }
368}
369
370#[cfg(feature = "timestamps")]
373fn now_realtime_ns() -> u64 {
374 let mut ts = libc::timespec {
375 tv_sec: 0,
376 tv_nsec: 0,
377 };
378 unsafe {
379 libc::clock_gettime(libc::CLOCK_REALTIME, &mut ts);
380 }
381 ts.tv_sec as u64 * 1_000_000_000 + ts.tv_nsec as u64
382}