use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use memcache::{Client, CommandError, MemcacheError, ServerError};
use async_trait::async_trait;
use crate::backend::{run_blocking, Backend, HealthStatus};
use crate::error::{BackendError, BackendErrorKind};
const MAX_MEMCACHED_TTL_SECS: u64 = 30 * 24 * 60 * 60;
const DEFAULT_MAX_ITEM_SIZE_BYTES: usize = 1024 * 1024;
const MAX_KEY_LEN: usize = 250;
const DEFAULT_OP_TIMEOUT: Duration = Duration::from_secs(1);
const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(2);
const DEFAULT_POOL_SIZE: u32 = 10;
fn validate_key(key: &str) -> Result<(), BackendError> {
if key.is_empty() {
return Err(BackendError::permanent("memcached key must not be empty"));
}
if key.len() > MAX_KEY_LEN {
return Err(BackendError::permanent(format!(
"memcached key is {} bytes (limit {MAX_KEY_LEN})",
key.len()
)));
}
if key.bytes().any(|b| b <= 0x20 || b >= 0x7f) {
return Err(BackendError::permanent(
"memcached key contains whitespace, control, or non-ASCII bytes — not \
representable in the memcached ASCII protocol (and rejected by cachekit-py, \
breaking cross-SDK key identity)",
));
}
Ok(())
}
fn memcached_err(e: MemcacheError) -> BackendError {
let kind = match &e {
MemcacheError::IOError(io) => match io.kind() {
std::io::ErrorKind::TimedOut | std::io::ErrorKind::WouldBlock => {
BackendErrorKind::Timeout
}
_ => BackendErrorKind::Transient,
},
MemcacheError::PoolError(_) => BackendErrorKind::Transient,
MemcacheError::ServerError(ServerError::Error(msg)) if msg.contains("object too large") => {
BackendErrorKind::Permanent
}
MemcacheError::ServerError(ServerError::Error(_)) => BackendErrorKind::Transient,
MemcacheError::ServerError(_) => BackendErrorKind::Permanent,
MemcacheError::CommandError(CommandError::ValueTooLarge) => BackendErrorKind::Permanent,
MemcacheError::ClientError(_)
| MemcacheError::CommandError(_)
| MemcacheError::ParseError(_)
| MemcacheError::BadURL(_) => BackendErrorKind::Permanent,
};
BackendError {
kind,
message: e.to_string(),
source: Some(Box::new(e)),
}
}
fn expiration_secs(ttl: Option<Duration>) -> u32 {
match ttl {
None => 0,
Some(d) => {
let secs = d.as_secs().clamp(1, MAX_MEMCACHED_TTL_SECS);
u32::try_from(secs).unwrap_or(u32::MAX)
}
}
}
pub struct MemcachedBackend {
client: Arc<Client>,
max_item_size_bytes: usize,
op_budget: Duration,
}
impl std::fmt::Debug for MemcachedBackend {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("MemcachedBackend").finish_non_exhaustive()
}
}
impl MemcachedBackend {
pub fn builder() -> MemcachedBackendBuilder {
MemcachedBackendBuilder::default()
}
#[cfg(not(feature = "unsync"))]
async fn run_op<T: Send + 'static>(
&self,
f: impl FnOnce() -> Result<T, BackendError> + Send + 'static,
) -> Result<T, BackendError> {
match tokio::time::timeout(self.op_budget, run_blocking(f)).await {
Ok(result) => result,
Err(_) => Err(BackendError::timeout(format!(
"memcached operation exceeded its {}ms budget (server hung or pool checkout \
retrying); the in-flight call was abandoned",
self.op_budget.as_millis()
))),
}
}
#[cfg(feature = "unsync")]
async fn run_op<T>(
&self,
f: impl FnOnce() -> Result<T, BackendError>,
) -> Result<T, BackendError> {
run_blocking(f).await
}
pub async fn refresh_ttl(
&self,
key: &str,
ttl: Option<Duration>,
) -> Result<bool, BackendError> {
validate_key(key)?;
let client = Arc::clone(&self.client);
let key = key.to_owned();
let expire = expiration_secs(ttl);
self.run_op(move || client.touch(&key, expire).map_err(memcached_err))
.await
}
}
#[cfg(not(target_arch = "wasm32"))]
#[cfg_attr(not(feature = "unsync"), async_trait)]
#[cfg_attr(feature = "unsync", async_trait(?Send))]
impl Backend for MemcachedBackend {
async fn get(&self, key: &str) -> Result<Option<Vec<u8>>, BackendError> {
validate_key(key)?;
let client = Arc::clone(&self.client);
let key = key.to_owned();
self.run_op(move || client.get::<Vec<u8>>(&key).map_err(memcached_err))
.await
}
async fn set(
&self,
key: &str,
value: Vec<u8>,
ttl: Option<Duration>,
) -> Result<(), BackendError> {
validate_key(key)?;
if self.max_item_size_bytes > 0 && value.len() > self.max_item_size_bytes {
return Err(BackendError::permanent(format!(
"value is {} bytes, exceeding the memcached max item size of {} bytes; \
enable compression, use a larger-payload backend (Redis/SaaS/File), or \
raise both the server's -I limit and the builder's max_item_size_bytes",
value.len(),
self.max_item_size_bytes,
)));
}
let client = Arc::clone(&self.client);
let key = key.to_owned();
let expire = expiration_secs(ttl);
self.run_op(move || {
client
.set(&key, value.as_slice(), expire)
.map_err(memcached_err)
})
.await
}
async fn delete(&self, key: &str) -> Result<bool, BackendError> {
validate_key(key)?;
let client = Arc::clone(&self.client);
let key = key.to_owned();
self.run_op(move || client.delete(&key).map_err(memcached_err))
.await
}
async fn exists(&self, key: &str) -> Result<bool, BackendError> {
validate_key(key)?;
let client = Arc::clone(&self.client);
let key = key.to_owned();
self.run_op(move || {
Ok(client
.get::<Vec<u8>>(&key)
.map_err(memcached_err)?
.is_some())
})
.await
}
async fn health(&self) -> Result<HealthStatus, BackendError> {
let start = std::time::Instant::now();
let client = Arc::clone(&self.client);
let versions = self
.run_op(move || client.version().map_err(memcached_err))
.await?;
let latency = start.elapsed();
let mut details = HashMap::new();
if let Some((_, version)) = versions.first() {
details.insert("version".to_string(), version.clone());
}
Ok(HealthStatus {
is_healthy: true,
latency_ms: latency.as_secs_f64() * 1000.0,
backend_type: "memcached".to_string(),
details,
})
}
}
#[derive(Default)]
#[must_use]
pub struct MemcachedBackendBuilder {
url: Option<String>,
max_item_size_bytes: Option<usize>,
timeout: Option<Duration>,
connect_timeout: Option<Duration>,
pool_size: Option<u32>,
}
impl MemcachedBackendBuilder {
pub fn url(mut self, url: impl Into<String>) -> Self {
self.url = Some(url.into());
self
}
pub fn max_item_size_bytes(mut self, bytes: usize) -> Self {
self.max_item_size_bytes = Some(bytes);
self
}
pub fn timeout(mut self, timeout: Duration) -> Self {
self.timeout = Some(timeout);
self
}
pub fn connect_timeout(mut self, timeout: Duration) -> Self {
self.connect_timeout = Some(timeout);
self
}
pub fn pool_size(mut self, size: u32) -> Self {
self.pool_size = Some(size.max(1));
self
}
fn normalized_url(url: &str) -> String {
if let Some(rest) = url.strip_prefix("tcp://") {
return format!("memcache://{rest}");
}
if url.contains("://") {
return url.to_string();
}
format!("memcache://{url}")
}
fn with_connection_params(url: &str, timeout: Duration) -> String {
let mut url = url.to_string();
let mut sep = if url.contains('?') { '&' } else { '?' };
let has_param = |u: &str, name: &str| {
u.contains(&format!("?{name}=")) || u.contains(&format!("&{name}="))
};
if !has_param(&url, "protocol") {
url.push(sep);
url.push_str("protocol=ascii");
sep = '&';
}
if !has_param(&url, "timeout") {
url.push(sep);
url.push_str(&format!("timeout={}", timeout.as_secs_f64()));
}
url
}
pub async fn connect(self) -> Result<MemcachedBackend, crate::error::CachekitError> {
use crate::error::CachekitError;
let url = self
.url
.filter(|u| !u.is_empty())
.ok_or_else(|| CachekitError::Config("url is required".to_string()))?;
let timeout = self.timeout.unwrap_or(DEFAULT_OP_TIMEOUT);
let connect_timeout = self.connect_timeout.unwrap_or(DEFAULT_CONNECT_TIMEOUT);
let pool_size = self.pool_size.unwrap_or(DEFAULT_POOL_SIZE);
let url = Self::with_connection_params(&Self::normalized_url(&url), timeout);
let op_budget = timeout * 3 + connect_timeout;
let build = move || {
let client = Client::builder()
.add_server(url)
.map_err(memcached_err)?
.with_max_pool_size(pool_size)
.with_min_idle_conns(0)
.with_connection_timeout(connect_timeout)
.with_read_timeout(timeout)
.with_write_timeout(timeout)
.build()
.map_err(memcached_err)?;
client.version().map_err(memcached_err)?;
Ok(client)
};
#[cfg(not(feature = "unsync"))]
let client = {
let connect_budget = (connect_timeout + timeout) * pool_size.min(4) + timeout;
match tokio::time::timeout(connect_budget, run_blocking(build)).await {
Ok(result) => result?,
Err(_) => {
return Err(BackendError::timeout(format!(
"memcached connect exceeded its {}ms budget",
connect_budget.as_millis()
))
.into())
}
}
};
#[cfg(feature = "unsync")]
let client = run_blocking(build).await?;
Ok(MemcachedBackend {
client: Arc::new(client),
max_item_size_bytes: self
.max_item_size_bytes
.unwrap_or(DEFAULT_MAX_ITEM_SIZE_BYTES),
op_budget,
})
}
}
#[cfg(test)]
#[allow(clippy::expect_used, clippy::unwrap_used)] mod tests {
use memcache::ClientError;
use super::*;
#[test]
fn error_mapping_classifies_by_retryability() {
let refused = std::io::Error::new(std::io::ErrorKind::ConnectionRefused, "refused");
assert_eq!(
memcached_err(MemcacheError::IOError(refused)).kind,
BackendErrorKind::Transient
);
let timed_out = std::io::Error::new(std::io::ErrorKind::TimedOut, "slow");
assert_eq!(
memcached_err(MemcacheError::IOError(timed_out)).kind,
BackendErrorKind::Timeout
);
let would_block = std::io::Error::new(std::io::ErrorKind::WouldBlock, "recv timeout");
assert_eq!(
memcached_err(MemcacheError::IOError(would_block)).kind,
BackendErrorKind::Timeout
);
assert_eq!(
memcached_err(MemcacheError::ServerError(ServerError::Error(
"SERVER_ERROR out of memory".into()
)))
.kind,
BackendErrorKind::Transient
);
assert_eq!(
memcached_err(MemcacheError::ServerError(ServerError::Error(
"SERVER_ERROR object too large for cache".into()
)))
.kind,
BackendErrorKind::Permanent
);
assert_eq!(
memcached_err(MemcacheError::CommandError(CommandError::ValueTooLarge)).kind,
BackendErrorKind::Permanent
);
assert_eq!(
memcached_err(MemcacheError::ClientError(ClientError::KeyTooLong)).kind,
BackendErrorKind::Permanent
);
}
#[test]
fn key_validation_rejects_protocol_metacharacters() {
assert!(validate_key("evil\r\nflush_all\r\n").is_err());
assert!(validate_key("evil\nkey").is_err());
assert!(validate_key("evil key").is_err());
assert!(validate_key("evil\tkey").is_err());
assert!(validate_key("evil\x00key").is_err());
assert!(validate_key("evil\x7fkey").is_err());
assert!(validate_key("").is_err());
assert!(validate_key(&"k".repeat(251)).is_err());
assert!(validate_key("café").is_err());
assert!(validate_key("键").is_err());
assert!(validate_key("ns:app:func:m.f:args:abc123:v1").is_ok());
assert!(validate_key(&"k".repeat(250)).is_ok());
}
#[test]
fn expiration_clamps_to_memcached_ceiling() {
assert_eq!(expiration_secs(None), 0);
assert_eq!(expiration_secs(Some(Duration::from_millis(10))), 1);
assert_eq!(expiration_secs(Some(Duration::from_secs(60))), 60);
let clamped = expiration_secs(Some(Duration::from_secs(MAX_MEMCACHED_TTL_SECS + 1)));
assert_eq!(u64::from(clamped), MAX_MEMCACHED_TTL_SECS);
}
#[test]
fn url_normalization() {
assert_eq!(
MemcachedBackendBuilder::normalized_url("tcp://h:11211"),
"memcache://h:11211"
);
assert_eq!(
MemcachedBackendBuilder::normalized_url("h:11211"),
"memcache://h:11211"
);
assert_eq!(
MemcachedBackendBuilder::normalized_url("memcache://h:11211?protocol=ascii"),
"memcache://h:11211?protocol=ascii"
);
}
#[test]
fn connection_params_pin_ascii_and_per_connection_timeout() {
assert_eq!(
MemcachedBackendBuilder::with_connection_params(
"memcache://h:11211",
Duration::from_secs(1)
),
"memcache://h:11211?protocol=ascii&timeout=1"
);
assert_eq!(
MemcachedBackendBuilder::with_connection_params(
"memcache://h:11211?protocol=binary&timeout=5",
Duration::from_secs(1)
),
"memcache://h:11211?protocol=binary&timeout=5"
);
assert_eq!(
MemcachedBackendBuilder::with_connection_params(
"memcache://h:11211?tcp_nodelay=true",
Duration::from_millis(1500)
),
"memcache://h:11211?tcp_nodelay=true&protocol=ascii&timeout=1.5"
);
assert_eq!(
MemcachedBackendBuilder::with_connection_params(
"memcache://h:11211?connect_timeout=5",
Duration::from_secs(1)
),
"memcache://h:11211?connect_timeout=5&protocol=ascii&timeout=1"
);
}
}