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
141#[derive(Debug, thiserror::Error)]
149#[non_exhaustive]
150pub enum Error {
151 #[error("connection closed")]
153 ConnectionClosed,
154
155 #[error("memcache error: {0}")]
157 Memcache(String),
158
159 #[error("unexpected response")]
161 UnexpectedResponse,
162
163 #[error("protocol error: {0}")]
165 Protocol(#[from] memcache_proto::ParseError),
166
167 #[error("io error: {0}")]
169 Io(#[from] io::Error),
170
171 #[error("all connections failed")]
173 AllConnectionsFailed,
174
175 #[error("no pending operations")]
177 NoPending,
178
179 #[error("too many in-flight operations")]
183 TooManyInFlight,
184
185 #[error("key too long (max 250 bytes)")]
190 KeyTooLong,
191
192 #[error("value too long (max 1048576 bytes)")]
195 ValueTooLong,
196}
197
198pub const MAX_KEY_LEN: usize = 250;
200
201pub const MAX_VALUE_LEN: usize = 1024 * 1024;
204
205#[inline]
206fn validate_key(key: &[u8]) -> Result<(), Error> {
207 if key.len() > MAX_KEY_LEN {
208 Err(Error::KeyTooLong)
209 } else {
210 Ok(())
211 }
212}
213
214#[inline]
215fn validate_value(value: &[u8]) -> Result<(), Error> {
216 if value.len() > MAX_VALUE_LEN {
217 Err(Error::ValueTooLong)
218 } else {
219 Ok(())
220 }
221}
222
223#[inline]
224fn validate_value_len(value_len: usize) -> Result<(), Error> {
225 if value_len > MAX_VALUE_LEN {
226 Err(Error::ValueTooLong)
227 } else {
228 Ok(())
229 }
230}
231
232#[derive(Debug, Clone)]
236pub struct Value {
237 pub data: Bytes,
239 pub flags: u32,
241}
242
243#[derive(Debug, Clone)]
245pub struct GetValue {
246 pub key: Bytes,
248 pub data: Bytes,
250 pub flags: u32,
252 pub cas: Option<u64>,
254}
255
256#[derive(Debug, Clone, Copy, PartialEq, Eq)]
260#[non_exhaustive]
261pub enum CommandType {
262 Get,
263 Set,
264 Delete,
265 Other,
266}
267
268#[derive(Debug, Clone)]
270pub struct CommandResult {
271 pub command: CommandType,
273 pub latency_ns: u64,
275 pub hit: Option<bool>,
277 pub success: bool,
279 pub ttfb_ns: Option<u64>,
281 pub tx_bytes: u32,
283 pub rx_bytes: u32,
285}
286
287#[cfg(feature = "metrics")]
292pub struct ClientMetrics {
293 pub latency: histogram::Histogram,
295 pub get_latency: histogram::Histogram,
297 pub set_latency: histogram::Histogram,
299 pub del_latency: histogram::Histogram,
301 pub requests: u64,
303 pub errors: u64,
305 pub hits: u64,
307 pub misses: u64,
309}
310
311#[cfg(feature = "metrics")]
312impl ClientMetrics {
313 fn new() -> Self {
314 Self {
315 latency: histogram::Histogram::new(7, 64).unwrap(),
316 get_latency: histogram::Histogram::new(7, 64).unwrap(),
317 set_latency: histogram::Histogram::new(7, 64).unwrap(),
318 del_latency: histogram::Histogram::new(7, 64).unwrap(),
319 requests: 0,
320 errors: 0,
321 hits: 0,
322 misses: 0,
323 }
324 }
325
326 fn record(&mut self, result: &CommandResult) {
327 self.requests += 1;
328 let _ = self.latency.increment(result.latency_ns);
329
330 if !result.success {
331 self.errors += 1;
332 }
333
334 match result.command {
335 CommandType::Get => {
336 let _ = self.get_latency.increment(result.latency_ns);
337 match result.hit {
338 Some(true) => self.hits += 1,
339 Some(false) => self.misses += 1,
340 None => {}
341 }
342 }
343 CommandType::Set => {
344 let _ = self.set_latency.increment(result.latency_ns);
345 }
346 CommandType::Delete => {
347 let _ = self.del_latency.increment(result.latency_ns);
348 }
349 _ => {}
350 }
351 }
352}
353
354enum PendingOpKind {
357 Get,
358 Set,
359 Delete,
360}
361
362struct PendingOp {
363 kind: PendingOpKind,
364 send_ts: u64,
365 start: Option<Instant>,
366 user_data: u64,
367 tx_bytes: u32,
368}
369
370#[non_exhaustive]
372pub enum CompletedOp {
373 Get {
375 result: Result<Option<Value>, Error>,
376 user_data: u64,
377 latency_ns: u64,
378 },
379 Set {
381 result: Result<(), Error>,
382 user_data: u64,
383 latency_ns: u64,
384 },
385 Delete {
387 result: Result<bool, Error>,
388 user_data: u64,
389 latency_ns: u64,
390 },
391}
392
393pub struct ClientBuilder {
397 conn: ConnCtx,
398 on_result: Option<ResultCallback>,
399 max_in_flight: usize,
400 #[cfg(feature = "timestamps")]
401 use_kernel_ts: bool,
402 #[cfg(feature = "metrics")]
403 with_metrics: bool,
404}
405
406impl ClientBuilder {
407 pub(crate) fn new(conn: ConnCtx) -> Self {
408 Self {
409 conn,
410 on_result: None,
411 max_in_flight: usize::MAX,
412 #[cfg(feature = "timestamps")]
413 use_kernel_ts: false,
414 #[cfg(feature = "metrics")]
415 with_metrics: false,
416 }
417 }
418
419 pub fn max_in_flight(mut self, n: usize) -> Self {
424 self.max_in_flight = n;
425 self
426 }
427
428 pub fn on_result<F: Fn(&CommandResult) + 'static>(mut self, f: F) -> Self {
430 self.on_result = Some(Box::new(f));
431 self
432 }
433
434 #[cfg(feature = "timestamps")]
436 pub fn kernel_timestamps(mut self, enabled: bool) -> Self {
437 self.use_kernel_ts = enabled;
438 self
439 }
440
441 #[cfg(feature = "metrics")]
443 pub fn with_metrics(mut self) -> Self {
444 self.with_metrics = true;
445 self
446 }
447
448 pub fn build(self) -> Client {
450 Client {
451 conn: self.conn,
452 on_result: self.on_result,
453 pending: VecDeque::with_capacity(16),
454 last_rx_bytes: Cell::new(0),
455 max_in_flight: self.max_in_flight,
456 #[cfg(feature = "timestamps")]
457 use_kernel_ts: self.use_kernel_ts,
458 #[cfg(feature = "metrics")]
459 metrics: if self.with_metrics {
460 Some(ClientMetrics::new())
461 } else {
462 None
463 },
464 }
465 }
466}
467
468pub struct Client {
476 conn: ConnCtx,
477 on_result: Option<ResultCallback>,
478 pending: VecDeque<PendingOp>,
479 last_rx_bytes: Cell<u32>,
480 max_in_flight: usize,
483 #[cfg(feature = "timestamps")]
484 use_kernel_ts: bool,
485 #[cfg(feature = "metrics")]
486 metrics: Option<ClientMetrics>,
487}
488
489impl Client {
490 pub fn new(conn: ConnCtx) -> Self {
494 Self {
495 conn,
496 on_result: None,
497 pending: VecDeque::new(),
498 last_rx_bytes: Cell::new(0),
499 max_in_flight: usize::MAX,
500 #[cfg(feature = "timestamps")]
501 use_kernel_ts: false,
502 #[cfg(feature = "metrics")]
503 metrics: None,
504 }
505 }
506
507 pub fn builder(conn: ConnCtx) -> ClientBuilder {
509 ClientBuilder::new(conn)
510 }
511
512 pub fn conn(&self) -> ConnCtx {
514 self.conn
515 }
516
517 #[cfg(feature = "metrics")]
519 pub fn metrics(&self) -> Option<&ClientMetrics> {
520 self.metrics.as_ref()
521 }
522
523 #[cfg(feature = "metrics")]
525 pub fn metrics_mut(&mut self) -> Option<&mut ClientMetrics> {
526 self.metrics.as_mut()
527 }
528
529 #[inline]
532 fn is_instrumented(&self) -> bool {
533 if self.on_result.is_some() {
534 return true;
535 }
536 #[cfg(feature = "metrics")]
537 if self.metrics.is_some() {
538 return true;
539 }
540 false
541 }
542
543 #[cfg(feature = "timestamps")]
544 #[inline]
545 fn send_timestamp(&self) -> u64 {
546 if self.use_kernel_ts {
547 now_realtime_ns()
548 } else {
549 0
550 }
551 }
552
553 #[cfg(not(feature = "timestamps"))]
554 #[inline]
555 fn send_timestamp(&self) -> u64 {
556 0
557 }
558
559 #[cfg(feature = "timestamps")]
560 #[inline]
561 fn finish_timing(&self, send_ts: u64, start: Instant) -> u64 {
562 if self.use_kernel_ts {
563 let recv_ts = self.conn.recv_timestamp();
564 if recv_ts > 0 && recv_ts > send_ts {
565 return recv_ts - send_ts;
566 }
567 }
568 start.elapsed().as_nanos() as u64
569 }
570
571 #[cfg(not(feature = "timestamps"))]
572 #[inline]
573 fn finish_timing(&self, _send_ts: u64, start: Instant) -> u64 {
574 start.elapsed().as_nanos() as u64
575 }
576
577 fn record(&mut self, result: &CommandResult) {
578 if let Some(ref cb) = self.on_result {
579 cb(result);
580 }
581 #[cfg(feature = "metrics")]
582 if let Some(ref mut m) = self.metrics {
583 m.record(result);
584 }
585 }
586
587 #[inline]
590 fn timing_start(&self) -> (u64, Option<Instant>) {
591 #[cfg(feature = "timestamps")]
592 {
593 if self.is_instrumented() {
594 (self.send_timestamp(), Some(Instant::now()))
595 } else {
596 (0, None)
597 }
598 }
599 #[cfg(not(feature = "timestamps"))]
600 {
601 if self.on_result.is_some() {
603 (0, Some(Instant::now()))
604 } else {
605 (0, None)
606 }
607 }
608 }
609
610 #[cfg(feature = "timestamps")]
611 #[inline]
612 fn compute_ttfb(&self, send_ts: u64) -> Option<u64> {
613 if self.use_kernel_ts {
614 let recv_ts = self.conn.recv_timestamp();
615 if recv_ts > 0 && recv_ts > send_ts {
616 return Some(recv_ts - send_ts);
617 }
618 }
619 None
620 }
621
622 #[cfg(not(feature = "timestamps"))]
623 #[inline]
624 fn compute_ttfb(&self, _send_ts: u64) -> Option<u64> {
625 None
626 }
627
628 pub fn pending_count(&self) -> usize {
630 self.pending.len()
631 }
632
633 #[inline]
637 fn check_in_flight(&self) -> Result<(), Error> {
638 if self.pending.len() >= self.max_in_flight {
639 Err(Error::TooManyInFlight)
640 } else {
641 Ok(())
642 }
643 }
644
645 pub fn fire_get(&mut self, key: &[u8], user_data: u64) -> Result<(), Error> {
647 self.check_in_flight()?;
648 let encoded = encode_request(&McRequest::get(key))?;
649 let tx_bytes = encoded.len() as u32;
650 self.conn.send_nowait(&encoded)?;
651 let (send_ts, start) = self.timing_start();
652 self.pending.push_back(PendingOp {
653 kind: PendingOpKind::Get,
654 send_ts,
655 start,
656 user_data,
657 tx_bytes,
658 });
659 Ok(())
660 }
661
662 pub fn fire_set(
664 &mut self,
665 key: &[u8],
666 value: &[u8],
667 flags: u32,
668 exptime: u32,
669 user_data: u64,
670 ) -> Result<(), Error> {
671 self.check_in_flight()?;
672 let encoded = encode_set(key, value, flags, exptime)?;
673 let tx_bytes = encoded.len() as u32;
674 self.conn.send_nowait(&encoded)?;
675 let (send_ts, start) = self.timing_start();
676 self.pending.push_back(PendingOp {
677 kind: PendingOpKind::Set,
678 send_ts,
679 start,
680 user_data,
681 tx_bytes,
682 });
683 Ok(())
684 }
685
686 pub fn fire_set_with_guard<G: SendGuard>(
688 &mut self,
689 key: &[u8],
690 guard: G,
691 flags: u32,
692 exptime: u32,
693 user_data: u64,
694 ) -> Result<(), Error> {
695 self.check_in_flight()?;
696 let (_, value_len) = guard.as_ptr_len();
697 let prefix = encode_set_guard_prefix(key, value_len as usize, flags, exptime)?;
698 let tx_bytes = (prefix.len() + value_len as usize + 2) as u32;
699 self.conn.send_parts().build(move |b| {
700 b.copy(&prefix)
701 .guard(GuardBox::new(guard))
702 .copy(b"\r\n")
703 .submit()
704 })?;
705 let (send_ts, start) = self.timing_start();
706 self.pending.push_back(PendingOp {
707 kind: PendingOpKind::Set,
708 send_ts,
709 start,
710 user_data,
711 tx_bytes,
712 });
713 Ok(())
714 }
715
716 pub fn fire_delete(&mut self, key: &[u8], user_data: u64) -> Result<(), Error> {
718 self.check_in_flight()?;
719 let encoded = encode_request(&McRequest::delete(key))?;
720 let tx_bytes = encoded.len() as u32;
721 self.conn.send_nowait(&encoded)?;
722 let (send_ts, start) = self.timing_start();
723 self.pending.push_back(PendingOp {
724 kind: PendingOpKind::Delete,
725 send_ts,
726 start,
727 user_data,
728 tx_bytes,
729 });
730 Ok(())
731 }
732
733 pub async fn recv(&mut self) -> Result<CompletedOp, Error> {
737 let pending = self.pending.pop_front().ok_or(Error::NoPending)?;
738
739 let ttfb_ns = self.compute_ttfb(pending.send_ts);
741
742 let response = match self.read_response().await {
743 Ok(v) => v,
744 Err(e) => {
745 self.pending.clear();
749 return Err(e);
750 }
751 };
752 let latency_ns = match pending.start {
753 Some(start) => self.finish_timing(pending.send_ts, start),
754 None => 0,
755 };
756 let rx_bytes = self.last_rx_bytes.get();
757 let tx_bytes = pending.tx_bytes;
758
759 let op = match pending.kind {
760 PendingOpKind::Get => {
761 let result = match check_error_bytes(&response) {
763 Err(e) => Err(e),
764 Ok(()) => match response {
765 McResponseBytes::Values(mut values) => {
766 if values.is_empty() {
767 Ok(None)
768 } else {
769 let v = values.swap_remove(0);
770 Ok(Some(Value {
771 data: v.data,
772 flags: v.flags,
773 }))
774 }
775 }
776 _ => Err(Error::UnexpectedResponse),
777 },
778 };
779 let (success, hit) = match &result {
780 Ok(Some(_)) => (true, Some(true)),
781 Ok(None) => (true, Some(false)),
782 Err(_) => (false, None),
783 };
784 self.record(&CommandResult {
785 command: CommandType::Get,
786 latency_ns,
787 hit,
788 success,
789 ttfb_ns,
790 tx_bytes,
791 rx_bytes,
792 });
793 CompletedOp::Get {
794 result,
795 user_data: pending.user_data,
796 latency_ns,
797 }
798 }
799 PendingOpKind::Set => {
800 let result = match check_error_bytes(&response) {
801 Err(e) => Err(e),
802 Ok(()) => match response {
803 McResponseBytes::Stored => Ok(()),
804 _ => Err(Error::UnexpectedResponse),
805 },
806 };
807 self.record(&CommandResult {
808 command: CommandType::Set,
809 latency_ns,
810 hit: None,
811 success: result.is_ok(),
812 ttfb_ns,
813 tx_bytes,
814 rx_bytes,
815 });
816 CompletedOp::Set {
817 result,
818 user_data: pending.user_data,
819 latency_ns,
820 }
821 }
822 PendingOpKind::Delete => {
823 let result = match check_error_bytes(&response) {
824 Err(e) => Err(e),
825 Ok(()) => match response {
826 McResponseBytes::Deleted => Ok(true),
827 McResponseBytes::NotFound => Ok(false),
828 _ => Err(Error::UnexpectedResponse),
829 },
830 };
831 self.record(&CommandResult {
832 command: CommandType::Delete,
833 latency_ns,
834 hit: None,
835 success: result.is_ok(),
836 ttfb_ns,
837 tx_bytes,
838 rx_bytes,
839 });
840 CompletedOp::Delete {
841 result,
842 user_data: pending.user_data,
843 latency_ns,
844 }
845 }
846 };
847
848 Ok(op)
849 }
850
851 pub(crate) async fn read_response(&self) -> Result<McResponseBytes, Error> {
866 let mut result: Option<Result<McResponseBytes, Error>> = None;
867 let n = self
868 .conn
869 .with_bytes(|bytes| {
870 let len = bytes.len();
871 match McResponseBytes::parse(bytes) {
872 Ok((response, consumed)) => {
873 result = Some(Ok(response));
874 ParseResult::Consumed(consumed)
875 }
876 Err(e) if e.is_incomplete() => ParseResult::Consumed(0),
877 Err(e) => {
878 result = Some(Err(Error::Protocol(e)));
879 ParseResult::Consumed(len)
880 }
881 }
882 })
883 .await;
884 self.last_rx_bytes.set(n as u32);
885 if n == 0 {
886 return result.unwrap_or(Err(Error::ConnectionClosed));
887 }
888 let r = result.unwrap();
889 if matches!(r, Err(Error::Protocol(_))) {
890 self.conn.close();
891 }
892 r
893 }
894
895 async fn execute(&self, encoded: &[u8]) -> Result<McResponseBytes, Error> {
898 self.conn.send(encoded)?;
899 let response = self.read_response().await?;
900 check_error_bytes(&response)?;
901 Ok(response)
902 }
903
904 pub async fn get(&mut self, key: impl AsRef<[u8]>) -> Result<Option<Value>, Error> {
908 let key = key.as_ref();
909 let encoded = encode_request(&McRequest::get(key))?;
910
911 if !self.is_instrumented() {
912 let response = self.execute(&encoded).await?;
913 return match response {
914 McResponseBytes::Values(mut values) => {
915 if values.is_empty() {
916 Ok(None)
917 } else {
918 let v = values.swap_remove(0);
919 Ok(Some(Value {
920 data: v.data,
921 flags: v.flags,
922 }))
923 }
924 }
925 _ => Err(Error::UnexpectedResponse),
926 };
927 }
928
929 let tx_bytes = encoded.len() as u32;
930 let send_ts = self.send_timestamp();
931 let start = Instant::now();
932 let response = self.execute(&encoded).await;
933 let latency_ns = self.finish_timing(send_ts, start);
934 let rx_bytes = self.last_rx_bytes.get();
935
936 let result = match response {
937 Ok(McResponseBytes::Values(mut values)) => {
938 if values.is_empty() {
939 Ok(None)
940 } else {
941 let v = values.swap_remove(0);
942 Ok(Some(Value {
943 data: v.data,
944 flags: v.flags,
945 }))
946 }
947 }
948 Ok(_) => Err(Error::UnexpectedResponse),
949 Err(e) => Err(e),
950 };
951
952 let (success, hit) = match &result {
953 Ok(Some(_)) => (true, Some(true)),
954 Ok(None) => (true, Some(false)),
955 Err(_) => (false, None),
956 };
957 self.record(&CommandResult {
958 command: CommandType::Get,
959 latency_ns,
960 hit,
961 success,
962 ttfb_ns: None,
963 tx_bytes,
964 rx_bytes,
965 });
966 result
967 }
968
969 pub async fn gets(&mut self, keys: &[&[u8]]) -> Result<Vec<GetValue>, Error> {
971 if keys.is_empty() {
972 return Ok(Vec::new());
973 }
974 let encoded = encode_request(&McRequest::gets(keys))?;
975 let response = self.execute(&encoded).await?;
976 match response {
977 McResponseBytes::Values(values) => Ok(values
978 .into_iter()
979 .map(|v| GetValue {
980 key: v.key,
981 data: v.data,
982 flags: v.flags,
983 cas: v.cas,
984 })
985 .collect()),
986 _ => Err(Error::UnexpectedResponse),
987 }
988 }
989
990 pub async fn set(
992 &mut self,
993 key: impl AsRef<[u8]>,
994 value: impl AsRef<[u8]>,
995 ) -> Result<(), Error> {
996 self.set_with_options(key, value, 0, 0).await
997 }
998
999 pub async fn set_with_options(
1001 &mut self,
1002 key: impl AsRef<[u8]>,
1003 value: impl AsRef<[u8]>,
1004 flags: u32,
1005 exptime: u32,
1006 ) -> Result<(), Error> {
1007 let key = key.as_ref();
1008 let value = value.as_ref();
1009 let encoded = encode_set(key, value, flags, exptime)?;
1010
1011 if !self.is_instrumented() {
1012 let response = self.execute(&encoded).await?;
1013 return match response {
1014 McResponseBytes::Stored => Ok(()),
1015 _ => Err(Error::UnexpectedResponse),
1016 };
1017 }
1018
1019 let tx_bytes = encoded.len() as u32;
1020 let send_ts = self.send_timestamp();
1021 let start = Instant::now();
1022 let response = self.execute(&encoded).await;
1023 let latency_ns = self.finish_timing(send_ts, start);
1024 let rx_bytes = self.last_rx_bytes.get();
1025
1026 let result = match response {
1027 Ok(McResponseBytes::Stored) => Ok(()),
1028 Ok(_) => Err(Error::UnexpectedResponse),
1029 Err(e) => Err(e),
1030 };
1031 self.record(&CommandResult {
1032 command: CommandType::Set,
1033 latency_ns,
1034 hit: None,
1035 success: result.is_ok(),
1036 ttfb_ns: None,
1037 tx_bytes,
1038 rx_bytes,
1039 });
1040 result
1041 }
1042
1043 pub async fn add(
1046 &mut self,
1047 key: impl AsRef<[u8]>,
1048 value: impl AsRef<[u8]>,
1049 ) -> Result<bool, Error> {
1050 let key = key.as_ref();
1051 let value = value.as_ref();
1052 let encoded = encode_add(key, value)?;
1053 let response = self.execute(&encoded).await?;
1054 match response {
1055 McResponseBytes::Stored => Ok(true),
1056 McResponseBytes::NotStored => Ok(false),
1057 _ => Err(Error::UnexpectedResponse),
1058 }
1059 }
1060
1061 pub async fn replace(
1064 &mut self,
1065 key: impl AsRef<[u8]>,
1066 value: impl AsRef<[u8]>,
1067 ) -> Result<bool, Error> {
1068 let key = key.as_ref();
1069 let value = value.as_ref();
1070 let encoded = encode_request(&McRequest::Replace {
1071 key,
1072 value,
1073 flags: 0,
1074 exptime: 0,
1075 })?;
1076 let response = self.execute(&encoded).await?;
1077 match response {
1078 McResponseBytes::Stored => Ok(true),
1079 McResponseBytes::NotStored => Ok(false),
1080 _ => Err(Error::UnexpectedResponse),
1081 }
1082 }
1083
1084 pub async fn incr(&mut self, key: impl AsRef<[u8]>, delta: u64) -> Result<Option<u64>, Error> {
1087 let key = key.as_ref();
1088 let encoded = encode_request(&McRequest::incr(key, delta))?;
1089 let response = self.execute(&encoded).await?;
1090 match response {
1091 McResponseBytes::Numeric(val) => Ok(Some(val)),
1092 McResponseBytes::NotFound => Ok(None),
1093 _ => Err(Error::UnexpectedResponse),
1094 }
1095 }
1096
1097 pub async fn decr(&mut self, key: impl AsRef<[u8]>, delta: u64) -> Result<Option<u64>, Error> {
1100 let key = key.as_ref();
1101 let encoded = encode_request(&McRequest::decr(key, delta))?;
1102 let response = self.execute(&encoded).await?;
1103 match response {
1104 McResponseBytes::Numeric(val) => Ok(Some(val)),
1105 McResponseBytes::NotFound => Ok(None),
1106 _ => Err(Error::UnexpectedResponse),
1107 }
1108 }
1109
1110 pub async fn append(
1113 &mut self,
1114 key: impl AsRef<[u8]>,
1115 value: impl AsRef<[u8]>,
1116 ) -> Result<bool, Error> {
1117 let key = key.as_ref();
1118 let value = value.as_ref();
1119 let encoded = encode_request(&McRequest::append(key, value))?;
1120 let response = self.execute(&encoded).await?;
1121 match response {
1122 McResponseBytes::Stored => Ok(true),
1123 McResponseBytes::NotStored => Ok(false),
1124 _ => Err(Error::UnexpectedResponse),
1125 }
1126 }
1127
1128 pub async fn prepend(
1131 &mut self,
1132 key: impl AsRef<[u8]>,
1133 value: impl AsRef<[u8]>,
1134 ) -> Result<bool, Error> {
1135 let key = key.as_ref();
1136 let value = value.as_ref();
1137 let encoded = encode_request(&McRequest::prepend(key, value))?;
1138 let response = self.execute(&encoded).await?;
1139 match response {
1140 McResponseBytes::Stored => Ok(true),
1141 McResponseBytes::NotStored => Ok(false),
1142 _ => Err(Error::UnexpectedResponse),
1143 }
1144 }
1145
1146 pub async fn cas(
1150 &mut self,
1151 key: impl AsRef<[u8]>,
1152 value: impl AsRef<[u8]>,
1153 cas_unique: u64,
1154 ) -> Result<bool, Error> {
1155 let key = key.as_ref();
1156 let value = value.as_ref();
1157 let encoded = encode_request(&McRequest::cas(key, value, cas_unique))?;
1158 let response = self.execute(&encoded).await?;
1159 match response {
1160 McResponseBytes::Stored => Ok(true),
1161 McResponseBytes::Exists => Ok(false),
1162 McResponseBytes::NotFound => Err(Error::Memcache("NOT_FOUND".into())),
1163 _ => Err(Error::UnexpectedResponse),
1164 }
1165 }
1166
1167 pub async fn delete(&mut self, key: impl AsRef<[u8]>) -> Result<bool, Error> {
1169 let key = key.as_ref();
1170 let encoded = encode_request(&McRequest::delete(key))?;
1171
1172 if !self.is_instrumented() {
1173 let response = self.execute(&encoded).await?;
1174 return match response {
1175 McResponseBytes::Deleted => Ok(true),
1176 McResponseBytes::NotFound => Ok(false),
1177 _ => Err(Error::UnexpectedResponse),
1178 };
1179 }
1180
1181 let tx_bytes = encoded.len() as u32;
1182 let send_ts = self.send_timestamp();
1183 let start = Instant::now();
1184 let response = self.execute(&encoded).await;
1185 let latency_ns = self.finish_timing(send_ts, start);
1186 let rx_bytes = self.last_rx_bytes.get();
1187
1188 let result = match response {
1189 Ok(McResponseBytes::Deleted) => Ok(true),
1190 Ok(McResponseBytes::NotFound) => Ok(false),
1191 Ok(_) => Err(Error::UnexpectedResponse),
1192 Err(e) => Err(e),
1193 };
1194 self.record(&CommandResult {
1195 command: CommandType::Delete,
1196 latency_ns,
1197 hit: None,
1198 success: result.is_ok(),
1199 ttfb_ns: None,
1200 tx_bytes,
1201 rx_bytes,
1202 });
1203 result
1204 }
1205
1206 pub async fn flush_all(&mut self) -> Result<(), Error> {
1208 let encoded = encode_request(&McRequest::flush_all())?;
1209 let response = self.execute(&encoded).await?;
1210 match response {
1211 McResponseBytes::Ok => Ok(()),
1212 _ => Err(Error::UnexpectedResponse),
1213 }
1214 }
1215
1216 pub async fn version(&mut self) -> Result<Box<str>, Error> {
1218 let encoded = encode_request(&McRequest::version())?;
1219 let response = self.execute(&encoded).await?;
1220 match response {
1221 McResponseBytes::Version(v) => Ok(Box::from(String::from_utf8_lossy(v.as_ref()))),
1222 _ => Err(Error::UnexpectedResponse),
1223 }
1224 }
1225
1226 pub async fn set_with_guard<G: SendGuard>(
1231 &mut self,
1232 key: &[u8],
1233 guard: G,
1234 flags: u32,
1235 exptime: u32,
1236 ) -> Result<(), Error> {
1237 if !self.is_instrumented() {
1238 let (_, value_len) = guard.as_ptr_len();
1239 let prefix = encode_set_guard_prefix(key, value_len as usize, flags, exptime)?;
1240
1241 self.conn.send_parts().build(move |b| {
1242 b.copy(&prefix)
1243 .guard(GuardBox::new(guard))
1244 .copy(b"\r\n")
1245 .submit()
1246 })?;
1247
1248 let response = self.read_response().await?;
1249 check_error_bytes(&response)?;
1250 return match response {
1251 McResponseBytes::Stored => Ok(()),
1252 _ => Err(Error::UnexpectedResponse),
1253 };
1254 }
1255
1256 let (_, value_len) = guard.as_ptr_len();
1257 let prefix = encode_set_guard_prefix(key, value_len as usize, flags, exptime)?;
1258 let tx_bytes = (prefix.len() + value_len as usize + 2) as u32;
1259
1260 let send_ts = self.send_timestamp();
1261 let start = Instant::now();
1262
1263 self.conn.send_parts().build(move |b| {
1264 b.copy(&prefix)
1265 .guard(GuardBox::new(guard))
1266 .copy(b"\r\n")
1267 .submit()
1268 })?;
1269
1270 let response = self.read_response().await;
1271 let latency_ns = self.finish_timing(send_ts, start);
1272 let rx_bytes = self.last_rx_bytes.get();
1273
1274 let result = match response {
1275 Ok(ref r) => {
1276 check_error_bytes(r)?;
1277 match r {
1278 McResponseBytes::Stored => Ok(()),
1279 _ => Err(Error::UnexpectedResponse),
1280 }
1281 }
1282 Err(e) => Err(e),
1283 };
1284 self.record(&CommandResult {
1285 command: CommandType::Set,
1286 latency_ns,
1287 hit: None,
1288 success: result.is_ok(),
1289 ttfb_ns: None,
1290 tx_bytes,
1291 rx_bytes,
1292 });
1293 result
1294 }
1295}
1296
1297fn encode_set_guard_prefix(
1308 key: &[u8],
1309 value_len: usize,
1310 flags: u32,
1311 exptime: u32,
1312) -> Result<Vec<u8>, Error> {
1313 use std::io::Write;
1314 validate_key(key)?;
1315 validate_value_len(value_len)?;
1316 let mut buf = Vec::with_capacity(32 + key.len());
1317 buf.extend_from_slice(b"set ");
1318 buf.extend_from_slice(key);
1319 write!(buf, " {} {} {}\r\n", flags, exptime, value_len).unwrap();
1320 Ok(buf)
1321}
1322
1323fn validate_request(req: &McRequest<'_>) -> Result<(), Error> {
1327 match req {
1328 McRequest::Get { key }
1329 | McRequest::Incr { key, .. }
1330 | McRequest::Decr { key, .. }
1331 | McRequest::Delete { key } => validate_key(key),
1332 McRequest::Gets { keys } => {
1333 for k in keys.iter() {
1334 validate_key(k)?;
1335 }
1336 Ok(())
1337 }
1338 McRequest::Set { key, value, .. }
1339 | McRequest::Add { key, value, .. }
1340 | McRequest::Replace { key, value, .. }
1341 | McRequest::Append { key, value }
1342 | McRequest::Prepend { key, value }
1343 | McRequest::Cas { key, value, .. } => {
1344 validate_key(key)?;
1345 validate_value(value)
1346 }
1347 McRequest::FlushAll | McRequest::Version | McRequest::Quit => Ok(()),
1348 }
1349}
1350
1351pub(crate) fn encode_request(req: &McRequest<'_>) -> Result<Vec<u8>, Error> {
1354 validate_request(req)?;
1355 let size = match req {
1356 McRequest::Get { key } => 6 + key.len(),
1357 McRequest::Gets { keys } => 6 + keys.iter().map(|k| 1 + k.len()).sum::<usize>(),
1358 McRequest::Set { key, value, .. } | McRequest::Add { key, value, .. } => {
1359 41 + key.len() + value.len()
1360 }
1361 McRequest::Replace { key, value, .. } => 45 + key.len() + value.len(),
1362 McRequest::Incr { key, .. } | McRequest::Decr { key, .. } => 27 + key.len(),
1363 McRequest::Append { key, value } => 44 + key.len() + value.len(),
1364 McRequest::Prepend { key, value } => 45 + key.len() + value.len(),
1365 McRequest::Cas { key, value, .. } => 61 + key.len() + value.len(),
1366 McRequest::Delete { key } => 9 + key.len(),
1367 McRequest::FlushAll => 11,
1368 McRequest::Version => 9,
1369 McRequest::Quit => 6,
1370 };
1371 let mut buf = vec![0u8; size];
1372 let len = req.encode(&mut buf);
1373 buf.truncate(len);
1374 Ok(buf)
1375}
1376
1377pub(crate) fn encode_set(
1379 key: &[u8],
1380 value: &[u8],
1381 flags: u32,
1382 exptime: u32,
1383) -> Result<Vec<u8>, Error> {
1384 encode_request(&McRequest::Set {
1385 key,
1386 value,
1387 flags,
1388 exptime,
1389 })
1390}
1391
1392pub(crate) fn encode_add(key: &[u8], value: &[u8]) -> Result<Vec<u8>, Error> {
1394 encode_request(&McRequest::Add {
1395 key,
1396 value,
1397 flags: 0,
1398 exptime: 0,
1399 })
1400}
1401
1402pub(crate) fn check_error_bytes(response: &McResponseBytes) -> Result<(), Error> {
1404 match response {
1405 McResponseBytes::Error => Err(Error::Memcache("ERROR".into())),
1406 McResponseBytes::ClientError(msg) => Err(Error::Memcache(format!(
1407 "CLIENT_ERROR {}",
1408 String::from_utf8_lossy(msg)
1409 ))),
1410 McResponseBytes::ServerError(msg) => Err(Error::Memcache(format!(
1411 "SERVER_ERROR {}",
1412 String::from_utf8_lossy(msg)
1413 ))),
1414 _ => Ok(()),
1415 }
1416}
1417
1418#[cfg(feature = "timestamps")]
1421fn now_realtime_ns() -> u64 {
1422 let mut ts = libc::timespec {
1423 tv_sec: 0,
1424 tv_nsec: 0,
1425 };
1426 unsafe {
1427 libc::clock_gettime(libc::CLOCK_REALTIME, &mut ts);
1428 }
1429 ts.tv_sec as u64 * 1_000_000_000 + ts.tv_nsec as u64
1430}
1431
1432#[cfg(test)]
1433mod tests {
1434 use super::*;
1435
1436 #[test]
1437 fn validate_key_accepts_max_len() {
1438 let key = vec![b'k'; MAX_KEY_LEN];
1439 assert!(validate_key(&key).is_ok());
1440 }
1441
1442 #[test]
1443 fn validate_key_rejects_oversized() {
1444 let key = vec![b'k'; MAX_KEY_LEN + 1];
1445 assert!(matches!(validate_key(&key), Err(Error::KeyTooLong)));
1446 }
1447
1448 #[test]
1449 fn validate_value_accepts_max_len() {
1450 let v = vec![0u8; MAX_VALUE_LEN];
1451 assert!(validate_value(&v).is_ok());
1452 }
1453
1454 #[test]
1455 fn validate_value_rejects_oversized() {
1456 let v = vec![0u8; MAX_VALUE_LEN + 1];
1457 assert!(matches!(validate_value(&v), Err(Error::ValueTooLong)));
1458 }
1459
1460 #[test]
1461 fn encode_request_get_rejects_long_key() {
1462 let key = vec![b'k'; MAX_KEY_LEN + 1];
1463 let r = encode_request(&McRequest::get(&key));
1464 assert!(matches!(r, Err(Error::KeyTooLong)));
1465 }
1466
1467 #[test]
1468 fn encode_request_set_rejects_long_value() {
1469 let key = b"k";
1470 let value = vec![0u8; MAX_VALUE_LEN + 1];
1471 let r = encode_request(&McRequest::Set {
1472 key,
1473 value: &value,
1474 flags: 0,
1475 exptime: 0,
1476 });
1477 assert!(matches!(r, Err(Error::ValueTooLong)));
1478 }
1479
1480 #[test]
1481 fn encode_request_gets_rejects_any_long_key() {
1482 let ok = vec![b'k'; MAX_KEY_LEN];
1483 let bad = vec![b'k'; MAX_KEY_LEN + 1];
1484 let keys: &[&[u8]] = &[&ok, &bad];
1485 let r = encode_request(&McRequest::gets(keys));
1486 assert!(matches!(r, Err(Error::KeyTooLong)));
1487 }
1488
1489 #[test]
1490 fn encode_set_guard_prefix_rejects_long_key() {
1491 let key = vec![b'k'; MAX_KEY_LEN + 1];
1492 let r = encode_set_guard_prefix(&key, 16, 0, 0);
1493 assert!(matches!(r, Err(Error::KeyTooLong)));
1494 }
1495
1496 #[test]
1497 fn encode_set_guard_prefix_rejects_long_value_len() {
1498 let r = encode_set_guard_prefix(b"k", MAX_VALUE_LEN + 1, 0, 0);
1499 assert!(matches!(r, Err(Error::ValueTooLong)));
1500 }
1501
1502 #[test]
1503 fn encode_set_guard_prefix_accepts_max_value_len() {
1504 let r = encode_set_guard_prefix(b"k", MAX_VALUE_LEN, 0, 0);
1505 assert!(r.is_ok());
1506 }
1507
1508 #[test]
1509 fn encode_request_passes_through_at_caps() {
1510 let key = vec![b'k'; MAX_KEY_LEN];
1511 let value = vec![0u8; MAX_VALUE_LEN];
1512 let r = encode_request(&McRequest::Set {
1513 key: &key,
1514 value: &value,
1515 flags: 0,
1516 exptime: 0,
1517 });
1518 assert!(r.is_ok());
1519 let buf = r.unwrap();
1520 assert!(buf.starts_with(b"set "));
1522 }
1523}