1pub mod pool;
125pub mod sharded;
126pub use pool::{Pool, PoolConfig};
127pub use sharded::{ShardedClient, ShardedConfig};
128
129use std::cell::Cell;
130use std::collections::VecDeque;
131use std::io;
132use std::time::Instant;
133
134use bytes::Bytes;
135use memcache_proto::{Request as McRequest, ResponseBytes as McResponseBytes};
136use ringline::{ConnCtx, GuardBox, ParseResult, SendGuard};
137
138type ResultCallback = Box<dyn Fn(&CommandResult)>;
140
141const MAX_FLUSH_GUARDS: usize = 8;
143const MAX_FLUSH_IOVECS: usize = 2 * MAX_FLUSH_GUARDS + 1;
148const DEFAULT_ZC_THRESHOLD: u32 = 4096;
152
153#[derive(Debug, thiserror::Error)]
161#[non_exhaustive]
162pub enum Error {
163 #[error("connection closed")]
165 ConnectionClosed,
166
167 #[error("memcache error: {0}")]
169 Memcache(String),
170
171 #[error("unexpected response")]
173 UnexpectedResponse,
174
175 #[error("protocol error: {0}")]
177 Protocol(#[from] memcache_proto::ParseError),
178
179 #[error("io error: {0}")]
181 Io(#[from] io::Error),
182
183 #[error("all connections failed")]
185 AllConnectionsFailed,
186
187 #[error("no pending operations")]
189 NoPending,
190
191 #[error("too many in-flight operations")]
195 TooManyInFlight,
196
197 #[error("key too long (max 250 bytes)")]
202 KeyTooLong,
203}
204
205pub const MAX_KEY_LEN: usize = 250;
207
208#[inline]
209fn validate_key(key: &[u8]) -> Result<(), Error> {
210 if key.len() > MAX_KEY_LEN {
211 Err(Error::KeyTooLong)
212 } else {
213 Ok(())
214 }
215}
216
217#[derive(Debug, Clone)]
221pub struct Value {
222 pub data: Bytes,
224 pub flags: u32,
226}
227
228#[derive(Debug, Clone)]
230pub struct GetValue {
231 pub key: Bytes,
233 pub data: Bytes,
235 pub flags: u32,
237 pub cas: Option<u64>,
239}
240
241#[derive(Debug, Clone, Copy, PartialEq, Eq)]
245#[non_exhaustive]
246pub enum CommandType {
247 Get,
248 Set,
249 Delete,
250 Other,
251}
252
253#[derive(Debug, Clone)]
255pub struct CommandResult {
256 pub command: CommandType,
258 pub latency_ns: u64,
260 pub hit: Option<bool>,
262 pub success: bool,
264 pub ttfb_ns: Option<u64>,
266 pub tx_bytes: u32,
268 pub rx_bytes: u32,
270}
271
272#[cfg(feature = "metrics")]
277pub struct ClientMetrics {
278 pub latency: histogram::Histogram,
280 pub get_latency: histogram::Histogram,
282 pub set_latency: histogram::Histogram,
284 pub del_latency: histogram::Histogram,
286 pub requests: u64,
288 pub errors: u64,
290 pub hits: u64,
292 pub misses: u64,
294}
295
296#[cfg(feature = "metrics")]
297impl ClientMetrics {
298 fn new() -> Self {
299 Self {
300 latency: histogram::Histogram::new(7, 64).unwrap(),
301 get_latency: histogram::Histogram::new(7, 64).unwrap(),
302 set_latency: histogram::Histogram::new(7, 64).unwrap(),
303 del_latency: histogram::Histogram::new(7, 64).unwrap(),
304 requests: 0,
305 errors: 0,
306 hits: 0,
307 misses: 0,
308 }
309 }
310
311 fn record(&mut self, result: &CommandResult) {
312 self.requests += 1;
313 let _ = self.latency.increment(result.latency_ns);
314
315 if !result.success {
316 self.errors += 1;
317 }
318
319 match result.command {
320 CommandType::Get => {
321 let _ = self.get_latency.increment(result.latency_ns);
322 match result.hit {
323 Some(true) => self.hits += 1,
324 Some(false) => self.misses += 1,
325 None => {}
326 }
327 }
328 CommandType::Set => {
329 let _ = self.set_latency.increment(result.latency_ns);
330 }
331 CommandType::Delete => {
332 let _ = self.del_latency.increment(result.latency_ns);
333 }
334 _ => {}
335 }
336 }
337}
338
339enum PendingOpKind {
342 Get,
343 Set,
344 Delete,
345}
346
347struct PendingOp {
348 kind: PendingOpKind,
349 send_ts: u64,
350 start: Option<Instant>,
351 user_data: u64,
352 tx_bytes: u32,
353}
354
355#[non_exhaustive]
357pub enum CompletedOp {
358 Get {
360 result: Result<Option<Value>, Error>,
361 user_data: u64,
362 latency_ns: u64,
363 },
364 Set {
366 result: Result<(), Error>,
367 user_data: u64,
368 latency_ns: u64,
369 },
370 Delete {
372 result: Result<bool, Error>,
373 user_data: u64,
374 latency_ns: u64,
375 },
376}
377
378pub struct ClientBuilder {
382 conn: ConnCtx,
383 on_result: Option<ResultCallback>,
384 max_batch_size: usize,
385 max_in_flight: usize,
386 zc_threshold: u32,
387 #[cfg(feature = "timestamps")]
388 use_kernel_ts: bool,
389 #[cfg(feature = "metrics")]
390 with_metrics: bool,
391}
392
393impl ClientBuilder {
394 pub(crate) fn new(conn: ConnCtx) -> Self {
395 Self {
396 conn,
397 on_result: None,
398 max_batch_size: 1,
399 max_in_flight: usize::MAX,
400 zc_threshold: DEFAULT_ZC_THRESHOLD,
401 #[cfg(feature = "timestamps")]
402 use_kernel_ts: false,
403 #[cfg(feature = "metrics")]
404 with_metrics: false,
405 }
406 }
407
408 pub fn zc_threshold(mut self, bytes: u32) -> Self {
417 self.zc_threshold = bytes;
418 self
419 }
420
421 pub fn max_in_flight(mut self, n: usize) -> Self {
426 self.max_in_flight = n;
427 self
428 }
429
430 pub fn on_result<F: Fn(&CommandResult) + 'static>(mut self, f: F) -> Self {
432 self.on_result = Some(Box::new(f));
433 self
434 }
435
436 #[cfg(feature = "timestamps")]
438 pub fn kernel_timestamps(mut self, enabled: bool) -> Self {
439 self.use_kernel_ts = enabled;
440 self
441 }
442
443 pub fn max_batch_size(mut self, n: usize) -> Self {
452 assert!(n > 0, "max_batch_size must be >= 1");
453 self.max_batch_size = n;
454 self
455 }
456
457 #[cfg(feature = "metrics")]
459 pub fn with_metrics(mut self) -> Self {
460 self.with_metrics = true;
461 self
462 }
463
464 pub fn build(self) -> Client {
466 Client {
467 conn: self.conn,
468 on_result: self.on_result,
469 pending: VecDeque::with_capacity(16),
470 last_rx_bytes: Cell::new(0),
471 write_buf: Vec::new(),
472 write_guards: Vec::new(),
473 flushed_count: 0,
474 max_batch_size: self.max_batch_size,
475 max_in_flight: self.max_in_flight,
476 zc_threshold: self.zc_threshold,
477 buffered_ops: 0,
478 encode_buf: Vec::new(),
479 #[cfg(feature = "timestamps")]
480 use_kernel_ts: self.use_kernel_ts,
481 #[cfg(feature = "metrics")]
482 metrics: if self.with_metrics {
483 Some(ClientMetrics::new())
484 } else {
485 None
486 },
487 }
488 }
489}
490
491pub struct Client {
499 conn: ConnCtx,
500 on_result: Option<ResultCallback>,
501 pending: VecDeque<PendingOp>,
502 last_rx_bytes: Cell<u32>,
503 write_buf: Vec<u8>,
507 write_guards: Vec<(usize, GuardBox)>,
511 flushed_count: usize,
513 max_batch_size: usize,
516 max_in_flight: usize,
519 zc_threshold: u32,
523 buffered_ops: usize,
525 encode_buf: Vec<u8>,
529 #[cfg(feature = "timestamps")]
530 use_kernel_ts: bool,
531 #[cfg(feature = "metrics")]
532 metrics: Option<ClientMetrics>,
533}
534
535impl Client {
536 pub fn new(conn: ConnCtx) -> Self {
540 Self {
541 conn,
542 on_result: None,
543 pending: VecDeque::new(),
544 last_rx_bytes: Cell::new(0),
545 write_buf: Vec::new(),
546 write_guards: Vec::new(),
547 flushed_count: 0,
548 max_batch_size: 1,
549 max_in_flight: usize::MAX,
550 zc_threshold: DEFAULT_ZC_THRESHOLD,
551 buffered_ops: 0,
552 encode_buf: Vec::new(),
553 #[cfg(feature = "timestamps")]
554 use_kernel_ts: false,
555 #[cfg(feature = "metrics")]
556 metrics: None,
557 }
558 }
559
560 pub fn builder(conn: ConnCtx) -> ClientBuilder {
562 ClientBuilder::new(conn)
563 }
564
565 pub fn conn(&self) -> ConnCtx {
567 self.conn
568 }
569
570 #[cfg(feature = "metrics")]
572 pub fn metrics(&self) -> Option<&ClientMetrics> {
573 self.metrics.as_ref()
574 }
575
576 #[cfg(feature = "metrics")]
578 pub fn metrics_mut(&mut self) -> Option<&mut ClientMetrics> {
579 self.metrics.as_mut()
580 }
581
582 #[inline]
585 fn is_instrumented(&self) -> bool {
586 if self.on_result.is_some() {
587 return true;
588 }
589 #[cfg(feature = "metrics")]
590 if self.metrics.is_some() {
591 return true;
592 }
593 false
594 }
595
596 #[cfg(feature = "timestamps")]
597 #[inline]
598 fn send_timestamp(&self) -> u64 {
599 if self.use_kernel_ts {
600 now_realtime_ns()
601 } else {
602 0
603 }
604 }
605
606 #[cfg(not(feature = "timestamps"))]
607 #[inline]
608 fn send_timestamp(&self) -> u64 {
609 0
610 }
611
612 #[cfg(feature = "timestamps")]
613 #[inline]
614 fn finish_timing(&self, send_ts: u64, start: Instant) -> u64 {
615 if self.use_kernel_ts {
616 let recv_ts = self.conn.recv_timestamp();
617 if recv_ts > 0 && recv_ts > send_ts {
618 return recv_ts - send_ts;
619 }
620 }
621 start.elapsed().as_nanos() as u64
622 }
623
624 #[cfg(not(feature = "timestamps"))]
625 #[inline]
626 fn finish_timing(&self, _send_ts: u64, start: Instant) -> u64 {
627 start.elapsed().as_nanos() as u64
628 }
629
630 fn record(&mut self, result: &CommandResult) {
631 if let Some(ref cb) = self.on_result {
632 cb(result);
633 }
634 #[cfg(feature = "metrics")]
635 if let Some(ref mut m) = self.metrics {
636 m.record(result);
637 }
638 }
639
640 #[inline]
643 fn timing_start(&self) -> (u64, Option<Instant>) {
644 #[cfg(feature = "timestamps")]
645 {
646 if self.is_instrumented() {
647 (self.send_timestamp(), Some(Instant::now()))
648 } else {
649 (0, None)
650 }
651 }
652 #[cfg(not(feature = "timestamps"))]
653 {
654 if self.on_result.is_some() {
656 (0, Some(Instant::now()))
657 } else {
658 (0, None)
659 }
660 }
661 }
662
663 #[inline]
671 fn timing_start_buffered(&self) -> (u64, Option<Instant>) {
672 if self.buffered_ops >= 1 {
673 (0, None)
674 } else {
675 self.timing_start()
676 }
677 }
678
679 #[cfg(feature = "timestamps")]
680 #[inline]
681 fn compute_ttfb(&self, send_ts: u64) -> Option<u64> {
682 if self.use_kernel_ts {
683 let recv_ts = self.conn.recv_timestamp();
684 if recv_ts > 0 && recv_ts > send_ts {
685 return Some(recv_ts - send_ts);
686 }
687 }
688 None
689 }
690
691 #[cfg(not(feature = "timestamps"))]
692 #[inline]
693 fn compute_ttfb(&self, _send_ts: u64) -> Option<u64> {
694 None
695 }
696
697 pub fn pending_count(&self) -> usize {
699 self.pending.len()
700 }
701
702 fn pre_flush_if_needed(&mut self, has_guard: bool) -> Result<(), Error> {
707 if self.pending.len() >= self.max_in_flight {
708 return Err(Error::TooManyInFlight);
709 }
710 if self.buffered_ops == 0 {
711 return Ok(());
712 }
713 let next_guards = self.write_guards.len() + usize::from(has_guard);
714 let next_parts = 2 * next_guards + 1;
715 if next_guards > MAX_FLUSH_GUARDS || next_parts > MAX_FLUSH_IOVECS {
716 self.flush()?;
717 }
718 Ok(())
719 }
720
721 fn post_flush_if_needed(&mut self) -> Result<(), Error> {
723 if self.buffered_ops >= self.max_batch_size {
724 self.flush()?;
725 }
726 Ok(())
727 }
728
729 pub fn flush(&mut self) -> Result<(), Error> {
735 if self.write_buf.is_empty() && self.write_guards.is_empty() {
736 self.buffered_ops = 0;
737 return Ok(());
738 }
739
740 let send_outcome: Result<(), Error> = if self.write_guards.is_empty() {
747 self.conn
748 .send_nowait(&self.write_buf)
749 .map(|_| ())
750 .map_err(Error::from)
751 } else {
752 use ringline::SendPart;
753 let mut parts: Vec<SendPart<'_>> = Vec::with_capacity(2 * MAX_FLUSH_GUARDS + 1);
754 let mut pos = 0;
755 for (offset, guard) in self.write_guards.drain(..) {
756 if offset > pos {
757 parts.push(SendPart::Copy(&self.write_buf[pos..offset]));
758 }
759 parts.push(SendPart::Guard(guard));
760 pos = offset;
761 }
762 if pos < self.write_buf.len() {
763 parts.push(SendPart::Copy(&self.write_buf[pos..]));
764 }
765 self.conn
766 .send_parts()
767 .submit_batch(parts)
768 .map(|_| ())
769 .map_err(Error::from)
770 };
771
772 if let Err(e) = send_outcome {
773 self.write_buf.clear();
779 self.write_guards.clear();
780 self.buffered_ops = 0;
781 self.pending.truncate(self.flushed_count);
782 return Err(e);
783 }
784
785 if self.buffered_ops >= 1 {
789 let (send_ts, start) = self.timing_start();
790 for pending in self.pending.iter_mut().skip(self.flushed_count) {
791 pending.send_ts = send_ts;
792 pending.start = start;
793 }
794 }
795 self.flushed_count = self.pending.len();
796
797 self.write_buf.clear();
798 self.write_guards.clear();
799 self.buffered_ops = 0;
800
801 Ok(())
802 }
803
804 pub fn fire_get(&mut self, key: &[u8], user_data: u64) -> Result<(), Error> {
806 self.pre_flush_if_needed(false)?;
807 let tx_bytes;
808 if self.max_batch_size == 1 && self.write_guards.is_empty() && self.write_buf.is_empty() {
809 self.encode_buf.clear();
811 encode_request_into(&McRequest::get(key), &mut self.encode_buf)?;
812 tx_bytes = self.encode_buf.len() as u32;
813 self.conn.send_nowait(&self.encode_buf)?;
814 } else {
815 let before = self.write_buf.len();
816 encode_request_into(&McRequest::get(key), &mut self.write_buf)?;
817 tx_bytes = (self.write_buf.len() - before) as u32;
818 self.buffered_ops += 1;
819 }
820 let (send_ts, start) = self.timing_start_buffered();
821 self.pending.push_back(PendingOp {
822 kind: PendingOpKind::Get,
823 send_ts,
824 start,
825 user_data,
826 tx_bytes,
827 });
828 if self.max_batch_size == 1 && self.write_guards.is_empty() && self.write_buf.is_empty() {
834 self.flushed_count += 1;
835 }
836 self.post_flush_if_needed()?;
837 Ok(())
838 }
839
840 pub fn fire_set(
842 &mut self,
843 key: &[u8],
844 value: &[u8],
845 flags: u32,
846 exptime: u32,
847 user_data: u64,
848 ) -> Result<(), Error> {
849 self.pre_flush_if_needed(false)?;
850 let req = McRequest::Set {
851 key,
852 value,
853 flags,
854 exptime,
855 };
856 let tx_bytes;
857 if self.max_batch_size == 1 && self.write_guards.is_empty() && self.write_buf.is_empty() {
858 self.encode_buf.clear();
860 encode_request_into(&req, &mut self.encode_buf)?;
861 tx_bytes = self.encode_buf.len() as u32;
862 self.conn.send_nowait(&self.encode_buf)?;
863 } else {
864 let before = self.write_buf.len();
865 encode_request_into(&req, &mut self.write_buf)?;
866 tx_bytes = (self.write_buf.len() - before) as u32;
867 self.buffered_ops += 1;
868 }
869 let (send_ts, start) = self.timing_start_buffered();
870 self.pending.push_back(PendingOp {
871 kind: PendingOpKind::Set,
872 send_ts,
873 start,
874 user_data,
875 tx_bytes,
876 });
877 if self.max_batch_size == 1 && self.write_guards.is_empty() && self.write_buf.is_empty() {
883 self.flushed_count += 1;
884 }
885 self.post_flush_if_needed()?;
886 Ok(())
887 }
888
889 pub fn fire_set_with_guard<G: SendGuard>(
894 &mut self,
895 key: &[u8],
896 guard: G,
897 flags: u32,
898 exptime: u32,
899 user_data: u64,
900 ) -> Result<(), Error> {
901 let (ptr, value_len) = guard.as_ptr_len();
902 if self.zc_threshold != 0 && value_len < self.zc_threshold {
906 let value = unsafe { core::slice::from_raw_parts(ptr, value_len as usize) };
912 return self.fire_set(key, value, flags, exptime, user_data);
913 }
914 self.pre_flush_if_needed(true)?;
915 let before = self.write_buf.len();
918 append_set_guard_prefix(&mut self.write_buf, key, value_len as usize, flags, exptime)?;
919 self.write_guards
920 .push((self.write_buf.len(), GuardBox::new(guard)));
921 self.write_buf.extend_from_slice(b"\r\n");
922 let tx_bytes = (self.write_buf.len() - before + value_len as usize) as u32;
923 self.buffered_ops += 1;
924 let (send_ts, start) = self.timing_start_buffered();
925 self.pending.push_back(PendingOp {
926 kind: PendingOpKind::Set,
927 send_ts,
928 start,
929 user_data,
930 tx_bytes,
931 });
932 self.post_flush_if_needed()?;
933 Ok(())
934 }
935
936 pub fn fire_delete(&mut self, key: &[u8], user_data: u64) -> Result<(), Error> {
938 self.pre_flush_if_needed(false)?;
939 let tx_bytes;
940 if self.max_batch_size == 1 && self.write_guards.is_empty() && self.write_buf.is_empty() {
941 self.encode_buf.clear();
943 encode_request_into(&McRequest::delete(key), &mut self.encode_buf)?;
944 tx_bytes = self.encode_buf.len() as u32;
945 self.conn.send_nowait(&self.encode_buf)?;
946 } else {
947 let before = self.write_buf.len();
948 encode_request_into(&McRequest::delete(key), &mut self.write_buf)?;
949 tx_bytes = (self.write_buf.len() - before) as u32;
950 self.buffered_ops += 1;
951 }
952 let (send_ts, start) = self.timing_start_buffered();
953 self.pending.push_back(PendingOp {
954 kind: PendingOpKind::Delete,
955 send_ts,
956 start,
957 user_data,
958 tx_bytes,
959 });
960 if self.max_batch_size == 1 && self.write_guards.is_empty() && self.write_buf.is_empty() {
966 self.flushed_count += 1;
967 }
968 self.post_flush_if_needed()?;
969 Ok(())
970 }
971
972 pub async fn recv(&mut self) -> Result<CompletedOp, Error> {
976 self.flush()?;
978
979 let pending = self.pending.pop_front().ok_or(Error::NoPending)?;
980 self.flushed_count = self.flushed_count.saturating_sub(1);
981
982 let response = match self.read_response().await {
983 Ok(v) => v,
984 Err(e) => {
985 self.pending.clear();
991 self.flushed_count = 0;
992 return Err(e);
993 }
994 };
995 let ttfb_ns = self.compute_ttfb(pending.send_ts);
999 let latency_ns = match pending.start {
1000 Some(start) => self.finish_timing(pending.send_ts, start),
1001 None => 0,
1002 };
1003 let rx_bytes = self.last_rx_bytes.get();
1004 let tx_bytes = pending.tx_bytes;
1005
1006 let op = match pending.kind {
1007 PendingOpKind::Get => {
1008 let result = match check_error_bytes(&response) {
1010 Err(e) => Err(e),
1011 Ok(()) => match response {
1012 McResponseBytes::Values(mut values) => {
1013 if values.is_empty() {
1014 Ok(None)
1015 } else {
1016 let v = values.swap_remove(0);
1017 Ok(Some(Value {
1018 data: v.data,
1019 flags: v.flags,
1020 }))
1021 }
1022 }
1023 _ => Err(Error::UnexpectedResponse),
1024 },
1025 };
1026 let (success, hit) = match &result {
1027 Ok(Some(_)) => (true, Some(true)),
1028 Ok(None) => (true, Some(false)),
1029 Err(_) => (false, None),
1030 };
1031 self.record(&CommandResult {
1032 command: CommandType::Get,
1033 latency_ns,
1034 hit,
1035 success,
1036 ttfb_ns,
1037 tx_bytes,
1038 rx_bytes,
1039 });
1040 CompletedOp::Get {
1041 result,
1042 user_data: pending.user_data,
1043 latency_ns,
1044 }
1045 }
1046 PendingOpKind::Set => {
1047 let result = match check_error_bytes(&response) {
1048 Err(e) => Err(e),
1049 Ok(()) => match response {
1050 McResponseBytes::Stored => Ok(()),
1051 _ => Err(Error::UnexpectedResponse),
1052 },
1053 };
1054 self.record(&CommandResult {
1055 command: CommandType::Set,
1056 latency_ns,
1057 hit: None,
1058 success: result.is_ok(),
1059 ttfb_ns,
1060 tx_bytes,
1061 rx_bytes,
1062 });
1063 CompletedOp::Set {
1064 result,
1065 user_data: pending.user_data,
1066 latency_ns,
1067 }
1068 }
1069 PendingOpKind::Delete => {
1070 let result = match check_error_bytes(&response) {
1071 Err(e) => Err(e),
1072 Ok(()) => match response {
1073 McResponseBytes::Deleted => Ok(true),
1074 McResponseBytes::NotFound => Ok(false),
1075 _ => Err(Error::UnexpectedResponse),
1076 },
1077 };
1078 self.record(&CommandResult {
1079 command: CommandType::Delete,
1080 latency_ns,
1081 hit: None,
1082 success: result.is_ok(),
1083 ttfb_ns,
1084 tx_bytes,
1085 rx_bytes,
1086 });
1087 CompletedOp::Delete {
1088 result,
1089 user_data: pending.user_data,
1090 latency_ns,
1091 }
1092 }
1093 };
1094
1095 Ok(op)
1096 }
1097
1098 pub(crate) async fn read_response(&self) -> Result<McResponseBytes, Error> {
1113 let mut result: Option<Result<McResponseBytes, Error>> = None;
1114 let n = self
1115 .conn
1116 .with_bytes(|bytes| {
1117 let len = bytes.len();
1118 match McResponseBytes::parse(bytes) {
1119 Ok((response, consumed)) => {
1120 result = Some(Ok(response));
1121 ParseResult::Consumed(consumed)
1122 }
1123 Err(e) if e.is_incomplete() => ParseResult::Consumed(0),
1124 Err(e) => {
1125 result = Some(Err(Error::Protocol(e)));
1126 ParseResult::Consumed(len)
1127 }
1128 }
1129 })
1130 .await;
1131 self.last_rx_bytes.set(n as u32);
1132 if n == 0 {
1133 return result.unwrap_or(Err(Error::ConnectionClosed));
1134 }
1135 let r = result.unwrap();
1136 if matches!(r, Err(Error::Protocol(_))) {
1137 self.conn.close();
1138 }
1139 r
1140 }
1141
1142 async fn execute(&self, encoded: &[u8]) -> Result<McResponseBytes, Error> {
1145 self.conn.send(encoded)?;
1146 let response = self.read_response().await?;
1147 check_error_bytes(&response)?;
1148 Ok(response)
1149 }
1150
1151 pub async fn get(&mut self, key: impl AsRef<[u8]>) -> Result<Option<Value>, Error> {
1155 let key = key.as_ref();
1156 self.encode_buf.clear();
1157 encode_request_into(&McRequest::get(key), &mut self.encode_buf)?;
1158
1159 if !self.is_instrumented() {
1160 let response = self.execute(&self.encode_buf).await?;
1161 return match response {
1162 McResponseBytes::Values(mut values) => {
1163 if values.is_empty() {
1164 Ok(None)
1165 } else {
1166 let v = values.swap_remove(0);
1167 Ok(Some(Value {
1168 data: v.data,
1169 flags: v.flags,
1170 }))
1171 }
1172 }
1173 _ => Err(Error::UnexpectedResponse),
1174 };
1175 }
1176
1177 let tx_bytes = self.encode_buf.len() as u32;
1178 let send_ts = self.send_timestamp();
1179 let start = Instant::now();
1180 let response = self.execute(&self.encode_buf).await;
1181 let latency_ns = self.finish_timing(send_ts, start);
1182 let rx_bytes = self.last_rx_bytes.get();
1183
1184 let result = match response {
1185 Ok(McResponseBytes::Values(mut values)) => {
1186 if values.is_empty() {
1187 Ok(None)
1188 } else {
1189 let v = values.swap_remove(0);
1190 Ok(Some(Value {
1191 data: v.data,
1192 flags: v.flags,
1193 }))
1194 }
1195 }
1196 Ok(_) => Err(Error::UnexpectedResponse),
1197 Err(e) => Err(e),
1198 };
1199
1200 let (success, hit) = match &result {
1201 Ok(Some(_)) => (true, Some(true)),
1202 Ok(None) => (true, Some(false)),
1203 Err(_) => (false, None),
1204 };
1205 self.record(&CommandResult {
1206 command: CommandType::Get,
1207 latency_ns,
1208 hit,
1209 success,
1210 ttfb_ns: None,
1211 tx_bytes,
1212 rx_bytes,
1213 });
1214 result
1215 }
1216
1217 pub async fn gets(&mut self, keys: &[&[u8]]) -> Result<Vec<GetValue>, Error> {
1219 if keys.is_empty() {
1220 return Ok(Vec::new());
1221 }
1222 let encoded = encode_request(&McRequest::gets(keys))?;
1223 let response = self.execute(&encoded).await?;
1224 match response {
1225 McResponseBytes::Values(values) => Ok(values
1226 .into_iter()
1227 .map(|v| GetValue {
1228 key: v.key,
1229 data: v.data,
1230 flags: v.flags,
1231 cas: v.cas,
1232 })
1233 .collect()),
1234 _ => Err(Error::UnexpectedResponse),
1235 }
1236 }
1237
1238 pub async fn set(
1240 &mut self,
1241 key: impl AsRef<[u8]>,
1242 value: impl AsRef<[u8]>,
1243 ) -> Result<(), Error> {
1244 self.set_with_options(key, value, 0, 0).await
1245 }
1246
1247 pub async fn set_with_options(
1249 &mut self,
1250 key: impl AsRef<[u8]>,
1251 value: impl AsRef<[u8]>,
1252 flags: u32,
1253 exptime: u32,
1254 ) -> Result<(), Error> {
1255 let key = key.as_ref();
1256 let value = value.as_ref();
1257 self.encode_buf.clear();
1258 encode_request_into(
1259 &McRequest::Set {
1260 key,
1261 value,
1262 flags,
1263 exptime,
1264 },
1265 &mut self.encode_buf,
1266 )?;
1267
1268 if !self.is_instrumented() {
1269 let response = self.execute(&self.encode_buf).await?;
1270 return match response {
1271 McResponseBytes::Stored => Ok(()),
1272 _ => Err(Error::UnexpectedResponse),
1273 };
1274 }
1275
1276 let tx_bytes = self.encode_buf.len() as u32;
1277 let send_ts = self.send_timestamp();
1278 let start = Instant::now();
1279 let response = self.execute(&self.encode_buf).await;
1280 let latency_ns = self.finish_timing(send_ts, start);
1281 let rx_bytes = self.last_rx_bytes.get();
1282
1283 let result = match response {
1284 Ok(McResponseBytes::Stored) => Ok(()),
1285 Ok(_) => Err(Error::UnexpectedResponse),
1286 Err(e) => Err(e),
1287 };
1288 self.record(&CommandResult {
1289 command: CommandType::Set,
1290 latency_ns,
1291 hit: None,
1292 success: result.is_ok(),
1293 ttfb_ns: None,
1294 tx_bytes,
1295 rx_bytes,
1296 });
1297 result
1298 }
1299
1300 pub async fn add(
1303 &mut self,
1304 key: impl AsRef<[u8]>,
1305 value: impl AsRef<[u8]>,
1306 ) -> Result<bool, Error> {
1307 let key = key.as_ref();
1308 let value = value.as_ref();
1309 let encoded = encode_add(key, value)?;
1310 let response = self.execute(&encoded).await?;
1311 match response {
1312 McResponseBytes::Stored => Ok(true),
1313 McResponseBytes::NotStored => Ok(false),
1314 _ => Err(Error::UnexpectedResponse),
1315 }
1316 }
1317
1318 pub async fn replace(
1321 &mut self,
1322 key: impl AsRef<[u8]>,
1323 value: impl AsRef<[u8]>,
1324 ) -> Result<bool, Error> {
1325 let key = key.as_ref();
1326 let value = value.as_ref();
1327 let encoded = encode_request(&McRequest::Replace {
1328 key,
1329 value,
1330 flags: 0,
1331 exptime: 0,
1332 })?;
1333 let response = self.execute(&encoded).await?;
1334 match response {
1335 McResponseBytes::Stored => Ok(true),
1336 McResponseBytes::NotStored => Ok(false),
1337 _ => Err(Error::UnexpectedResponse),
1338 }
1339 }
1340
1341 pub async fn incr(&mut self, key: impl AsRef<[u8]>, delta: u64) -> Result<Option<u64>, Error> {
1344 let key = key.as_ref();
1345 let encoded = encode_request(&McRequest::incr(key, delta))?;
1346 let response = self.execute(&encoded).await?;
1347 match response {
1348 McResponseBytes::Numeric(val) => Ok(Some(val)),
1349 McResponseBytes::NotFound => Ok(None),
1350 _ => Err(Error::UnexpectedResponse),
1351 }
1352 }
1353
1354 pub async fn decr(&mut self, key: impl AsRef<[u8]>, delta: u64) -> Result<Option<u64>, Error> {
1357 let key = key.as_ref();
1358 let encoded = encode_request(&McRequest::decr(key, delta))?;
1359 let response = self.execute(&encoded).await?;
1360 match response {
1361 McResponseBytes::Numeric(val) => Ok(Some(val)),
1362 McResponseBytes::NotFound => Ok(None),
1363 _ => Err(Error::UnexpectedResponse),
1364 }
1365 }
1366
1367 pub async fn append(
1370 &mut self,
1371 key: impl AsRef<[u8]>,
1372 value: impl AsRef<[u8]>,
1373 ) -> Result<bool, Error> {
1374 let key = key.as_ref();
1375 let value = value.as_ref();
1376 let encoded = encode_request(&McRequest::append(key, value))?;
1377 let response = self.execute(&encoded).await?;
1378 match response {
1379 McResponseBytes::Stored => Ok(true),
1380 McResponseBytes::NotStored => Ok(false),
1381 _ => Err(Error::UnexpectedResponse),
1382 }
1383 }
1384
1385 pub async fn prepend(
1388 &mut self,
1389 key: impl AsRef<[u8]>,
1390 value: impl AsRef<[u8]>,
1391 ) -> Result<bool, Error> {
1392 let key = key.as_ref();
1393 let value = value.as_ref();
1394 let encoded = encode_request(&McRequest::prepend(key, value))?;
1395 let response = self.execute(&encoded).await?;
1396 match response {
1397 McResponseBytes::Stored => Ok(true),
1398 McResponseBytes::NotStored => Ok(false),
1399 _ => Err(Error::UnexpectedResponse),
1400 }
1401 }
1402
1403 pub async fn cas(
1407 &mut self,
1408 key: impl AsRef<[u8]>,
1409 value: impl AsRef<[u8]>,
1410 cas_unique: u64,
1411 ) -> Result<bool, Error> {
1412 let key = key.as_ref();
1413 let value = value.as_ref();
1414 let encoded = encode_request(&McRequest::cas(key, value, cas_unique))?;
1415 let response = self.execute(&encoded).await?;
1416 match response {
1417 McResponseBytes::Stored => Ok(true),
1418 McResponseBytes::Exists => Ok(false),
1419 McResponseBytes::NotFound => Err(Error::Memcache("NOT_FOUND".into())),
1420 _ => Err(Error::UnexpectedResponse),
1421 }
1422 }
1423
1424 pub async fn delete(&mut self, key: impl AsRef<[u8]>) -> Result<bool, Error> {
1426 let key = key.as_ref();
1427 self.encode_buf.clear();
1428 encode_request_into(&McRequest::delete(key), &mut self.encode_buf)?;
1429
1430 if !self.is_instrumented() {
1431 let response = self.execute(&self.encode_buf).await?;
1432 return match response {
1433 McResponseBytes::Deleted => Ok(true),
1434 McResponseBytes::NotFound => Ok(false),
1435 _ => Err(Error::UnexpectedResponse),
1436 };
1437 }
1438
1439 let tx_bytes = self.encode_buf.len() as u32;
1440 let send_ts = self.send_timestamp();
1441 let start = Instant::now();
1442 let response = self.execute(&self.encode_buf).await;
1443 let latency_ns = self.finish_timing(send_ts, start);
1444 let rx_bytes = self.last_rx_bytes.get();
1445
1446 let result = match response {
1447 Ok(McResponseBytes::Deleted) => Ok(true),
1448 Ok(McResponseBytes::NotFound) => Ok(false),
1449 Ok(_) => Err(Error::UnexpectedResponse),
1450 Err(e) => Err(e),
1451 };
1452 self.record(&CommandResult {
1453 command: CommandType::Delete,
1454 latency_ns,
1455 hit: None,
1456 success: result.is_ok(),
1457 ttfb_ns: None,
1458 tx_bytes,
1459 rx_bytes,
1460 });
1461 result
1462 }
1463
1464 pub async fn flush_all(&mut self) -> Result<(), Error> {
1466 let encoded = encode_request(&McRequest::flush_all())?;
1467 let response = self.execute(&encoded).await?;
1468 match response {
1469 McResponseBytes::Ok => Ok(()),
1470 _ => Err(Error::UnexpectedResponse),
1471 }
1472 }
1473
1474 pub async fn version(&mut self) -> Result<Box<str>, Error> {
1476 let encoded = encode_request(&McRequest::version())?;
1477 let response = self.execute(&encoded).await?;
1478 match response {
1479 McResponseBytes::Version(v) => Ok(Box::from(String::from_utf8_lossy(v.as_ref()))),
1480 _ => Err(Error::UnexpectedResponse),
1481 }
1482 }
1483
1484 pub async fn set_with_guard<G: SendGuard>(
1489 &mut self,
1490 key: &[u8],
1491 guard: G,
1492 flags: u32,
1493 exptime: u32,
1494 ) -> Result<(), Error> {
1495 if !self.is_instrumented() {
1496 let (_, value_len) = guard.as_ptr_len();
1497 self.encode_buf.clear();
1498 append_set_guard_prefix(
1499 &mut self.encode_buf,
1500 key,
1501 value_len as usize,
1502 flags,
1503 exptime,
1504 )?;
1505
1506 let prefix: &[u8] = &self.encode_buf;
1507 self.conn.send_parts().build(move |b| {
1508 b.copy(prefix)
1509 .guard(GuardBox::new(guard))
1510 .copy(b"\r\n")
1511 .submit()
1512 })?;
1513
1514 let response = self.read_response().await?;
1515 check_error_bytes(&response)?;
1516 return match response {
1517 McResponseBytes::Stored => Ok(()),
1518 _ => Err(Error::UnexpectedResponse),
1519 };
1520 }
1521
1522 let (_, value_len) = guard.as_ptr_len();
1523 self.encode_buf.clear();
1524 append_set_guard_prefix(
1525 &mut self.encode_buf,
1526 key,
1527 value_len as usize,
1528 flags,
1529 exptime,
1530 )?;
1531 let tx_bytes = (self.encode_buf.len() + value_len as usize + 2) as u32;
1532
1533 let send_ts = self.send_timestamp();
1534 let start = Instant::now();
1535
1536 let prefix: &[u8] = &self.encode_buf;
1537 self.conn.send_parts().build(move |b| {
1538 b.copy(prefix)
1539 .guard(GuardBox::new(guard))
1540 .copy(b"\r\n")
1541 .submit()
1542 })?;
1543
1544 let response = self.read_response().await;
1545 let latency_ns = self.finish_timing(send_ts, start);
1546 let rx_bytes = self.last_rx_bytes.get();
1547
1548 let result = match response {
1549 Ok(ref r) => {
1550 check_error_bytes(r)?;
1551 match r {
1552 McResponseBytes::Stored => Ok(()),
1553 _ => Err(Error::UnexpectedResponse),
1554 }
1555 }
1556 Err(e) => Err(e),
1557 };
1558 self.record(&CommandResult {
1559 command: CommandType::Set,
1560 latency_ns,
1561 hit: None,
1562 success: result.is_ok(),
1563 ttfb_ns: None,
1564 tx_bytes,
1565 rx_bytes,
1566 });
1567 result
1568 }
1569}
1570
1571fn append_set_guard_prefix(
1582 buf: &mut Vec<u8>,
1583 key: &[u8],
1584 value_len: usize,
1585 flags: u32,
1586 exptime: u32,
1587) -> Result<(), Error> {
1588 validate_key(key)?;
1589 let mut itoa_buf = itoa::Buffer::new();
1590 buf.extend_from_slice(b"set ");
1591 buf.extend_from_slice(key);
1592 buf.push(b' ');
1593 buf.extend_from_slice(itoa_buf.format(flags).as_bytes());
1594 buf.push(b' ');
1595 buf.extend_from_slice(itoa_buf.format(exptime).as_bytes());
1596 buf.push(b' ');
1597 buf.extend_from_slice(itoa_buf.format(value_len).as_bytes());
1598 buf.extend_from_slice(b"\r\n");
1599 Ok(())
1600}
1601
1602fn validate_request(req: &McRequest<'_>) -> Result<(), Error> {
1609 match req {
1610 McRequest::Get { key }
1611 | McRequest::Incr { key, .. }
1612 | McRequest::Decr { key, .. }
1613 | McRequest::Delete { key } => validate_key(key),
1614 McRequest::Gets { keys } => {
1615 for k in keys.iter() {
1616 validate_key(k)?;
1617 }
1618 Ok(())
1619 }
1620 McRequest::Set { key, .. }
1621 | McRequest::Add { key, .. }
1622 | McRequest::Replace { key, .. }
1623 | McRequest::Append { key, .. }
1624 | McRequest::Prepend { key, .. }
1625 | McRequest::Cas { key, .. } => validate_key(key),
1626 McRequest::FlushAll | McRequest::Version | McRequest::Quit => Ok(()),
1627 }
1628}
1629
1630pub(crate) fn encode_request_into(req: &McRequest<'_>, buf: &mut Vec<u8>) -> Result<(), Error> {
1634 validate_request(req)?;
1635 let size = match req {
1636 McRequest::Get { key } => 6 + key.len(),
1637 McRequest::Gets { keys } => 6 + keys.iter().map(|k| 1 + k.len()).sum::<usize>(),
1638 McRequest::Set { key, value, .. } | McRequest::Add { key, value, .. } => {
1639 41 + key.len() + value.len()
1640 }
1641 McRequest::Replace { key, value, .. } => 45 + key.len() + value.len(),
1642 McRequest::Incr { key, .. } | McRequest::Decr { key, .. } => 27 + key.len(),
1643 McRequest::Append { key, value } => 44 + key.len() + value.len(),
1644 McRequest::Prepend { key, value } => 45 + key.len() + value.len(),
1645 McRequest::Cas { key, value, .. } => 61 + key.len() + value.len(),
1646 McRequest::Delete { key } => 9 + key.len(),
1647 McRequest::FlushAll => 11,
1648 McRequest::Version => 9,
1649 McRequest::Quit => 6,
1650 };
1651 let start = buf.len();
1652 buf.resize(start + size, 0);
1653 let len = req.encode(&mut buf[start..]);
1654 buf.truncate(start + len);
1655 Ok(())
1656}
1657
1658pub(crate) fn encode_request(req: &McRequest<'_>) -> Result<Vec<u8>, Error> {
1662 let mut buf = Vec::new();
1663 encode_request_into(req, &mut buf)?;
1664 Ok(buf)
1665}
1666
1667#[cfg(test)]
1670pub(crate) fn encode_set(
1671 key: &[u8],
1672 value: &[u8],
1673 flags: u32,
1674 exptime: u32,
1675) -> Result<Vec<u8>, Error> {
1676 encode_request(&McRequest::Set {
1677 key,
1678 value,
1679 flags,
1680 exptime,
1681 })
1682}
1683
1684pub(crate) fn encode_add(key: &[u8], value: &[u8]) -> Result<Vec<u8>, Error> {
1686 encode_request(&McRequest::Add {
1687 key,
1688 value,
1689 flags: 0,
1690 exptime: 0,
1691 })
1692}
1693
1694pub(crate) fn check_error_bytes(response: &McResponseBytes) -> Result<(), Error> {
1696 match response {
1697 McResponseBytes::Error => Err(Error::Memcache("ERROR".into())),
1698 McResponseBytes::ClientError(msg) => Err(Error::Memcache(format!(
1699 "CLIENT_ERROR {}",
1700 String::from_utf8_lossy(msg)
1701 ))),
1702 McResponseBytes::ServerError(msg) => Err(Error::Memcache(format!(
1703 "SERVER_ERROR {}",
1704 String::from_utf8_lossy(msg)
1705 ))),
1706 _ => Ok(()),
1707 }
1708}
1709
1710#[cfg(feature = "timestamps")]
1713fn now_realtime_ns() -> u64 {
1714 let mut ts = libc::timespec {
1715 tv_sec: 0,
1716 tv_nsec: 0,
1717 };
1718 unsafe {
1719 libc::clock_gettime(libc::CLOCK_REALTIME, &mut ts);
1720 }
1721 ts.tv_sec as u64 * 1_000_000_000 + ts.tv_nsec as u64
1722}
1723
1724#[cfg(test)]
1725mod encode_tests {
1726 use super::*;
1727
1728 const KEY: &[u8] = b"user:123456789";
1731
1732 #[test]
1733 fn golden_encode_request_get() {
1734 let encoded = encode_request(&McRequest::get(KEY)).unwrap();
1735 assert_eq!(encoded, b"get user:123456789\r\n");
1736 }
1737
1738 #[test]
1739 fn golden_encode_request_delete() {
1740 let encoded = encode_request(&McRequest::delete(KEY)).unwrap();
1741 assert_eq!(encoded, b"delete user:123456789\r\n");
1742 }
1743
1744 #[test]
1745 fn golden_encode_request_set() {
1746 let encoded = encode_set(KEY, b"hello world value", 42, 7200).unwrap();
1747 assert_eq!(
1748 encoded,
1749 &b"set user:123456789 42 7200 17\r\nhello world value\r\n"[..]
1750 );
1751 }
1752
1753 #[test]
1754 fn golden_encode_request_set_zero_fields() {
1755 let encoded = encode_set(KEY, b"v", 0, 0).unwrap();
1756 assert_eq!(encoded, &b"set user:123456789 0 0 1\r\nv\r\n"[..]);
1757 }
1758
1759 #[test]
1760 fn golden_set_guard_prefix() {
1761 let mut prefix = Vec::new();
1764 append_set_guard_prefix(&mut prefix, KEY, 1024, 42, 7200).unwrap();
1765 assert_eq!(prefix, &b"set user:123456789 42 7200 1024\r\n"[..]);
1766 }
1767
1768 #[test]
1769 fn encode_into_appends_without_clobbering() {
1770 let mut buf = b"EXISTING".to_vec();
1773 encode_request_into(&McRequest::get(KEY), &mut buf).unwrap();
1774 assert_eq!(buf, &b"EXISTINGget user:123456789\r\n"[..]);
1775 }
1776
1777 #[test]
1778 fn encode_into_error_leaves_buf_untouched() {
1779 let mut buf = b"EXISTING".to_vec();
1780 let long_key = vec![b'k'; MAX_KEY_LEN + 1];
1781 assert!(encode_request_into(&McRequest::get(&long_key), &mut buf).is_err());
1782 assert_eq!(buf, b"EXISTING");
1783 assert!(append_set_guard_prefix(&mut buf, &long_key, 16, 0, 0).is_err());
1784 assert_eq!(buf, b"EXISTING");
1785 }
1786
1787 #[test]
1797 fn coalesced_buf_equals_concatenation_of_individual_encodings() {
1798 let mut coalesced = Vec::new();
1801 encode_request_into(&McRequest::get(b"key1"), &mut coalesced).unwrap();
1802 encode_request_into(
1803 &McRequest::Set {
1804 key: b"key2",
1805 value: b"val",
1806 flags: 0,
1807 exptime: 0,
1808 },
1809 &mut coalesced,
1810 )
1811 .unwrap();
1812 encode_request_into(&McRequest::delete(b"key3"), &mut coalesced).unwrap();
1813
1814 let mut expected = Vec::new();
1817 expected.extend_from_slice(&encode_request(&McRequest::get(b"key1")).unwrap());
1818 expected.extend_from_slice(
1819 &encode_request(&McRequest::Set {
1820 key: b"key2",
1821 value: b"val",
1822 flags: 0,
1823 exptime: 0,
1824 })
1825 .unwrap(),
1826 );
1827 expected.extend_from_slice(&encode_request(&McRequest::delete(b"key3")).unwrap());
1828
1829 assert_eq!(coalesced, expected);
1830 assert_eq!(
1832 coalesced,
1833 &b"get key1\r\nset key2 0 0 3\r\nval\r\ndelete key3\r\n"[..]
1834 );
1835 }
1836
1837 #[test]
1838 fn coalesced_guard_set_byte_layout() {
1839 let mut write_buf = Vec::new();
1843 let value: &[u8] = b"hello";
1844 append_set_guard_prefix(&mut write_buf, KEY, value.len(), 42, 7200).unwrap();
1845 let guard_offset = write_buf.len();
1846 write_buf.extend_from_slice(b"\r\n");
1847
1848 let mut wire = Vec::new();
1851 wire.extend_from_slice(&write_buf[..guard_offset]);
1852 wire.extend_from_slice(value);
1853 wire.extend_from_slice(&write_buf[guard_offset..]);
1854
1855 assert_eq!(
1856 wire,
1857 &b"set user:123456789 42 7200 5\r\nhello\r\n"[..],
1858 "coalesced guard SET must match a standard SET on the wire"
1859 );
1860
1861 assert_eq!(write_buf.len() + value.len(), wire.len());
1865 }
1866
1867 #[test]
1868 fn two_coalesced_guard_sets_interleave_correctly() {
1869 let v1: &[u8] = b"aa";
1873 let v2: &[u8] = b"bbbb";
1874
1875 let mut write_buf = Vec::new();
1876 let mut guards: Vec<usize> = Vec::new();
1877
1878 append_set_guard_prefix(&mut write_buf, b"k1", v1.len(), 0, 0).unwrap();
1879 guards.push(write_buf.len());
1880 write_buf.extend_from_slice(b"\r\n");
1881
1882 append_set_guard_prefix(&mut write_buf, b"k2", v2.len(), 0, 0).unwrap();
1883 guards.push(write_buf.len());
1884 write_buf.extend_from_slice(b"\r\n");
1885
1886 let values = [v1, v2];
1888 let mut wire = Vec::new();
1889 let mut pos = 0;
1890 for (i, &offset) in guards.iter().enumerate() {
1891 if offset > pos {
1892 wire.extend_from_slice(&write_buf[pos..offset]);
1893 }
1894 wire.extend_from_slice(values[i]);
1895 pos = offset;
1896 }
1897 if pos < write_buf.len() {
1898 wire.extend_from_slice(&write_buf[pos..]);
1899 }
1900
1901 assert_eq!(wire, &b"set k1 0 0 2\r\naa\r\nset k2 0 0 4\r\nbbbb\r\n"[..]);
1902 }
1903}
1904
1905#[cfg(test)]
1906mod tests {
1907 use super::*;
1908
1909 #[test]
1910 fn validate_key_accepts_max_len() {
1911 let key = vec![b'k'; MAX_KEY_LEN];
1912 assert!(validate_key(&key).is_ok());
1913 }
1914
1915 #[test]
1916 fn validate_key_rejects_oversized() {
1917 let key = vec![b'k'; MAX_KEY_LEN + 1];
1918 assert!(matches!(validate_key(&key), Err(Error::KeyTooLong)));
1919 }
1920
1921 #[test]
1922 fn encode_request_get_rejects_long_key() {
1923 let key = vec![b'k'; MAX_KEY_LEN + 1];
1924 let r = encode_request(&McRequest::get(&key));
1925 assert!(matches!(r, Err(Error::KeyTooLong)));
1926 }
1927
1928 #[test]
1929 fn encode_request_set_passes_large_value() {
1930 let key = b"k";
1934 let value = vec![0u8; 2 * 1024 * 1024];
1935 let r = encode_request(&McRequest::Set {
1936 key,
1937 value: &value,
1938 flags: 0,
1939 exptime: 0,
1940 });
1941 assert!(r.is_ok());
1942 assert!(r.unwrap().starts_with(b"set "));
1943 }
1944
1945 #[test]
1946 fn encode_request_gets_rejects_any_long_key() {
1947 let ok = vec![b'k'; MAX_KEY_LEN];
1948 let bad = vec![b'k'; MAX_KEY_LEN + 1];
1949 let keys: &[&[u8]] = &[&ok, &bad];
1950 let r = encode_request(&McRequest::gets(keys));
1951 assert!(matches!(r, Err(Error::KeyTooLong)));
1952 }
1953
1954 #[test]
1955 fn set_guard_prefix_rejects_long_key() {
1956 let key = vec![b'k'; MAX_KEY_LEN + 1];
1957 let r = append_set_guard_prefix(&mut Vec::new(), &key, 16, 0, 0);
1958 assert!(matches!(r, Err(Error::KeyTooLong)));
1959 }
1960
1961 #[test]
1962 fn set_guard_prefix_passes_large_value_len() {
1963 let r = append_set_guard_prefix(&mut Vec::new(), b"k", 2 * 1024 * 1024, 0, 0);
1965 assert!(r.is_ok());
1966 }
1967
1968 #[test]
1969 fn encode_request_passes_through_at_key_cap() {
1970 let key = vec![b'k'; MAX_KEY_LEN];
1971 let value = vec![0u8; 1024];
1972 let r = encode_request(&McRequest::Set {
1973 key: &key,
1974 value: &value,
1975 flags: 0,
1976 exptime: 0,
1977 });
1978 assert!(r.is_ok());
1979 let buf = r.unwrap();
1980 assert!(buf.starts_with(b"set "));
1982 }
1983}
1984
1985#[cfg(test)]
1986mod zc_threshold_tests {
1987 use super::*;
1988 use ringline::{ConnCtx, RegionId, SendGuard};
1989
1990 const KEY: &[u8] = b"user:123456789";
1991
1992 struct VecGuard(Vec<u8>);
1995
1996 impl SendGuard for VecGuard {
1997 fn as_ptr_len(&self) -> (*const u8, u32) {
1998 (self.0.as_ptr(), self.0.len() as u32)
1999 }
2000 fn region(&self) -> RegionId {
2001 RegionId::UNREGISTERED
2002 }
2003 }
2004
2005 fn test_client(max_batch_size: usize, zc_threshold: u32) -> Client {
2008 let conn = ConnCtx::for_test(0, 0);
2009 Client::builder(conn)
2010 .max_batch_size(max_batch_size)
2011 .zc_threshold(zc_threshold)
2012 .build()
2013 }
2014
2015 #[test]
2016 fn small_guard_set_is_byte_identical_to_plain_set() {
2017 let value = vec![0xABu8; 64];
2020
2021 let mut guarded = test_client(4, 4096);
2022 guarded
2023 .fire_set_with_guard(KEY, VecGuard(value.clone()), 7, 42, 1)
2024 .unwrap();
2025
2026 let mut plain = test_client(4, 4096);
2027 plain.fire_set(KEY, &value, 7, 42, 1).unwrap();
2028
2029 assert_eq!(guarded.write_buf, plain.write_buf);
2030 assert!(guarded.write_guards.is_empty());
2031 }
2032
2033 #[test]
2034 fn small_value_takes_copy_path() {
2035 let mut client = test_client(4, 4096);
2037 client
2038 .fire_set_with_guard(KEY, VecGuard(vec![0u8; 64]), 0, 0, 1)
2039 .unwrap();
2040 assert!(client.write_guards.is_empty());
2041 assert_eq!(client.buffered_ops, 1);
2042 assert!(!client.write_buf.is_empty());
2043 }
2044
2045 #[test]
2046 fn large_value_takes_guard_path() {
2047 let mut client = test_client(4, 4096);
2049 client
2050 .fire_set_with_guard(KEY, VecGuard(vec![0u8; 8192]), 0, 0, 1)
2051 .unwrap();
2052 assert_eq!(client.write_guards.len(), 1);
2053 }
2054
2055 #[test]
2056 fn threshold_zero_disables_fold() {
2057 let mut client = test_client(4, 0);
2059 client
2060 .fire_set_with_guard(KEY, VecGuard(vec![0u8; 64]), 0, 0, 1)
2061 .unwrap();
2062 assert_eq!(client.write_guards.len(), 1);
2063 }
2064
2065 #[test]
2066 fn guard_batch_accumulates_max_flush_guards_before_flush() {
2067 let mut client = test_client(16, 0);
2072 for i in 0..MAX_FLUSH_GUARDS {
2073 client
2074 .fire_set_with_guard(KEY, VecGuard(vec![0u8; 64]), 0, 0, i as u64)
2075 .unwrap();
2076 }
2077 assert_eq!(
2078 client.write_guards.len(),
2079 MAX_FLUSH_GUARDS,
2080 "a full guard batch must accumulate without a premature flush"
2081 );
2082 }
2083
2084 #[test]
2085 fn batch_of_small_guard_sets_coalesces() {
2086 const N: usize = 8;
2090 let mut client = test_client(N, 4096);
2091 for i in 0..N - 1 {
2092 client
2093 .fire_set_with_guard(KEY, VecGuard(vec![0u8; 64]), 0, 0, i as u64)
2094 .unwrap();
2095 assert!(
2096 client.write_guards.is_empty(),
2097 "op {i} must use the copy path, not the guard path"
2098 );
2099 assert_eq!(client.buffered_ops, i + 1);
2100 }
2101 assert_eq!(client.buffered_ops, N - 1);
2102 assert!(client.write_guards.is_empty());
2103 }
2104}