lease readouts: a pool of their own, a refusal over it, the arm's promises pinned, the request type named for the lease
Readouts no longer take a lease's connection slot: the accept loop admits max_connections + READOUT_CONNECTIONS, a lease holds one of max_connections after it is read, a readout one of READOUT_CONNECTIONS, and a readout over its limit is refused Busy rather than dropped. Socket tests hold the service lock and a metered cell across readouts (no asking, no week wait, no lock), flood the readout pool beside leases, and pin the refusal. The new public request type is LeaseUsageReadoutRequest so it no longer shares a name with the kit's EnvelopeKind::UsageReadoutRequest; wire bytes and fixtures are unchanged. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
parent
254cf83d9d
commit
51982b7f63
4 changed files with 213 additions and 9 deletions
|
|
@ -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<QuotaCell>,
|
||||
tokens: Arc<dyn TokenSource>,
|
||||
units: Arc<dyn PeerUnits>,
|
||||
/// Connections holding or waiting on a lease, at most `max_connections`.
|
||||
leases: Arc<AtomicUsize>,
|
||||
/// Connections being answered a readout, at most `READOUT_CONNECTIONS`.
|
||||
readouts: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
/// 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<zeroize::Zeroizing<String>> {
|
||||
panic!("a readout read the token")
|
||||
}
|
||||
}
|
||||
struct NoUnits;
|
||||
impl PeerUnits for NoUnits {
|
||||
fn invocation_in(&self, _: BorrowedFd<'_>, _: &str) -> Option<PeerUnit> {
|
||||
panic!("a readout looked up the peer's unit")
|
||||
}
|
||||
}
|
||||
|
||||
fn lender(
|
||||
quota: &Arc<QuotaCell>,
|
||||
dir: &std::path::Path,
|
||||
max_connections: usize,
|
||||
) -> (Arc<Mutex<TransportService>>, 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<LeaseMessage> {
|
||||
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();
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
Loading…
Reference in a new issue