pub mod actions;
pub mod cluster;
pub mod pubsub;
pub mod resp3;
pub mod resp_value;
use crate::llm::action_helper::call_llm;
use crate::llm::actions::protocol_trait::ActionResult;
use crate::llm::ollama_client::OllamaClient;
use crate::logging::emit::Log;
use crate::protocol::Event;
use crate::server::connection::ConnectionId;
use crate::state::app_state::AppState;
use crate::utils::redis_keys::{key_hash_slot, slot_keys};
use crate::utils::resp::{scan_resp2_frame, RespLength, RespLimits, RespScan};
use actions::{
encode_error, encode_integer, RedisProtocol, ACCEPT_SUBSCRIPTION, DELIVER_MESSAGE,
REDIS_COMMAND_EVENT, REDIS_PUBLISH_EVENT, REDIS_SUBSCRIBE_EVENT,
};
use anyhow::Result;
use pubsub::{Kind, Outbox, OutboxReceiver, PubSubHub};
use redis_protocol::resp2::decode::decode;
use redis_protocol::resp2::types::OwnedFrame as Frame;
use resp_value::Resp;
use std::net::SocketAddr;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::mpsc;
use tracing::{debug, error, trace, warn};
pub const MAX_PENDING_FRAME_BYTES: usize = 64 * 1024 * 1024;
const RESP_LIMITS: RespLimits = RespLimits::with_max_bulk_len(MAX_PENDING_FRAME_BYTES);
const FIRST_COMMAND_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300);
const IDLE_BETWEEN_COMMANDS_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300);
const MAX_CONNECTIONS: usize = crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS;
const CONNECTION_CAP_REFUSAL: &[u8] = b"-ERR max number of clients reached\r\n";
pub struct RedisShared {
pub hub: PubSubHub,
pub cluster: Option<std::sync::Mutex<cluster::ClusterState>>,
}
impl RedisShared {
fn cluster(&self) -> Option<std::sync::MutexGuard<'_, cluster::ClusterState>> {
self.cluster
.as_ref()
.map(|c| c.lock().unwrap_or_else(|p| p.into_inner()))
}
}
#[derive(Debug, Default, Clone)]
pub struct RedisServerOptions {
pub first_byte_timeout_secs: Option<u64>,
pub idle_timeout_secs: Option<u64>,
pub cluster_nodes: Option<serde_json::Value>,
}
pub struct RedisServer {
llm_client: OllamaClient,
app_state: Arc<AppState>,
#[allow(dead_code)]
status_tx: mpsc::UnboundedSender<String>,
server_id: Option<crate::state::ServerId>,
}
impl RedisServer {
pub fn new(
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
server_id: Option<crate::state::ServerId>,
) -> Self {
Self {
llm_client,
app_state,
status_tx,
server_id,
}
}
pub async fn spawn_with_llm_actions(
listen_addr: SocketAddr,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
_send_first: bool,
server_id: crate::state::ServerId,
options: RedisServerOptions,
) -> Result<SocketAddr> {
let first_byte_timeout = options
.first_byte_timeout_secs
.map(std::time::Duration::from_secs)
.unwrap_or(FIRST_COMMAND_READ_TIMEOUT);
let idle_timeout = options
.idle_timeout_secs
.map(std::time::Duration::from_secs)
.unwrap_or(IDLE_BETWEEN_COMMANDS_TIMEOUT);
let listener = TcpListener::bind(listen_addr).await?;
let actual_addr = listener.local_addr()?;
let cluster = options
.cluster_nodes
.as_ref()
.map(|nodes| cluster::ClusterState::from_param(nodes, actual_addr))
.transpose()
.map_err(|e| anyhow::anyhow!("invalid cluster_nodes: {e:#}"))?;
let shared = Arc::new(RedisShared {
hub: PubSubHub::default(),
cluster: cluster.map(std::sync::Mutex::new),
});
app_state
.register_server_handle(server_id, shared.clone())
.await;
let log = Log::new(Some(&status_tx));
match shared.cluster() {
Some(c) => log.info(format!(
"Redis server listening on {} (cluster mode: {} of {} nodes)",
actual_addr,
c.me().addr(),
c.nodes.len()
)),
None => log.info(format!("Redis server listening on {}", actual_addr)),
}
let server = Arc::new(RedisServer::new(
llm_client,
app_state.clone(),
status_tx.clone(),
Some(server_id),
));
let status_tx_clone = status_tx.clone();
let task_registrar = app_state.clone();
let limiter = crate::server::accept_bounded::ConnectionLimiter::new(MAX_CONNECTIONS);
let accept_handle = tokio::spawn(async move {
loop {
match crate::server::accept_bounded::accept_bounded(
&listener,
&limiter,
CONNECTION_CAP_REFUSAL,
"Redis",
Some(&status_tx),
)
.await
{
Ok((stream, addr, permit)) => {
Log::new(Some(&status_tx)).info(format!("Redis connection from {}", addr));
let connection_id =
ConnectionId::new(app_state.get_next_unified_id().await);
let local_addr_conn = stream.local_addr().unwrap_or(actual_addr);
if let Some(server_id) = server.server_id {
use crate::state::server::{
ConnectionState as ServerConnectionState, ConnectionStatus,
ProtocolConnectionInfo,
};
let now = crate::utils::clock::Instant::now();
let conn_state = ServerConnectionState {
id: connection_id,
remote_addr: addr,
local_addr: local_addr_conn,
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(),
};
server
.app_state
.add_connection_to_server(server_id, conn_state)
.await;
}
let handler = RedisHandler {
connection_id,
llm_client: server.llm_client.clone(),
app_state: server.app_state.clone(),
status_tx: status_tx.clone(),
server_id: server.server_id,
first_byte_timeout,
idle_timeout,
shared: shared.clone(),
};
let conn_handle = tokio::spawn(async move {
let _permit = permit;
if let Err(e) = handler.handle_connection(stream).await {
error!("Redis connection error: {:?}", e);
}
});
if let Some(server_id) = server.server_id {
server
.app_state
.register_server_task(server_id, conn_handle)
.await;
}
}
Err(e) => {
Log::new(Some(&status_tx)).error(format!("Redis accept error: {}", e));
}
}
}
});
task_registrar
.register_server_task(server_id, accept_handle)
.await;
let _ = status_tx_clone.send("__UPDATE_UI__".to_string());
Ok(actual_addr)
}
}
struct RedisHandler {
connection_id: ConnectionId,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
server_id: Option<crate::state::ServerId>,
first_byte_timeout: std::time::Duration,
idle_timeout: std::time::Duration,
shared: Arc<RedisShared>,
}
struct AbortWriter(tokio::task::AbortHandle);
impl Drop for AbortWriter {
fn drop(&mut self) {
self.0.abort();
}
}
#[derive(Default)]
struct ModelAnswer {
output: Vec<u8>,
close: bool,
accept: bool,
deliver: bool,
failed: bool,
}
struct Conn<'a> {
protocol: Arc<RedisProtocol>,
outbox: Outbox,
write_half: Arc<tokio::sync::Mutex<tokio::io::WriteHalf<TcpStream>>>,
log: Log<'a>,
asking: bool,
}
enum Flow {
Continue,
Close,
}
const SUBSCRIBED_MODE_COMMANDS: &[&str] = &[
"SUBSCRIBE",
"PSUBSCRIBE",
"SSUBSCRIBE",
"UNSUBSCRIBE",
"PUNSUBSCRIBE",
"SUNSUBSCRIBE",
"PING",
"QUIT",
"RESET",
];
fn wrong_arity(name: &str) -> String {
format!(
"ERR wrong number of arguments for '{}' command",
name.to_ascii_lowercase()
)
}
impl RedisHandler {
async fn handle_connection(self, stream: TcpStream) -> Result<()> {
let server_id = self.server_id;
let connection_id = self.connection_id;
let app_state = self.app_state.clone();
let status_tx = self.status_tx.clone();
let shared = self.shared.clone();
let result = self.run(stream).await;
shared.hub.remove_connection(connection_id.as_u32());
if let Some(server_id) = server_id {
app_state
.remove_peer_handle(server_id, connection_id.as_u32())
.await;
app_state
.close_connection_on_server(server_id, connection_id)
.await;
let _ = status_tx.send("__UPDATE_UI__".to_string());
}
result
}
async fn refuse(&self, conn: &Conn<'_>, message: &str) {
let _ = conn.outbox.send_reply(encode_error(message)).await;
let _ = conn.write_half.lock().await.shutdown().await;
}
async fn write_reply(&self, conn: &Conn<'_>, reply: Vec<u8>) -> Result<()> {
conn.outbox.send_reply(reply).await?;
Ok(())
}
async fn writer_task(
mut out: OutboxReceiver,
outbox: Outbox,
write_half: Arc<tokio::sync::Mutex<tokio::io::WriteHalf<TcpStream>>>,
app_state: Arc<AppState>,
server_id: Option<crate::state::ServerId>,
connection_id: ConnectionId,
) {
while let Some(frame) = out.rx.recv().await {
let len = frame.bytes.len();
let result = if len == 0 {
Ok(())
} else {
let mut writer = write_half.lock().await;
match writer.write_all(&frame.bytes).await {
Ok(()) => writer.flush().await,
Err(e) => Err(e),
}
};
if frame.counted {
out.queued.fetch_sub(len, Ordering::AcqRel);
}
let failed = result.is_err();
if !failed && len > 0 {
if let Some(server_id) = server_id {
app_state
.update_connection_stats(
server_id,
connection_id,
None,
Some(len as u64),
None,
Some(1),
)
.await;
}
}
if let Some(ack) = frame.ack {
let _ = ack.send(result);
}
if failed {
outbox.mark_killed();
break;
}
}
}
async fn run(self, stream: TcpStream) -> Result<()> {
let first_byte_timeout = self.first_byte_timeout;
let idle_timeout = self.idle_timeout;
let protocol = Arc::new(RedisProtocol::for_connection(
self.connection_id,
self.app_state.clone(),
self.status_tx.clone(),
self.shared.clone(),
));
let (mut read_half, write_half) = tokio::io::split(stream);
let write_half = Arc::new(tokio::sync::Mutex::new(write_half));
if let Some(server_id) = self.server_id {
let peer_rx = crate::server::peer_support::register_peer_channel(
&self.app_state,
server_id,
self.connection_id.as_u32(),
)
.await;
crate::server::peer_support::spawn_peer_command_task(
peer_rx,
protocol.clone(),
self.app_state.clone(),
server_id,
self.connection_id.as_u32(),
write_half.clone(),
self.status_tx.clone(),
);
let _ = self.status_tx.send("__UPDATE_UI__".to_string());
}
let (outbox, outbox_rx) = Outbox::new();
let writer = tokio::spawn(Self::writer_task(
outbox_rx,
outbox.clone(),
write_half.clone(),
self.app_state.clone(),
self.server_id,
self.connection_id,
));
let _writer_guard = AbortWriter(writer.abort_handle());
if let Some(server_id) = self.server_id {
self.app_state.register_server_task(server_id, writer).await;
}
let mut conn = Conn {
protocol,
outbox,
write_half,
log: Log::new(Some(&self.status_tx)),
asking: false,
};
let mut buffer = Vec::new();
let mut answered_one = false;
loop {
let read_timeout = if answered_one {
idle_timeout
} else {
first_byte_timeout
};
let mut chunk = vec![0u8; 4096];
let read = tokio::select! {
read = tokio::time::timeout(read_timeout, read_half.read(&mut chunk)) => read,
_ = conn.outbox.closing() => {
conn.log.warn(format!(
"Redis connection {} closed: it stopped reading while {} bytes of \
pub/sub messages were queued for it",
self.connection_id,
conn.outbox.queued_bytes()
));
return Ok(());
}
};
let read = match read {
Ok(read) => read,
Err(_) => {
conn.log.info(format!(
"Redis client {} sent nothing for {}s; closing idle connection",
self.connection_id,
read_timeout.as_secs()
));
return Ok(());
}
};
let n = match read {
Ok(0) => {
debug!("Redis client disconnected");
return Ok(());
}
Ok(n) => n,
Err(e) => {
error!("Redis read error: {}", e);
return Err(e.into());
}
};
buffer.extend_from_slice(&chunk[..n]);
if let Some(server_id) = self.server_id {
self.app_state
.update_connection_stats(
server_id,
self.connection_id,
Some(n as u64),
None,
Some(1),
None,
)
.await;
}
if buffer.len() > MAX_PENDING_FRAME_BYTES {
conn.log.error(format!(
"Redis client {} exceeded the {} byte frame buffer limit without completing a frame; closing",
self.connection_id, MAX_PENDING_FRAME_BYTES
));
self.refuse(&conn, "ERR Protocol error: invalid multibulk length")
.await;
return Ok(());
}
let mut offset = 0;
while offset < buffer.len() {
match scan_resp2_frame(&buffer[offset..], &RESP_LIMITS) {
RespScan::Complete { .. } => {}
RespScan::Incomplete => break,
RespScan::TooDeep { limit } => {
conn.log.warn(format!(
"Redis connection {} decision=fail_closed_resp_too_deep: frame nests \
arrays deeper than {} levels; refusing and closing",
self.connection_id, limit
));
self.refuse(&conn, "ERR Protocol error: nesting too deep")
.await;
return Ok(());
}
RespScan::TooLong(kind) => {
let (what, reply) = match kind {
RespLength::Aggregate => (
"array length",
"ERR Protocol error: invalid multibulk length",
),
RespLength::Bulk => {
("bulk length", "ERR Protocol error: invalid bulk length")
}
};
conn.log.warn(format!(
"Redis connection {} decision=fail_closed_resp_too_long: declared {} \
exceeds the frame limit; refusing and closing",
self.connection_id, what
));
self.refuse(&conn, reply).await;
return Ok(());
}
RespScan::Malformed => {
error!(
"Redis connection {} sent a frame that is not RESP2; closing",
self.connection_id
);
return Err(anyhow::anyhow!("undecodable RESP frame"));
}
}
match decode(&buffer[offset..]) {
Ok(Some((frame, consumed))) => {
answered_one = true;
trace!("Redis frame: {:?}", frame);
offset += consumed;
match self.handle_frame(&mut conn, &frame).await? {
Flow::Continue => {}
Flow::Close => return Ok(()),
}
}
Ok(None) => {
break;
}
Err(e) => {
error!("Redis decode error: {:?}", e);
return Err(e.into());
}
}
}
buffer.drain(..offset);
}
}
async fn handle_frame(&self, conn: &mut Conn<'_>, frame: &Frame) -> Result<Flow> {
let args = frame_args(frame);
let name = args
.first()
.map(|a| String::from_utf8_lossy(a).to_ascii_uppercase())
.unwrap_or_default();
let conn_key = self.connection_id.as_u32();
let proto = conn.protocol.protocol_version();
let asking = std::mem::take(&mut conn.asking);
let route = self.shared.cluster().map(|cluster| {
let slots: Vec<u16> = slot_keys(&args).iter().map(|k| key_hash_slot(k)).collect();
cluster.route(&slots, asking)
});
if let Some(route) = route {
if let Some(error) = route.error() {
debug!(
"Redis connection {} decision=cluster_redirect command={} {}",
self.connection_id, name, error
);
self.write_reply(conn, encode_error(&error)).await?;
return Ok(Flow::Continue);
}
}
if proto == 2
&& self.shared.hub.subscription_count(conn_key) > 0
&& !SUBSCRIBED_MODE_COMMANDS.contains(&name.as_str())
{
self.write_reply(
conn,
encode_error(&format!(
"ERR Can't execute '{}': only (P|S)SUBSCRIBE / (P|S)UNSUBSCRIBE / PING / QUIT / RESET are allowed in this context",
name.to_ascii_lowercase()
)),
)
.await?;
return Ok(Flow::Continue);
}
let command_str = frame_to_command_string(frame);
let command_for_log = crate::utils::truncate_for_log(&command_str, 512);
conn.log
.debug(format!("Redis command: {}", command_for_log));
let subscribed_resp2 = proto == 2 && self.shared.hub.subscription_count(conn_key) > 0;
match name.as_str() {
"HELLO" => {
return self
.hello(conn, frame, &command_str, &command_for_log)
.await
}
"RESET" => {
self.shared.hub.remove_connection(conn_key);
conn.protocol.set_protocol_version(2);
self.write_reply(conn, b"+RESET\r\n".to_vec()).await?;
return Ok(Flow::Continue);
}
"PING" if subscribed_resp2 => {
if args.len() > 2 {
self.write_reply(conn, encode_error(&wrong_arity("ping")))
.await?;
} else {
let payload = args.get(1).cloned().unwrap_or_default();
let reply = Resp::Array(vec![Resp::bulk("pong"), Resp::Bulk(payload)]);
self.write_reply(conn, reply.encode(2)).await?;
}
return Ok(Flow::Continue);
}
"QUIT" if subscribed_resp2 => {
self.write_reply(conn, b"+OK\r\n".to_vec()).await?;
let _ = conn.write_half.lock().await.shutdown().await;
return Ok(Flow::Close);
}
"SUBSCRIBE" | "PSUBSCRIBE" | "SSUBSCRIBE" => {
let kind = match name.as_str() {
"SUBSCRIBE" => Kind::Channel,
"PSUBSCRIBE" => Kind::Pattern,
_ => Kind::Shard,
};
return self
.subscribe(conn, kind, &name, &args, &command_for_log)
.await;
}
"UNSUBSCRIBE" | "PUNSUBSCRIBE" | "SUNSUBSCRIBE" => {
let kind = match name.as_str() {
"UNSUBSCRIBE" => Kind::Channel,
"PUNSUBSCRIBE" => Kind::Pattern,
_ => Kind::Shard,
};
self.shared
.hub
.unsubscribe(conn_key, &conn.outbox, proto, kind, &args[1..]);
self.write_reply(conn, Vec::new()).await?;
return Ok(Flow::Continue);
}
"PUBLISH" | "SPUBLISH" => {
return self.publish(conn, &name, &args, &command_for_log).await;
}
"PUBSUB" => {
let reply = self.pubsub_command(&args, proto);
self.write_reply(conn, reply).await?;
return Ok(Flow::Continue);
}
_ => {}
}
if self.shared.cluster.is_some() {
if let Some(reply) = self.cluster_command(conn, &name, &args, proto) {
self.write_reply(conn, reply).await?;
return Ok(Flow::Continue);
}
}
let event = Event::new(
&REDIS_COMMAND_EVENT,
serde_json::json!({
"command": command_str,
"protocol": proto,
}),
);
let answer = self.ask_model(conn, &event, &command_for_log).await;
if answer.accept || answer.deliver {
warn!(
"Redis connection {}: redis_accept_subscription / redis_deliver_message only \
answer redis_subscribe / redis_publish events; ignored for '{}'",
self.connection_id, command_for_log
);
}
let mut response = answer.output;
if response.is_empty() && !answer.close {
warn!(
"Redis connection {} decision=fail_closed_no_action command='{}': \
no response action produced, replying with an error",
self.connection_id, command_for_log
);
response = encode_error("ERR no response produced for this command");
}
if !response.is_empty() {
self.write_reply(conn, response).await?;
}
if answer.close {
let _ = conn.write_half.lock().await.shutdown().await;
return Ok(Flow::Close);
}
Ok(Flow::Continue)
}
async fn ask_model(
&self,
conn: &Conn<'_>,
event: &Event,
command_for_log: &str,
) -> ModelAnswer {
let server_id = self
.server_id
.unwrap_or_else(|| crate::state::ServerId::new(0));
let llm_result = call_llm(
&self.llm_client,
&self.app_state,
server_id,
Some(self.connection_id),
event,
conn.protocol.as_ref(),
)
.await;
let mut answer = ModelAnswer::default();
match llm_result {
Ok(execution_result) => {
let mut stack: Vec<ActionResult> = execution_result.protocol_results;
stack.reverse();
while let Some(result) = stack.pop() {
match result {
ActionResult::Output(bytes) => answer.output.extend_from_slice(&bytes),
ActionResult::Multiple(items) => {
for inner in items.into_iter().rev() {
stack.push(inner);
}
}
ActionResult::CloseConnection => {
debug!("Redis closing connection");
answer.close = true;
}
ActionResult::Custom { name, .. } if name == ACCEPT_SUBSCRIPTION => {
answer.accept = true;
}
ActionResult::Custom { name, .. } if name == DELIVER_MESSAGE => {
answer.deliver = true;
}
ActionResult::Custom { name, .. } => {
warn!(
"Redis: action result '{}' is Custom, not Output; client will receive nothing for it",
name
);
}
_ => {
}
}
}
}
Err(e) => {
let overloaded = crate::llm::is_overload_error(&e);
let message = redis_error_message(&e, overloaded);
let decision = if overloaded {
"fail_closed_llm_overloaded"
} else {
"fail_closed_llm_error"
};
conn.log.warn(format!(
"Redis connection {} decision={} command='{}' replying -{}: {}",
self.connection_id, decision, command_for_log, message, e
));
answer.output = encode_error(&message);
answer.failed = true;
}
}
answer
}
async fn hello(
&self,
conn: &mut Conn<'_>,
frame: &Frame,
command_str: &str,
command_for_log: &str,
) -> Result<Flow> {
let args = hello_args(frame).unwrap_or_default();
let mode = if self.shared.cluster.is_some() {
"cluster"
} else {
"standalone"
};
let id = u64::from(self.connection_id.as_u32());
let hello = match resp3::parse_hello(&args) {
Err(message) => {
self.write_reply(conn, encode_error(&message)).await?;
return Ok(Flow::Continue);
}
Ok(hello) => hello,
};
let target = hello.protover.unwrap_or(conn.protocol.protocol_version());
if hello.auth.is_none() {
conn.protocol.set_protocol_version(target);
conn.log.info(format!(
"Redis connection {} speaks RESP{} (HELLO)",
self.connection_id, target
));
self.write_reply(conn, resp3::hello_reply(target, id, mode))
.await?;
return Ok(Flow::Continue);
}
let event = Event::new(
&REDIS_COMMAND_EVENT,
serde_json::json!({
"command": command_str,
"protocol": conn.protocol.protocol_version(),
}),
);
let answer = self.ask_model(conn, &event, command_for_log).await;
let response = if answer.output.starts_with(b"-") {
conn.log.info(format!(
"Redis connection {} decision=model_reject: HELLO with AUTH refused; staying on \
RESP{}",
self.connection_id,
conn.protocol.protocol_version()
));
answer.output
} else if answer.output.is_empty() && !answer.close {
warn!(
"Redis connection {} decision=fail_closed_no_action command='{}': no answer to \
HELLO with AUTH; refusing",
self.connection_id, command_for_log
);
encode_error("ERR no response produced for this command")
} else {
conn.protocol.set_protocol_version(target);
resp3::hello_reply(target, id, mode)
};
self.write_reply(conn, response).await?;
if answer.close {
let _ = conn.write_half.lock().await.shutdown().await;
return Ok(Flow::Close);
}
Ok(Flow::Continue)
}
async fn subscribe(
&self,
conn: &mut Conn<'_>,
kind: Kind,
name: &str,
args: &[Vec<u8>],
command_for_log: &str,
) -> Result<Flow> {
let conn_key = self.connection_id.as_u32();
if args.len() < 2 {
self.write_reply(conn, encode_error(&wrong_arity(name)))
.await?;
return Ok(Flow::Continue);
}
let names = &args[1..];
if let Err(message) = self.shared.hub.check_subscribe(conn_key, kind, names) {
conn.log.warn(format!(
"Redis connection {} decision=fail_closed_subscription_bound: {}",
self.connection_id, message
));
self.write_reply(conn, encode_error(&message)).await?;
return Ok(Flow::Continue);
}
let event = Event::new(
&REDIS_SUBSCRIBE_EVENT,
serde_json::json!({
"command": name,
"channels": names.iter().map(|n| String::from_utf8_lossy(n).to_string()).collect::<Vec<_>>(),
"subscriptions": self.shared.hub.subscription_count(conn_key),
"protocol": conn.protocol.protocol_version(),
}),
);
let answer = self.ask_model(conn, &event, command_for_log).await;
if answer.accept && !answer.failed {
if !answer.output.is_empty() {
warn!(
"Redis connection {}: the subscription was accepted, so the reply verbs \
sent alongside redis_accept_subscription are dropped (the confirmations \
are the reply)",
self.connection_id
);
}
match self.shared.hub.subscribe(
conn_key,
&conn.outbox,
&conn.protocol.protover_cell(),
kind,
names,
) {
Ok(()) => {
debug!(
"Redis connection {} decision=model_accept command='{}'",
self.connection_id, command_for_log
);
self.write_reply(conn, Vec::new()).await?;
}
Err(message) => {
conn.log.warn(format!(
"Redis connection {} decision=fail_closed_subscription_bound: {}",
self.connection_id, message
));
self.write_reply(conn, encode_error(&message)).await?;
}
}
} else if !answer.output.is_empty() {
if !answer.failed {
debug!(
"Redis connection {} decision=model_reject command='{}': not subscribed",
self.connection_id, command_for_log
);
}
self.write_reply(conn, answer.output).await?;
} else if !answer.close {
warn!(
"Redis connection {} decision=fail_closed_no_action command='{}': no decision \
on the subscription; refusing it",
self.connection_id, command_for_log
);
self.write_reply(
conn,
encode_error("ERR no response produced for this command"),
)
.await?;
}
if answer.close {
let _ = conn.write_half.lock().await.shutdown().await;
return Ok(Flow::Close);
}
Ok(Flow::Continue)
}
async fn publish(
&self,
conn: &mut Conn<'_>,
name: &str,
args: &[Vec<u8>],
command_for_log: &str,
) -> Result<Flow> {
if args.len() != 3 {
self.write_reply(conn, encode_error(&wrong_arity(name)))
.await?;
return Ok(Flow::Continue);
}
let shard = name == "SPUBLISH";
let (channel, message) = (&args[1], &args[2]);
let kind = if shard { Kind::Shard } else { Kind::Channel };
if let Err(error) =
pubsub::check_name(kind, channel).and_then(|()| pubsub::check_message(message))
{
conn.log.warn(format!(
"Redis connection {} decision=fail_closed_message_bound: {}",
self.connection_id, error
));
self.write_reply(conn, encode_error(&error)).await?;
return Ok(Flow::Continue);
}
let event = Event::new(
&REDIS_PUBLISH_EVENT,
serde_json::json!({
"command": name,
"channel": String::from_utf8_lossy(channel),
"message": String::from_utf8_lossy(message),
"receivers": self.shared.hub.receivers(channel, shard),
"protocol": conn.protocol.protocol_version(),
}),
);
let answer = self.ask_model(conn, &event, command_for_log).await;
if answer.deliver && !answer.failed {
if !answer.output.is_empty() {
warn!(
"Redis connection {}: the message was delivered, so the reply verbs sent \
alongside redis_deliver_message are dropped (the receiver count is the \
reply)",
self.connection_id
);
}
let receivers = self.shared.hub.publish(channel, message, shard);
debug!(
"Redis connection {} decision=model_deliver command='{}' receivers={}",
self.connection_id, command_for_log, receivers
);
self.write_reply(conn, encode_integer(receivers as i64))
.await?;
} else if !answer.output.is_empty() {
if !answer.failed {
debug!(
"Redis connection {} decision=model_withhold command='{}': answered without \
delivering",
self.connection_id, command_for_log
);
}
self.write_reply(conn, answer.output).await?;
} else if !answer.close {
warn!(
"Redis connection {} decision=fail_closed_no_action command='{}': no decision \
on the message; nothing delivered",
self.connection_id, command_for_log
);
self.write_reply(
conn,
encode_error("ERR no response produced for this command"),
)
.await?;
}
if answer.close {
let _ = conn.write_half.lock().await.shutdown().await;
return Ok(Flow::Close);
}
Ok(Flow::Continue)
}
fn pubsub_command(&self, args: &[Vec<u8>], proto: u8) -> Vec<u8> {
let sub = args
.get(1)
.map(|a| String::from_utf8_lossy(a).to_ascii_uppercase())
.unwrap_or_default();
let hub = &self.shared.hub;
let channels = |shard: bool| -> Vec<u8> {
if args.len() > 3 {
return encode_error(&wrong_arity(&format!("pubsub|{}", sub.to_lowercase())));
}
let pattern = args.get(2).map(|p| p.as_slice());
if let Some(p) = pattern {
if let Err(e) = pubsub::check_name(Kind::Pattern, p) {
return encode_error(&e);
}
}
Resp::Array(
hub.active_channels(shard, pattern)
.into_iter()
.map(Resp::Bulk)
.collect(),
)
.encode(proto)
};
let numsub = |shard: bool| -> Vec<u8> {
Resp::Map(
hub.numsub(shard, &args[2..])
.into_iter()
.map(|(c, n)| (Resp::Bulk(c), Resp::Int(n as i64)))
.collect(),
)
.encode(proto)
};
match sub.as_str() {
"CHANNELS" => channels(false),
"SHARDCHANNELS" => channels(true),
"NUMSUB" => numsub(false),
"SHARDNUMSUB" => numsub(true),
"NUMPAT" if args.len() == 2 => encode_integer(hub.numpat() as i64),
"NUMPAT" => encode_error(&wrong_arity("pubsub|numpat")),
"" => encode_error(&wrong_arity("pubsub")),
other => encode_error(&format!(
"ERR unknown subcommand '{}'. Try PUBSUB HELP.",
crate::utils::truncate_for_log(&other.to_lowercase(), 64)
)),
}
}
fn cluster_command(
&self,
conn: &mut Conn<'_>,
name: &str,
args: &[Vec<u8>],
proto: u8,
) -> Option<Vec<u8>> {
match name {
"ASKING" => {
conn.asking = true;
return Some(b"+OK\r\n".to_vec());
}
"READONLY" | "READWRITE" => return Some(b"+OK\r\n".to_vec()),
"CLUSTER" => {}
_ => return None,
}
let sub = args
.get(1)
.map(|a| String::from_utf8_lossy(a).to_ascii_uppercase())
.unwrap_or_default();
let cluster = self.shared.cluster()?;
let reply = match (sub.as_str(), args.len()) {
("SLOTS", 2) => cluster.slots_reply().encode(proto),
("SHARDS", 2) => cluster.shards_reply().encode(proto),
("NODES", 2) => cluster.nodes_reply().encode(proto),
("INFO", 2) => cluster.info_reply().encode(proto),
("MYID", 2) => Resp::bulk(&cluster.me().id).encode(proto),
("KEYSLOT", 3) => encode_integer(i64::from(key_hash_slot(&args[2]))),
("SLOTS" | "SHARDS" | "NODES" | "INFO" | "MYID" | "KEYSLOT", _) => encode_error(
&wrong_arity(&format!("cluster|{}", sub.to_ascii_lowercase())),
),
_ => return None,
};
Some(reply)
}
}
fn frame_args(frame: &Frame) -> Vec<Vec<u8>> {
match frame {
Frame::Array(parts) => parts
.iter()
.map(|f| match f {
Frame::BulkString(b) | Frame::SimpleString(b) => b.clone(),
Frame::Integer(i) => i.to_string().into_bytes(),
other => format!("{other:?}").into_bytes(),
})
.collect(),
other => vec![format!("{other:?}").into_bytes()],
}
}
fn hello_args(frame: &Frame) -> Option<Vec<String>> {
let Frame::Array(parts) = frame else {
return None;
};
let text = |f: &Frame| match f {
Frame::BulkString(b) | Frame::SimpleString(b) => String::from_utf8_lossy(b).to_string(),
Frame::Integer(i) => i.to_string(),
other => format!("{other:?}"),
};
let first = parts.first().map(text)?;
first
.eq_ignore_ascii_case("HELLO")
.then(|| parts[1..].iter().map(text).collect())
}
fn frame_to_command_string(frame: &Frame) -> String {
match frame {
Frame::Array(frames) => {
let parts: Vec<String> = frames
.iter()
.map(|f| match f {
Frame::BulkString(bytes) => String::from_utf8_lossy(bytes).to_string(),
Frame::SimpleString(bytes) => String::from_utf8_lossy(bytes).to_string(),
Frame::Integer(i) => i.to_string(),
_ => format!("{:?}", f),
})
.collect();
parts.join(" ")
}
_ => format!("{:?}", frame),
}
}
fn redis_error_message(err: &anyhow::Error, overloaded: bool) -> String {
let text = crate::utils::WireFailure::classify(err).prefixed_text();
if overloaded {
format!("LOADING {text}")
} else {
format!("ERR {text}")
}
}