1pub mod pool;
27pub use pool::{Pool, PoolConfig};
28
29use std::io;
30use std::time::Instant;
31
32use ping_proto::{Request as PingRequest, Response as PingResponse};
33use ringline::{ConnCtx, ParseResult};
34
35#[derive(Debug, thiserror::Error)]
39pub enum Error {
40 #[error("connection closed")]
42 ConnectionClosed,
43
44 #[error("unexpected response")]
46 UnexpectedResponse,
47
48 #[error("protocol error: {0}")]
50 Protocol(#[from] ping_proto::ParseError),
51
52 #[error("io error: {0}")]
54 Io(#[from] io::Error),
55
56 #[error("all connections failed")]
58 AllConnectionsFailed,
59}
60
61#[derive(Debug, Clone, Copy, PartialEq, Eq)]
65pub enum CommandType {
66 Ping,
67}
68
69#[derive(Debug, Clone)]
71pub struct CommandResult {
72 pub command: CommandType,
74 pub latency_ns: u64,
76 pub success: bool,
78 pub ttfb_ns: Option<u64>,
80}
81
82#[cfg(feature = "metrics")]
87pub struct ClientMetrics {
88 pub latency: histogram::Histogram,
90 pub requests: u64,
92 pub errors: u64,
94}
95
96#[cfg(feature = "metrics")]
97impl ClientMetrics {
98 fn new() -> Self {
99 Self {
100 latency: histogram::Histogram::new(7, 64).unwrap(),
101 requests: 0,
102 errors: 0,
103 }
104 }
105
106 fn record(&mut self, result: &CommandResult) {
107 self.requests += 1;
108 let _ = self.latency.increment(result.latency_ns);
109
110 if !result.success {
111 self.errors += 1;
112 }
113 }
114}
115
116type ResultCallback = Box<dyn Fn(&CommandResult)>;
119
120pub struct ClientBuilder {
122 conn: ConnCtx,
123 on_result: Option<ResultCallback>,
124 #[cfg(feature = "timestamps")]
125 use_kernel_ts: bool,
126 #[cfg(feature = "metrics")]
127 with_metrics: bool,
128}
129
130impl ClientBuilder {
131 pub(crate) fn new(conn: ConnCtx) -> Self {
132 Self {
133 conn,
134 on_result: None,
135 #[cfg(feature = "timestamps")]
136 use_kernel_ts: false,
137 #[cfg(feature = "metrics")]
138 with_metrics: false,
139 }
140 }
141
142 pub fn on_result<F: Fn(&CommandResult) + 'static>(mut self, f: F) -> Self {
144 self.on_result = Some(Box::new(f));
145 self
146 }
147
148 #[cfg(feature = "timestamps")]
150 pub fn kernel_timestamps(mut self, enabled: bool) -> Self {
151 self.use_kernel_ts = enabled;
152 self
153 }
154
155 #[cfg(feature = "metrics")]
157 pub fn with_metrics(mut self) -> Self {
158 self.with_metrics = true;
159 self
160 }
161
162 pub fn build(self) -> Client {
164 Client {
165 conn: self.conn,
166 on_result: self.on_result,
167 #[cfg(feature = "timestamps")]
168 use_kernel_ts: self.use_kernel_ts,
169 #[cfg(feature = "metrics")]
170 metrics: if self.with_metrics {
171 Some(ClientMetrics::new())
172 } else {
173 None
174 },
175 }
176 }
177}
178
179pub struct Client {
187 conn: ConnCtx,
188 on_result: Option<ResultCallback>,
189 #[cfg(feature = "timestamps")]
190 use_kernel_ts: bool,
191 #[cfg(feature = "metrics")]
192 metrics: Option<ClientMetrics>,
193}
194
195impl Client {
196 pub fn new(conn: ConnCtx) -> Self {
200 Self {
201 conn,
202 on_result: None,
203 #[cfg(feature = "timestamps")]
204 use_kernel_ts: false,
205 #[cfg(feature = "metrics")]
206 metrics: None,
207 }
208 }
209
210 pub fn builder(conn: ConnCtx) -> ClientBuilder {
212 ClientBuilder::new(conn)
213 }
214
215 pub fn conn(&self) -> ConnCtx {
217 self.conn
218 }
219
220 #[cfg(feature = "metrics")]
222 pub fn metrics(&self) -> Option<&ClientMetrics> {
223 self.metrics.as_ref()
224 }
225
226 #[cfg(feature = "metrics")]
228 pub fn metrics_mut(&mut self) -> Option<&mut ClientMetrics> {
229 self.metrics.as_mut()
230 }
231
232 #[inline]
235 fn is_instrumented(&self) -> bool {
236 if self.on_result.is_some() {
237 return true;
238 }
239 #[cfg(feature = "metrics")]
240 if self.metrics.is_some() {
241 return true;
242 }
243 false
244 }
245
246 #[cfg(feature = "timestamps")]
247 #[inline]
248 fn send_timestamp(&self) -> u64 {
249 if self.use_kernel_ts {
250 now_realtime_ns()
251 } else {
252 0
253 }
254 }
255
256 #[cfg(not(feature = "timestamps"))]
257 #[inline]
258 fn send_timestamp(&self) -> u64 {
259 0
260 }
261
262 #[cfg(feature = "timestamps")]
263 #[inline]
264 fn finish_timing(&self, send_ts: u64, start: Instant) -> u64 {
265 if self.use_kernel_ts {
266 let recv_ts = self.conn.recv_timestamp();
267 if recv_ts > 0 && recv_ts > send_ts {
268 return recv_ts - send_ts;
269 }
270 }
271 start.elapsed().as_nanos() as u64
272 }
273
274 #[cfg(not(feature = "timestamps"))]
275 #[inline]
276 fn finish_timing(&self, _send_ts: u64, start: Instant) -> u64 {
277 start.elapsed().as_nanos() as u64
278 }
279
280 fn record(&mut self, result: &CommandResult) {
281 if let Some(ref cb) = self.on_result {
282 cb(result);
283 }
284 #[cfg(feature = "metrics")]
285 if let Some(ref mut m) = self.metrics {
286 m.record(result);
287 }
288 }
289
290 pub(crate) async fn read_response(&self) -> Result<PingResponse, Error> {
294 let mut result: Option<Result<PingResponse, Error>> = None;
295 let n = self
296 .conn
297 .with_data(|data| match PingResponse::parse(data) {
298 Ok((response, consumed)) => {
299 result = Some(Ok(response));
300 ParseResult::Consumed(consumed)
301 }
302 Err(ping_proto::ParseError::Incomplete) => ParseResult::Consumed(0),
303 Err(e) => {
304 result = Some(Err(Error::Protocol(e)));
305 ParseResult::Consumed(data.len())
306 }
307 })
308 .await;
309 if n == 0 {
310 return result.unwrap_or(Err(Error::ConnectionClosed));
311 }
312 result.unwrap()
313 }
314
315 async fn execute(&self, encoded: &[u8]) -> Result<PingResponse, Error> {
317 self.conn.send(encoded)?;
318 self.read_response().await
319 }
320
321 pub async fn ping(&mut self) -> Result<(), Error> {
325 let mut buf = [0u8; 6];
326 let len = PingRequest::Ping.encode(&mut buf);
327
328 if !self.is_instrumented() {
329 let response = self.execute(&buf[..len]).await?;
330 return match response {
331 PingResponse::Pong => Ok(()),
332 #[allow(unreachable_patterns)]
333 _ => Err(Error::UnexpectedResponse),
334 };
335 }
336
337 let send_ts = self.send_timestamp();
338 let start = Instant::now();
339 let response = self.execute(&buf[..len]).await;
340 let latency_ns = self.finish_timing(send_ts, start);
341
342 let result = match response {
343 Ok(PingResponse::Pong) => Ok(()),
344 Ok(_) => Err(Error::UnexpectedResponse),
345 Err(e) => Err(e),
346 };
347 self.record(&CommandResult {
348 command: CommandType::Ping,
349 latency_ns,
350 success: result.is_ok(),
351 ttfb_ns: None,
352 });
353 result
354 }
355}
356
357#[cfg(feature = "timestamps")]
360fn now_realtime_ns() -> u64 {
361 let mut ts = libc::timespec {
362 tv_sec: 0,
363 tv_nsec: 0,
364 };
365 unsafe {
366 libc::clock_gettime(libc::CLOCK_REALTIME, &mut ts);
367 }
368 ts.tv_sec as u64 * 1_000_000_000 + ts.tv_nsec as u64
369}