pub mod actions;
pub mod wire;
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, EventType};
use crate::server::connection::ConnectionId;
use crate::state::app_state::AppState;
use anyhow::Result;
use std::net::SocketAddr;
use std::pin::Pin;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::sync::{mpsc, Mutex, Notify};
use wire::Command;
pub use wire::{MAX_JOB_BYTES, MAX_LINE_BYTES};
const FIRST_COMMAND_TIMEOUT: Duration = Duration::from_secs(300);
const IDLE_TIMEOUT: Duration = Duration::from_secs(300);
const MAX_CONNECTIONS: usize = crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS;
const CONNECTION_CAP_REFUSAL: &[u8] = b"OUT_OF_MEMORY\r\n";
const UNAVAILABLE_REPLY: &[u8] = b"INTERNAL_ERROR\r\n";
const OVERLOADED_REPLY: &[u8] = b"OUT_OF_MEMORY\r\n";
pub struct BeanstalkdServer;
impl BeanstalkdServer {
pub async fn spawn_with_llm_actions(
listen_addr: SocketAddr,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
server_id: crate::state::ServerId,
first_byte_timeout_secs: Option<u64>,
idle_timeout_secs: Option<u64>,
) -> Result<SocketAddr> {
let first_timeout = first_byte_timeout_secs
.map(Duration::from_secs)
.unwrap_or(FIRST_COMMAND_TIMEOUT);
let idle_timeout = idle_timeout_secs
.map(Duration::from_secs)
.unwrap_or(IDLE_TIMEOUT);
let listener = TcpListener::bind(listen_addr).await?;
let local_addr = listener.local_addr()?;
Log::new(Some(&status_tx)).info(format!("Beanstalkd server listening on {}", local_addr));
let protocol = Arc::new(actions::BeanstalkdProtocol::new());
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,
"Beanstalkd",
Some(&status_tx),
)
.await
{
Ok((socket, peer_addr, permit)) => {
let connection_id =
ConnectionId::new(app_state.get_next_unified_id().await);
use crate::state::server::{
ConnectionState as ServerConnectionState, ConnectionStatus,
ProtocolConnectionInfo,
};
let now = crate::utils::clock::Instant::now();
app_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;
let _ = status_tx.send("__UPDATE_UI__".to_string());
Log::new(Some(&status_tx))
.info(format!("Beanstalkd client connected from {}", peer_addr));
let session = Session {
peer_addr,
llm_client: llm_client.clone(),
app_state: app_state.clone(),
status_tx: status_tx.clone(),
server_id,
protocol: protocol.clone(),
connection_id,
first_timeout,
idle_timeout,
};
let task_owner = app_state.clone();
task_owner
.spawn_server_task(server_id, async move {
let _permit = permit;
session.run(socket).await
})
.await;
}
Err(e) => {
Log::new(Some(&status_tx)).error(format!("Beanstalkd accept error: {}", e));
break;
}
}
}
});
task_registrar
.register_server_task(server_id, accept_handle)
.await;
Ok(local_addr)
}
}
#[derive(Default)]
struct WriteWatch {
writes: AtomicU64,
notify: Notify,
}
impl WriteWatch {
fn count(&self) -> u64 {
self.writes.load(Ordering::SeqCst)
}
}
struct WatchedWriter<W> {
inner: W,
watch: Arc<WriteWatch>,
}
impl<W: AsyncWrite + Unpin> AsyncWrite for WatchedWriter<W> {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<std::io::Result<usize>> {
let polled = Pin::new(&mut self.inner).poll_write(cx, buf);
if let Poll::Ready(Ok(n)) = &polled {
if *n > 0 {
self.watch.writes.fetch_add(1, Ordering::SeqCst);
self.watch.notify.notify_waiters();
}
}
polled
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
Pin::new(&mut self.inner).poll_flush(cx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
Pin::new(&mut self.inner).poll_shutdown(cx)
}
}
type Writer = Arc<Mutex<WatchedWriter<tokio::io::WriteHalf<tokio::net::TcpStream>>>>;
struct Session {
peer_addr: SocketAddr,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
server_id: crate::state::ServerId,
protocol: Arc<actions::BeanstalkdProtocol>,
connection_id: ConnectionId,
first_timeout: Duration,
idle_timeout: Duration,
}
struct Tubes {
used: String,
watched: Vec<String>,
}
enum Next {
Continue,
Wait,
Close,
}
impl Session {
async fn run(self, socket: tokio::net::TcpStream) {
let (mut reader, write_half) = tokio::io::split(socket);
let watch = Arc::new(WriteWatch::default());
let writer: Writer = Arc::new(Mutex::new(WatchedWriter {
inner: write_half,
watch: watch.clone(),
}));
let peer_rx = crate::server::peer_support::register_peer_channel(
&self.app_state,
self.server_id,
self.connection_id.as_u32(),
)
.await;
crate::server::peer_support::spawn_peer_command_task(
peer_rx,
self.protocol.clone(),
self.app_state.clone(),
self.server_id,
self.connection_id.as_u32(),
writer.clone(),
self.status_tx.clone(),
);
{
let mut framer = Framer::new(&mut reader);
self.session(&mut framer, &writer, &watch).await;
}
self.app_state
.remove_peer_handle(self.server_id, self.connection_id.as_u32())
.await;
let _ = writer.lock().await.shutdown().await;
linger(&mut reader).await;
self.app_state
.update_connection_status(
self.server_id,
self.connection_id,
crate::state::server::ConnectionStatus::Closed,
)
.await;
let _ = self.status_tx.send("__UPDATE_UI__".to_string());
}
async fn write(&self, writer: &Writer, data: &[u8]) -> std::io::Result<()> {
{
let mut w = writer.lock().await;
w.write_all(data).await?;
w.flush().await?;
}
self.record_sent(data.len()).await;
Ok(())
}
async fn record_sent(&self, n: usize) {
self.app_state
.update_connection_stats(
self.server_id,
self.connection_id,
None,
Some(n as u64),
None,
Some(1),
)
.await;
}
async fn record_received(&self, n: usize) {
self.app_state
.update_connection_stats(
self.server_id,
self.connection_id,
Some(n as u64),
None,
Some(1),
None,
)
.await;
}
async fn session<R: AsyncRead + Unpin>(
&self,
framer: &mut Framer<'_, R>,
writer: &Writer,
watch: &Arc<WriteWatch>,
) {
let log = Log::new(Some(&self.status_tx));
let mut tubes = Tubes {
used: "default".to_string(),
watched: vec!["default".to_string()],
};
let mut answered_one = false;
loop {
let timeout = if answered_one {
self.idle_timeout
} else {
self.first_timeout
};
let (line, n) = match framer.next_line(timeout).await {
LineRead::Line(line, n) => (line, n),
LineRead::Eof => {
log.info(format!("Beanstalkd client {} disconnected", self.peer_addr));
return;
}
LineRead::TimedOut => {
log.info(format!(
"Beanstalkd client {} sent no command within {}s; closing",
self.peer_addr,
timeout.as_secs()
));
return;
}
LineRead::TooLong => {
log.warn(format!(
"Beanstalkd command from {} exceeded {} bytes \
decision=fail_closed_line_too_long",
self.peer_addr, MAX_LINE_BYTES
));
let _ = self.write(writer, b"BAD_FORMAT\r\n").await;
return;
}
LineRead::Failed(e) => {
log.error(format!(
"Beanstalkd read error from {}: {}",
self.peer_addr, e
));
return;
}
};
self.record_received(n).await;
log.trace(format!(
"Beanstalkd command from {}: {:?}",
self.peer_addr, line
));
answered_one = true;
let command = wire::parse_command(&line);
let fixed: Option<Vec<u8>> = match &command {
Command::Use(tube) => {
tubes.used = tube.clone();
Some(format!("USING {tube}\r\n").into_bytes())
}
Command::Watch(tube) => {
if !tubes.watched.contains(tube) {
tubes.watched.push(tube.clone());
}
Some(format!("WATCHING {}\r\n", tubes.watched.len()).into_bytes())
}
Command::Ignore(tube) => {
if tubes.watched.len() == 1 && tubes.watched[0] == *tube {
Some(b"NOT_IGNORED\r\n".to_vec())
} else {
tubes.watched.retain(|t| t != tube);
Some(format!("WATCHING {}\r\n", tubes.watched.len()).into_bytes())
}
}
Command::ListTubeUsed => Some(format!("USING {}\r\n", tubes.used).into_bytes()),
Command::ListTubesWatched => Some(
wire::render_tube_list(&tubes.watched)
.map(String::into_bytes)
.unwrap_or_else(|_| UNAVAILABLE_REPLY.to_vec()),
),
Command::Quit => return,
Command::BadFormat => Some(b"BAD_FORMAT\r\n".to_vec()),
Command::Unknown => Some(b"UNKNOWN_COMMAND\r\n".to_vec()),
_ => None,
};
if let Some(reply) = fixed {
if self.write(writer, &reply).await.is_err() {
return;
}
continue;
}
let mut body: Option<Vec<u8>> = None;
if let Command::Put { bytes, .. } = &command {
let declared = *bytes;
match declared {
Some(n) if n <= MAX_JOB_BYTES as u64 => {
match framer.read_exact(n as usize + 2, self.idle_timeout).await {
Ok(mut data) => {
self.record_received(data.len()).await;
if !data.ends_with(b"\r\n") {
log.warn(format!(
"Beanstalkd put from {}: body not followed by CRLF \
decision=fail_closed_expected_crlf",
self.peer_addr
));
if self.write(writer, b"EXPECTED_CRLF\r\n").await.is_err() {
return;
}
continue;
}
data.truncate(n as usize);
body = Some(data);
}
Err(reason) => {
log.info(format!(
"Beanstalkd put from {}: body never arrived ({reason}); \
closing",
self.peer_addr
));
return;
}
}
}
_ => {
log.warn(format!(
"Beanstalkd put from {} declared {} bytes, over max-job-size {} \
decision=fail_closed_job_too_big",
self.peer_addr,
declared
.map(|n| n.to_string())
.unwrap_or_else(|| "more than 2^64".to_string()),
MAX_JOB_BYTES
));
if self.write(writer, b"JOB_TOO_BIG\r\n").await.is_err() {
return;
}
match declared.map(|n| n.saturating_add(2)) {
Some(skip) if skip <= wire::MAX_DISCARD_BYTES => {
if !framer.discard(skip, self.idle_timeout).await {
return;
}
continue;
}
_ => return,
}
}
}
}
match self
.answer_with_model(&command, body.as_deref(), &tubes, writer)
.await
{
Next::Continue => {}
Next::Close => return,
Next::Wait => {
if !self.wait_in_reserve(&command, framer, writer, watch).await {
return;
}
}
}
}
}
async fn wait_in_reserve<R: AsyncRead + Unpin>(
&self,
command: &Command,
framer: &mut Framer<'_, R>,
writer: &Writer,
watch: &Arc<WriteWatch>,
) -> bool {
let log = Log::new(Some(&self.status_tx));
let deadline = match command {
Command::ReserveWithTimeout(secs) => {
Some(tokio::time::Instant::now() + Duration::from_secs(u64::from(*secs)))
}
_ => None,
};
let start = watch.count();
loop {
let notified = watch.notify.notified();
tokio::pin!(notified);
notified.as_mut().enable();
if watch.count() != start {
log.info(format!(
"Beanstalkd reserve from {} answered by a reply written to the connection",
self.peer_addr
));
return true;
}
let timer = async {
match deadline {
Some(at) => tokio::time::sleep_until(at).await,
None => std::future::pending::<()>().await,
}
};
tokio::select! {
read = framer.fill() => match read {
Ok(0) => {
log.info(format!(
"Beanstalkd client {} hung up while waiting in reserve",
self.peer_addr
));
return false;
}
Ok(n) => {
log.warn(format!(
"Beanstalkd client {} sent {n} bytes while its reserve was waiting; \
the reserve is abandoned and the command is read",
self.peer_addr
));
return true;
}
Err(e) => {
log.error(format!("Beanstalkd read error from {}: {}", self.peer_addr, e));
return false;
}
},
_ = timer => {
let wrote = {
let mut w = writer.lock().await;
if watch.count() != start {
false
} else {
if w.write_all(b"TIMED_OUT\r\n").await.is_err()
|| w.flush().await.is_err()
{
return false;
}
true
}
};
if wrote {
self.record_sent(11).await;
log.info(format!(
"Beanstalkd reserve from {} decision=model_wait_timed_out",
self.peer_addr
));
}
return true;
}
_ = &mut notified => continue,
}
}
}
async fn answer_with_model(
&self,
command: &Command,
body: Option<&[u8]>,
tubes: &Tubes,
writer: &Writer,
) -> Next {
let log = Log::new(Some(&self.status_tx));
let Some((event_type, data)) = event_for(command, body, tubes) else {
return Next::Continue;
};
let event = Event::new(event_type, data);
let result = match call_llm(
&self.llm_client,
&self.app_state,
self.server_id,
Some(self.connection_id),
&event,
self.protocol.as_ref(),
)
.await
{
Ok(result) => result,
Err(e) => {
let (category, reply) = match crate::utils::WireFailure::classify(&e) {
crate::utils::WireFailure::Overloaded => ("overloaded", OVERLOADED_REPLY),
crate::utils::WireFailure::Unavailable => ("unavailable", UNAVAILABLE_REPLY),
};
log.warn(format!(
"Beanstalkd {:?} from {} decision=fail_closed_llm_error category={}",
command, self.peer_addr, category
));
log.debug(format!("Beanstalkd LLM call failed: {}", e));
return match self.write(writer, reply).await {
Ok(()) => Next::Continue,
Err(_) => Next::Close,
};
}
};
for message in &result.messages {
log.info(message);
}
let mut replies: Vec<Vec<u8>> = Vec::new();
let mut close = false;
let mut wait = false;
let mut stack = result.protocol_results;
stack.reverse();
while let Some(item) = stack.pop() {
match item {
ActionResult::Output(bytes) => replies.push(bytes),
ActionResult::CloseConnection => close = true,
ActionResult::WaitForMore => wait = true,
ActionResult::Multiple(items) => stack.extend(items.into_iter().rev()),
_ => {}
}
}
let is_reserve = matches!(command, Command::Reserve | Command::ReserveWithTimeout(_));
let Some(reply) = replies.first() else {
if wait && is_reserve {
log.info(format!(
"Beanstalkd {:?} from {} decision=model_wait",
command, self.peer_addr
));
return if close { Next::Close } else { Next::Wait };
}
let decision = if wait {
"fail_closed_mismatched_reply"
} else {
"model_silent"
};
log.warn(format!(
"Beanstalkd {:?} from {} decision={} ({} failed action(s)); answering \
INTERNAL_ERROR",
command,
self.peer_addr,
decision,
result.failures.len()
));
return match self.write(writer, UNAVAILABLE_REPLY).await {
Ok(()) if !close => Next::Continue,
_ => Next::Close,
};
};
if !wire::reply_fits(command, reply) {
let (word, _) = wire::reply_head(reply);
log.warn(format!(
"Beanstalkd {:?} from {} decision=fail_closed_mismatched_reply reply={}; \
answering INTERNAL_ERROR",
command, self.peer_addr, word
));
return match self.write(writer, UNAVAILABLE_REPLY).await {
Ok(()) if !close => Next::Continue,
_ => Next::Close,
};
}
if replies.len() > 1 {
log.warn(format!(
"Beanstalkd {:?}: the model produced {} replies to one command; sending the first",
command,
replies.len()
));
}
let (word, _) = wire::reply_head(reply);
let decision = if wire::is_refusal(reply) {
"model_reject"
} else {
"model_answer"
};
log.info(format!(
"Beanstalkd {:?} from {} decision={} reply={}",
command, self.peer_addr, decision, word
));
if self.write(writer, reply).await.is_err() || close {
return Next::Close;
}
Next::Continue
}
}
fn event_for(
command: &Command,
body: Option<&[u8]>,
tubes: &Tubes,
) -> Option<(&'static EventType, serde_json::Value)> {
use actions::{
BEANSTALKD_JOB_COMMAND_EVENT, BEANSTALKD_PUT_EVENT, BEANSTALKD_RESERVE_EVENT,
BEANSTALKD_STATS_EVENT,
};
use serde_json::json;
let used = &tubes.used;
let job = |name: &str, id: Option<u64>| {
let mut data = json!({"command": name, "tube": used});
if let Some(id) = id {
data["job_id"] = json!(id);
}
(&*BEANSTALKD_JOB_COMMAND_EVENT, data)
};
Some(match command {
Command::Put {
priority,
delay,
ttr,
..
} => {
let body = body.unwrap_or_default();
(
&*BEANSTALKD_PUT_EVENT,
json!({
"tube": used,
"priority": priority,
"delay": delay,
"ttr": ttr,
"body": String::from_utf8_lossy(body),
"body_bytes": body.len(),
}),
)
}
Command::Reserve => (
&*BEANSTALKD_RESERVE_EVENT,
json!({
"tubes": tubes.watched,
"answer_with": actions::reserve_answer_with(&tubes.watched),
}),
),
Command::ReserveWithTimeout(secs) => (
&*BEANSTALKD_RESERVE_EVENT,
json!({
"tubes": tubes.watched,
"timeout_secs": secs,
"answer_with": actions::reserve_answer_with(&tubes.watched),
}),
),
Command::ReserveJob(id) => job("reserve-job", Some(*id)),
Command::Delete(id) => job("delete", Some(*id)),
Command::Release {
id,
priority,
delay,
} => {
let (event, mut data) = job("release", Some(*id));
data["priority"] = json!(priority);
data["delay"] = json!(delay);
(event, data)
}
Command::Bury { id, priority } => {
let (event, mut data) = job("bury", Some(*id));
data["priority"] = json!(priority);
(event, data)
}
Command::Touch(id) => job("touch", Some(*id)),
Command::Peek(id) => job("peek", Some(*id)),
Command::PeekReady => job("peek-ready", None),
Command::PeekDelayed => job("peek-delayed", None),
Command::PeekBuried => job("peek-buried", None),
Command::Kick(bound) => {
let (event, mut data) = job("kick", None);
data["bound"] = json!(bound);
(event, data)
}
Command::KickJob(id) => job("kick-job", Some(*id)),
Command::PauseTube { tube, delay } => (
&*BEANSTALKD_JOB_COMMAND_EVENT,
json!({"command": "pause-tube", "tube": tube, "delay": delay}),
),
Command::Stats => (
&*BEANSTALKD_STATS_EVENT,
json!({"scope": "server", "answer_with": actions::stats_answer_with("server")}),
),
Command::StatsTube(tube) => (
&*BEANSTALKD_STATS_EVENT,
json!({
"scope": "tube",
"tube": tube,
"answer_with": actions::stats_answer_with("tube"),
}),
),
Command::StatsJob(id) => (
&*BEANSTALKD_STATS_EVENT,
json!({
"scope": "job",
"job_id": id,
"answer_with": actions::stats_answer_with("job"),
}),
),
Command::ListTubes => (
&*BEANSTALKD_STATS_EVENT,
json!({"scope": "tubes", "answer_with": actions::stats_answer_with("tubes")}),
),
_ => return None,
})
}
const LINGER_TIME: Duration = Duration::from_secs(2);
const LINGER_BYTES: usize = 64 * 1024;
async fn linger<R: AsyncRead + Unpin>(reader: &mut R) {
let deadline = tokio::time::Instant::now() + LINGER_TIME;
let mut sink = [0u8; 4096];
let mut drained = 0usize;
while drained < LINGER_BYTES {
match tokio::time::timeout_at(deadline, reader.read(&mut sink)).await {
Ok(Ok(n)) if n > 0 => drained += n,
_ => break,
}
}
}
enum LineRead {
Line(String, usize),
Eof,
TimedOut,
TooLong,
Failed(std::io::Error),
}
struct Framer<'a, R> {
reader: &'a mut R,
pending: Vec<u8>,
chunk: Vec<u8>,
}
impl<'a, R: AsyncRead + Unpin> Framer<'a, R> {
fn new(reader: &'a mut R) -> Self {
Self {
reader,
pending: Vec::new(),
chunk: vec![0u8; 4096],
}
}
async fn fill(&mut self) -> std::io::Result<usize> {
let n = self.reader.read(&mut self.chunk).await?;
self.pending.extend_from_slice(&self.chunk[..n]);
Ok(n)
}
async fn next_line(&mut self, timeout: Duration) -> LineRead {
loop {
if let Some(idx) = self.pending.iter().position(|b| *b == b'\n') {
let consumed = idx + 1;
if consumed > MAX_LINE_BYTES {
return LineRead::TooLong;
}
let line: Vec<u8> = self.pending.drain(..consumed).collect();
let text = String::from_utf8_lossy(&line[..idx]);
return LineRead::Line(text.trim_end_matches('\r').to_string(), consumed);
}
if self.pending.len() >= MAX_LINE_BYTES {
return LineRead::TooLong;
}
match tokio::time::timeout(timeout, self.fill()).await {
Err(_) => return LineRead::TimedOut,
Ok(Ok(0)) => {
if self.pending.is_empty() {
return LineRead::Eof;
}
let line = std::mem::take(&mut self.pending);
let consumed = line.len();
let text = String::from_utf8_lossy(&line);
return LineRead::Line(text.trim_end_matches('\r').to_string(), consumed);
}
Ok(Ok(_)) => {}
Ok(Err(e)) => return LineRead::Failed(e),
}
}
}
async fn read_exact(&mut self, n: usize, timeout: Duration) -> Result<Vec<u8>, String> {
while self.pending.len() < n {
match tokio::time::timeout(timeout, self.fill()).await {
Err(_) => return Err(format!("no bytes for {}s", timeout.as_secs())),
Ok(Ok(0)) => return Err("EOF".to_string()),
Ok(Ok(_)) => {}
Ok(Err(e)) => return Err(e.to_string()),
}
}
Ok(self.pending.drain(..n).collect())
}
async fn discard(&mut self, mut n: u64, timeout: Duration) -> bool {
loop {
let take = (self.pending.len() as u64).min(n) as usize;
self.pending.drain(..take);
n -= take as u64;
if n == 0 {
return true;
}
match tokio::time::timeout(timeout, self.fill()).await {
Ok(Ok(k)) if k > 0 => {}
_ => return false,
}
}
}
}