Hardswap Invariants in tests + Quota fixes

This commit is contained in:
Alexey
2026-09-20 00:28:52 +03:00
parent 89dacbd17e
commit d706b3f3ba
66 changed files with 1821 additions and 505 deletions
+29
View File
@@ -90,6 +90,35 @@ async fn stop_runtime(
generation.stop_background_tasks().await;
}
#[tokio::test]
async fn user_revocation_interrupts_live_session_before_periodic_cleanup() {
let (runtime, generation, listener) = live_runtime().await;
let bootstrap = issue_bootstrap(&runtime);
let (_, token) = create_session(&listener, &runtime, &bootstrap).await;
let session_hash = token_hash(&token);
let session = runtime
.get_session(session_hash, "proxy.example.com")
.unwrap();
let polling = Arc::clone(&session);
let poll = tokio::spawn(async move { polling.poll_down(0).await });
tokio::task::yield_now().await;
let mutation = generation.proxy_shared.delete_user("alice");
assert!(mutation.cancelled >= 1);
assert!(matches!(
tokio::time::timeout(Duration::from_millis(250), poll)
.await
.unwrap()
.unwrap(),
Err(ManagerError::Closed)
));
assert!(runtime
.get_session(session_hash, "proxy.example.com")
.is_err());
stop_runtime(runtime, generation).await;
}
#[tokio::test]
async fn pause_preserves_decoy_retry_and_exact_session_replay() {
let (runtime, generation, listener) = live_runtime().await;
+2
View File
@@ -45,6 +45,7 @@ pub(super) async fn run_upgraded(
UpgradeDeadlineLease::deadline,
);
let upgraded = tokio::select! {
biased;
_ = cancellation.cancelled() => return,
result = tokio::time::timeout_at(deadline, on_upgrade) => result,
};
@@ -137,6 +138,7 @@ async fn run_multiplex(
let down = session.poll_down_websocket(cursor);
tokio::pin!(down);
let event = tokio::select! {
biased;
_ = cancellation.cancelled() => return Err(()),
_ = tokio::time::sleep_until(open_deadline.into()), if !active => return Err(()),
_ = tokio::time::sleep_until(next_ping.into()) => DriverEvent::Liveness,
+7
View File
@@ -20,6 +20,7 @@ pub(super) async fn read_message(
backpressure_timeout: Duration,
) -> Result<(Message, Option<WebSocketBudgetLease>), ()> {
tokio::select! {
biased;
_ = cancellation.cancelled() => return Err(()),
ready = socket.get_ref().readable() => ready.map_err(|_| ())?,
}
@@ -28,6 +29,7 @@ pub(super) async fn read_message(
Some(reserve_data(runtime, owner, maximum, cancellation, backpressure_timeout).await?);
}
let message = tokio::select! {
biased;
_ = cancellation.cancelled() => return Err(()),
message = socket.next() => message.ok_or(())?.map_err(|_| ())?,
};
@@ -59,6 +61,7 @@ pub(super) async fn reserve_data(
return Ok(budget);
}
tokio::select! {
biased;
_ = cancellation.cancelled() => return Err(()),
_ = notified => {}
}
@@ -114,6 +117,7 @@ pub(super) async fn process_lane(
Err(_) => return Err(()),
}
tokio::select! {
biased;
_ = cancellation.cancelled() => return Err(()),
_ = notified => {}
}
@@ -147,6 +151,7 @@ where
Err(_) => return Err(()),
}
tokio::select! {
biased;
_ = cancellation.cancelled() => return Err(()),
_ = notified => {}
}
@@ -163,6 +168,7 @@ pub(super) async fn send(
timeout: Duration,
) -> Result<(), ()> {
tokio::select! {
biased;
_ = cancellation.cancelled() => Err(()),
result = tokio::time::timeout(timeout, socket.send(message)) => {
result.map_err(|_| ())?.map_err(|_| ())
@@ -176,6 +182,7 @@ pub(super) async fn flush(
timeout: Duration,
) -> Result<(), ()> {
tokio::select! {
biased;
_ = cancellation.cancelled() => Err(()),
result = tokio::time::timeout(timeout, socket.flush()) => {
result.map_err(|_| ())?.map_err(|_| ())
+1
View File
@@ -39,6 +39,7 @@ pub(super) async fn run_lane(
let down = session.poll_down_websocket_lane(reservation.lane_identity(), cursor);
tokio::pin!(down);
let event = tokio::select! {
biased;
_ = cancellation.cancelled() => return Err(()),
_ = tokio::time::sleep_until(open_deadline.into()), if !active => return Err(()),
_ = tokio::time::sleep_until(next_ping.into()) => DriverEvent::Liveness,
+38 -21
View File
@@ -116,19 +116,9 @@ impl WebProcessRuntime {
.record_rejection(WebRejectionReason::BootstrapCapacity);
return Err(ManagerError::Limit);
}
if !allow_rate(
&mut state.bootstrap_rate,
now,
self.limits.new_bootstraps_per_minute,
self.limits.new_bootstraps_burst,
) {
self.record_limit_hit();
self.telemetry
.record_rejection(WebRejectionReason::BootstrapRate);
return Err(ManagerError::Limit);
}
if state.bootstraps.len() >= self.limits.max_bootstraps_global
&& !evict_oldest_unused_bootstrap(&mut state)
let global_capacity_full = state.bootstraps.len() >= self.limits.max_bootstraps_global;
if global_capacity_full
&& !state.bootstraps.values().any(|bootstrap| !bootstrap.used)
{
self.record_limit_hit();
self.telemetry
@@ -152,6 +142,23 @@ impl WebProcessRuntime {
let Some(user_registration) = user_publication.take_registration() else {
return Err(ManagerError::Closed);
};
if !allow_rate(
&mut state.bootstrap_rate,
now,
self.limits.new_bootstraps_per_minute,
self.limits.new_bootstraps_burst,
) {
self.record_limit_hit();
self.telemetry
.record_rejection(WebRejectionReason::BootstrapRate);
return Err(ManagerError::Limit);
}
let evicted_bootstrap = global_capacity_full
.then(|| evict_oldest_unused_bootstrap(&mut state))
.flatten();
if global_capacity_full && evicted_bootstrap.is_none() {
return Err(ManagerError::Limit);
}
let trace_session_id = self.trace.next_session_id();
let bridge_diagnostics_enabled = config.web.debug.bridge_diagnostics_enabled();
let (user_agent, user_agent_id) = bounded_user_agent(user_agent);
@@ -196,6 +203,7 @@ impl WebProcessRuntime {
*state.bootstraps_per_ip.entry(client_ip).or_insert(0) += 1;
user_publication.commit();
drop(state);
drop(evicted_bootstrap);
if recovery {
self.telemetry
.record_bridge_recovery(WebBridgeRecoveryEvent::BootstrapIssued);
@@ -243,7 +251,11 @@ impl WebProcessRuntime {
.lock()
.bootstraps
.get(&hash)
.filter(|entry| entry.profile.host == host && now <= entry.expires_at)
.filter(|entry| {
entry.profile.host == host
&& now <= entry.expires_at
&& !entry.user_registration.is_cancelled()
})
.map(|entry| {
(
entry.trace_session_id,
@@ -263,20 +275,23 @@ impl WebProcessRuntime {
host: &str,
) -> std::result::Result<Arc<WebSession>, ManagerError> {
let state = self.state.lock();
if let Some(session) = state
let session = state
.sessions
.get(&hash)
.cloned()
.filter(|session| session.matches_host(host))
{
return Ok(session);
}
.filter(|session| session.matches_host(host));
let retired_carrier = state
.closed_tokens
.get(&hash)
.filter(|closed| closed.host == host)
.map(|closed| closed.carrier);
drop(state);
if let Some(session) = session {
if session.close_if_cancelled() {
return Err(ManagerError::Closed);
}
return Ok(session);
}
if let Some(carrier) = retired_carrier {
self.telemetry.record_session_observation(
carrier,
@@ -294,14 +309,16 @@ impl WebProcessRuntime {
profile: &WebRuntimeProfile,
) -> Option<Arc<WebSession>> {
let expected_profile = profile_key(profile);
self.state
let session = self
.state
.lock()
.sessions
.get(&hash)
.filter(|session| {
session.matches_host(host) && session.profile_key() == expected_profile
})
.cloned()
.cloned();
session.filter(|session| !session.close_if_cancelled())
}
/// Closes a live token and accepts bounded tombstone retries.
+6 -5
View File
@@ -307,8 +307,8 @@ pub(super) fn allow_rate(state: &mut RateState, now: Instant, per_minute: u32, b
true
}
/// Evicts the oldest unused bootstrap while preserving used retry state.
pub(super) fn evict_oldest_unused_bootstrap(state: &mut ManagerState) -> bool {
/// Detaches the oldest unused bootstrap while preserving used retry state.
pub(super) fn evict_oldest_unused_bootstrap(state: &mut ManagerState) -> Option<Bootstrap> {
let Some(hash) = state
.bootstraps
.iter()
@@ -316,10 +316,11 @@ pub(super) fn evict_oldest_unused_bootstrap(state: &mut ManagerState) -> bool {
.min_by_key(|(_, bootstrap)| bootstrap.issued_at)
.map(|(hash, _)| *hash)
else {
return false;
return None;
};
remove_bootstrap_locked(state, hash);
true
let bootstrap = state.bootstraps.remove(&hash)?;
decrement_map(&mut state.bootstraps_per_ip, &bootstrap.issuance_ip);
Some(bootstrap)
}
/// Removes expired bootstrap and closed-token entries while the manager lock is held.
+8
View File
@@ -20,6 +20,13 @@ impl WebSession {
completion: StreamCompletion,
retain_reservation_on_reject: bool,
) -> bool {
if self.close_if_cancelled() {
completion
.retain_rejected
.store(retain_reservation_on_reject, Ordering::Release);
drop(completion);
return false;
}
let stream = completion.stream;
let peer_port = completion.peer_port;
let Some(manager) = self.manager.upgrade() else {
@@ -60,6 +67,7 @@ impl WebSession {
);
let logical_stream = WebLogicalStream::new(Arc::clone(&session), stream);
tokio::select! {
biased;
_ = cancel.cancelled() => {}
_ = run_stream(
Arc::clone(&session),
+37 -12
View File
@@ -26,9 +26,14 @@ impl WebSession {
if !self.carrier().is_multiplexed() {
return Err(ManagerError::Protocol);
}
if self.close_if_cancelled() {
return Err(ManagerError::Closed);
}
let (epoch, healthy) = {
let mut state = self.state.lock();
if state.closed {
if state.closed || self.cancel.is_cancelled() {
drop(state);
self.close_if_cancelled();
return Err(ManagerError::Closed);
}
if let Some(unacked) = &state.unacked {
@@ -92,6 +97,11 @@ impl WebSession {
notified.as_mut().enable();
{
let mut state = self.state.lock();
if self.cancel.is_cancelled() {
drop(state);
self.close_if_cancelled();
return Err(ManagerError::Closed);
}
if state.down_epoch != epoch {
return Ok(PollResult {
body: Bytes::new(),
@@ -129,18 +139,33 @@ impl WebSession {
notified.await;
}
};
match tokio::time::timeout(deadline, poll).await {
Ok(result) => result,
Err(_) => {
let mut state = self.state.lock();
if state.down_epoch == epoch {
state.activity.touch_progress(Instant::now());
tokio::select! {
biased;
_ = self.cancel.cancelled() => {
self.close_if_cancelled();
Err(ManagerError::Closed)
}
result = tokio::time::timeout(deadline, poll) => match result {
Ok(result) => result,
Err(_) => {
if self.close_if_cancelled() {
return Err(ManagerError::Closed);
}
let mut state = self.state.lock();
if self.cancel.is_cancelled() {
drop(state);
self.close_if_cancelled();
return Err(ManagerError::Closed);
}
if state.down_epoch == epoch {
state.activity.touch_progress(Instant::now());
}
Ok(PollResult {
body: Bytes::new(),
next_cursor: cursor,
lane_closed: false,
})
}
Ok(PollResult {
body: Bytes::new(),
next_cursor: cursor,
lane_closed: false,
})
}
}
}
+7
View File
@@ -22,6 +22,9 @@ impl WebSession {
if self.carrier() != WebCarrier::HttpsLanes || lane_id > frame::MAX_STREAM_ID {
return Err(ManagerError::Protocol);
}
if self.close_if_cancelled() {
return Err(ManagerError::Closed);
}
let frames = match frame::parse_all(body, &self.limits) {
Ok(frames) => frames,
Err(_) => {
@@ -175,6 +178,10 @@ impl WebSession {
if matches!(result, Err(ManagerError::Backpressure)) {
return result;
}
if matches!(result, Err(ManagerError::Closed)) && self.close_if_cancelled() {
drop(opened);
return result;
}
if result.is_err() {
self.close(SessionCloseReason::Protocol);
drop(opened);
+79 -41
View File
@@ -42,9 +42,14 @@ impl WebSession {
if !self.carrier().uses_lanes() || lane_id > frame::MAX_STREAM_ID {
return Err(ManagerError::Protocol);
}
if self.close_if_cancelled() {
return Err(ManagerError::Closed);
}
let lane_ready = if let Some(expected_instance) = expected_instance {
let state = self.state.lock();
if state.closed {
if state.closed || self.cancel.is_cancelled() {
drop(state);
self.close_if_cancelled();
return Err(ManagerError::Closed);
}
state
@@ -63,7 +68,9 @@ impl WebSession {
}
let (instance, epoch, notify, healthy) = {
let mut state = self.state.lock();
if state.closed {
if state.closed || self.cancel.is_cancelled() {
drop(state);
self.close_if_cancelled();
return Err(ManagerError::Closed);
}
let (acknowledged, replay) = {
@@ -175,7 +182,9 @@ impl WebSession {
notified.as_mut().enable();
{
let mut state = self.state.lock();
if state.closed {
if state.closed || self.cancel.is_cancelled() {
drop(state);
self.close_if_cancelled();
return Err(ManagerError::Closed);
}
let carrier_health_eligible = lane_id != 0
@@ -245,55 +254,72 @@ impl WebSession {
notified.await;
}
};
match tokio::time::timeout(deadline, poll).await {
Ok(result) => result,
Err(_) => {
let mut state = self.state.lock();
if state.closed {
return Err(ManagerError::Closed);
}
if !state.carrier_lanes.contains_key(&lane_id) {
return Ok(PollResult {
body: Bytes::new(),
next_cursor: cursor,
lane_closed: true,
});
}
if lane_id != 0
&& !state.streams.contains_key(&lane_id)
&& state.closed_streams.contains(&lane_id)
{
return Ok(PollResult {
body: Bytes::new(),
next_cursor: cursor,
lane_closed: true,
});
}
if let Some(lane) = state.carrier_lanes.get(&lane_id) {
if lane.instance != instance {
tokio::select! {
biased;
_ = self.cancel.cancelled() => {
self.close_if_cancelled();
Err(ManagerError::Closed)
}
result = tokio::time::timeout(deadline, poll) => match result {
Ok(result) => result,
Err(_) => {
if self.close_if_cancelled() {
return Err(ManagerError::Closed);
}
let mut state = self.state.lock();
if state.closed || self.cancel.is_cancelled() {
drop(state);
self.close_if_cancelled();
return Err(ManagerError::Closed);
}
if !state.carrier_lanes.contains_key(&lane_id) {
return Ok(PollResult {
body: Bytes::new(),
next_cursor: cursor,
lane_closed: true,
});
}
if lane.down_epoch == epoch {
state.activity.touch_progress(Instant::now());
if lane_id != 0
&& !state.streams.contains_key(&lane_id)
&& state.closed_streams.contains(&lane_id)
{
return Ok(PollResult {
body: Bytes::new(),
next_cursor: cursor,
lane_closed: true,
});
}
if let Some(lane) = state.carrier_lanes.get(&lane_id) {
if lane.instance != instance {
return Ok(PollResult {
body: Bytes::new(),
next_cursor: cursor,
lane_closed: true,
});
}
if lane.down_epoch == epoch {
state.activity.touch_progress(Instant::now());
}
}
Ok(PollResult {
body: Bytes::new(),
next_cursor: cursor,
lane_closed: false,
})
}
Ok(PollResult {
body: Bytes::new(),
next_cursor: cursor,
lane_closed: false,
})
}
}
}
async fn wait_for_lane_open(&self, lane_id: u32, cursor: u64) -> Result<bool, ManagerError> {
if self.close_if_cancelled() {
return Err(ManagerError::Closed);
}
let wait = {
let mut state = self.state.lock();
if state.closed {
if state.closed || self.cancel.is_cancelled() {
drop(state);
self.close_if_cancelled();
return Err(ManagerError::Closed);
}
if state.carrier_lanes.contains_key(&lane_id) {
@@ -334,7 +360,9 @@ impl WebSession {
notified.as_mut().enable();
{
let state = self.state.lock();
if state.closed {
if state.closed || self.cancel.is_cancelled() {
drop(state);
self.close_if_cancelled();
return Err(ManagerError::Closed);
}
if state.carrier_lanes.contains_key(&lane_id)
@@ -346,14 +374,24 @@ impl WebSession {
}
notified.await;
}
})
.await;
});
let opened = tokio::select! {
biased;
_ = self.cancel.cancelled() => {
drop(wait);
self.close_if_cancelled();
return Err(ManagerError::Closed);
}
opened = opened => opened,
};
drop(wait);
match opened {
Ok(result) => result,
Err(_) => {
let state = self.state.lock();
if state.closed {
if state.closed || self.cancel.is_cancelled() {
drop(state);
self.close_if_cancelled();
Err(ManagerError::Closed)
} else {
Ok(state.carrier_lanes.contains_key(&lane_id)
+11 -13
View File
@@ -121,6 +121,15 @@ impl CarrierSupersedeCompletion<'_> {
}
impl WebSession {
/// Closes a bearer as soon as its process-owned user registration is revoked.
pub(crate) fn close_if_cancelled(&self) -> bool {
if !self.cancel.is_cancelled() {
return false;
}
self.close(SessionCloseReason::UserDisabled);
true
}
/// Closes carrier state while relay tasks retain their admission until exit.
pub(crate) fn close(&self, reason: SessionCloseReason) -> SessionCloseOutcome {
let mut state = self.state.lock();
@@ -221,19 +230,8 @@ impl WebSession {
/// Atomically closes a session only when reconnect grace is still due.
pub(crate) fn close_if_due(&self, now: Instant) -> bool {
if self.cancel.is_cancelled() {
let released = {
let mut state = self.state.lock();
if state.closed || state.close_requested.is_some() {
None
} else {
Some(self.release_on_close_locked(&mut state, SessionCloseReason::UserDisabled))
}
};
if let Some(released) = released {
self.finish_close(released);
return true;
}
if self.close_if_cancelled() {
return true;
}
let healthy = {
let mut state = self.state.lock();
+3
View File
@@ -34,6 +34,9 @@ impl WebSession {
&self,
state: &SessionState,
) -> Result<(), crate::web::manager::ManagerError> {
if self.cancel.is_cancelled() {
return Err(crate::web::manager::ManagerError::Closed);
}
if state.negotiation_phase == SessionNegotiationPhase::Uncommitted
&& self
.carrier_deadline_at
+7
View File
@@ -69,6 +69,9 @@ impl WebSession {
if !self.carrier().is_multiplexed() {
return Err(ManagerError::Protocol);
}
if self.close_if_cancelled() {
return Err(ManagerError::Closed);
}
if self
.up_active
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
@@ -157,6 +160,10 @@ impl WebSession {
if matches!(result, Err(ManagerError::Backpressure)) {
return result;
}
if matches!(result, Err(ManagerError::Closed)) && self.close_if_cancelled() {
drop(opened);
return result;
}
if result.is_err() {
self.close(SessionCloseReason::Protocol);
drop(opened);
+24 -3
View File
@@ -39,8 +39,12 @@ pub(crate) struct WebSocketProbeReservation {
impl WebSocketProbeReservation {
/// Binds the admitted process connection to the future commit acknowledgement.
pub(crate) fn bind(&mut self, owner: u64) -> Result<(), ManagerError> {
if self.session.close_if_cancelled() {
return Err(ManagerError::Closed);
}
let mut state = self.session.state.lock();
if state.closed
|| self.session.cancel.is_cancelled()
|| !state.websocket_probe_claimed
|| state.websocket_commit_ack_owner.is_some()
{
@@ -85,8 +89,12 @@ impl WebSocketLaneReservation {
if self.phase != WebSocketLaneReservationPhase::Reserved {
return Err(ManagerError::Concurrent);
}
if self.session.close_if_cancelled() {
return Err(ManagerError::Closed);
}
let mut state = self.session.state.lock();
if state.closed
|| self.session.cancel.is_cancelled()
|| state
.carrier_lanes
.get(&self.claim.lane.lane_id)
@@ -113,8 +121,12 @@ impl WebSocketLaneReservation {
{
return Err(ManagerError::Protocol);
}
if self.session.close_if_cancelled() {
return Err(ManagerError::Closed);
}
let mut state = self.session.state.lock();
if state
if self.session.cancel.is_cancelled()
|| state
.carrier_lanes
.get(&self.claim.lane.lane_id)
.is_none_or(|lane| lane.instance != self.claim.lane.instance)
@@ -174,8 +186,11 @@ impl WebSession {
self: &Arc<Self>,
acknowledge_commit: bool,
) -> Result<Option<WebSocketProbeReservation>, ManagerError> {
if self.close_if_cancelled() {
return Err(ManagerError::Closed);
}
let mut state = self.state.lock();
if state.closed {
if state.closed || self.cancel.is_cancelled() {
return Err(ManagerError::Closed);
}
self.ensure_carrier_active_locked(&state)?;
@@ -216,8 +231,11 @@ impl WebSession {
{
return Err(ManagerError::Protocol);
}
if self.close_if_cancelled() {
return Err(ManagerError::Closed);
}
let mut state = self.state.lock();
if state.closed {
if state.closed || self.cancel.is_cancelled() {
return Err(ManagerError::Closed);
}
if state.active_peer_ports.len() >= self.profile.max_streams_per_session
@@ -306,6 +324,9 @@ impl WebSession {
{
return Err(ManagerError::Protocol);
}
if self.close_if_cancelled() {
return Err(ManagerError::Closed);
}
let lane_id = reservation.lane_id();
let frames = frame::parse_all(body, &self.limits).map_err(|_| ManagerError::Protocol)?;
if frames