diff --git a/crates/claude-connector/src/lease.rs b/crates/claude-connector/src/lease.rs index 1749f6f..a049bba 100644 --- a/crates/claude-connector/src/lease.rs +++ b/crates/claude-connector/src/lease.rs @@ -53,7 +53,7 @@ pub enum LeaseMessage { ClaudeLeaseGrant(LeaseGranted), ClaudeLeaseRefusal(LeaseRefused), /// Asks for the plan's usage windows. Opens no lease. - ClaudeUsageReadoutRequest(UsageReadoutRequest), + ClaudeUsageReadoutRequest(LeaseUsageReadoutRequest), /// The answer to a readout request: the cell's windows as of now, no fd. ClaudeUsageReadout(UsageReadout), } @@ -66,7 +66,7 @@ pub struct LeaseRequest { /// A readout request names nothing: the answer is the same for every asker. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct UsageReadoutRequest {} +pub struct LeaseUsageReadoutRequest {} #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct LeaseGranted { @@ -395,8 +395,17 @@ struct Lender { quota: Arc, tokens: Arc, units: Arc, + /// Connections holding or waiting on a lease, at most `max_connections`. + leases: Arc, + /// Connections being answered a readout, at most `READOUT_CONNECTIONS`. + readouts: Arc, } +/// How many usage readouts are served at once, beside the leases. A readout +/// holds nothing, so a few suffice, and a flood of them cannot take a lease's +/// slot: the accept loop admits `max_connections` plus this many. +const READOUT_CONNECTIONS: usize = 4; + /// Serves leases on `listener` until it fails. Each connection is one lease or /// one usage readout. A lease: request, then grant or refusal, then (after a /// grant) the hold. A readout: request, then one frame, then close; its reach @@ -422,11 +431,14 @@ pub fn serve_leases( quota, tokens, units, + leases: Arc::new(AtomicUsize::new(0)), + readouts: Arc::new(AtomicUsize::new(0)), }); let active = Arc::new(AtomicUsize::new(0)); loop { let (stream, _) = listener.accept()?; - let Some(permit) = Permit::acquire(&active, limits.max_connections) else { + let admitted = limits.max_connections.saturating_add(READOUT_CONNECTIONS); + let Some(permit) = Permit::acquire(&active, admitted) else { continue; }; let lender = lender.clone(); @@ -447,10 +459,22 @@ impl Lender { let request_deadline = Instant::now() + self.limits.socket_timeout; match read_message(&mut stream, request_deadline)? { LeaseMessage::ClaudeLeaseRequest(request) => { + // A lease over the limit is dropped unanswered; readouts + // never count against it. + let Some(_slot) = Permit::acquire(&self.leases, self.limits.max_connections) else { + return Ok(()); + }; let peer = peer_credentials(&stream)?; self.lend(stream, peer, request) } LeaseMessage::ClaudeUsageReadoutRequest(_) => { + let Some(_slot) = Permit::acquire(&self.readouts, READOUT_CONNECTIONS) else { + return refuse( + &mut stream, + LeaseRefusalKind::Busy, + "too many usage readouts at once", + ); + }; stream.write_all(&encode_frame(&readout(&self.quota, unix_ms()))) } _ => Err(io::Error::from(io::ErrorKind::InvalidData)), @@ -662,6 +686,186 @@ fn refuse(stream: &mut UnixStream, kind: LeaseRefusalKind, detail: &str) -> io:: mod tests { use super::*; + struct NoTokens; + impl TokenSource for NoTokens { + fn read(&self) -> Option> { + panic!("a readout read the token") + } + } + struct NoUnits; + impl PeerUnits for NoUnits { + fn invocation_in(&self, _: BorrowedFd<'_>, _: &str) -> Option { + panic!("a readout looked up the peer's unit") + } + } + + fn lender( + quota: &Arc, + dir: &std::path::Path, + max_connections: usize, + ) -> (Arc>, Lender) { + let service = Arc::new(Mutex::new( + TransportService::open( + &dir.join("replay.cc"), + "claude-code", + [model_connector::CallerAdmission::new( + "lease-test-caller", + "lease-test-connection-key", + 1, + ["model".to_string()], + 1, + 65_536, + 32_768, + ) + .unwrap()], + 300_000, + ) + .unwrap(), + )); + let lender = Lender { + callers: Vec::new(), + limits: LeaseLimits { + socket_timeout: Duration::from_secs(2), + max_lease_ms: 3_600_000, + max_connections, + }, + service: service.clone(), + quota: quota.clone(), + tokens: Arc::new(NoTokens), + units: Arc::new(NoUnits), + leases: Arc::new(AtomicUsize::new(0)), + readouts: Arc::new(AtomicUsize::new(0)), + }; + (service, lender) + } + + /// Serves `lender`'s parts on a socket in `dir`; the path comes back. + fn serve(lender: Lender, dir: &std::path::Path) -> std::path::PathBuf { + let socket = dir.join("lease.sock"); + let listener = UnixListener::bind(&socket).unwrap(); + let (service, quota) = (lender.service.clone(), lender.quota.clone()); + thread::spawn(move || { + let _ = serve_leases( + &listener, + lender.callers, + lender.limits, + service, + quota, + lender.tokens, + lender.units, + ); + }); + socket + } + + fn ask(socket: &std::path::Path, message: &LeaseMessage) -> Option { + let mut stream = UnixStream::connect(socket).ok()?; + stream + .set_read_timeout(Some(Duration::from_secs(3))) + .unwrap(); + stream.write_all(&encode_frame(message)).ok()?; + let mut raw = Vec::new(); + let _ = stream.read_to_end(&mut raw); + decode_frame(&raw) + } + + fn readout_request() -> LeaseMessage { + LeaseMessage::ClaudeUsageReadoutRequest(LeaseUsageReadoutRequest {}) + } + + fn scratch(name: &str) -> std::path::PathBuf { + let dir = std::env::temp_dir().join(format!("{name}-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + dir + } + + #[test] + fn a_socket_readout_never_asks_the_meter_waits_for_a_week_or_takes_the_service_lock() { + let dir = scratch("readout-arm"); + let quota = Arc::new(QuotaCell::new()); + quota.attach_meter(); + let (service, lender) = lender("a, &dir, 8); + let socket = serve(lender, &dir); + // Held throughout: a readout that takes it hangs past the read timeout. + let guard = service.lock().unwrap(); + for _ in 0..20 { + let started = Instant::now(); + let Some(LeaseMessage::ClaudeUsageReadout(answer)) = ask(&socket, &readout_request()) + else { + panic!("no readout"); + }; + assert!(started.elapsed() < Duration::from_secs(1)); + assert!(answer.plan_quota.is_empty()); + } + drop(guard); + assert_eq!(quota.asks_so_far(), 0); + assert!(!quota.wanted_at(unix_ms(), Duration::from_secs(15 * 60))); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn a_flood_of_readouts_at_their_limit_leaves_every_lease_request_answered() { + let dir = scratch("readout-flood"); + let quota = Arc::new(QuotaCell::new()); + // Fewer lease slots than flooding readers: the readers must not need them. + let (_service, lender) = lender("a, &dir, 2); + let socket = serve(lender, &dir); + let stop = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let flood: Vec<_> = (0..READOUT_CONNECTIONS) + .map(|_| { + let (socket, stop) = (socket.clone(), stop.clone()); + thread::spawn(move || { + while !stop.load(Ordering::Acquire) { + let _ = ask(&socket, &readout_request()); + } + }) + }) + .collect(); + // No caller is configured, so each lease is answered UnknownCaller: + // an answer is all that is asked. + let lease = LeaseMessage::ClaudeLeaseRequest(LeaseRequest { + expires_at_unix_ms: unix_ms() + 60_000, + }); + let unanswered = (0..40) + .filter(|_| { + !matches!( + ask(&socket, &lease), + Some(LeaseMessage::ClaudeLeaseRefusal(LeaseRefused { + kind: LeaseRefusalKind::UnknownCaller, + .. + })) + ) + }) + .count(); + stop.store(true, Ordering::Release); + for client in flood { + client.join().unwrap(); + } + assert_eq!(unanswered, 0); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn a_readout_over_the_readout_limit_is_refused_busy_not_dropped() { + let dir = scratch("readout-limit"); + let quota = Arc::new(QuotaCell::new()); + let (_service, lender) = lender("a, &dir, 8); + let held: Vec<_> = (0..READOUT_CONNECTIONS) + .map(|_| Permit::acquire(&lender.readouts, READOUT_CONNECTIONS).unwrap()) + .collect(); + let (mut client, server) = UnixStream::pair().unwrap(); + client.write_all(&encode_frame(&readout_request())).unwrap(); + lender.serve_one(server).unwrap(); + let mut raw = Vec::new(); + client.read_to_end(&mut raw).unwrap(); + let Some(LeaseMessage::ClaudeLeaseRefusal(refusal)) = decode_frame(&raw) else { + panic!("no refusal"); + }; + assert_eq!(refusal.kind, LeaseRefusalKind::Busy); + drop(held); + let _ = std::fs::remove_dir_all(&dir); + } + #[test] fn a_readout_of_an_empty_metered_cell_does_not_wait() { let quota = QuotaCell::new(); diff --git a/crates/claude-connector/src/lib.rs b/crates/claude-connector/src/lib.rs index 52f957f..88de37b 100644 --- a/crates/claude-connector/src/lib.rs +++ b/crates/claude-connector/src/lib.rs @@ -29,8 +29,8 @@ pub use launch::ChildEnv; #[cfg(target_os = "linux")] pub use lease::{ LeaseCaller, LeaseGranted, LeaseLimits, LeaseMessage, LeaseRefusalKind, LeaseRefused, - LeaseRequest, MAX_LEASE_FRAME_BYTES, PeerCredentials, PeerUnit, PeerUnits, UsageReadoutRequest, - decode_frame, encode_frame, peer_credentials, serve_leases, user_uid, + LeaseRequest, LeaseUsageReadoutRequest, MAX_LEASE_FRAME_BYTES, PeerCredentials, PeerUnit, + PeerUnits, decode_frame, encode_frame, peer_credentials, serve_leases, user_uid, }; #[cfg(unix)] pub use process::StdProcessPort; diff --git a/crates/claude-connector/tests/lease.rs b/crates/claude-connector/tests/lease.rs index 0316b9f..61e5a5e 100644 --- a/crates/claude-connector/tests/lease.rs +++ b/crates/claude-connector/tests/lease.rs @@ -15,7 +15,7 @@ use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use claude_connector::{ LeaseCaller, LeaseGranted, LeaseLimits, LeaseMessage, LeaseRefusalKind, LeaseRefused, - LeaseRequest, PeerUnit, PeerUnits, QuotaCell, TokenSource, UsageReadoutRequest, decode_frame, + LeaseRequest, PeerUnit, PeerUnits, QuotaCell, TokenSource, LeaseUsageReadoutRequest, decode_frame, encode_frame, serve_leases, }; use model_connector::{CallerAdmission, LeaseTurn, PlanQuota, TransportService, UsageReadout}; @@ -658,7 +658,7 @@ fn read_usage(harness: &Harness) -> Answer { .unwrap(); stream .write_all(&encode_frame(&LeaseMessage::ClaudeUsageReadoutRequest( - UsageReadoutRequest {}, + LeaseUsageReadoutRequest {}, ))) .unwrap(); let (raw, fds) = receive(&stream); diff --git a/crates/claude-connector/tests/lease_wire_golden.rs b/crates/claude-connector/tests/lease_wire_golden.rs index c713889..83820e3 100644 --- a/crates/claude-connector/tests/lease_wire_golden.rs +++ b/crates/claude-connector/tests/lease_wire_golden.rs @@ -8,7 +8,7 @@ use std::path::{Path, PathBuf}; use claude_connector::{ - LeaseGranted, LeaseMessage, LeaseRefusalKind, LeaseRefused, LeaseRequest, UsageReadoutRequest, + LeaseGranted, LeaseMessage, LeaseRefusalKind, LeaseRefused, LeaseRequest, LeaseUsageReadoutRequest, decode_frame, encode_frame, }; use model_connector::{PlanQuota, PlanQuotaSnapshot, UsageReadout}; @@ -107,7 +107,7 @@ fn the_usage_readout_frames_are_pinned() { status: status.to_string(), observed_at_unix_ms: 1_793_300_000_000, }; - let request = LeaseMessage::ClaudeUsageReadoutRequest(UsageReadoutRequest {}); + let request = LeaseMessage::ClaudeUsageReadoutRequest(LeaseUsageReadoutRequest {}); let answer = LeaseMessage::ClaudeUsageReadout(UsageReadout::new( 1_793_300_100_000, PlanQuotaSnapshot {