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 = 4;
143const MAX_FLUSH_IOVECS: usize = 8;
145const DEFAULT_ZC_THRESHOLD: u32 = 4096;
149
150#[derive(Debug, thiserror::Error)]
158#[non_exhaustive]
159pub enum Error {
160 #[error("connection closed")]
162 ConnectionClosed,
163
164 #[error("memcache error: {0}")]
166 Memcache(String),
167
168 #[error("unexpected response")]
170 UnexpectedResponse,
171
172 #[error("protocol error: {0}")]
174 Protocol(#[from] memcache_proto::ParseError),
175
176 #[error("io error: {0}")]
178 Io(#[from] io::Error),
179
180 #[error("all connections failed")]
182 AllConnectionsFailed,
183
184 #[error("no pending operations")]
186 NoPending,
187
188 #[error("too many in-flight operations")]
192 TooManyInFlight,
193
194 #[error("key too long (max 250 bytes)")]
199 KeyTooLong,
200
201 #[error("value too long (max 1048576 bytes)")]
204 ValueTooLong,
205}
206
207pub const MAX_KEY_LEN: usize = 250;
209
210pub const MAX_VALUE_LEN: usize = 1024 * 1024;
213
214#[inline]
215fn validate_key(key: &[u8]) -> Result<(), Error> {
216 if key.len() > MAX_KEY_LEN {
217 Err(Error::KeyTooLong)
218 } else {
219 Ok(())
220 }
221}
222
223#[inline]
224fn validate_value(value: &[u8]) -> Result<(), Error> {
225 if value.len() > MAX_VALUE_LEN {
226 Err(Error::ValueTooLong)
227 } else {
228 Ok(())
229 }
230}
231
232#[inline]
233fn validate_value_len(value_len: usize) -> Result<(), Error> {
234 if value_len > MAX_VALUE_LEN {
235 Err(Error::ValueTooLong)
236 } else {
237 Ok(())
238 }
239}
240
241#[derive(Debug, Clone)]
245pub struct Value {
246 pub data: Bytes,
248 pub flags: u32,
250}
251
252#[derive(Debug, Clone)]
254pub struct GetValue {
255 pub key: Bytes,
257 pub data: Bytes,
259 pub flags: u32,
261 pub cas: Option<u64>,
263}
264
265#[derive(Debug, Clone, Copy, PartialEq, Eq)]
269#[non_exhaustive]
270pub enum CommandType {
271 Get,
272 Set,
273 Delete,
274 Other,
275}
276
277#[derive(Debug, Clone)]
279pub struct CommandResult {
280 pub command: CommandType,
282 pub latency_ns: u64,
284 pub hit: Option<bool>,
286 pub success: bool,
288 pub ttfb_ns: Option<u64>,
290 pub tx_bytes: u32,
292 pub rx_bytes: u32,
294}
295
296#[cfg(feature = "metrics")]
301pub struct ClientMetrics {
302 pub latency: histogram::Histogram,
304 pub get_latency: histogram::Histogram,
306 pub set_latency: histogram::Histogram,
308 pub del_latency: histogram::Histogram,
310 pub requests: u64,
312 pub errors: u64,
314 pub hits: u64,
316 pub misses: u64,
318}
319
320#[cfg(feature = "metrics")]
321impl ClientMetrics {
322 fn new() -> Self {
323 Self {
324 latency: histogram::Histogram::new(7, 64).unwrap(),
325 get_latency: histogram::Histogram::new(7, 64).unwrap(),
326 set_latency: histogram::Histogram::new(7, 64).unwrap(),
327 del_latency: histogram::Histogram::new(7, 64).unwrap(),
328 requests: 0,
329 errors: 0,
330 hits: 0,
331 misses: 0,
332 }
333 }
334
335 fn record(&mut self, result: &CommandResult) {
336 self.requests += 1;
337 let _ = self.latency.increment(result.latency_ns);
338
339 if !result.success {
340 self.errors += 1;
341 }
342
343 match result.command {
344 CommandType::Get => {
345 let _ = self.get_latency.increment(result.latency_ns);
346 match result.hit {
347 Some(true) => self.hits += 1,
348 Some(false) => self.misses += 1,
349 None => {}
350 }
351 }
352 CommandType::Set => {
353 let _ = self.set_latency.increment(result.latency_ns);
354 }
355 CommandType::Delete => {
356 let _ = self.del_latency.increment(result.latency_ns);
357 }
358 _ => {}
359 }
360 }
361}
362
363enum PendingOpKind {
366 Get,
367 Set,
368 Delete,
369}
370
371struct PendingOp {
372 kind: PendingOpKind,
373 send_ts: u64,
374 start: Option<Instant>,
375 user_data: u64,
376 tx_bytes: u32,
377}
378
379#[non_exhaustive]
381pub enum CompletedOp {
382 Get {
384 result: Result<Option<Value>, Error>,
385 user_data: u64,
386 latency_ns: u64,
387 },
388 Set {
390 result: Result<(), Error>,
391 user_data: u64,
392 latency_ns: u64,
393 },
394 Delete {
396 result: Result<bool, Error>,
397 user_data: u64,
398 latency_ns: u64,
399 },
400}
401
402pub struct ClientBuilder {
406 conn: ConnCtx,
407 on_result: Option<ResultCallback>,
408 max_batch_size: usize,
409 max_in_flight: usize,
410 zc_threshold: u32,
411 #[cfg(feature = "timestamps")]
412 use_kernel_ts: bool,
413 #[cfg(feature = "metrics")]
414 with_metrics: bool,
415}
416
417impl ClientBuilder {
418 pub(crate) fn new(conn: ConnCtx) -> Self {
419 Self {
420 conn,
421 on_result: None,
422 max_batch_size: 1,
423 max_in_flight: usize::MAX,
424 zc_threshold: DEFAULT_ZC_THRESHOLD,
425 #[cfg(feature = "timestamps")]
426 use_kernel_ts: false,
427 #[cfg(feature = "metrics")]
428 with_metrics: false,
429 }
430 }
431
432 pub fn zc_threshold(mut self, bytes: u32) -> Self {
441 self.zc_threshold = bytes;
442 self
443 }
444
445 pub fn max_in_flight(mut self, n: usize) -> Self {
450 self.max_in_flight = n;
451 self
452 }
453
454 pub fn on_result<F: Fn(&CommandResult) + 'static>(mut self, f: F) -> Self {
456 self.on_result = Some(Box::new(f));
457 self
458 }
459
460 #[cfg(feature = "timestamps")]
462 pub fn kernel_timestamps(mut self, enabled: bool) -> Self {
463 self.use_kernel_ts = enabled;
464 self
465 }
466
467 pub fn max_batch_size(mut self, n: usize) -> Self {
476 assert!(n > 0, "max_batch_size must be >= 1");
477 self.max_batch_size = n;
478 self
479 }
480
481 #[cfg(feature = "metrics")]
483 pub fn with_metrics(mut self) -> Self {
484 self.with_metrics = true;
485 self
486 }
487
488 pub fn build(self) -> Client {
490 Client {
491 conn: self.conn,
492 on_result: self.on_result,
493 pending: VecDeque::with_capacity(16),
494 last_rx_bytes: Cell::new(0),
495 write_buf: Vec::new(),
496 write_guards: Vec::new(),
497 flushed_count: 0,
498 max_batch_size: self.max_batch_size,
499 max_in_flight: self.max_in_flight,
500 zc_threshold: self.zc_threshold,
501 buffered_ops: 0,
502 encode_buf: Vec::new(),
503 #[cfg(feature = "timestamps")]
504 use_kernel_ts: self.use_kernel_ts,
505 #[cfg(feature = "metrics")]
506 metrics: if self.with_metrics {
507 Some(ClientMetrics::new())
508 } else {
509 None
510 },
511 }
512 }
513}
514
515pub struct Client {
523 conn: ConnCtx,
524 on_result: Option<ResultCallback>,
525 pending: VecDeque<PendingOp>,
526 last_rx_bytes: Cell<u32>,
527 write_buf: Vec<u8>,
531 write_guards: Vec<(usize, GuardBox)>,
535 flushed_count: usize,
537 max_batch_size: usize,
540 max_in_flight: usize,
543 zc_threshold: u32,
547 buffered_ops: usize,
549 encode_buf: Vec<u8>,
553 #[cfg(feature = "timestamps")]
554 use_kernel_ts: bool,
555 #[cfg(feature = "metrics")]
556 metrics: Option<ClientMetrics>,
557}
558
559impl Client {
560 pub fn new(conn: ConnCtx) -> Self {
564 Self {
565 conn,
566 on_result: None,
567 pending: VecDeque::new(),
568 last_rx_bytes: Cell::new(0),
569 write_buf: Vec::new(),
570 write_guards: Vec::new(),
571 flushed_count: 0,
572 max_batch_size: 1,
573 max_in_flight: usize::MAX,
574 zc_threshold: DEFAULT_ZC_THRESHOLD,
575 buffered_ops: 0,
576 encode_buf: Vec::new(),
577 #[cfg(feature = "timestamps")]
578 use_kernel_ts: false,
579 #[cfg(feature = "metrics")]
580 metrics: None,
581 }
582 }
583
584 pub fn builder(conn: ConnCtx) -> ClientBuilder {
586 ClientBuilder::new(conn)
587 }
588
589 pub fn conn(&self) -> ConnCtx {
591 self.conn
592 }
593
594 #[cfg(feature = "metrics")]
596 pub fn metrics(&self) -> Option<&ClientMetrics> {
597 self.metrics.as_ref()
598 }
599
600 #[cfg(feature = "metrics")]
602 pub fn metrics_mut(&mut self) -> Option<&mut ClientMetrics> {
603 self.metrics.as_mut()
604 }
605
606 #[inline]
609 fn is_instrumented(&self) -> bool {
610 if self.on_result.is_some() {
611 return true;
612 }
613 #[cfg(feature = "metrics")]
614 if self.metrics.is_some() {
615 return true;
616 }
617 false
618 }
619
620 #[cfg(feature = "timestamps")]
621 #[inline]
622 fn send_timestamp(&self) -> u64 {
623 if self.use_kernel_ts {
624 now_realtime_ns()
625 } else {
626 0
627 }
628 }
629
630 #[cfg(not(feature = "timestamps"))]
631 #[inline]
632 fn send_timestamp(&self) -> u64 {
633 0
634 }
635
636 #[cfg(feature = "timestamps")]
637 #[inline]
638 fn finish_timing(&self, send_ts: u64, start: Instant) -> u64 {
639 if self.use_kernel_ts {
640 let recv_ts = self.conn.recv_timestamp();
641 if recv_ts > 0 && recv_ts > send_ts {
642 return recv_ts - send_ts;
643 }
644 }
645 start.elapsed().as_nanos() as u64
646 }
647
648 #[cfg(not(feature = "timestamps"))]
649 #[inline]
650 fn finish_timing(&self, _send_ts: u64, start: Instant) -> u64 {
651 start.elapsed().as_nanos() as u64
652 }
653
654 fn record(&mut self, result: &CommandResult) {
655 if let Some(ref cb) = self.on_result {
656 cb(result);
657 }
658 #[cfg(feature = "metrics")]
659 if let Some(ref mut m) = self.metrics {
660 m.record(result);
661 }
662 }
663
664 #[inline]
667 fn timing_start(&self) -> (u64, Option<Instant>) {
668 #[cfg(feature = "timestamps")]
669 {
670 if self.is_instrumented() {
671 (self.send_timestamp(), Some(Instant::now()))
672 } else {
673 (0, None)
674 }
675 }
676 #[cfg(not(feature = "timestamps"))]
677 {
678 if self.on_result.is_some() {
680 (0, Some(Instant::now()))
681 } else {
682 (0, None)
683 }
684 }
685 }
686
687 #[inline]
695 fn timing_start_buffered(&self) -> (u64, Option<Instant>) {
696 if self.buffered_ops > 1 {
697 (0, None)
698 } else {
699 self.timing_start()
700 }
701 }
702
703 #[cfg(feature = "timestamps")]
704 #[inline]
705 fn compute_ttfb(&self, send_ts: u64) -> Option<u64> {
706 if self.use_kernel_ts {
707 let recv_ts = self.conn.recv_timestamp();
708 if recv_ts > 0 && recv_ts > send_ts {
709 return Some(recv_ts - send_ts);
710 }
711 }
712 None
713 }
714
715 #[cfg(not(feature = "timestamps"))]
716 #[inline]
717 fn compute_ttfb(&self, _send_ts: u64) -> Option<u64> {
718 None
719 }
720
721 pub fn pending_count(&self) -> usize {
723 self.pending.len()
724 }
725
726 fn pre_flush_if_needed(&mut self, has_guard: bool) -> Result<(), Error> {
731 if self.pending.len() >= self.max_in_flight {
732 return Err(Error::TooManyInFlight);
733 }
734 if self.buffered_ops == 0 {
735 return Ok(());
736 }
737 let next_guards = self.write_guards.len() + usize::from(has_guard);
738 let next_parts = 2 * next_guards + 1;
739 if next_guards > MAX_FLUSH_GUARDS || next_parts > MAX_FLUSH_IOVECS {
740 self.flush()?;
741 }
742 Ok(())
743 }
744
745 fn post_flush_if_needed(&mut self) -> Result<(), Error> {
747 if self.buffered_ops >= self.max_batch_size {
748 self.flush()?;
749 }
750 Ok(())
751 }
752
753 pub fn flush(&mut self) -> Result<(), Error> {
759 if self.write_buf.is_empty() && self.write_guards.is_empty() {
760 self.buffered_ops = 0;
761 return Ok(());
762 }
763
764 let send_outcome: Result<(), Error> = if self.write_guards.is_empty() {
771 self.conn
772 .send_nowait(&self.write_buf)
773 .map(|_| ())
774 .map_err(Error::from)
775 } else {
776 use ringline::SendPart;
777 let mut parts: Vec<SendPart<'_>> = Vec::with_capacity(2 * MAX_FLUSH_GUARDS + 1);
778 let mut pos = 0;
779 for (offset, guard) in self.write_guards.drain(..) {
780 if offset > pos {
781 parts.push(SendPart::Copy(&self.write_buf[pos..offset]));
782 }
783 parts.push(SendPart::Guard(guard));
784 pos = offset;
785 }
786 if pos < self.write_buf.len() {
787 parts.push(SendPart::Copy(&self.write_buf[pos..]));
788 }
789 self.conn
790 .send_parts()
791 .submit_batch(parts)
792 .map(|_| ())
793 .map_err(Error::from)
794 };
795
796 if let Err(e) = send_outcome {
797 self.write_buf.clear();
803 self.write_guards.clear();
804 self.buffered_ops = 0;
805 self.pending.truncate(self.flushed_count);
806 return Err(e);
807 }
808
809 if self.buffered_ops > 1 {
813 let (send_ts, start) = self.timing_start();
814 for pending in self.pending.iter_mut().skip(self.flushed_count) {
815 pending.send_ts = send_ts;
816 pending.start = start;
817 }
818 }
819 self.flushed_count = self.pending.len();
820
821 self.write_buf.clear();
822 self.write_guards.clear();
823 self.buffered_ops = 0;
824
825 Ok(())
826 }
827
828 pub fn fire_get(&mut self, key: &[u8], user_data: u64) -> Result<(), Error> {
830 self.pre_flush_if_needed(false)?;
831 let tx_bytes;
832 if self.max_batch_size == 1 && self.write_guards.is_empty() && self.write_buf.is_empty() {
833 self.encode_buf.clear();
835 encode_request_into(&McRequest::get(key), &mut self.encode_buf)?;
836 tx_bytes = self.encode_buf.len() as u32;
837 self.conn.send_nowait(&self.encode_buf)?;
838 } else {
839 let before = self.write_buf.len();
840 encode_request_into(&McRequest::get(key), &mut self.write_buf)?;
841 tx_bytes = (self.write_buf.len() - before) as u32;
842 self.buffered_ops += 1;
843 }
844 let (send_ts, start) = self.timing_start_buffered();
845 self.pending.push_back(PendingOp {
846 kind: PendingOpKind::Get,
847 send_ts,
848 start,
849 user_data,
850 tx_bytes,
851 });
852 self.post_flush_if_needed()?;
853 Ok(())
854 }
855
856 pub fn fire_set(
858 &mut self,
859 key: &[u8],
860 value: &[u8],
861 flags: u32,
862 exptime: u32,
863 user_data: u64,
864 ) -> Result<(), Error> {
865 self.pre_flush_if_needed(false)?;
866 let req = McRequest::Set {
867 key,
868 value,
869 flags,
870 exptime,
871 };
872 let tx_bytes;
873 if self.max_batch_size == 1 && self.write_guards.is_empty() && self.write_buf.is_empty() {
874 self.encode_buf.clear();
876 encode_request_into(&req, &mut self.encode_buf)?;
877 tx_bytes = self.encode_buf.len() as u32;
878 self.conn.send_nowait(&self.encode_buf)?;
879 } else {
880 let before = self.write_buf.len();
881 encode_request_into(&req, &mut self.write_buf)?;
882 tx_bytes = (self.write_buf.len() - before) as u32;
883 self.buffered_ops += 1;
884 }
885 let (send_ts, start) = self.timing_start_buffered();
886 self.pending.push_back(PendingOp {
887 kind: PendingOpKind::Set,
888 send_ts,
889 start,
890 user_data,
891 tx_bytes,
892 });
893 self.post_flush_if_needed()?;
894 Ok(())
895 }
896
897 pub fn fire_set_with_guard<G: SendGuard>(
902 &mut self,
903 key: &[u8],
904 guard: G,
905 flags: u32,
906 exptime: u32,
907 user_data: u64,
908 ) -> Result<(), Error> {
909 let (ptr, value_len) = guard.as_ptr_len();
910 if self.zc_threshold != 0 && value_len < self.zc_threshold {
914 let value = unsafe { core::slice::from_raw_parts(ptr, value_len as usize) };
920 return self.fire_set(key, value, flags, exptime, user_data);
921 }
922 self.pre_flush_if_needed(true)?;
923 let before = self.write_buf.len();
926 append_set_guard_prefix(&mut self.write_buf, key, value_len as usize, flags, exptime)?;
927 self.write_guards
928 .push((self.write_buf.len(), GuardBox::new(guard)));
929 self.write_buf.extend_from_slice(b"\r\n");
930 let tx_bytes = (self.write_buf.len() - before + value_len as usize) as u32;
931 self.buffered_ops += 1;
932 let (send_ts, start) = self.timing_start_buffered();
933 self.pending.push_back(PendingOp {
934 kind: PendingOpKind::Set,
935 send_ts,
936 start,
937 user_data,
938 tx_bytes,
939 });
940 self.post_flush_if_needed()?;
941 Ok(())
942 }
943
944 pub fn fire_delete(&mut self, key: &[u8], user_data: u64) -> Result<(), Error> {
946 self.pre_flush_if_needed(false)?;
947 let tx_bytes;
948 if self.max_batch_size == 1 && self.write_guards.is_empty() && self.write_buf.is_empty() {
949 self.encode_buf.clear();
951 encode_request_into(&McRequest::delete(key), &mut self.encode_buf)?;
952 tx_bytes = self.encode_buf.len() as u32;
953 self.conn.send_nowait(&self.encode_buf)?;
954 } else {
955 let before = self.write_buf.len();
956 encode_request_into(&McRequest::delete(key), &mut self.write_buf)?;
957 tx_bytes = (self.write_buf.len() - before) as u32;
958 self.buffered_ops += 1;
959 }
960 let (send_ts, start) = self.timing_start_buffered();
961 self.pending.push_back(PendingOp {
962 kind: PendingOpKind::Delete,
963 send_ts,
964 start,
965 user_data,
966 tx_bytes,
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 ttfb_ns = self.compute_ttfb(pending.send_ts);
984
985 let response = match self.read_response().await {
986 Ok(v) => v,
987 Err(e) => {
988 self.pending.clear();
994 self.flushed_count = 0;
995 return Err(e);
996 }
997 };
998 let latency_ns = match pending.start {
999 Some(start) => self.finish_timing(pending.send_ts, start),
1000 None => 0,
1001 };
1002 let rx_bytes = self.last_rx_bytes.get();
1003 let tx_bytes = pending.tx_bytes;
1004
1005 let op = match pending.kind {
1006 PendingOpKind::Get => {
1007 let result = match check_error_bytes(&response) {
1009 Err(e) => Err(e),
1010 Ok(()) => match response {
1011 McResponseBytes::Values(mut values) => {
1012 if values.is_empty() {
1013 Ok(None)
1014 } else {
1015 let v = values.swap_remove(0);
1016 Ok(Some(Value {
1017 data: v.data,
1018 flags: v.flags,
1019 }))
1020 }
1021 }
1022 _ => Err(Error::UnexpectedResponse),
1023 },
1024 };
1025 let (success, hit) = match &result {
1026 Ok(Some(_)) => (true, Some(true)),
1027 Ok(None) => (true, Some(false)),
1028 Err(_) => (false, None),
1029 };
1030 self.record(&CommandResult {
1031 command: CommandType::Get,
1032 latency_ns,
1033 hit,
1034 success,
1035 ttfb_ns,
1036 tx_bytes,
1037 rx_bytes,
1038 });
1039 CompletedOp::Get {
1040 result,
1041 user_data: pending.user_data,
1042 latency_ns,
1043 }
1044 }
1045 PendingOpKind::Set => {
1046 let result = match check_error_bytes(&response) {
1047 Err(e) => Err(e),
1048 Ok(()) => match response {
1049 McResponseBytes::Stored => Ok(()),
1050 _ => Err(Error::UnexpectedResponse),
1051 },
1052 };
1053 self.record(&CommandResult {
1054 command: CommandType::Set,
1055 latency_ns,
1056 hit: None,
1057 success: result.is_ok(),
1058 ttfb_ns,
1059 tx_bytes,
1060 rx_bytes,
1061 });
1062 CompletedOp::Set {
1063 result,
1064 user_data: pending.user_data,
1065 latency_ns,
1066 }
1067 }
1068 PendingOpKind::Delete => {
1069 let result = match check_error_bytes(&response) {
1070 Err(e) => Err(e),
1071 Ok(()) => match response {
1072 McResponseBytes::Deleted => Ok(true),
1073 McResponseBytes::NotFound => Ok(false),
1074 _ => Err(Error::UnexpectedResponse),
1075 },
1076 };
1077 self.record(&CommandResult {
1078 command: CommandType::Delete,
1079 latency_ns,
1080 hit: None,
1081 success: result.is_ok(),
1082 ttfb_ns,
1083 tx_bytes,
1084 rx_bytes,
1085 });
1086 CompletedOp::Delete {
1087 result,
1088 user_data: pending.user_data,
1089 latency_ns,
1090 }
1091 }
1092 };
1093
1094 Ok(op)
1095 }
1096
1097 pub(crate) async fn read_response(&self) -> Result<McResponseBytes, Error> {
1112 let mut result: Option<Result<McResponseBytes, Error>> = None;
1113 let n = self
1114 .conn
1115 .with_bytes(|bytes| {
1116 let len = bytes.len();
1117 match McResponseBytes::parse(bytes) {
1118 Ok((response, consumed)) => {
1119 result = Some(Ok(response));
1120 ParseResult::Consumed(consumed)
1121 }
1122 Err(e) if e.is_incomplete() => ParseResult::Consumed(0),
1123 Err(e) => {
1124 result = Some(Err(Error::Protocol(e)));
1125 ParseResult::Consumed(len)
1126 }
1127 }
1128 })
1129 .await;
1130 self.last_rx_bytes.set(n as u32);
1131 if n == 0 {
1132 return result.unwrap_or(Err(Error::ConnectionClosed));
1133 }
1134 let r = result.unwrap();
1135 if matches!(r, Err(Error::Protocol(_))) {
1136 self.conn.close();
1137 }
1138 r
1139 }
1140
1141 async fn execute(&self, encoded: &[u8]) -> Result<McResponseBytes, Error> {
1144 self.conn.send(encoded)?;
1145 let response = self.read_response().await?;
1146 check_error_bytes(&response)?;
1147 Ok(response)
1148 }
1149
1150 pub async fn get(&mut self, key: impl AsRef<[u8]>) -> Result<Option<Value>, Error> {
1154 let key = key.as_ref();
1155 self.encode_buf.clear();
1156 encode_request_into(&McRequest::get(key), &mut self.encode_buf)?;
1157
1158 if !self.is_instrumented() {
1159 let response = self.execute(&self.encode_buf).await?;
1160 return match response {
1161 McResponseBytes::Values(mut values) => {
1162 if values.is_empty() {
1163 Ok(None)
1164 } else {
1165 let v = values.swap_remove(0);
1166 Ok(Some(Value {
1167 data: v.data,
1168 flags: v.flags,
1169 }))
1170 }
1171 }
1172 _ => Err(Error::UnexpectedResponse),
1173 };
1174 }
1175
1176 let tx_bytes = self.encode_buf.len() as u32;
1177 let send_ts = self.send_timestamp();
1178 let start = Instant::now();
1179 let response = self.execute(&self.encode_buf).await;
1180 let latency_ns = self.finish_timing(send_ts, start);
1181 let rx_bytes = self.last_rx_bytes.get();
1182
1183 let result = match response {
1184 Ok(McResponseBytes::Values(mut values)) => {
1185 if values.is_empty() {
1186 Ok(None)
1187 } else {
1188 let v = values.swap_remove(0);
1189 Ok(Some(Value {
1190 data: v.data,
1191 flags: v.flags,
1192 }))
1193 }
1194 }
1195 Ok(_) => Err(Error::UnexpectedResponse),
1196 Err(e) => Err(e),
1197 };
1198
1199 let (success, hit) = match &result {
1200 Ok(Some(_)) => (true, Some(true)),
1201 Ok(None) => (true, Some(false)),
1202 Err(_) => (false, None),
1203 };
1204 self.record(&CommandResult {
1205 command: CommandType::Get,
1206 latency_ns,
1207 hit,
1208 success,
1209 ttfb_ns: None,
1210 tx_bytes,
1211 rx_bytes,
1212 });
1213 result
1214 }
1215
1216 pub async fn gets(&mut self, keys: &[&[u8]]) -> Result<Vec<GetValue>, Error> {
1218 if keys.is_empty() {
1219 return Ok(Vec::new());
1220 }
1221 let encoded = encode_request(&McRequest::gets(keys))?;
1222 let response = self.execute(&encoded).await?;
1223 match response {
1224 McResponseBytes::Values(values) => Ok(values
1225 .into_iter()
1226 .map(|v| GetValue {
1227 key: v.key,
1228 data: v.data,
1229 flags: v.flags,
1230 cas: v.cas,
1231 })
1232 .collect()),
1233 _ => Err(Error::UnexpectedResponse),
1234 }
1235 }
1236
1237 pub async fn set(
1239 &mut self,
1240 key: impl AsRef<[u8]>,
1241 value: impl AsRef<[u8]>,
1242 ) -> Result<(), Error> {
1243 self.set_with_options(key, value, 0, 0).await
1244 }
1245
1246 pub async fn set_with_options(
1248 &mut self,
1249 key: impl AsRef<[u8]>,
1250 value: impl AsRef<[u8]>,
1251 flags: u32,
1252 exptime: u32,
1253 ) -> Result<(), Error> {
1254 let key = key.as_ref();
1255 let value = value.as_ref();
1256 self.encode_buf.clear();
1257 encode_request_into(
1258 &McRequest::Set {
1259 key,
1260 value,
1261 flags,
1262 exptime,
1263 },
1264 &mut self.encode_buf,
1265 )?;
1266
1267 if !self.is_instrumented() {
1268 let response = self.execute(&self.encode_buf).await?;
1269 return match response {
1270 McResponseBytes::Stored => Ok(()),
1271 _ => Err(Error::UnexpectedResponse),
1272 };
1273 }
1274
1275 let tx_bytes = self.encode_buf.len() as u32;
1276 let send_ts = self.send_timestamp();
1277 let start = Instant::now();
1278 let response = self.execute(&self.encode_buf).await;
1279 let latency_ns = self.finish_timing(send_ts, start);
1280 let rx_bytes = self.last_rx_bytes.get();
1281
1282 let result = match response {
1283 Ok(McResponseBytes::Stored) => Ok(()),
1284 Ok(_) => Err(Error::UnexpectedResponse),
1285 Err(e) => Err(e),
1286 };
1287 self.record(&CommandResult {
1288 command: CommandType::Set,
1289 latency_ns,
1290 hit: None,
1291 success: result.is_ok(),
1292 ttfb_ns: None,
1293 tx_bytes,
1294 rx_bytes,
1295 });
1296 result
1297 }
1298
1299 pub async fn add(
1302 &mut self,
1303 key: impl AsRef<[u8]>,
1304 value: impl AsRef<[u8]>,
1305 ) -> Result<bool, Error> {
1306 let key = key.as_ref();
1307 let value = value.as_ref();
1308 let encoded = encode_add(key, value)?;
1309 let response = self.execute(&encoded).await?;
1310 match response {
1311 McResponseBytes::Stored => Ok(true),
1312 McResponseBytes::NotStored => Ok(false),
1313 _ => Err(Error::UnexpectedResponse),
1314 }
1315 }
1316
1317 pub async fn replace(
1320 &mut self,
1321 key: impl AsRef<[u8]>,
1322 value: impl AsRef<[u8]>,
1323 ) -> Result<bool, Error> {
1324 let key = key.as_ref();
1325 let value = value.as_ref();
1326 let encoded = encode_request(&McRequest::Replace {
1327 key,
1328 value,
1329 flags: 0,
1330 exptime: 0,
1331 })?;
1332 let response = self.execute(&encoded).await?;
1333 match response {
1334 McResponseBytes::Stored => Ok(true),
1335 McResponseBytes::NotStored => Ok(false),
1336 _ => Err(Error::UnexpectedResponse),
1337 }
1338 }
1339
1340 pub async fn incr(&mut self, key: impl AsRef<[u8]>, delta: u64) -> Result<Option<u64>, Error> {
1343 let key = key.as_ref();
1344 let encoded = encode_request(&McRequest::incr(key, delta))?;
1345 let response = self.execute(&encoded).await?;
1346 match response {
1347 McResponseBytes::Numeric(val) => Ok(Some(val)),
1348 McResponseBytes::NotFound => Ok(None),
1349 _ => Err(Error::UnexpectedResponse),
1350 }
1351 }
1352
1353 pub async fn decr(&mut self, key: impl AsRef<[u8]>, delta: u64) -> Result<Option<u64>, Error> {
1356 let key = key.as_ref();
1357 let encoded = encode_request(&McRequest::decr(key, delta))?;
1358 let response = self.execute(&encoded).await?;
1359 match response {
1360 McResponseBytes::Numeric(val) => Ok(Some(val)),
1361 McResponseBytes::NotFound => Ok(None),
1362 _ => Err(Error::UnexpectedResponse),
1363 }
1364 }
1365
1366 pub async fn append(
1369 &mut self,
1370 key: impl AsRef<[u8]>,
1371 value: impl AsRef<[u8]>,
1372 ) -> Result<bool, Error> {
1373 let key = key.as_ref();
1374 let value = value.as_ref();
1375 let encoded = encode_request(&McRequest::append(key, value))?;
1376 let response = self.execute(&encoded).await?;
1377 match response {
1378 McResponseBytes::Stored => Ok(true),
1379 McResponseBytes::NotStored => Ok(false),
1380 _ => Err(Error::UnexpectedResponse),
1381 }
1382 }
1383
1384 pub async fn prepend(
1387 &mut self,
1388 key: impl AsRef<[u8]>,
1389 value: impl AsRef<[u8]>,
1390 ) -> Result<bool, Error> {
1391 let key = key.as_ref();
1392 let value = value.as_ref();
1393 let encoded = encode_request(&McRequest::prepend(key, value))?;
1394 let response = self.execute(&encoded).await?;
1395 match response {
1396 McResponseBytes::Stored => Ok(true),
1397 McResponseBytes::NotStored => Ok(false),
1398 _ => Err(Error::UnexpectedResponse),
1399 }
1400 }
1401
1402 pub async fn cas(
1406 &mut self,
1407 key: impl AsRef<[u8]>,
1408 value: impl AsRef<[u8]>,
1409 cas_unique: u64,
1410 ) -> Result<bool, Error> {
1411 let key = key.as_ref();
1412 let value = value.as_ref();
1413 let encoded = encode_request(&McRequest::cas(key, value, cas_unique))?;
1414 let response = self.execute(&encoded).await?;
1415 match response {
1416 McResponseBytes::Stored => Ok(true),
1417 McResponseBytes::Exists => Ok(false),
1418 McResponseBytes::NotFound => Err(Error::Memcache("NOT_FOUND".into())),
1419 _ => Err(Error::UnexpectedResponse),
1420 }
1421 }
1422
1423 pub async fn delete(&mut self, key: impl AsRef<[u8]>) -> Result<bool, Error> {
1425 let key = key.as_ref();
1426 self.encode_buf.clear();
1427 encode_request_into(&McRequest::delete(key), &mut self.encode_buf)?;
1428
1429 if !self.is_instrumented() {
1430 let response = self.execute(&self.encode_buf).await?;
1431 return match response {
1432 McResponseBytes::Deleted => Ok(true),
1433 McResponseBytes::NotFound => Ok(false),
1434 _ => Err(Error::UnexpectedResponse),
1435 };
1436 }
1437
1438 let tx_bytes = self.encode_buf.len() as u32;
1439 let send_ts = self.send_timestamp();
1440 let start = Instant::now();
1441 let response = self.execute(&self.encode_buf).await;
1442 let latency_ns = self.finish_timing(send_ts, start);
1443 let rx_bytes = self.last_rx_bytes.get();
1444
1445 let result = match response {
1446 Ok(McResponseBytes::Deleted) => Ok(true),
1447 Ok(McResponseBytes::NotFound) => Ok(false),
1448 Ok(_) => Err(Error::UnexpectedResponse),
1449 Err(e) => Err(e),
1450 };
1451 self.record(&CommandResult {
1452 command: CommandType::Delete,
1453 latency_ns,
1454 hit: None,
1455 success: result.is_ok(),
1456 ttfb_ns: None,
1457 tx_bytes,
1458 rx_bytes,
1459 });
1460 result
1461 }
1462
1463 pub async fn flush_all(&mut self) -> Result<(), Error> {
1465 let encoded = encode_request(&McRequest::flush_all())?;
1466 let response = self.execute(&encoded).await?;
1467 match response {
1468 McResponseBytes::Ok => Ok(()),
1469 _ => Err(Error::UnexpectedResponse),
1470 }
1471 }
1472
1473 pub async fn version(&mut self) -> Result<Box<str>, Error> {
1475 let encoded = encode_request(&McRequest::version())?;
1476 let response = self.execute(&encoded).await?;
1477 match response {
1478 McResponseBytes::Version(v) => Ok(Box::from(String::from_utf8_lossy(v.as_ref()))),
1479 _ => Err(Error::UnexpectedResponse),
1480 }
1481 }
1482
1483 pub async fn set_with_guard<G: SendGuard>(
1488 &mut self,
1489 key: &[u8],
1490 guard: G,
1491 flags: u32,
1492 exptime: u32,
1493 ) -> Result<(), Error> {
1494 if !self.is_instrumented() {
1495 let (_, value_len) = guard.as_ptr_len();
1496 self.encode_buf.clear();
1497 append_set_guard_prefix(
1498 &mut self.encode_buf,
1499 key,
1500 value_len as usize,
1501 flags,
1502 exptime,
1503 )?;
1504
1505 let prefix: &[u8] = &self.encode_buf;
1506 self.conn.send_parts().build(move |b| {
1507 b.copy(prefix)
1508 .guard(GuardBox::new(guard))
1509 .copy(b"\r\n")
1510 .submit()
1511 })?;
1512
1513 let response = self.read_response().await?;
1514 check_error_bytes(&response)?;
1515 return match response {
1516 McResponseBytes::Stored => Ok(()),
1517 _ => Err(Error::UnexpectedResponse),
1518 };
1519 }
1520
1521 let (_, value_len) = guard.as_ptr_len();
1522 self.encode_buf.clear();
1523 append_set_guard_prefix(
1524 &mut self.encode_buf,
1525 key,
1526 value_len as usize,
1527 flags,
1528 exptime,
1529 )?;
1530 let tx_bytes = (self.encode_buf.len() + value_len as usize + 2) as u32;
1531
1532 let send_ts = self.send_timestamp();
1533 let start = Instant::now();
1534
1535 let prefix: &[u8] = &self.encode_buf;
1536 self.conn.send_parts().build(move |b| {
1537 b.copy(prefix)
1538 .guard(GuardBox::new(guard))
1539 .copy(b"\r\n")
1540 .submit()
1541 })?;
1542
1543 let response = self.read_response().await;
1544 let latency_ns = self.finish_timing(send_ts, start);
1545 let rx_bytes = self.last_rx_bytes.get();
1546
1547 let result = match response {
1548 Ok(ref r) => {
1549 check_error_bytes(r)?;
1550 match r {
1551 McResponseBytes::Stored => Ok(()),
1552 _ => Err(Error::UnexpectedResponse),
1553 }
1554 }
1555 Err(e) => Err(e),
1556 };
1557 self.record(&CommandResult {
1558 command: CommandType::Set,
1559 latency_ns,
1560 hit: None,
1561 success: result.is_ok(),
1562 ttfb_ns: None,
1563 tx_bytes,
1564 rx_bytes,
1565 });
1566 result
1567 }
1568}
1569
1570fn append_set_guard_prefix(
1581 buf: &mut Vec<u8>,
1582 key: &[u8],
1583 value_len: usize,
1584 flags: u32,
1585 exptime: u32,
1586) -> Result<(), Error> {
1587 validate_key(key)?;
1588 validate_value_len(value_len)?;
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> {
1606 match req {
1607 McRequest::Get { key }
1608 | McRequest::Incr { key, .. }
1609 | McRequest::Decr { key, .. }
1610 | McRequest::Delete { key } => validate_key(key),
1611 McRequest::Gets { keys } => {
1612 for k in keys.iter() {
1613 validate_key(k)?;
1614 }
1615 Ok(())
1616 }
1617 McRequest::Set { key, value, .. }
1618 | McRequest::Add { key, value, .. }
1619 | McRequest::Replace { key, value, .. }
1620 | McRequest::Append { key, value }
1621 | McRequest::Prepend { key, value }
1622 | McRequest::Cas { key, value, .. } => {
1623 validate_key(key)?;
1624 validate_value(value)
1625 }
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
1667pub(crate) fn encode_set(
1669 key: &[u8],
1670 value: &[u8],
1671 flags: u32,
1672 exptime: u32,
1673) -> Result<Vec<u8>, Error> {
1674 encode_request(&McRequest::Set {
1675 key,
1676 value,
1677 flags,
1678 exptime,
1679 })
1680}
1681
1682pub(crate) fn encode_add(key: &[u8], value: &[u8]) -> Result<Vec<u8>, Error> {
1684 encode_request(&McRequest::Add {
1685 key,
1686 value,
1687 flags: 0,
1688 exptime: 0,
1689 })
1690}
1691
1692pub(crate) fn check_error_bytes(response: &McResponseBytes) -> Result<(), Error> {
1694 match response {
1695 McResponseBytes::Error => Err(Error::Memcache("ERROR".into())),
1696 McResponseBytes::ClientError(msg) => Err(Error::Memcache(format!(
1697 "CLIENT_ERROR {}",
1698 String::from_utf8_lossy(msg)
1699 ))),
1700 McResponseBytes::ServerError(msg) => Err(Error::Memcache(format!(
1701 "SERVER_ERROR {}",
1702 String::from_utf8_lossy(msg)
1703 ))),
1704 _ => Ok(()),
1705 }
1706}
1707
1708#[cfg(feature = "timestamps")]
1711fn now_realtime_ns() -> u64 {
1712 let mut ts = libc::timespec {
1713 tv_sec: 0,
1714 tv_nsec: 0,
1715 };
1716 unsafe {
1717 libc::clock_gettime(libc::CLOCK_REALTIME, &mut ts);
1718 }
1719 ts.tv_sec as u64 * 1_000_000_000 + ts.tv_nsec as u64
1720}
1721
1722#[cfg(test)]
1723mod encode_tests {
1724 use super::*;
1725
1726 const KEY: &[u8] = b"user:123456789";
1729
1730 #[test]
1731 fn golden_encode_request_get() {
1732 let encoded = encode_request(&McRequest::get(KEY)).unwrap();
1733 assert_eq!(encoded, b"get user:123456789\r\n");
1734 }
1735
1736 #[test]
1737 fn golden_encode_request_delete() {
1738 let encoded = encode_request(&McRequest::delete(KEY)).unwrap();
1739 assert_eq!(encoded, b"delete user:123456789\r\n");
1740 }
1741
1742 #[test]
1743 fn golden_encode_request_set() {
1744 let encoded = encode_set(KEY, b"hello world value", 42, 7200).unwrap();
1745 assert_eq!(
1746 encoded,
1747 &b"set user:123456789 42 7200 17\r\nhello world value\r\n"[..]
1748 );
1749 }
1750
1751 #[test]
1752 fn golden_encode_request_set_zero_fields() {
1753 let encoded = encode_set(KEY, b"v", 0, 0).unwrap();
1754 assert_eq!(encoded, &b"set user:123456789 0 0 1\r\nv\r\n"[..]);
1755 }
1756
1757 #[test]
1758 fn golden_set_guard_prefix() {
1759 let mut prefix = Vec::new();
1762 append_set_guard_prefix(&mut prefix, KEY, 1024, 42, 7200).unwrap();
1763 assert_eq!(prefix, &b"set user:123456789 42 7200 1024\r\n"[..]);
1764 }
1765
1766 #[test]
1767 fn encode_into_appends_without_clobbering() {
1768 let mut buf = b"EXISTING".to_vec();
1771 encode_request_into(&McRequest::get(KEY), &mut buf).unwrap();
1772 assert_eq!(buf, &b"EXISTINGget user:123456789\r\n"[..]);
1773 }
1774
1775 #[test]
1776 fn encode_into_error_leaves_buf_untouched() {
1777 let mut buf = b"EXISTING".to_vec();
1778 let long_key = vec![b'k'; MAX_KEY_LEN + 1];
1779 assert!(encode_request_into(&McRequest::get(&long_key), &mut buf).is_err());
1780 assert_eq!(buf, b"EXISTING");
1781 assert!(append_set_guard_prefix(&mut buf, &long_key, 16, 0, 0).is_err());
1782 assert_eq!(buf, b"EXISTING");
1783 }
1784
1785 #[test]
1795 fn coalesced_buf_equals_concatenation_of_individual_encodings() {
1796 let mut coalesced = Vec::new();
1799 encode_request_into(&McRequest::get(b"key1"), &mut coalesced).unwrap();
1800 encode_request_into(
1801 &McRequest::Set {
1802 key: b"key2",
1803 value: b"val",
1804 flags: 0,
1805 exptime: 0,
1806 },
1807 &mut coalesced,
1808 )
1809 .unwrap();
1810 encode_request_into(&McRequest::delete(b"key3"), &mut coalesced).unwrap();
1811
1812 let mut expected = Vec::new();
1815 expected.extend_from_slice(&encode_request(&McRequest::get(b"key1")).unwrap());
1816 expected.extend_from_slice(
1817 &encode_request(&McRequest::Set {
1818 key: b"key2",
1819 value: b"val",
1820 flags: 0,
1821 exptime: 0,
1822 })
1823 .unwrap(),
1824 );
1825 expected.extend_from_slice(&encode_request(&McRequest::delete(b"key3")).unwrap());
1826
1827 assert_eq!(coalesced, expected);
1828 assert_eq!(
1830 coalesced,
1831 &b"get key1\r\nset key2 0 0 3\r\nval\r\ndelete key3\r\n"[..]
1832 );
1833 }
1834
1835 #[test]
1836 fn coalesced_guard_set_byte_layout() {
1837 let mut write_buf = Vec::new();
1841 let value: &[u8] = b"hello";
1842 append_set_guard_prefix(&mut write_buf, KEY, value.len(), 42, 7200).unwrap();
1843 let guard_offset = write_buf.len();
1844 write_buf.extend_from_slice(b"\r\n");
1845
1846 let mut wire = Vec::new();
1849 wire.extend_from_slice(&write_buf[..guard_offset]);
1850 wire.extend_from_slice(value);
1851 wire.extend_from_slice(&write_buf[guard_offset..]);
1852
1853 assert_eq!(
1854 wire,
1855 &b"set user:123456789 42 7200 5\r\nhello\r\n"[..],
1856 "coalesced guard SET must match a standard SET on the wire"
1857 );
1858
1859 assert_eq!(write_buf.len() + value.len(), wire.len());
1863 }
1864
1865 #[test]
1866 fn two_coalesced_guard_sets_interleave_correctly() {
1867 let v1: &[u8] = b"aa";
1871 let v2: &[u8] = b"bbbb";
1872
1873 let mut write_buf = Vec::new();
1874 let mut guards: Vec<usize> = Vec::new();
1875
1876 append_set_guard_prefix(&mut write_buf, b"k1", v1.len(), 0, 0).unwrap();
1877 guards.push(write_buf.len());
1878 write_buf.extend_from_slice(b"\r\n");
1879
1880 append_set_guard_prefix(&mut write_buf, b"k2", v2.len(), 0, 0).unwrap();
1881 guards.push(write_buf.len());
1882 write_buf.extend_from_slice(b"\r\n");
1883
1884 let values = [v1, v2];
1886 let mut wire = Vec::new();
1887 let mut pos = 0;
1888 for (i, &offset) in guards.iter().enumerate() {
1889 if offset > pos {
1890 wire.extend_from_slice(&write_buf[pos..offset]);
1891 }
1892 wire.extend_from_slice(values[i]);
1893 pos = offset;
1894 }
1895 if pos < write_buf.len() {
1896 wire.extend_from_slice(&write_buf[pos..]);
1897 }
1898
1899 assert_eq!(wire, &b"set k1 0 0 2\r\naa\r\nset k2 0 0 4\r\nbbbb\r\n"[..]);
1900 }
1901}
1902
1903#[cfg(test)]
1904mod tests {
1905 use super::*;
1906
1907 #[test]
1908 fn validate_key_accepts_max_len() {
1909 let key = vec![b'k'; MAX_KEY_LEN];
1910 assert!(validate_key(&key).is_ok());
1911 }
1912
1913 #[test]
1914 fn validate_key_rejects_oversized() {
1915 let key = vec![b'k'; MAX_KEY_LEN + 1];
1916 assert!(matches!(validate_key(&key), Err(Error::KeyTooLong)));
1917 }
1918
1919 #[test]
1920 fn validate_value_accepts_max_len() {
1921 let v = vec![0u8; MAX_VALUE_LEN];
1922 assert!(validate_value(&v).is_ok());
1923 }
1924
1925 #[test]
1926 fn validate_value_rejects_oversized() {
1927 let v = vec![0u8; MAX_VALUE_LEN + 1];
1928 assert!(matches!(validate_value(&v), Err(Error::ValueTooLong)));
1929 }
1930
1931 #[test]
1932 fn encode_request_get_rejects_long_key() {
1933 let key = vec![b'k'; MAX_KEY_LEN + 1];
1934 let r = encode_request(&McRequest::get(&key));
1935 assert!(matches!(r, Err(Error::KeyTooLong)));
1936 }
1937
1938 #[test]
1939 fn encode_request_set_rejects_long_value() {
1940 let key = b"k";
1941 let value = vec![0u8; MAX_VALUE_LEN + 1];
1942 let r = encode_request(&McRequest::Set {
1943 key,
1944 value: &value,
1945 flags: 0,
1946 exptime: 0,
1947 });
1948 assert!(matches!(r, Err(Error::ValueTooLong)));
1949 }
1950
1951 #[test]
1952 fn encode_request_gets_rejects_any_long_key() {
1953 let ok = vec![b'k'; MAX_KEY_LEN];
1954 let bad = vec![b'k'; MAX_KEY_LEN + 1];
1955 let keys: &[&[u8]] = &[&ok, &bad];
1956 let r = encode_request(&McRequest::gets(keys));
1957 assert!(matches!(r, Err(Error::KeyTooLong)));
1958 }
1959
1960 #[test]
1961 fn set_guard_prefix_rejects_long_key() {
1962 let key = vec![b'k'; MAX_KEY_LEN + 1];
1963 let r = append_set_guard_prefix(&mut Vec::new(), &key, 16, 0, 0);
1964 assert!(matches!(r, Err(Error::KeyTooLong)));
1965 }
1966
1967 #[test]
1968 fn set_guard_prefix_rejects_long_value_len() {
1969 let r = append_set_guard_prefix(&mut Vec::new(), b"k", MAX_VALUE_LEN + 1, 0, 0);
1970 assert!(matches!(r, Err(Error::ValueTooLong)));
1971 }
1972
1973 #[test]
1974 fn set_guard_prefix_accepts_max_value_len() {
1975 let r = append_set_guard_prefix(&mut Vec::new(), b"k", MAX_VALUE_LEN, 0, 0);
1976 assert!(r.is_ok());
1977 }
1978
1979 #[test]
1980 fn encode_request_passes_through_at_caps() {
1981 let key = vec![b'k'; MAX_KEY_LEN];
1982 let value = vec![0u8; MAX_VALUE_LEN];
1983 let r = encode_request(&McRequest::Set {
1984 key: &key,
1985 value: &value,
1986 flags: 0,
1987 exptime: 0,
1988 });
1989 assert!(r.is_ok());
1990 let buf = r.unwrap();
1991 assert!(buf.starts_with(b"set "));
1993 }
1994}
1995
1996#[cfg(test)]
1997mod zc_threshold_tests {
1998 use super::*;
1999 use ringline::{ConnCtx, RegionId, SendGuard};
2000
2001 const KEY: &[u8] = b"user:123456789";
2002
2003 struct VecGuard(Vec<u8>);
2006
2007 impl SendGuard for VecGuard {
2008 fn as_ptr_len(&self) -> (*const u8, u32) {
2009 (self.0.as_ptr(), self.0.len() as u32)
2010 }
2011 fn region(&self) -> RegionId {
2012 RegionId::UNREGISTERED
2013 }
2014 }
2015
2016 fn test_client(max_batch_size: usize, zc_threshold: u32) -> Client {
2019 let conn = ConnCtx::for_test(0, 0);
2020 Client::builder(conn)
2021 .max_batch_size(max_batch_size)
2022 .zc_threshold(zc_threshold)
2023 .build()
2024 }
2025
2026 #[test]
2027 fn small_guard_set_is_byte_identical_to_plain_set() {
2028 let value = vec![0xABu8; 64];
2031
2032 let mut guarded = test_client(4, 4096);
2033 guarded
2034 .fire_set_with_guard(KEY, VecGuard(value.clone()), 7, 42, 1)
2035 .unwrap();
2036
2037 let mut plain = test_client(4, 4096);
2038 plain.fire_set(KEY, &value, 7, 42, 1).unwrap();
2039
2040 assert_eq!(guarded.write_buf, plain.write_buf);
2041 assert!(guarded.write_guards.is_empty());
2042 }
2043
2044 #[test]
2045 fn small_value_takes_copy_path() {
2046 let mut client = test_client(4, 4096);
2048 client
2049 .fire_set_with_guard(KEY, VecGuard(vec![0u8; 64]), 0, 0, 1)
2050 .unwrap();
2051 assert!(client.write_guards.is_empty());
2052 assert_eq!(client.buffered_ops, 1);
2053 assert!(!client.write_buf.is_empty());
2054 }
2055
2056 #[test]
2057 fn large_value_takes_guard_path() {
2058 let mut client = test_client(4, 4096);
2060 client
2061 .fire_set_with_guard(KEY, VecGuard(vec![0u8; 8192]), 0, 0, 1)
2062 .unwrap();
2063 assert_eq!(client.write_guards.len(), 1);
2064 }
2065
2066 #[test]
2067 fn threshold_zero_disables_fold() {
2068 let mut client = test_client(4, 0);
2070 client
2071 .fire_set_with_guard(KEY, VecGuard(vec![0u8; 64]), 0, 0, 1)
2072 .unwrap();
2073 assert_eq!(client.write_guards.len(), 1);
2074 }
2075
2076 #[test]
2077 fn batch_of_small_guard_sets_coalesces() {
2078 const N: usize = 8;
2082 let mut client = test_client(N, 4096);
2083 for i in 0..N - 1 {
2084 client
2085 .fire_set_with_guard(KEY, VecGuard(vec![0u8; 64]), 0, 0, i as u64)
2086 .unwrap();
2087 assert!(
2088 client.write_guards.is_empty(),
2089 "op {i} must use the copy path, not the guard path"
2090 );
2091 assert_eq!(client.buffered_ops, i + 1);
2092 }
2093 assert_eq!(client.buffered_ops, N - 1);
2094 assert!(client.write_guards.is_empty());
2095 }
2096}