pub mod pool;
pub mod sharded;
pub use pool::{Pool, PoolConfig};
pub use sharded::{ShardedClient, ShardedConfig};
use std::cell::Cell;
use std::collections::VecDeque;
use std::io;
use std::time::Instant;
use bytes::Bytes;
use memcache_proto::{Request as McRequest, ResponseBytes as McResponseBytes};
use ringline::{ConnCtx, GuardBox, ParseResult, SendGuard};
type ResultCallback = Box<dyn Fn(&CommandResult)>;
const MAX_FLUSH_GUARDS: usize = 8;
const MAX_FLUSH_IOVECS: usize = 2 * MAX_FLUSH_GUARDS + 1;
const DEFAULT_ZC_THRESHOLD: u32 = 4096;
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum Error {
#[error("connection closed")]
ConnectionClosed,
#[error("memcache error: {0}")]
Memcache(String),
#[error("unexpected response")]
UnexpectedResponse,
#[error("protocol error: {0}")]
Protocol(#[from] memcache_proto::ParseError),
#[error("io error: {0}")]
Io(#[from] io::Error),
#[error("all connections failed")]
AllConnectionsFailed,
#[error("no pending operations")]
NoPending,
#[error("too many in-flight operations")]
TooManyInFlight,
#[error("key too long (max 250 bytes)")]
KeyTooLong,
#[error("value too long (max 1048576 bytes)")]
ValueTooLong,
}
pub const MAX_KEY_LEN: usize = 250;
pub const MAX_VALUE_LEN: usize = 1024 * 1024;
#[inline]
fn validate_key(key: &[u8]) -> Result<(), Error> {
if key.len() > MAX_KEY_LEN {
Err(Error::KeyTooLong)
} else {
Ok(())
}
}
#[inline]
fn validate_value(value: &[u8]) -> Result<(), Error> {
if value.len() > MAX_VALUE_LEN {
Err(Error::ValueTooLong)
} else {
Ok(())
}
}
#[inline]
fn validate_value_len(value_len: usize) -> Result<(), Error> {
if value_len > MAX_VALUE_LEN {
Err(Error::ValueTooLong)
} else {
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct Value {
pub data: Bytes,
pub flags: u32,
}
#[derive(Debug, Clone)]
pub struct GetValue {
pub key: Bytes,
pub data: Bytes,
pub flags: u32,
pub cas: Option<u64>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum CommandType {
Get,
Set,
Delete,
Other,
}
#[derive(Debug, Clone)]
pub struct CommandResult {
pub command: CommandType,
pub latency_ns: u64,
pub hit: Option<bool>,
pub success: bool,
pub ttfb_ns: Option<u64>,
pub tx_bytes: u32,
pub rx_bytes: u32,
}
#[cfg(feature = "metrics")]
pub struct ClientMetrics {
pub latency: histogram::Histogram,
pub get_latency: histogram::Histogram,
pub set_latency: histogram::Histogram,
pub del_latency: histogram::Histogram,
pub requests: u64,
pub errors: u64,
pub hits: u64,
pub misses: u64,
}
#[cfg(feature = "metrics")]
impl ClientMetrics {
fn new() -> Self {
Self {
latency: histogram::Histogram::new(7, 64).unwrap(),
get_latency: histogram::Histogram::new(7, 64).unwrap(),
set_latency: histogram::Histogram::new(7, 64).unwrap(),
del_latency: histogram::Histogram::new(7, 64).unwrap(),
requests: 0,
errors: 0,
hits: 0,
misses: 0,
}
}
fn record(&mut self, result: &CommandResult) {
self.requests += 1;
let _ = self.latency.increment(result.latency_ns);
if !result.success {
self.errors += 1;
}
match result.command {
CommandType::Get => {
let _ = self.get_latency.increment(result.latency_ns);
match result.hit {
Some(true) => self.hits += 1,
Some(false) => self.misses += 1,
None => {}
}
}
CommandType::Set => {
let _ = self.set_latency.increment(result.latency_ns);
}
CommandType::Delete => {
let _ = self.del_latency.increment(result.latency_ns);
}
_ => {}
}
}
}
enum PendingOpKind {
Get,
Set,
Delete,
}
struct PendingOp {
kind: PendingOpKind,
send_ts: u64,
start: Option<Instant>,
user_data: u64,
tx_bytes: u32,
}
#[non_exhaustive]
pub enum CompletedOp {
Get {
result: Result<Option<Value>, Error>,
user_data: u64,
latency_ns: u64,
},
Set {
result: Result<(), Error>,
user_data: u64,
latency_ns: u64,
},
Delete {
result: Result<bool, Error>,
user_data: u64,
latency_ns: u64,
},
}
pub struct ClientBuilder {
conn: ConnCtx,
on_result: Option<ResultCallback>,
max_batch_size: usize,
max_in_flight: usize,
zc_threshold: u32,
#[cfg(feature = "timestamps")]
use_kernel_ts: bool,
#[cfg(feature = "metrics")]
with_metrics: bool,
}
impl ClientBuilder {
pub(crate) fn new(conn: ConnCtx) -> Self {
Self {
conn,
on_result: None,
max_batch_size: 1,
max_in_flight: usize::MAX,
zc_threshold: DEFAULT_ZC_THRESHOLD,
#[cfg(feature = "timestamps")]
use_kernel_ts: false,
#[cfg(feature = "metrics")]
with_metrics: false,
}
}
pub fn zc_threshold(mut self, bytes: u32) -> Self {
self.zc_threshold = bytes;
self
}
pub fn max_in_flight(mut self, n: usize) -> Self {
self.max_in_flight = n;
self
}
pub fn on_result<F: Fn(&CommandResult) + 'static>(mut self, f: F) -> Self {
self.on_result = Some(Box::new(f));
self
}
#[cfg(feature = "timestamps")]
pub fn kernel_timestamps(mut self, enabled: bool) -> Self {
self.use_kernel_ts = enabled;
self
}
pub fn max_batch_size(mut self, n: usize) -> Self {
assert!(n > 0, "max_batch_size must be >= 1");
self.max_batch_size = n;
self
}
#[cfg(feature = "metrics")]
pub fn with_metrics(mut self) -> Self {
self.with_metrics = true;
self
}
pub fn build(self) -> Client {
Client {
conn: self.conn,
on_result: self.on_result,
pending: VecDeque::with_capacity(16),
last_rx_bytes: Cell::new(0),
write_buf: Vec::new(),
write_guards: Vec::new(),
flushed_count: 0,
max_batch_size: self.max_batch_size,
max_in_flight: self.max_in_flight,
zc_threshold: self.zc_threshold,
buffered_ops: 0,
encode_buf: Vec::new(),
#[cfg(feature = "timestamps")]
use_kernel_ts: self.use_kernel_ts,
#[cfg(feature = "metrics")]
metrics: if self.with_metrics {
Some(ClientMetrics::new())
} else {
None
},
}
}
}
pub struct Client {
conn: ConnCtx,
on_result: Option<ResultCallback>,
pending: VecDeque<PendingOp>,
last_rx_bytes: Cell<u32>,
write_buf: Vec<u8>,
write_guards: Vec<(usize, GuardBox)>,
flushed_count: usize,
max_batch_size: usize,
max_in_flight: usize,
zc_threshold: u32,
buffered_ops: usize,
encode_buf: Vec<u8>,
#[cfg(feature = "timestamps")]
use_kernel_ts: bool,
#[cfg(feature = "metrics")]
metrics: Option<ClientMetrics>,
}
impl Client {
pub fn new(conn: ConnCtx) -> Self {
Self {
conn,
on_result: None,
pending: VecDeque::new(),
last_rx_bytes: Cell::new(0),
write_buf: Vec::new(),
write_guards: Vec::new(),
flushed_count: 0,
max_batch_size: 1,
max_in_flight: usize::MAX,
zc_threshold: DEFAULT_ZC_THRESHOLD,
buffered_ops: 0,
encode_buf: Vec::new(),
#[cfg(feature = "timestamps")]
use_kernel_ts: false,
#[cfg(feature = "metrics")]
metrics: None,
}
}
pub fn builder(conn: ConnCtx) -> ClientBuilder {
ClientBuilder::new(conn)
}
pub fn conn(&self) -> ConnCtx {
self.conn
}
#[cfg(feature = "metrics")]
pub fn metrics(&self) -> Option<&ClientMetrics> {
self.metrics.as_ref()
}
#[cfg(feature = "metrics")]
pub fn metrics_mut(&mut self) -> Option<&mut ClientMetrics> {
self.metrics.as_mut()
}
#[inline]
fn is_instrumented(&self) -> bool {
if self.on_result.is_some() {
return true;
}
#[cfg(feature = "metrics")]
if self.metrics.is_some() {
return true;
}
false
}
#[cfg(feature = "timestamps")]
#[inline]
fn send_timestamp(&self) -> u64 {
if self.use_kernel_ts {
now_realtime_ns()
} else {
0
}
}
#[cfg(not(feature = "timestamps"))]
#[inline]
fn send_timestamp(&self) -> u64 {
0
}
#[cfg(feature = "timestamps")]
#[inline]
fn finish_timing(&self, send_ts: u64, start: Instant) -> u64 {
if self.use_kernel_ts {
let recv_ts = self.conn.recv_timestamp();
if recv_ts > 0 && recv_ts > send_ts {
return recv_ts - send_ts;
}
}
start.elapsed().as_nanos() as u64
}
#[cfg(not(feature = "timestamps"))]
#[inline]
fn finish_timing(&self, _send_ts: u64, start: Instant) -> u64 {
start.elapsed().as_nanos() as u64
}
fn record(&mut self, result: &CommandResult) {
if let Some(ref cb) = self.on_result {
cb(result);
}
#[cfg(feature = "metrics")]
if let Some(ref mut m) = self.metrics {
m.record(result);
}
}
#[inline]
fn timing_start(&self) -> (u64, Option<Instant>) {
#[cfg(feature = "timestamps")]
{
if self.is_instrumented() {
(self.send_timestamp(), Some(Instant::now()))
} else {
(0, None)
}
}
#[cfg(not(feature = "timestamps"))]
{
if self.on_result.is_some() {
(0, Some(Instant::now()))
} else {
(0, None)
}
}
}
#[inline]
fn timing_start_buffered(&self) -> (u64, Option<Instant>) {
if self.buffered_ops >= 1 {
(0, None)
} else {
self.timing_start()
}
}
#[cfg(feature = "timestamps")]
#[inline]
fn compute_ttfb(&self, send_ts: u64) -> Option<u64> {
if self.use_kernel_ts {
let recv_ts = self.conn.recv_timestamp();
if recv_ts > 0 && recv_ts > send_ts {
return Some(recv_ts - send_ts);
}
}
None
}
#[cfg(not(feature = "timestamps"))]
#[inline]
fn compute_ttfb(&self, _send_ts: u64) -> Option<u64> {
None
}
pub fn pending_count(&self) -> usize {
self.pending.len()
}
fn pre_flush_if_needed(&mut self, has_guard: bool) -> Result<(), Error> {
if self.pending.len() >= self.max_in_flight {
return Err(Error::TooManyInFlight);
}
if self.buffered_ops == 0 {
return Ok(());
}
let next_guards = self.write_guards.len() + usize::from(has_guard);
let next_parts = 2 * next_guards + 1;
if next_guards > MAX_FLUSH_GUARDS || next_parts > MAX_FLUSH_IOVECS {
self.flush()?;
}
Ok(())
}
fn post_flush_if_needed(&mut self) -> Result<(), Error> {
if self.buffered_ops >= self.max_batch_size {
self.flush()?;
}
Ok(())
}
pub fn flush(&mut self) -> Result<(), Error> {
if self.write_buf.is_empty() && self.write_guards.is_empty() {
self.buffered_ops = 0;
return Ok(());
}
let send_outcome: Result<(), Error> = if self.write_guards.is_empty() {
self.conn
.send_nowait(&self.write_buf)
.map(|_| ())
.map_err(Error::from)
} else {
use ringline::SendPart;
let mut parts: Vec<SendPart<'_>> = Vec::with_capacity(2 * MAX_FLUSH_GUARDS + 1);
let mut pos = 0;
for (offset, guard) in self.write_guards.drain(..) {
if offset > pos {
parts.push(SendPart::Copy(&self.write_buf[pos..offset]));
}
parts.push(SendPart::Guard(guard));
pos = offset;
}
if pos < self.write_buf.len() {
parts.push(SendPart::Copy(&self.write_buf[pos..]));
}
self.conn
.send_parts()
.submit_batch(parts)
.map(|_| ())
.map_err(Error::from)
};
if let Err(e) = send_outcome {
self.write_buf.clear();
self.write_guards.clear();
self.buffered_ops = 0;
self.pending.truncate(self.flushed_count);
return Err(e);
}
if self.buffered_ops >= 1 {
let (send_ts, start) = self.timing_start();
for pending in self.pending.iter_mut().skip(self.flushed_count) {
pending.send_ts = send_ts;
pending.start = start;
}
}
self.flushed_count = self.pending.len();
self.write_buf.clear();
self.write_guards.clear();
self.buffered_ops = 0;
Ok(())
}
pub fn fire_get(&mut self, key: &[u8], user_data: u64) -> Result<(), Error> {
self.pre_flush_if_needed(false)?;
let tx_bytes;
if self.max_batch_size == 1 && self.write_guards.is_empty() && self.write_buf.is_empty() {
self.encode_buf.clear();
encode_request_into(&McRequest::get(key), &mut self.encode_buf)?;
tx_bytes = self.encode_buf.len() as u32;
self.conn.send_nowait(&self.encode_buf)?;
} else {
let before = self.write_buf.len();
encode_request_into(&McRequest::get(key), &mut self.write_buf)?;
tx_bytes = (self.write_buf.len() - before) as u32;
self.buffered_ops += 1;
}
let (send_ts, start) = self.timing_start_buffered();
self.pending.push_back(PendingOp {
kind: PendingOpKind::Get,
send_ts,
start,
user_data,
tx_bytes,
});
if self.max_batch_size == 1 && self.write_guards.is_empty() && self.write_buf.is_empty() {
self.flushed_count += 1;
}
self.post_flush_if_needed()?;
Ok(())
}
pub fn fire_set(
&mut self,
key: &[u8],
value: &[u8],
flags: u32,
exptime: u32,
user_data: u64,
) -> Result<(), Error> {
self.pre_flush_if_needed(false)?;
let req = McRequest::Set {
key,
value,
flags,
exptime,
};
let tx_bytes;
if self.max_batch_size == 1 && self.write_guards.is_empty() && self.write_buf.is_empty() {
self.encode_buf.clear();
encode_request_into(&req, &mut self.encode_buf)?;
tx_bytes = self.encode_buf.len() as u32;
self.conn.send_nowait(&self.encode_buf)?;
} else {
let before = self.write_buf.len();
encode_request_into(&req, &mut self.write_buf)?;
tx_bytes = (self.write_buf.len() - before) as u32;
self.buffered_ops += 1;
}
let (send_ts, start) = self.timing_start_buffered();
self.pending.push_back(PendingOp {
kind: PendingOpKind::Set,
send_ts,
start,
user_data,
tx_bytes,
});
if self.max_batch_size == 1 && self.write_guards.is_empty() && self.write_buf.is_empty() {
self.flushed_count += 1;
}
self.post_flush_if_needed()?;
Ok(())
}
pub fn fire_set_with_guard<G: SendGuard>(
&mut self,
key: &[u8],
guard: G,
flags: u32,
exptime: u32,
user_data: u64,
) -> Result<(), Error> {
let (ptr, value_len) = guard.as_ptr_len();
if self.zc_threshold != 0 && value_len < self.zc_threshold {
let value = unsafe { core::slice::from_raw_parts(ptr, value_len as usize) };
return self.fire_set(key, value, flags, exptime, user_data);
}
self.pre_flush_if_needed(true)?;
let before = self.write_buf.len();
append_set_guard_prefix(&mut self.write_buf, key, value_len as usize, flags, exptime)?;
self.write_guards
.push((self.write_buf.len(), GuardBox::new(guard)));
self.write_buf.extend_from_slice(b"\r\n");
let tx_bytes = (self.write_buf.len() - before + value_len as usize) as u32;
self.buffered_ops += 1;
let (send_ts, start) = self.timing_start_buffered();
self.pending.push_back(PendingOp {
kind: PendingOpKind::Set,
send_ts,
start,
user_data,
tx_bytes,
});
self.post_flush_if_needed()?;
Ok(())
}
pub fn fire_delete(&mut self, key: &[u8], user_data: u64) -> Result<(), Error> {
self.pre_flush_if_needed(false)?;
let tx_bytes;
if self.max_batch_size == 1 && self.write_guards.is_empty() && self.write_buf.is_empty() {
self.encode_buf.clear();
encode_request_into(&McRequest::delete(key), &mut self.encode_buf)?;
tx_bytes = self.encode_buf.len() as u32;
self.conn.send_nowait(&self.encode_buf)?;
} else {
let before = self.write_buf.len();
encode_request_into(&McRequest::delete(key), &mut self.write_buf)?;
tx_bytes = (self.write_buf.len() - before) as u32;
self.buffered_ops += 1;
}
let (send_ts, start) = self.timing_start_buffered();
self.pending.push_back(PendingOp {
kind: PendingOpKind::Delete,
send_ts,
start,
user_data,
tx_bytes,
});
if self.max_batch_size == 1 && self.write_guards.is_empty() && self.write_buf.is_empty() {
self.flushed_count += 1;
}
self.post_flush_if_needed()?;
Ok(())
}
pub async fn recv(&mut self) -> Result<CompletedOp, Error> {
self.flush()?;
let pending = self.pending.pop_front().ok_or(Error::NoPending)?;
self.flushed_count = self.flushed_count.saturating_sub(1);
let response = match self.read_response().await {
Ok(v) => v,
Err(e) => {
self.pending.clear();
self.flushed_count = 0;
return Err(e);
}
};
let ttfb_ns = self.compute_ttfb(pending.send_ts);
let latency_ns = match pending.start {
Some(start) => self.finish_timing(pending.send_ts, start),
None => 0,
};
let rx_bytes = self.last_rx_bytes.get();
let tx_bytes = pending.tx_bytes;
let op = match pending.kind {
PendingOpKind::Get => {
let result = match check_error_bytes(&response) {
Err(e) => Err(e),
Ok(()) => match response {
McResponseBytes::Values(mut values) => {
if values.is_empty() {
Ok(None)
} else {
let v = values.swap_remove(0);
Ok(Some(Value {
data: v.data,
flags: v.flags,
}))
}
}
_ => Err(Error::UnexpectedResponse),
},
};
let (success, hit) = match &result {
Ok(Some(_)) => (true, Some(true)),
Ok(None) => (true, Some(false)),
Err(_) => (false, None),
};
self.record(&CommandResult {
command: CommandType::Get,
latency_ns,
hit,
success,
ttfb_ns,
tx_bytes,
rx_bytes,
});
CompletedOp::Get {
result,
user_data: pending.user_data,
latency_ns,
}
}
PendingOpKind::Set => {
let result = match check_error_bytes(&response) {
Err(e) => Err(e),
Ok(()) => match response {
McResponseBytes::Stored => Ok(()),
_ => Err(Error::UnexpectedResponse),
},
};
self.record(&CommandResult {
command: CommandType::Set,
latency_ns,
hit: None,
success: result.is_ok(),
ttfb_ns,
tx_bytes,
rx_bytes,
});
CompletedOp::Set {
result,
user_data: pending.user_data,
latency_ns,
}
}
PendingOpKind::Delete => {
let result = match check_error_bytes(&response) {
Err(e) => Err(e),
Ok(()) => match response {
McResponseBytes::Deleted => Ok(true),
McResponseBytes::NotFound => Ok(false),
_ => Err(Error::UnexpectedResponse),
},
};
self.record(&CommandResult {
command: CommandType::Delete,
latency_ns,
hit: None,
success: result.is_ok(),
ttfb_ns,
tx_bytes,
rx_bytes,
});
CompletedOp::Delete {
result,
user_data: pending.user_data,
latency_ns,
}
}
};
Ok(op)
}
pub(crate) async fn read_response(&self) -> Result<McResponseBytes, Error> {
let mut result: Option<Result<McResponseBytes, Error>> = None;
let n = self
.conn
.with_bytes(|bytes| {
let len = bytes.len();
match McResponseBytes::parse(bytes) {
Ok((response, consumed)) => {
result = Some(Ok(response));
ParseResult::Consumed(consumed)
}
Err(e) if e.is_incomplete() => ParseResult::Consumed(0),
Err(e) => {
result = Some(Err(Error::Protocol(e)));
ParseResult::Consumed(len)
}
}
})
.await;
self.last_rx_bytes.set(n as u32);
if n == 0 {
return result.unwrap_or(Err(Error::ConnectionClosed));
}
let r = result.unwrap();
if matches!(r, Err(Error::Protocol(_))) {
self.conn.close();
}
r
}
async fn execute(&self, encoded: &[u8]) -> Result<McResponseBytes, Error> {
self.conn.send(encoded)?;
let response = self.read_response().await?;
check_error_bytes(&response)?;
Ok(response)
}
pub async fn get(&mut self, key: impl AsRef<[u8]>) -> Result<Option<Value>, Error> {
let key = key.as_ref();
self.encode_buf.clear();
encode_request_into(&McRequest::get(key), &mut self.encode_buf)?;
if !self.is_instrumented() {
let response = self.execute(&self.encode_buf).await?;
return match response {
McResponseBytes::Values(mut values) => {
if values.is_empty() {
Ok(None)
} else {
let v = values.swap_remove(0);
Ok(Some(Value {
data: v.data,
flags: v.flags,
}))
}
}
_ => Err(Error::UnexpectedResponse),
};
}
let tx_bytes = self.encode_buf.len() as u32;
let send_ts = self.send_timestamp();
let start = Instant::now();
let response = self.execute(&self.encode_buf).await;
let latency_ns = self.finish_timing(send_ts, start);
let rx_bytes = self.last_rx_bytes.get();
let result = match response {
Ok(McResponseBytes::Values(mut values)) => {
if values.is_empty() {
Ok(None)
} else {
let v = values.swap_remove(0);
Ok(Some(Value {
data: v.data,
flags: v.flags,
}))
}
}
Ok(_) => Err(Error::UnexpectedResponse),
Err(e) => Err(e),
};
let (success, hit) = match &result {
Ok(Some(_)) => (true, Some(true)),
Ok(None) => (true, Some(false)),
Err(_) => (false, None),
};
self.record(&CommandResult {
command: CommandType::Get,
latency_ns,
hit,
success,
ttfb_ns: None,
tx_bytes,
rx_bytes,
});
result
}
pub async fn gets(&mut self, keys: &[&[u8]]) -> Result<Vec<GetValue>, Error> {
if keys.is_empty() {
return Ok(Vec::new());
}
let encoded = encode_request(&McRequest::gets(keys))?;
let response = self.execute(&encoded).await?;
match response {
McResponseBytes::Values(values) => Ok(values
.into_iter()
.map(|v| GetValue {
key: v.key,
data: v.data,
flags: v.flags,
cas: v.cas,
})
.collect()),
_ => Err(Error::UnexpectedResponse),
}
}
pub async fn set(
&mut self,
key: impl AsRef<[u8]>,
value: impl AsRef<[u8]>,
) -> Result<(), Error> {
self.set_with_options(key, value, 0, 0).await
}
pub async fn set_with_options(
&mut self,
key: impl AsRef<[u8]>,
value: impl AsRef<[u8]>,
flags: u32,
exptime: u32,
) -> Result<(), Error> {
let key = key.as_ref();
let value = value.as_ref();
self.encode_buf.clear();
encode_request_into(
&McRequest::Set {
key,
value,
flags,
exptime,
},
&mut self.encode_buf,
)?;
if !self.is_instrumented() {
let response = self.execute(&self.encode_buf).await?;
return match response {
McResponseBytes::Stored => Ok(()),
_ => Err(Error::UnexpectedResponse),
};
}
let tx_bytes = self.encode_buf.len() as u32;
let send_ts = self.send_timestamp();
let start = Instant::now();
let response = self.execute(&self.encode_buf).await;
let latency_ns = self.finish_timing(send_ts, start);
let rx_bytes = self.last_rx_bytes.get();
let result = match response {
Ok(McResponseBytes::Stored) => Ok(()),
Ok(_) => Err(Error::UnexpectedResponse),
Err(e) => Err(e),
};
self.record(&CommandResult {
command: CommandType::Set,
latency_ns,
hit: None,
success: result.is_ok(),
ttfb_ns: None,
tx_bytes,
rx_bytes,
});
result
}
pub async fn add(
&mut self,
key: impl AsRef<[u8]>,
value: impl AsRef<[u8]>,
) -> Result<bool, Error> {
let key = key.as_ref();
let value = value.as_ref();
let encoded = encode_add(key, value)?;
let response = self.execute(&encoded).await?;
match response {
McResponseBytes::Stored => Ok(true),
McResponseBytes::NotStored => Ok(false),
_ => Err(Error::UnexpectedResponse),
}
}
pub async fn replace(
&mut self,
key: impl AsRef<[u8]>,
value: impl AsRef<[u8]>,
) -> Result<bool, Error> {
let key = key.as_ref();
let value = value.as_ref();
let encoded = encode_request(&McRequest::Replace {
key,
value,
flags: 0,
exptime: 0,
})?;
let response = self.execute(&encoded).await?;
match response {
McResponseBytes::Stored => Ok(true),
McResponseBytes::NotStored => Ok(false),
_ => Err(Error::UnexpectedResponse),
}
}
pub async fn incr(&mut self, key: impl AsRef<[u8]>, delta: u64) -> Result<Option<u64>, Error> {
let key = key.as_ref();
let encoded = encode_request(&McRequest::incr(key, delta))?;
let response = self.execute(&encoded).await?;
match response {
McResponseBytes::Numeric(val) => Ok(Some(val)),
McResponseBytes::NotFound => Ok(None),
_ => Err(Error::UnexpectedResponse),
}
}
pub async fn decr(&mut self, key: impl AsRef<[u8]>, delta: u64) -> Result<Option<u64>, Error> {
let key = key.as_ref();
let encoded = encode_request(&McRequest::decr(key, delta))?;
let response = self.execute(&encoded).await?;
match response {
McResponseBytes::Numeric(val) => Ok(Some(val)),
McResponseBytes::NotFound => Ok(None),
_ => Err(Error::UnexpectedResponse),
}
}
pub async fn append(
&mut self,
key: impl AsRef<[u8]>,
value: impl AsRef<[u8]>,
) -> Result<bool, Error> {
let key = key.as_ref();
let value = value.as_ref();
let encoded = encode_request(&McRequest::append(key, value))?;
let response = self.execute(&encoded).await?;
match response {
McResponseBytes::Stored => Ok(true),
McResponseBytes::NotStored => Ok(false),
_ => Err(Error::UnexpectedResponse),
}
}
pub async fn prepend(
&mut self,
key: impl AsRef<[u8]>,
value: impl AsRef<[u8]>,
) -> Result<bool, Error> {
let key = key.as_ref();
let value = value.as_ref();
let encoded = encode_request(&McRequest::prepend(key, value))?;
let response = self.execute(&encoded).await?;
match response {
McResponseBytes::Stored => Ok(true),
McResponseBytes::NotStored => Ok(false),
_ => Err(Error::UnexpectedResponse),
}
}
pub async fn cas(
&mut self,
key: impl AsRef<[u8]>,
value: impl AsRef<[u8]>,
cas_unique: u64,
) -> Result<bool, Error> {
let key = key.as_ref();
let value = value.as_ref();
let encoded = encode_request(&McRequest::cas(key, value, cas_unique))?;
let response = self.execute(&encoded).await?;
match response {
McResponseBytes::Stored => Ok(true),
McResponseBytes::Exists => Ok(false),
McResponseBytes::NotFound => Err(Error::Memcache("NOT_FOUND".into())),
_ => Err(Error::UnexpectedResponse),
}
}
pub async fn delete(&mut self, key: impl AsRef<[u8]>) -> Result<bool, Error> {
let key = key.as_ref();
self.encode_buf.clear();
encode_request_into(&McRequest::delete(key), &mut self.encode_buf)?;
if !self.is_instrumented() {
let response = self.execute(&self.encode_buf).await?;
return match response {
McResponseBytes::Deleted => Ok(true),
McResponseBytes::NotFound => Ok(false),
_ => Err(Error::UnexpectedResponse),
};
}
let tx_bytes = self.encode_buf.len() as u32;
let send_ts = self.send_timestamp();
let start = Instant::now();
let response = self.execute(&self.encode_buf).await;
let latency_ns = self.finish_timing(send_ts, start);
let rx_bytes = self.last_rx_bytes.get();
let result = match response {
Ok(McResponseBytes::Deleted) => Ok(true),
Ok(McResponseBytes::NotFound) => Ok(false),
Ok(_) => Err(Error::UnexpectedResponse),
Err(e) => Err(e),
};
self.record(&CommandResult {
command: CommandType::Delete,
latency_ns,
hit: None,
success: result.is_ok(),
ttfb_ns: None,
tx_bytes,
rx_bytes,
});
result
}
pub async fn flush_all(&mut self) -> Result<(), Error> {
let encoded = encode_request(&McRequest::flush_all())?;
let response = self.execute(&encoded).await?;
match response {
McResponseBytes::Ok => Ok(()),
_ => Err(Error::UnexpectedResponse),
}
}
pub async fn version(&mut self) -> Result<Box<str>, Error> {
let encoded = encode_request(&McRequest::version())?;
let response = self.execute(&encoded).await?;
match response {
McResponseBytes::Version(v) => Ok(Box::from(String::from_utf8_lossy(v.as_ref()))),
_ => Err(Error::UnexpectedResponse),
}
}
pub async fn set_with_guard<G: SendGuard>(
&mut self,
key: &[u8],
guard: G,
flags: u32,
exptime: u32,
) -> Result<(), Error> {
if !self.is_instrumented() {
let (_, value_len) = guard.as_ptr_len();
self.encode_buf.clear();
append_set_guard_prefix(
&mut self.encode_buf,
key,
value_len as usize,
flags,
exptime,
)?;
let prefix: &[u8] = &self.encode_buf;
self.conn.send_parts().build(move |b| {
b.copy(prefix)
.guard(GuardBox::new(guard))
.copy(b"\r\n")
.submit()
})?;
let response = self.read_response().await?;
check_error_bytes(&response)?;
return match response {
McResponseBytes::Stored => Ok(()),
_ => Err(Error::UnexpectedResponse),
};
}
let (_, value_len) = guard.as_ptr_len();
self.encode_buf.clear();
append_set_guard_prefix(
&mut self.encode_buf,
key,
value_len as usize,
flags,
exptime,
)?;
let tx_bytes = (self.encode_buf.len() + value_len as usize + 2) as u32;
let send_ts = self.send_timestamp();
let start = Instant::now();
let prefix: &[u8] = &self.encode_buf;
self.conn.send_parts().build(move |b| {
b.copy(prefix)
.guard(GuardBox::new(guard))
.copy(b"\r\n")
.submit()
})?;
let response = self.read_response().await;
let latency_ns = self.finish_timing(send_ts, start);
let rx_bytes = self.last_rx_bytes.get();
let result = match response {
Ok(ref r) => {
check_error_bytes(r)?;
match r {
McResponseBytes::Stored => Ok(()),
_ => Err(Error::UnexpectedResponse),
}
}
Err(e) => Err(e),
};
self.record(&CommandResult {
command: CommandType::Set,
latency_ns,
hit: None,
success: result.is_ok(),
ttfb_ns: None,
tx_bytes,
rx_bytes,
});
result
}
}
fn append_set_guard_prefix(
buf: &mut Vec<u8>,
key: &[u8],
value_len: usize,
flags: u32,
exptime: u32,
) -> Result<(), Error> {
validate_key(key)?;
validate_value_len(value_len)?;
let mut itoa_buf = itoa::Buffer::new();
buf.extend_from_slice(b"set ");
buf.extend_from_slice(key);
buf.push(b' ');
buf.extend_from_slice(itoa_buf.format(flags).as_bytes());
buf.push(b' ');
buf.extend_from_slice(itoa_buf.format(exptime).as_bytes());
buf.push(b' ');
buf.extend_from_slice(itoa_buf.format(value_len).as_bytes());
buf.extend_from_slice(b"\r\n");
Ok(())
}
fn validate_request(req: &McRequest<'_>) -> Result<(), Error> {
match req {
McRequest::Get { key }
| McRequest::Incr { key, .. }
| McRequest::Decr { key, .. }
| McRequest::Delete { key } => validate_key(key),
McRequest::Gets { keys } => {
for k in keys.iter() {
validate_key(k)?;
}
Ok(())
}
McRequest::Set { key, value, .. }
| McRequest::Add { key, value, .. }
| McRequest::Replace { key, value, .. }
| McRequest::Append { key, value }
| McRequest::Prepend { key, value }
| McRequest::Cas { key, value, .. } => {
validate_key(key)?;
validate_value(value)
}
McRequest::FlushAll | McRequest::Version | McRequest::Quit => Ok(()),
}
}
pub(crate) fn encode_request_into(req: &McRequest<'_>, buf: &mut Vec<u8>) -> Result<(), Error> {
validate_request(req)?;
let size = match req {
McRequest::Get { key } => 6 + key.len(),
McRequest::Gets { keys } => 6 + keys.iter().map(|k| 1 + k.len()).sum::<usize>(),
McRequest::Set { key, value, .. } | McRequest::Add { key, value, .. } => {
41 + key.len() + value.len()
}
McRequest::Replace { key, value, .. } => 45 + key.len() + value.len(),
McRequest::Incr { key, .. } | McRequest::Decr { key, .. } => 27 + key.len(),
McRequest::Append { key, value } => 44 + key.len() + value.len(),
McRequest::Prepend { key, value } => 45 + key.len() + value.len(),
McRequest::Cas { key, value, .. } => 61 + key.len() + value.len(),
McRequest::Delete { key } => 9 + key.len(),
McRequest::FlushAll => 11,
McRequest::Version => 9,
McRequest::Quit => 6,
};
let start = buf.len();
buf.resize(start + size, 0);
let len = req.encode(&mut buf[start..]);
buf.truncate(start + len);
Ok(())
}
pub(crate) fn encode_request(req: &McRequest<'_>) -> Result<Vec<u8>, Error> {
let mut buf = Vec::new();
encode_request_into(req, &mut buf)?;
Ok(buf)
}
#[cfg(test)]
pub(crate) fn encode_set(
key: &[u8],
value: &[u8],
flags: u32,
exptime: u32,
) -> Result<Vec<u8>, Error> {
encode_request(&McRequest::Set {
key,
value,
flags,
exptime,
})
}
pub(crate) fn encode_add(key: &[u8], value: &[u8]) -> Result<Vec<u8>, Error> {
encode_request(&McRequest::Add {
key,
value,
flags: 0,
exptime: 0,
})
}
pub(crate) fn check_error_bytes(response: &McResponseBytes) -> Result<(), Error> {
match response {
McResponseBytes::Error => Err(Error::Memcache("ERROR".into())),
McResponseBytes::ClientError(msg) => Err(Error::Memcache(format!(
"CLIENT_ERROR {}",
String::from_utf8_lossy(msg)
))),
McResponseBytes::ServerError(msg) => Err(Error::Memcache(format!(
"SERVER_ERROR {}",
String::from_utf8_lossy(msg)
))),
_ => Ok(()),
}
}
#[cfg(feature = "timestamps")]
fn now_realtime_ns() -> u64 {
let mut ts = libc::timespec {
tv_sec: 0,
tv_nsec: 0,
};
unsafe {
libc::clock_gettime(libc::CLOCK_REALTIME, &mut ts);
}
ts.tv_sec as u64 * 1_000_000_000 + ts.tv_nsec as u64
}
#[cfg(test)]
mod encode_tests {
use super::*;
const KEY: &[u8] = b"user:123456789";
#[test]
fn golden_encode_request_get() {
let encoded = encode_request(&McRequest::get(KEY)).unwrap();
assert_eq!(encoded, b"get user:123456789\r\n");
}
#[test]
fn golden_encode_request_delete() {
let encoded = encode_request(&McRequest::delete(KEY)).unwrap();
assert_eq!(encoded, b"delete user:123456789\r\n");
}
#[test]
fn golden_encode_request_set() {
let encoded = encode_set(KEY, b"hello world value", 42, 7200).unwrap();
assert_eq!(
encoded,
&b"set user:123456789 42 7200 17\r\nhello world value\r\n"[..]
);
}
#[test]
fn golden_encode_request_set_zero_fields() {
let encoded = encode_set(KEY, b"v", 0, 0).unwrap();
assert_eq!(encoded, &b"set user:123456789 0 0 1\r\nv\r\n"[..]);
}
#[test]
fn golden_set_guard_prefix() {
let mut prefix = Vec::new();
append_set_guard_prefix(&mut prefix, KEY, 1024, 42, 7200).unwrap();
assert_eq!(prefix, &b"set user:123456789 42 7200 1024\r\n"[..]);
}
#[test]
fn encode_into_appends_without_clobbering() {
let mut buf = b"EXISTING".to_vec();
encode_request_into(&McRequest::get(KEY), &mut buf).unwrap();
assert_eq!(buf, &b"EXISTINGget user:123456789\r\n"[..]);
}
#[test]
fn encode_into_error_leaves_buf_untouched() {
let mut buf = b"EXISTING".to_vec();
let long_key = vec![b'k'; MAX_KEY_LEN + 1];
assert!(encode_request_into(&McRequest::get(&long_key), &mut buf).is_err());
assert_eq!(buf, b"EXISTING");
assert!(append_set_guard_prefix(&mut buf, &long_key, 16, 0, 0).is_err());
assert_eq!(buf, b"EXISTING");
}
#[test]
fn coalesced_buf_equals_concatenation_of_individual_encodings() {
let mut coalesced = Vec::new();
encode_request_into(&McRequest::get(b"key1"), &mut coalesced).unwrap();
encode_request_into(
&McRequest::Set {
key: b"key2",
value: b"val",
flags: 0,
exptime: 0,
},
&mut coalesced,
)
.unwrap();
encode_request_into(&McRequest::delete(b"key3"), &mut coalesced).unwrap();
let mut expected = Vec::new();
expected.extend_from_slice(&encode_request(&McRequest::get(b"key1")).unwrap());
expected.extend_from_slice(
&encode_request(&McRequest::Set {
key: b"key2",
value: b"val",
flags: 0,
exptime: 0,
})
.unwrap(),
);
expected.extend_from_slice(&encode_request(&McRequest::delete(b"key3")).unwrap());
assert_eq!(coalesced, expected);
assert_eq!(
coalesced,
&b"get key1\r\nset key2 0 0 3\r\nval\r\ndelete key3\r\n"[..]
);
}
#[test]
fn coalesced_guard_set_byte_layout() {
let mut write_buf = Vec::new();
let value: &[u8] = b"hello";
append_set_guard_prefix(&mut write_buf, KEY, value.len(), 42, 7200).unwrap();
let guard_offset = write_buf.len();
write_buf.extend_from_slice(b"\r\n");
let mut wire = Vec::new();
wire.extend_from_slice(&write_buf[..guard_offset]);
wire.extend_from_slice(value);
wire.extend_from_slice(&write_buf[guard_offset..]);
assert_eq!(
wire,
&b"set user:123456789 42 7200 5\r\nhello\r\n"[..],
"coalesced guard SET must match a standard SET on the wire"
);
assert_eq!(write_buf.len() + value.len(), wire.len());
}
#[test]
fn two_coalesced_guard_sets_interleave_correctly() {
let v1: &[u8] = b"aa";
let v2: &[u8] = b"bbbb";
let mut write_buf = Vec::new();
let mut guards: Vec<usize> = Vec::new();
append_set_guard_prefix(&mut write_buf, b"k1", v1.len(), 0, 0).unwrap();
guards.push(write_buf.len());
write_buf.extend_from_slice(b"\r\n");
append_set_guard_prefix(&mut write_buf, b"k2", v2.len(), 0, 0).unwrap();
guards.push(write_buf.len());
write_buf.extend_from_slice(b"\r\n");
let values = [v1, v2];
let mut wire = Vec::new();
let mut pos = 0;
for (i, &offset) in guards.iter().enumerate() {
if offset > pos {
wire.extend_from_slice(&write_buf[pos..offset]);
}
wire.extend_from_slice(values[i]);
pos = offset;
}
if pos < write_buf.len() {
wire.extend_from_slice(&write_buf[pos..]);
}
assert_eq!(wire, &b"set k1 0 0 2\r\naa\r\nset k2 0 0 4\r\nbbbb\r\n"[..]);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn validate_key_accepts_max_len() {
let key = vec![b'k'; MAX_KEY_LEN];
assert!(validate_key(&key).is_ok());
}
#[test]
fn validate_key_rejects_oversized() {
let key = vec![b'k'; MAX_KEY_LEN + 1];
assert!(matches!(validate_key(&key), Err(Error::KeyTooLong)));
}
#[test]
fn validate_value_accepts_max_len() {
let v = vec![0u8; MAX_VALUE_LEN];
assert!(validate_value(&v).is_ok());
}
#[test]
fn validate_value_rejects_oversized() {
let v = vec![0u8; MAX_VALUE_LEN + 1];
assert!(matches!(validate_value(&v), Err(Error::ValueTooLong)));
}
#[test]
fn encode_request_get_rejects_long_key() {
let key = vec![b'k'; MAX_KEY_LEN + 1];
let r = encode_request(&McRequest::get(&key));
assert!(matches!(r, Err(Error::KeyTooLong)));
}
#[test]
fn encode_request_set_rejects_long_value() {
let key = b"k";
let value = vec![0u8; MAX_VALUE_LEN + 1];
let r = encode_request(&McRequest::Set {
key,
value: &value,
flags: 0,
exptime: 0,
});
assert!(matches!(r, Err(Error::ValueTooLong)));
}
#[test]
fn encode_request_gets_rejects_any_long_key() {
let ok = vec![b'k'; MAX_KEY_LEN];
let bad = vec![b'k'; MAX_KEY_LEN + 1];
let keys: &[&[u8]] = &[&ok, &bad];
let r = encode_request(&McRequest::gets(keys));
assert!(matches!(r, Err(Error::KeyTooLong)));
}
#[test]
fn set_guard_prefix_rejects_long_key() {
let key = vec![b'k'; MAX_KEY_LEN + 1];
let r = append_set_guard_prefix(&mut Vec::new(), &key, 16, 0, 0);
assert!(matches!(r, Err(Error::KeyTooLong)));
}
#[test]
fn set_guard_prefix_rejects_long_value_len() {
let r = append_set_guard_prefix(&mut Vec::new(), b"k", MAX_VALUE_LEN + 1, 0, 0);
assert!(matches!(r, Err(Error::ValueTooLong)));
}
#[test]
fn set_guard_prefix_accepts_max_value_len() {
let r = append_set_guard_prefix(&mut Vec::new(), b"k", MAX_VALUE_LEN, 0, 0);
assert!(r.is_ok());
}
#[test]
fn encode_request_passes_through_at_caps() {
let key = vec![b'k'; MAX_KEY_LEN];
let value = vec![0u8; MAX_VALUE_LEN];
let r = encode_request(&McRequest::Set {
key: &key,
value: &value,
flags: 0,
exptime: 0,
});
assert!(r.is_ok());
let buf = r.unwrap();
assert!(buf.starts_with(b"set "));
}
}
#[cfg(test)]
mod zc_threshold_tests {
use super::*;
use ringline::{ConnCtx, RegionId, SendGuard};
const KEY: &[u8] = b"user:123456789";
struct VecGuard(Vec<u8>);
impl SendGuard for VecGuard {
fn as_ptr_len(&self) -> (*const u8, u32) {
(self.0.as_ptr(), self.0.len() as u32)
}
fn region(&self) -> RegionId {
RegionId::UNREGISTERED
}
}
fn test_client(max_batch_size: usize, zc_threshold: u32) -> Client {
let conn = ConnCtx::for_test(0, 0);
Client::builder(conn)
.max_batch_size(max_batch_size)
.zc_threshold(zc_threshold)
.build()
}
#[test]
fn small_guard_set_is_byte_identical_to_plain_set() {
let value = vec![0xABu8; 64];
let mut guarded = test_client(4, 4096);
guarded
.fire_set_with_guard(KEY, VecGuard(value.clone()), 7, 42, 1)
.unwrap();
let mut plain = test_client(4, 4096);
plain.fire_set(KEY, &value, 7, 42, 1).unwrap();
assert_eq!(guarded.write_buf, plain.write_buf);
assert!(guarded.write_guards.is_empty());
}
#[test]
fn small_value_takes_copy_path() {
let mut client = test_client(4, 4096);
client
.fire_set_with_guard(KEY, VecGuard(vec![0u8; 64]), 0, 0, 1)
.unwrap();
assert!(client.write_guards.is_empty());
assert_eq!(client.buffered_ops, 1);
assert!(!client.write_buf.is_empty());
}
#[test]
fn large_value_takes_guard_path() {
let mut client = test_client(4, 4096);
client
.fire_set_with_guard(KEY, VecGuard(vec![0u8; 8192]), 0, 0, 1)
.unwrap();
assert_eq!(client.write_guards.len(), 1);
}
#[test]
fn threshold_zero_disables_fold() {
let mut client = test_client(4, 0);
client
.fire_set_with_guard(KEY, VecGuard(vec![0u8; 64]), 0, 0, 1)
.unwrap();
assert_eq!(client.write_guards.len(), 1);
}
#[test]
fn guard_batch_accumulates_max_flush_guards_before_flush() {
let mut client = test_client(16, 0);
for i in 0..MAX_FLUSH_GUARDS {
client
.fire_set_with_guard(KEY, VecGuard(vec![0u8; 64]), 0, 0, i as u64)
.unwrap();
}
assert_eq!(
client.write_guards.len(),
MAX_FLUSH_GUARDS,
"a full guard batch must accumulate without a premature flush"
);
}
#[test]
fn batch_of_small_guard_sets_coalesces() {
const N: usize = 8;
let mut client = test_client(N, 4096);
for i in 0..N - 1 {
client
.fire_set_with_guard(KEY, VecGuard(vec![0u8; 64]), 0, 0, i as u64)
.unwrap();
assert!(
client.write_guards.is_empty(),
"op {i} must use the copy path, not the guard path"
);
assert_eq!(client.buffered_ops, i + 1);
}
assert_eq!(client.buffered_ops, N - 1);
assert!(client.write_guards.is_empty());
}
}