pub mod actions;
pub mod protocol;
pub use actions::MemcachedProtocol;
use crate::llm::action_helper::call_llm;
use crate::llm::ollama_client::OllamaClient;
use crate::logging::emit::Log;
use crate::protocol::{Event, SpawnContext};
use crate::server::connection::ConnectionId;
use crate::state::app_state::AppState;
use crate::state::ServerId;
use anyhow::{Context, Result};
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::mpsc;
use tracing::{debug, warn};
use actions::{
MEMCACHED_ARITHMETIC_EVENT, MEMCACHED_DELETE_EVENT, MEMCACHED_FLUSH_ALL_EVENT,
MEMCACHED_GET_EVENT, MEMCACHED_STATS_EVENT, MEMCACHED_STORE_EVENT, MEMCACHED_TOUCH_EVENT,
MEMCACHED_UNKNOWN_COMMAND_EVENT, MEMCACHED_VERSION_EVENT,
};
use protocol::{Command, Parsed};
const MAX_BUFFERED: usize = protocol::MAX_VALUE_LEN + protocol::MAX_COMMAND_LINE + 16;
const LINGER_AFTER_REFUSAL: std::time::Duration = std::time::Duration::from_secs(2);
async fn linger_discard<R: tokio::io::AsyncRead + Unpin>(reader: &mut R, chunk: &mut [u8]) {
let deadline = tokio::time::Instant::now() + LINGER_AFTER_REFUSAL;
while let Ok(Ok(n)) = tokio::time::timeout_at(deadline, reader.read(chunk)).await {
if n == 0 {
break;
}
}
}
const FIRST_COMMAND_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
const IDLE_BETWEEN_COMMANDS_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(120);
const MAX_CONNECTIONS: usize = crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS;
const CONNECTION_CAP_REFUSAL: &[u8] = b"SERVER_ERROR too many connections\r\n";
const VALUE_TOO_LARGE_REFUSAL: &str = "SERVER_ERROR object too large for cache\r\n";
pub struct MemcachedServer;
impl MemcachedServer {
pub async fn spawn_with_llm_actions(ctx: SpawnContext) -> Result<SocketAddr> {
let listen_addr = ctx.legacy_listen_addr();
let SpawnContext {
llm_client,
state,
status_tx,
server_id,
..
} = ctx;
let listener = TcpListener::bind(listen_addr)
.await
.with_context(|| format!("Memcached failed to bind {}", listen_addr))?;
let actual_addr = listener.local_addr()?;
Log::new(Some(&status_tx)).info(format!("Memcached server listening on {}", actual_addr));
let task_registrar = state.clone();
let limiter = crate::server::accept_bounded::ConnectionLimiter::new(MAX_CONNECTIONS);
let accept_handle = tokio::spawn(async move {
loop {
let (stream, peer_addr, permit) =
match crate::server::accept_bounded::accept_bounded(
&listener,
&limiter,
CONNECTION_CAP_REFUSAL,
"Memcached",
Some(&status_tx),
)
.await
{
Ok(triple) => triple,
Err(e) => {
Log::new(Some(&status_tx))
.error(format!("Memcached accept error: {}", e));
continue;
}
};
let connection_id = ConnectionId::new(state.get_next_unified_id().await);
let local_addr = stream.local_addr().unwrap_or(actual_addr);
Self::track_connection(&state, server_id, connection_id, peer_addr, local_addr)
.await;
Log::new(Some(&status_tx)).info(format!(
"Memcached connection {} from {}",
connection_id, peer_addr
));
let _ = status_tx.send("__UPDATE_UI__".to_string());
let llm = llm_client.clone();
let st = state.clone();
let tx = status_tx.clone();
let registrar = state.clone();
let conn_handle = tokio::spawn(async move {
let _permit = permit;
if let Err(e) = Self::handle_connection(
stream,
peer_addr,
connection_id,
server_id,
llm,
st.clone(),
tx,
)
.await
{
debug!("Memcached connection {} ended: {}", connection_id, e);
}
st.close_connection_on_server(server_id, connection_id)
.await;
});
registrar.register_server_task(server_id, conn_handle).await;
}
});
task_registrar
.register_server_task(server_id, accept_handle)
.await;
Ok(actual_addr)
}
async fn track_connection(
state: &Arc<AppState>,
server_id: ServerId,
connection_id: ConnectionId,
peer_addr: SocketAddr,
local_addr: SocketAddr,
) {
use crate::state::server::{
ConnectionState as ServerConnectionState, ConnectionStatus, ProtocolConnectionInfo,
};
let now = crate::utils::clock::Instant::now();
state
.add_connection_to_server(
server_id,
ServerConnectionState {
id: connection_id,
remote_addr: peer_addr,
local_addr,
bytes_sent: 0,
bytes_received: 0,
packets_sent: 0,
packets_received: 0,
last_activity: now,
status: ConnectionStatus::Active,
status_changed_at: now,
protocol_info: ProtocolConnectionInfo::empty(),
},
)
.await;
}
#[allow(clippy::too_many_arguments)]
async fn handle_connection(
stream: TcpStream,
peer_addr: SocketAddr,
connection_id: ConnectionId,
server_id: ServerId,
llm_client: OllamaClient,
state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
) -> Result<()> {
let (reader, write_half) = tokio::io::split(stream);
let write_half = Arc::new(tokio::sync::Mutex::new(write_half));
let handler = Arc::new(actions::MemcachedProtocol::new());
let peer_rx = crate::server::peer_support::register_peer_channel(
&state,
server_id,
connection_id.as_u32(),
)
.await;
crate::server::peer_support::spawn_peer_command_task(
peer_rx,
handler.clone(),
state.clone(),
server_id,
connection_id.as_u32(),
write_half.clone(),
status_tx.clone(),
);
let result = Self::run_connection(
reader,
&write_half,
&handler,
peer_addr,
connection_id,
server_id,
&llm_client,
&state,
&status_tx,
)
.await;
state
.remove_peer_handle(server_id, connection_id.as_u32())
.await;
result
}
#[allow(clippy::too_many_arguments)]
async fn run_connection<R, W>(
mut reader: R,
write_half: &Arc<tokio::sync::Mutex<W>>,
handler: &Arc<actions::MemcachedProtocol>,
peer_addr: SocketAddr,
connection_id: ConnectionId,
server_id: ServerId,
llm_client: &OllamaClient,
state: &Arc<AppState>,
status_tx: &mpsc::UnboundedSender<String>,
) -> Result<()>
where
R: tokio::io::AsyncRead + Unpin,
W: tokio::io::AsyncWrite + Unpin,
{
let log = Log::new(Some(status_tx));
let mut buffer: Vec<u8> = Vec::with_capacity(4096);
let mut chunk = vec![0u8; 8192];
let mut answered_one = false;
loop {
loop {
match protocol::parse_command(&buffer) {
Parsed::Incomplete => break,
Parsed::Invalid { message, consumed } => {
let dropped = consumed.min(buffer.len());
buffer.drain(..dropped);
warn!(
"Memcached {} sent an unusable command: {}",
peer_addr, message
);
Self::count_received(state, server_id, connection_id, dropped).await;
let line = format!("CLIENT_ERROR {}\r\n", message);
{
let mut writer = write_half.lock().await;
writer.write_all(line.as_bytes()).await?;
writer.flush().await?;
}
Self::count_sent(state, server_id, connection_id, line.len()).await;
}
Parsed::Fatal { message, too_large } => {
let received = buffer.len();
buffer.clear();
let line = if too_large {
warn!(
"Memcached {} decision=fail_closed_value_too_large: {}; the \
cap is {} bytes, closing before the data block is read",
peer_addr,
message,
protocol::MAX_VALUE_LEN
);
VALUE_TOO_LARGE_REFUSAL.to_string()
} else {
warn!(
"Memcached {} decision=connection_closed_unframable: {}; \
closing, because the data block's extent is unknown and \
skipping it would put the peer's payload on the command path",
peer_addr, message
);
format!("CLIENT_ERROR {}\r\n", message)
};
Self::count_received(state, server_id, connection_id, received).await;
{
let mut writer = write_half.lock().await;
writer.write_all(line.as_bytes()).await?;
writer.flush().await?;
let _ = writer.shutdown().await;
}
Self::count_sent(state, server_id, connection_id, line.len()).await;
if too_large {
linger_discard(&mut reader, &mut chunk).await;
}
return Ok(());
}
Parsed::Complete { command, consumed } => {
buffer.drain(..consumed);
answered_one = true;
Self::count_received(state, server_id, connection_id, consumed).await;
if matches!(command, Command::Quit) {
debug!("Memcached {} sent quit", peer_addr);
return Ok(());
}
let suppress = protocol::is_noreply(&command);
let event = Self::event_for(&command);
let outcome = call_llm(
llm_client,
state,
server_id,
Some(connection_id),
&event,
handler.as_ref(),
)
.await;
let (reply, close) = match outcome {
Ok(execution) => {
for message in &execution.messages {
log.info(message);
}
let mut bytes = Vec::new();
let mut close = false;
let chosen =
execution.chosen_reply(Self::preferred_actions(&command));
let mut replies = 0usize;
for (index, result) in execution.protocol_results.iter().enumerate()
{
for (n, out) in result.get_all_output().iter().enumerate() {
replies += 1;
if Some(index) == chosen && n == 0 {
bytes.extend_from_slice(out);
}
}
close |= result.closes_connection();
}
if replies > 1 {
warn!(
"Memcached {} decision=duplicate_response_dropped \
command={}: {} extra repl{} not sent",
peer_addr,
Self::command_name(&command),
replies - 1,
if replies == 2 { "y" } else { "ies" }
);
log.warn(format!(
"Memcached {} decision=duplicate_response_dropped ({})",
peer_addr,
replies - 1
));
}
if let Command::Retrieval { keys, .. } = &command {
if let Some((kept, dropped)) =
protocol::retain_requested_values(&bytes, keys)
{
if !dropped.is_empty() {
log.warn(format!(
"Memcached {} decision=unrequested_key_dropped: \
the answer carried {:?}, which the client did \
not ask for (asked: {:?})",
peer_addr, dropped, keys
));
bytes = kept;
}
}
}
if bytes.is_empty() && !close {
warn!(
"Memcached {} decision=fail_closed_no_action command={}: \
no action produced",
peer_addr,
Self::command_name(&command)
);
bytes.extend_from_slice(
b"SERVER_ERROR no response was produced\r\n",
);
}
(bytes, close)
}
Err(e) => {
let overloaded = crate::llm::is_overload_error(&e);
let decision = if overloaded {
"fail_closed_llm_overloaded"
} else {
"fail_closed_llm_error"
};
let text = crate::utils::WireFailure::classify(&e).prefixed_text();
let reply = format!("SERVER_ERROR {text}\r\n");
log.warn(format!(
"Memcached {} decision={} command={} replying {}: {}",
peer_addr,
decision,
Self::command_name(&command),
reply.trim(),
e
));
(reply.into_bytes(), false)
}
};
if !suppress && !reply.is_empty() {
log.trace(format!(
"Memcached -> {}: {}",
peer_addr,
String::from_utf8_lossy(&reply)
));
{
let mut writer = write_half.lock().await;
writer.write_all(&reply).await?;
writer.flush().await?;
}
Self::count_sent(state, server_id, connection_id, reply.len()).await;
}
if close {
debug!("Memcached closing {} on model request", peer_addr);
return Ok(());
}
}
}
}
if buffer.len() > MAX_BUFFERED {
warn!(
"Memcached {} decision=fail_closed_buffer_too_large: buffered {} bytes \
without a complete command; closing",
peer_addr,
buffer.len()
);
let _ = write_half
.lock()
.await
.write_all(b"SERVER_ERROR command too large\r\n")
.await;
return Ok(());
}
let read_timeout = if answered_one {
IDLE_BETWEEN_COMMANDS_TIMEOUT
} else {
FIRST_COMMAND_READ_TIMEOUT
};
let n = match tokio::time::timeout(read_timeout, reader.read(&mut chunk)).await {
Ok(read) => read?,
Err(_) => {
debug!(
"Memcached {} sent nothing for {}s; closing idle connection",
peer_addr,
read_timeout.as_secs()
);
return Ok(());
}
};
if n == 0 {
debug!("Memcached {} closed the connection", peer_addr);
return Ok(());
}
buffer.extend_from_slice(&chunk[..n]);
}
}
async fn count_received(
state: &Arc<AppState>,
server_id: ServerId,
connection_id: ConnectionId,
bytes: usize,
) {
state
.update_connection_stats(
server_id,
connection_id,
Some(bytes as u64),
None,
Some(1),
None,
)
.await;
}
async fn count_sent(
state: &Arc<AppState>,
server_id: ServerId,
connection_id: ConnectionId,
bytes: usize,
) {
state
.update_connection_stats(
server_id,
connection_id,
None,
Some(bytes as u64),
None,
Some(1),
)
.await;
}
fn preferred_actions(command: &Command) -> &'static [&'static str] {
match command {
Command::Retrieval { .. } => &["send_memcached_values"],
Command::Stats { .. } => &["send_memcached_stats"],
Command::Version => &["send_memcached_version"],
Command::Arithmetic { .. } => &["send_memcached_number"],
_ => &[],
}
}
fn command_name(command: &Command) -> &'static str {
match command {
Command::Retrieval { command, .. } => command,
Command::Storage { command, .. } => command,
Command::Delete { .. } => "delete",
Command::Arithmetic { command, .. } => command,
Command::Touch { .. } => "touch",
Command::Stats { .. } => "stats",
Command::Version => "version",
Command::FlushAll { .. } => "flush_all",
Command::Quit => "quit",
Command::Unknown { .. } => "unknown",
}
}
fn event_for(command: &Command) -> Event {
match command {
Command::Retrieval { command, keys } => Event::new(
&MEMCACHED_GET_EVENT,
serde_json::json!({
"command": command,
"keys": keys,
"answer_with": format!(
"send_memcached_values with one entry for each of {:?} that the cache \
holds, and none for a key it does not hold; an empty values list is \
a miss. Never include a key that is not in that list",
keys
),
}),
),
Command::Storage {
command,
key,
flags,
exptime,
bytes,
cas_unique,
data,
..
} => {
let (value, encoding) = match std::str::from_utf8(data) {
Ok(s) => (s.to_string(), "utf8"),
Err(_) => (hex::encode(data), "hex"),
};
Event::new(
&MEMCACHED_STORE_EVENT,
serde_json::json!({
"command": command,
"key": key,
"flags": flags,
"exptime": exptime,
"bytes": bytes,
"cas_unique": cas_unique,
"value": value,
"value_encoding": encoding,
}),
)
}
Command::Delete { key, .. } => {
Event::new(&MEMCACHED_DELETE_EVENT, serde_json::json!({ "key": key }))
}
Command::Arithmetic {
command,
key,
delta,
..
} => Event::new(
&MEMCACHED_ARITHMETIC_EVENT,
serde_json::json!({ "command": command, "key": key, "delta": delta }),
),
Command::Touch { key, exptime, .. } => Event::new(
&MEMCACHED_TOUCH_EVENT,
serde_json::json!({ "key": key, "exptime": exptime }),
),
Command::Stats { argument } => Event::new(
&MEMCACHED_STATS_EVENT,
serde_json::json!({
"argument": argument,
"answer_with": "send_memcached_stats with a stats object under memcached's \
own stat names - version, curr_items, total_items, bytes, \
uptime, pid, curr_connections - holding what your \
instruction says (a version it names goes in version, an \
item count in curr_items)",
}),
),
Command::Version => Event::new(&MEMCACHED_VERSION_EVENT, serde_json::json!({})),
Command::FlushAll { delay, .. } => Event::new(
&MEMCACHED_FLUSH_ALL_EVENT,
serde_json::json!({ "delay": delay }),
),
Command::Unknown { line } => Event::new(
&MEMCACHED_UNKNOWN_COMMAND_EVENT,
serde_json::json!({ "line": line }),
),
Command::Quit => {
Event::new(
&MEMCACHED_UNKNOWN_COMMAND_EVENT,
serde_json::json!({ "line": "quit" }),
)
}
}
}
}