use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::{Duration, Instant};
use anyhow::{Context, Result};
use futures::{SinkExt, StreamExt};
use serde_json::json;
use tokio::net::{UnixListener, UnixStream};
use tokio::task::{JoinError, JoinSet};
use tokio_util::codec::{Framed, LinesCodec, LinesCodecError};
use tokio_util::sync::CancellationToken;
use super::accept::{AcceptBackoff, ConnectionSource};
use super::lifecycle;
use super::paths;
use super::protocol::{DaemonEnvelope, DaemonReply, StatusReport, DAEMON_SERVICE, MAX_LINE_BYTES};
use super::registry::ServiceRegistry;
use super::service::ServiceStream;
use super::single_instance;
const DRAIN_TIMEOUT: Duration = Duration::from_secs(5);
const IDLE_TIMEOUT: Duration = Duration::from_secs(60);
const ENV_STREAM_TICK: &str = "OMNI_DEV_DAEMON_STREAM_TICK";
const DEFAULT_STREAM_TICK: Duration = Duration::from_secs(10);
pub(crate) fn stream_tick() -> Duration {
duration_secs_from_env(ENV_STREAM_TICK, DEFAULT_STREAM_TICK)
}
pub(crate) fn duration_secs_from_env(var: &str, default: Duration) -> Duration {
duration_secs_from_raw(std::env::var(var).ok(), default)
}
fn duration_secs_from_raw(raw: Option<String>, default: Duration) -> Duration {
raw.and_then(|s| s.trim().parse::<u64>().ok())
.filter(|&secs| secs > 0)
.map_or(default, Duration::from_secs)
}
#[derive(Debug, Clone)]
pub struct DaemonOptions {
pub socket_path: PathBuf,
}
pub async fn run(registry: ServiceRegistry, opts: DaemonOptions) -> Result<()> {
run_with_shutdown(Arc::new(registry), opts, CancellationToken::new()).await
}
pub async fn run_with_shutdown(
registry: Arc<ServiceRegistry>,
opts: DaemonOptions,
shutdown: CancellationToken,
) -> Result<()> {
lifecycle::raise_nofile_limit();
if let Some(parent) = opts.socket_path.parent() {
paths::ensure_dir_0700(parent)?;
}
tighten_daemon_log(&opts.socket_path);
paths::check_socket_path_len(&opts.socket_path)?;
let (listener, socket_activated) = acquire_listener(&opts.socket_path).await?;
tracing::info!("daemon listening on {}", opts.socket_path.display());
lifecycle::install_signal_handlers(shutdown.clone());
let mut conns: JoinSet<()> = JoinSet::new();
accept_loop(&listener, ®istry, &shutdown, &mut conns).await;
drop(listener);
if !socket_activated {
remove_socket(&opts.socket_path);
}
drain_connections(&mut conns, DRAIN_TIMEOUT).await;
tracing::info!("daemon shutting down; draining services");
registry.shutdown_all().await;
Ok(())
}
async fn accept_loop<S: ConnectionSource>(
source: &S,
registry: &Arc<ServiceRegistry>,
shutdown: &CancellationToken,
conns: &mut JoinSet<()>,
) {
let mut backoff = AcceptBackoff::new();
loop {
tokio::select! {
() = shutdown.cancelled() => break,
accepted = source.accept() => match accepted {
Ok(stream) => {
if let Some(recovery) = backoff.on_success(Instant::now()) {
registry.credit_accept_outage(recovery.outage);
if recovery.announce {
tracing::info!(
failures = recovery.failures,
outage = ?recovery.duration,
"daemon accept recovered"
);
}
}
conns.spawn(handle_connection(
stream,
registry.clone(),
shutdown.clone(),
IDLE_TIMEOUT,
));
}
Err(e) => {
let step = backoff.on_failure(Instant::now());
registry.credit_accept_outage(step.outage);
if step.log {
tracing::warn!(
"daemon accept error: {e} (backing off; further errors are not logged for 30s)"
);
}
tokio::select! {
() = shutdown.cancelled() => break,
() = tokio::time::sleep(step.delay) => {}
}
}
},
joined = conns.join_next(), if !conns.is_empty() => {
if let Some(result) = joined {
note_reaped(result);
}
}
}
}
}
async fn acquire_listener(socket_path: &Path) -> Result<(UnixListener, bool)> {
#[cfg(target_os = "macos")]
if let Some(listener) = super::launchd::launchd_listener("Listener")? {
tracing::info!("daemon adopting launchd-activated control socket");
return Ok((listener, true));
}
#[cfg(target_os = "linux")]
if let Some(listener) = super::systemd::systemd_listener()? {
tracing::info!("daemon adopting systemd-activated control socket");
return Ok((listener, true));
}
let listener = single_instance::bind_or_reclaim(socket_path).await?;
Ok((listener, false))
}
fn tighten_daemon_log(socket_path: &Path) {
let log_path = paths::log_path_for_socket(socket_path);
if !log_path.exists() {
return;
}
if let Err(e) = paths::set_file_0600(&log_path) {
tracing::warn!("failed to tighten {} to 0600: {e}", log_path.display());
}
}
fn remove_socket(path: &Path) {
if let Err(e) = std::fs::remove_file(path) {
if e.kind() != std::io::ErrorKind::NotFound {
tracing::warn!("failed to remove socket {}: {e}", path.display());
}
}
}
fn note_reaped(result: Result<(), JoinError>) {
if let Err(e) = result {
if e.is_panic() {
tracing::warn!("daemon connection task panicked: {e}");
}
}
}
async fn drain_connections(conns: &mut JoinSet<()>, timeout: Duration) {
let count = conns.len();
if count == 0 {
return;
}
tracing::info!("draining {count} in-flight connection(s)");
let drain = async {
while let Some(result) = conns.join_next().await {
note_reaped(result);
}
};
if tokio::time::timeout(timeout, drain).await.is_err() {
tracing::warn!(
"timed out draining connections after {timeout:?}; aborting {} straggler(s)",
conns.len()
);
conns.abort_all();
while conns.join_next().await.is_some() {}
}
}
async fn handle_connection(
stream: UnixStream,
registry: Arc<ServiceRegistry>,
shutdown: CancellationToken,
idle_timeout: Duration,
) {
let mut framed = Framed::new(stream, LinesCodec::new_with_max_length(MAX_LINE_BYTES));
loop {
let line = match tokio::time::timeout(idle_timeout, framed.next()).await {
Ok(Some(line)) => line,
Ok(None) => break,
Err(_elapsed) => {
tracing::debug!(?idle_timeout, "closing an idle control-socket connection");
break;
}
};
let line = match line {
Ok(line) => line,
Err(e) => {
let msg = match e {
LinesCodecError::MaxLineLengthExceeded => {
format!("request line exceeds the {MAX_LINE_BYTES}-byte limit")
}
LinesCodecError::Io(io) => format!("read error: {io}"),
};
let _ = send_reply(&mut framed, DaemonReply::err(msg)).await;
break;
}
};
let envelope: DaemonEnvelope = match serde_json::from_str(&line) {
Ok(envelope) => envelope,
Err(e) => {
if !send_reply(
&mut framed,
DaemonReply::err(format!("invalid envelope: {e}")),
)
.await
{
break;
}
continue;
}
};
if let Some(name) = envelope.service.as_deref() {
if name != DAEMON_SERVICE {
if let Some(stream) = registry.subscribe(name, &envelope.op, &envelope.payload) {
run_stream(&mut framed, stream, &shutdown, name, &envelope.op).await;
return;
}
}
}
let reply = dispatch_envelope(envelope, ®istry, &shutdown).await;
if !send_reply(&mut framed, reply).await {
break;
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StreamEnd {
ClientCancel,
ClientEof,
ReadDecodeError,
DaemonShutdown,
WriteFailure,
}
fn inbound_end(line: Option<Result<String, LinesCodecError>>) -> StreamEnd {
match line {
Some(Ok(_)) => StreamEnd::ClientCancel,
None => StreamEnd::ClientEof,
Some(Err(_)) => StreamEnd::ReadDecodeError,
}
}
async fn run_stream(
framed: &mut Framed<UnixStream, LinesCodec>,
mut stream: Box<dyn ServiceStream>,
shutdown: &CancellationToken,
service: &str,
op: &str,
) -> StreamEnd {
tracing::debug!(service, op, "subscription_started");
let mut pushed = 0_u64;
let mut suppressed = 0_u64;
let mut last = stream.snapshot().await;
let end = if !send_reply(framed, DaemonReply::ok(last.clone())).await {
StreamEnd::WriteFailure
} else {
pushed += 1;
tracing::trace!(
service,
op,
trigger = "initial",
outcome = "pushed",
"subscription_sample"
);
let mut tick = tokio::time::interval(stream_tick());
tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
tick.tick().await;
loop {
let trigger = tokio::select! {
() = stream.changed() => "change_notification",
_ = tick.tick() => "periodic_tick",
line = framed.next() => {
if let Some(Err(error)) = &line {
tracing::debug!(service, op, %error, "subscription_read_failed");
}
break inbound_end(line);
},
() = shutdown.cancelled() => break StreamEnd::DaemonShutdown,
};
let snap = stream.snapshot().await;
if snap != last {
if !send_reply(framed, DaemonReply::ok(snap.clone())).await {
break StreamEnd::WriteFailure;
}
pushed = pushed.saturating_add(1);
last = snap;
tracing::trace!(
service,
op,
trigger,
outcome = "pushed",
"subscription_sample"
);
} else {
suppressed = suppressed.saturating_add(1);
tracing::trace!(
service,
op,
trigger,
outcome = "suppressed_identical",
"subscription_sample"
);
}
}
};
tracing::debug!(service, op, outcome = ?end, pushed, suppressed, "subscription_ended");
end
}
async fn send_reply(framed: &mut Framed<UnixStream, LinesCodec>, reply: DaemonReply) -> bool {
let encoded = match serde_json::to_string(&reply) {
Ok(encoded) => encoded,
Err(e) => {
tracing::warn!("failed to encode daemon reply: {e}");
return false;
}
};
if let Err(e) = framed.send(encoded).await {
tracing::warn!(error = %e, "daemon_client_write_failed");
return false;
}
true
}
async fn dispatch_envelope(
envelope: DaemonEnvelope,
registry: &ServiceRegistry,
shutdown: &CancellationToken,
) -> DaemonReply {
match envelope.service.as_deref() {
None | Some(DAEMON_SERVICE) => {
let reply = handle_builtin(&envelope.op, registry, shutdown).await;
if !reply.ok {
tracing::warn!(service = DAEMON_SERVICE, op = %envelope.op, error = ?reply.error, "daemon_dispatch_failed");
}
reply
}
Some(name) => {
let dispatch = registry.dispatch(name, &envelope.op, envelope.payload);
let result = match envelope.origin_invocation_id {
Some(origin) => crate::request_log::scope_origin_id(origin, dispatch).await,
None => dispatch.await,
};
match result {
Ok(payload) => DaemonReply::ok(payload),
Err(e) => {
tracing::warn!(service = name, op = %envelope.op, error = %format!("{e:#}"), "daemon_dispatch_failed");
DaemonReply::err(format!("{e:#}"))
}
}
}
}
}
async fn handle_builtin(
op: &str,
registry: &ServiceRegistry,
shutdown: &CancellationToken,
) -> DaemonReply {
match op {
"ping" => {
let mut payload = serde_json::Map::new();
payload.insert("pong".to_string(), json!(true));
payload.insert("version".to_string(), json!(crate::VERSION));
if let Ok(serde_json::Value::Object(prov)) =
serde_json::to_value(crate::build_info::provenance())
{
payload.extend(prov);
}
DaemonReply::ok(serde_json::Value::Object(payload))
}
"status" => {
let mut report = StatusReport::current(registry.statuses().await);
report.github_rate_limit = registry.github_rate_limit();
match serde_json::to_value(report) {
Ok(payload) => DaemonReply::ok(payload),
Err(e) => DaemonReply::err(format!("failed to encode status: {e}")),
}
}
"shutdown" => {
shutdown.cancel();
DaemonReply::ok(json!({ "stopping": true }))
}
other => DaemonReply::err(format!("unknown daemon op: {other}")),
}
}
pub fn resolve_socket(socket: Option<PathBuf>) -> Result<PathBuf> {
match socket {
Some(path) => Ok(path),
None => paths::socket_path().context("failed to resolve the default daemon socket path"),
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
#[tokio::test]
async fn subscriptions_report_all_terminal_causes() {
use tokio::io::AsyncWriteExt;
for expected in [
StreamEnd::ClientCancel,
StreamEnd::ClientEof,
StreamEnd::ReadDecodeError,
StreamEnd::DaemonShutdown,
StreamEnd::WriteFailure,
] {
let (mut client, server) = UnixStream::pair().unwrap();
let (_tx, rx) = watch::channel(0_u64);
let fake = FakeStream {
rx,
snap: Arc::new(StdMutex::new(json!({"n":0}))),
};
let token = CancellationToken::new();
let shutdown = token.clone();
let task = tokio::spawn(async move {
let mut framed = Framed::new(server, LinesCodec::new_with_max_length(64));
crate::test_support::capture_future_at(
tracing::Level::TRACE,
run_stream(
&mut framed,
Box::new(fake),
&shutdown,
"sessions",
"subscribe",
),
)
.await
});
if expected == StreamEnd::WriteFailure {
drop(client);
} else {
let mut reader = BufReader::new(&mut client);
let _ = read_reply(&mut reader).await;
drop(reader);
match expected {
StreamEnd::ClientCancel => {
client.write_all(b"cancel\n").await.unwrap();
}
StreamEnd::ClientEof => {
client.shutdown().await.unwrap();
}
StreamEnd::ReadDecodeError => {
client.write_all(&[b'x'; 65]).await.unwrap();
}
StreamEnd::DaemonShutdown => token.cancel(),
StreamEnd::WriteFailure => unreachable!(),
}
}
let (cause, logs) = tokio::time::timeout(Duration::from_secs(2), task)
.await
.unwrap()
.unwrap();
assert_eq!(cause, expected);
assert!(logs.contains("subscription_started"));
assert!(logs.contains("subscription_ended"));
assert!(logs.contains(&format!("outcome={expected:?}")));
if expected != StreamEnd::WriteFailure {
assert!(logs.contains("pushed=1"));
}
}
}
#[test]
fn inbound_subscription_end_reasons_are_distinct() {
assert_eq!(
super::inbound_end(Some(Ok("cancel".into()))),
super::StreamEnd::ClientCancel
);
assert_eq!(super::inbound_end(None), super::StreamEnd::ClientEof);
assert_eq!(
super::inbound_end(Some(Err(LinesCodecError::MaxLineLengthExceeded))),
super::StreamEnd::ReadDecodeError
);
}
use super::*;
#[test]
fn tighten_daemon_log_sets_0600_and_tolerates_absence() {
use std::os::unix::fs::PermissionsExt;
let dir = tempfile::tempdir().unwrap();
let socket = dir.path().join("daemon.sock");
tighten_daemon_log(&socket);
let log = paths::log_path_for_socket(&socket);
std::fs::write(&log, b"daemon listening on ...\n").unwrap();
std::fs::set_permissions(&log, std::fs::Permissions::from_mode(0o644)).unwrap();
tighten_daemon_log(&socket);
assert_eq!(
std::fs::metadata(&log).unwrap().permissions().mode() & 0o777,
0o600
);
}
#[tokio::test]
async fn drain_connections_returns_immediately_when_empty() {
let mut conns: JoinSet<()> = JoinSet::new();
drain_connections(&mut conns, Duration::from_secs(5)).await;
assert!(conns.is_empty());
}
#[tokio::test]
async fn drain_connections_awaits_completed_tasks() {
let mut conns: JoinSet<()> = JoinSet::new();
conns.spawn(async {});
drain_connections(&mut conns, Duration::from_secs(5)).await;
assert!(conns.is_empty());
}
#[tokio::test]
async fn drain_connections_times_out_and_aborts_stragglers() {
let mut conns: JoinSet<()> = JoinSet::new();
conns.spawn(std::future::pending::<()>());
drain_connections(&mut conns, Duration::from_millis(50)).await;
assert!(
conns.is_empty(),
"straggler should have been aborted and joined"
);
}
#[tokio::test]
async fn note_reaped_ignores_success_and_logs_panic() {
note_reaped(Ok(()));
let mut js: JoinSet<()> = JoinSet::new();
js.spawn(async { panic!("boom") });
let result = js.join_next().await.unwrap();
assert!(result.is_err());
note_reaped(result);
}
#[test]
fn duration_secs_from_raw_parses_seconds_and_falls_back_on_junk() {
let default = Duration::from_secs(10);
assert_eq!(duration_secs_from_raw(None, default), default);
assert_eq!(
duration_secs_from_raw(Some(String::new()), default),
default
);
assert_eq!(
duration_secs_from_raw(Some("garbage".to_string()), default),
default
);
assert_eq!(
duration_secs_from_raw(Some(" 0 ".to_string()), default),
default
);
assert_eq!(
duration_secs_from_raw(Some("30".to_string()), default),
Duration::from_secs(30)
);
assert_eq!(
duration_secs_from_raw(Some(" 5\n".to_string()), default),
Duration::from_secs(5)
);
}
use std::sync::Mutex as StdMutex;
use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::sync::watch;
struct FakeStream {
rx: watch::Receiver<u64>,
snap: Arc<StdMutex<serde_json::Value>>,
}
#[async_trait::async_trait]
impl ServiceStream for FakeStream {
async fn changed(&mut self) {
if self.rx.changed().await.is_err() {
std::future::pending::<()>().await;
}
}
async fn snapshot(&self) -> serde_json::Value {
self.snap.lock().unwrap().clone()
}
}
async fn read_reply<R: tokio::io::AsyncBufRead + Unpin>(reader: &mut R) -> DaemonReply {
let mut line = String::new();
let n = reader.read_line(&mut line).await.unwrap();
assert!(n > 0, "expected a reply line, got EOF");
serde_json::from_str(line.trim_end()).unwrap()
}
#[tokio::test]
async fn run_stream_pushes_initial_then_deltas_and_dedupes() {
let (client, server) = UnixStream::pair().unwrap();
let (tx, rx) = watch::channel(0u64);
let snap = Arc::new(StdMutex::new(json!({ "n": 0 })));
let fake = FakeStream {
rx,
snap: snap.clone(),
};
let shutdown = CancellationToken::new();
let server_shutdown = shutdown.clone();
let server_task = tokio::spawn(async move {
let mut framed = Framed::new(server, LinesCodec::new_with_max_length(MAX_LINE_BYTES));
run_stream(
&mut framed,
Box::new(fake),
&server_shutdown,
"fake",
"subscribe",
)
.await;
});
let mut reader = BufReader::new(client);
let initial = read_reply(&mut reader).await;
assert!(initial.ok);
assert_eq!(initial.payload, json!({ "n": 0 }));
tx.send(1).unwrap(); *snap.lock().unwrap() = json!({ "n": 1 });
tx.send(2).unwrap(); let delta = read_reply(&mut reader).await;
assert_eq!(delta.payload, json!({ "n": 1 }));
shutdown.cancel();
let mut tail = String::new();
let n = reader.read_line(&mut tail).await.unwrap();
assert_eq!(n, 0, "stream should close cleanly on shutdown");
server_task.await.unwrap();
}
#[tokio::test]
async fn run_stream_ends_when_client_sends_a_line() {
use tokio::io::AsyncWriteExt;
let (mut client, server) = UnixStream::pair().unwrap();
let (_tx, rx) = watch::channel(0u64);
let snap = Arc::new(StdMutex::new(json!({ "n": 0 })));
let fake = FakeStream { rx, snap };
let shutdown = CancellationToken::new();
let server_shutdown = shutdown.clone();
let server_task = tokio::spawn(async move {
let mut framed = Framed::new(server, LinesCodec::new_with_max_length(MAX_LINE_BYTES));
run_stream(
&mut framed,
Box::new(fake),
&server_shutdown,
"fake",
"subscribe",
)
.await;
});
let mut reader = BufReader::new(&mut client);
let _initial = read_reply(&mut reader).await;
drop(reader);
client.write_all(b"cancel\n").await.unwrap();
tokio::time::timeout(Duration::from_secs(2), server_task)
.await
.expect("run_stream should end after a client line")
.unwrap();
}
#[tokio::test]
async fn handle_connection_rejects_bad_envelope_then_serves_normal_op() {
use tokio::io::AsyncWriteExt;
let (client, server) = UnixStream::pair().unwrap();
let mut registry = ServiceRegistry::new();
registry.register(Arc::new(
crate::daemon::services::worktrees::WorktreesService::new(),
));
let shutdown = CancellationToken::new();
let task = tokio::spawn(handle_connection(
server,
Arc::new(registry),
shutdown,
IDLE_TIMEOUT,
));
let (read_half, mut write_half) = client.into_split();
let mut reader = BufReader::new(read_half);
write_half.write_all(b"not json\n").await.unwrap();
let bad = read_reply(&mut reader).await;
assert!(!bad.ok);
assert!(bad.error.unwrap().contains("invalid envelope"));
let env = serde_json::to_string(&DaemonEnvelope::service(
"worktrees",
"list",
serde_json::Value::Null,
))
.unwrap();
write_half.write_all(env.as_bytes()).await.unwrap();
write_half.write_all(b"\n").await.unwrap();
let listed = read_reply(&mut reader).await;
assert!(listed.ok);
assert!(listed.payload.get("windows").is_some());
drop(write_half);
drop(reader);
tokio::time::timeout(Duration::from_secs(2), task)
.await
.expect("handler should end after the client hangs up")
.unwrap();
}
#[tokio::test]
async fn handle_connection_enters_streaming_for_subscribe() {
use tokio::io::AsyncWriteExt;
let (client, server) = UnixStream::pair().unwrap();
let mut registry = ServiceRegistry::new();
registry.register(Arc::new(
crate::daemon::services::worktrees::WorktreesService::new(),
));
let shutdown = CancellationToken::new();
let task = tokio::spawn(handle_connection(
server,
Arc::new(registry),
shutdown.clone(),
IDLE_TIMEOUT,
));
let (read_half, mut write_half) = client.into_split();
let mut reader = BufReader::new(read_half);
let env = serde_json::to_string(&DaemonEnvelope::service(
"worktrees",
"subscribe",
serde_json::Value::Null,
))
.unwrap();
write_half.write_all(env.as_bytes()).await.unwrap();
write_half.write_all(b"\n").await.unwrap();
let initial = read_reply(&mut reader).await;
assert!(initial.ok);
assert_eq!(initial.payload, json!({ "repos": [], "show_closed": true }));
shutdown.cancel();
tokio::time::timeout(Duration::from_secs(2), task)
.await
.expect("shutdown should end the streaming handler")
.unwrap();
}
#[tokio::test]
async fn run_stream_ends_when_a_pushed_delta_cannot_be_written() {
let (mut client, server) = UnixStream::pair().unwrap();
let server = server.into_std().unwrap();
let hangup = server.try_clone().unwrap();
let server = UnixStream::from_std(server).unwrap();
let (tx, rx) = watch::channel(0u64);
let snap = Arc::new(StdMutex::new(json!({ "n": 0 })));
let fake = FakeStream {
rx,
snap: Arc::clone(&snap),
};
let shutdown = CancellationToken::new();
let task = tokio::spawn(async move {
let mut framed = Framed::new(server, LinesCodec::new_with_max_length(MAX_LINE_BYTES));
run_stream(&mut framed, Box::new(fake), &shutdown, "fake", "subscribe").await
});
let mut reader = BufReader::new(&mut client);
read_reply(&mut reader).await;
hangup.shutdown(std::net::Shutdown::Write).unwrap();
*snap.lock().unwrap() = json!({ "n": 1 });
tx.send(1).unwrap();
let end = tokio::time::timeout(Duration::from_secs(2), task)
.await
.expect("run_stream should end when a push cannot be written")
.unwrap();
assert_eq!(end, StreamEnd::WriteFailure);
}
#[tokio::test]
async fn run_stream_returns_when_initial_send_fails() {
let (client, server) = UnixStream::pair().unwrap();
drop(client);
let (_tx, rx) = watch::channel(0u64);
let fake = FakeStream {
rx,
snap: Arc::new(StdMutex::new(json!({ "n": 0 }))),
};
let shutdown = CancellationToken::new();
let mut framed = Framed::new(server, LinesCodec::new_with_max_length(MAX_LINE_BYTES));
tokio::time::timeout(
Duration::from_secs(2),
run_stream(&mut framed, Box::new(fake), &shutdown, "fake", "subscribe"),
)
.await
.expect("run_stream should return promptly when the initial send fails");
}
fn ping_line() -> &'static [u8] {
b"{\"op\":\"ping\"}\n"
}
#[tokio::test]
async fn handle_connection_closes_a_connection_that_never_sends() {
use tokio::io::AsyncReadExt;
let (mut client, server) = UnixStream::pair().unwrap();
let task = tokio::spawn(handle_connection(
server,
worktrees_registry(),
CancellationToken::new(),
Duration::from_millis(50),
));
tokio::time::timeout(Duration::from_secs(2), task)
.await
.expect("an idle connection should be closed")
.unwrap();
let mut buf = Vec::new();
assert_eq!(client.read_to_end(&mut buf).await.unwrap(), 0);
}
#[tokio::test]
async fn handle_connection_stays_open_while_requests_keep_arriving() {
use tokio::io::AsyncWriteExt;
let (client, server) = UnixStream::pair().unwrap();
let task = tokio::spawn(handle_connection(
server,
worktrees_registry(),
CancellationToken::new(),
Duration::from_millis(600),
));
let (read_half, mut write_half) = client.into_split();
let mut reader = BufReader::new(read_half);
for _ in 0..3 {
write_half.write_all(ping_line()).await.unwrap();
assert!(read_reply(&mut reader).await.ok);
tokio::time::sleep(Duration::from_millis(250)).await;
}
assert!(
!task.is_finished(),
"the connection was closed while active"
);
}
#[tokio::test]
async fn handle_connection_does_not_time_out_a_subscription() {
use tokio::io::AsyncWriteExt;
let (client, server) = UnixStream::pair().unwrap();
let shutdown = CancellationToken::new();
let task = tokio::spawn(handle_connection(
server,
worktrees_registry(),
shutdown.clone(),
Duration::from_millis(50),
));
let (read_half, mut write_half) = client.into_split();
let mut reader = BufReader::new(read_half);
let env = serde_json::to_string(&DaemonEnvelope::service(
"worktrees",
"subscribe",
serde_json::Value::Null,
))
.unwrap();
write_half.write_all(env.as_bytes()).await.unwrap();
write_half.write_all(b"\n").await.unwrap();
assert!(read_reply(&mut reader).await.ok);
tokio::time::sleep(Duration::from_millis(250)).await;
assert!(!task.is_finished(), "the idle timeout cut a subscription");
shutdown.cancel();
tokio::time::timeout(Duration::from_secs(2), task)
.await
.expect("shutdown still ends the stream")
.unwrap();
}
#[tokio::test]
async fn handle_connection_never_interrupts_a_request_in_progress() {
use crate::daemon::testutil::StubService;
use tokio::io::AsyncWriteExt;
let (client, server) = UnixStream::pair().unwrap();
let mut registry = ServiceRegistry::new();
registry.register(Arc::new(StubService::slow(
"slow",
Duration::from_millis(250),
)));
let task = tokio::spawn(handle_connection(
server,
Arc::new(registry),
CancellationToken::new(),
Duration::from_millis(50),
));
let (read_half, mut write_half) = client.into_split();
let mut reader = BufReader::new(read_half);
let env = serde_json::to_string(&DaemonEnvelope::service(
"slow",
"work",
serde_json::Value::Null,
))
.unwrap();
write_half.write_all(env.as_bytes()).await.unwrap();
write_half.write_all(b"\n").await.unwrap();
let reply = read_reply(&mut reader).await;
assert!(reply.ok);
assert_eq!(reply.payload, json!({ "done": true }));
drop(task);
}
fn emfile() -> std::io::Error {
std::io::Error::from_raw_os_error(24)
}
struct ScriptedSource {
script: StdMutex<std::collections::VecDeque<std::io::Result<UnixStream>>>,
calls: std::sync::atomic::AtomicUsize,
}
impl ScriptedSource {
fn new(script: Vec<std::io::Result<UnixStream>>) -> Self {
Self {
script: StdMutex::new(script.into()),
calls: std::sync::atomic::AtomicUsize::new(0),
}
}
fn calls(&self) -> usize {
self.calls.load(std::sync::atomic::Ordering::SeqCst)
}
}
#[async_trait::async_trait]
impl ConnectionSource for ScriptedSource {
async fn accept(&self) -> std::io::Result<UnixStream> {
self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
let next = self.script.lock().unwrap().pop_front();
match next {
Some(result) => result,
None => std::future::pending().await,
}
}
}
fn worktrees_registry() -> Arc<ServiceRegistry> {
let mut registry = ServiceRegistry::new();
registry.register(Arc::new(
crate::daemon::services::worktrees::WorktreesService::new(),
));
Arc::new(registry)
}
#[tokio::test]
async fn accept_loop_backs_off_on_errors_then_serves_the_next_connection() {
use tokio::io::AsyncWriteExt;
let (client, server) = UnixStream::pair().unwrap();
let source = ScriptedSource::new(vec![
Err(emfile()),
Err(emfile()),
Err(emfile()),
Ok(server),
]);
let registry = worktrees_registry();
let shutdown = CancellationToken::new();
let mut conns: JoinSet<()> = JoinSet::new();
let started = Instant::now();
let serving = async {
accept_loop(&source, ®istry, &shutdown, &mut conns).await;
};
let talking = async {
let (read_half, mut write_half) = client.into_split();
let mut reader = BufReader::new(read_half);
write_half.write_all(b"{\"op\":\"ping\"}\n").await.unwrap();
let reply = read_reply(&mut reader).await;
shutdown.cancel();
reply
};
let ((), reply) = tokio::time::timeout(Duration::from_secs(5), async {
tokio::join!(serving, talking)
})
.await
.expect("the loop should recover and serve");
assert!(reply.ok, "{reply:?}");
let elapsed = started.elapsed();
assert!(
elapsed >= Duration::from_millis(35),
"retried too fast: {elapsed:?}"
);
assert!(source.calls() <= 5, "spun: {} accept calls", source.calls());
}
#[tokio::test]
async fn accept_loop_credits_the_outage_to_the_services() {
use crate::daemon::testutil::StubService;
use tokio::io::AsyncWriteExt;
let (client, server) = UnixStream::pair().unwrap();
let source = ScriptedSource::new(vec![
Err(emfile()),
Err(emfile()),
Err(emfile()),
Ok(server),
]);
let recorder = Arc::new(StubService::new("recorder"));
let mut registry = ServiceRegistry::new();
registry.register(recorder.clone());
let registry = Arc::new(registry);
let shutdown = CancellationToken::new();
let mut conns: JoinSet<()> = JoinSet::new();
let talking = async {
let (read_half, mut write_half) = client.into_split();
let mut reader = BufReader::new(read_half);
write_half.write_all(b"{\"op\":\"ping\"}\n").await.unwrap();
read_reply(&mut reader).await;
shutdown.cancel();
};
tokio::time::timeout(Duration::from_secs(5), async {
tokio::join!(
accept_loop(&source, ®istry, &shutdown, &mut conns),
talking
)
})
.await
.expect("the loop should recover and serve");
let credited = recorder.credited();
assert_eq!(credited.len(), 4, "{credited:?}");
assert_eq!(credited[0], Duration::ZERO);
let total: Duration = credited.iter().sum();
assert!(total >= Duration::from_millis(35), "credited {total:?}");
assert!(total < Duration::from_secs(30), "credited {total:?}");
}
#[tokio::test]
async fn accept_loop_recovers_from_an_unlogged_burst_and_still_credits_it() {
use crate::daemon::testutil::StubService;
use tokio::io::AsyncWriteExt;
let (client_a, server_a) = UnixStream::pair().unwrap();
let (client_b, server_b) = UnixStream::pair().unwrap();
let source = ScriptedSource::new(vec![
Err(emfile()),
Ok(server_a),
Err(emfile()),
Ok(server_b),
]);
let recorder = Arc::new(StubService::new("recorder"));
let mut registry = ServiceRegistry::new();
registry.register(recorder.clone());
let registry = Arc::new(registry);
let shutdown = CancellationToken::new();
let mut conns: JoinSet<()> = JoinSet::new();
let talking = async {
for client in [client_a, client_b] {
let (read_half, mut write_half) = client.into_split();
let mut reader = BufReader::new(read_half);
write_half.write_all(b"{\"op\":\"ping\"}\n").await.unwrap();
assert!(read_reply(&mut reader).await.ok);
}
shutdown.cancel();
};
tokio::time::timeout(Duration::from_secs(5), async {
tokio::join!(
accept_loop(&source, ®istry, &shutdown, &mut conns),
talking
)
})
.await
.expect("the loop should serve both connections");
assert_eq!(recorder.credited().len(), 4, "{:?}", recorder.credited());
}
#[tokio::test]
async fn accept_loop_stops_promptly_when_shutdown_arrives_mid_backoff() {
let errors = (0..32).map(|_| Err(emfile())).collect();
let source = ScriptedSource::new(errors);
let registry = worktrees_registry();
let shutdown = CancellationToken::new();
let mut conns: JoinSet<()> = JoinSet::new();
let stopper = shutdown.clone();
let stop_later = async move {
tokio::time::sleep(Duration::from_millis(1400)).await;
stopper.cancel();
};
let started = Instant::now();
tokio::time::timeout(Duration::from_secs(5), async {
tokio::join!(
accept_loop(&source, ®istry, &shutdown, &mut conns),
stop_later
)
})
.await
.expect("shutdown must interrupt the backoff sleep");
assert!(
started.elapsed() < Duration::from_millis(1900),
"shutdown waited out a backoff sleep: {:?}",
started.elapsed()
);
}
}