mirror of https://github.com/telemt/telemt.git
Old Test Deletion
This commit is contained in:
parent
6f4356f72a
commit
2f9fddfa6f
|
|
@ -1,113 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use dashmap::DashMap;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
||||||
use tokio::sync::Barrier;
|
|
||||||
use tokio::time::{Duration, timeout};
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn blackhat_campaign_saturation_quota_race_with_queue_pressure_stays_fail_closed() {
|
|
||||||
let _guard = super::quota_user_lock_test_scope();
|
|
||||||
let _pressure_guard = super::relay_idle_pressure_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX);
|
|
||||||
for idx in 0..QUOTA_USER_LOCKS_MAX {
|
|
||||||
retained.push(quota_user_lock(&format!(
|
|
||||||
"middle-blackhat-held-{}-{idx}",
|
|
||||||
std::process::id()
|
|
||||||
)));
|
|
||||||
}
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
map.len(),
|
|
||||||
QUOTA_USER_LOCKS_MAX,
|
|
||||||
"precondition: bounded lock cache must be saturated"
|
|
||||||
);
|
|
||||||
|
|
||||||
let (tx, _rx) = mpsc::channel::<C2MeCommand>(1);
|
|
||||||
tx.send(C2MeCommand::Close)
|
|
||||||
.await
|
|
||||||
.expect("queue prefill should succeed");
|
|
||||||
|
|
||||||
let pressure_seq_before = relay_pressure_event_seq();
|
|
||||||
let pressure_errors = Arc::new(AtomicUsize::new(0));
|
|
||||||
let mut pressure_workers = Vec::new();
|
|
||||||
for _ in 0..16 {
|
|
||||||
let tx = tx.clone();
|
|
||||||
let pressure_errors = Arc::clone(&pressure_errors);
|
|
||||||
pressure_workers.push(tokio::spawn(async move {
|
|
||||||
if enqueue_c2me_command(&tx, C2MeCommand::Close).await.is_err() {
|
|
||||||
pressure_errors.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let user = format!("middle-blackhat-quota-race-{}", std::process::id());
|
|
||||||
let gate = Arc::new(Barrier::new(16));
|
|
||||||
|
|
||||||
let mut quota_workers = Vec::new();
|
|
||||||
for _ in 0..16u8 {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let user = user.clone();
|
|
||||||
let gate = Arc::clone(&gate);
|
|
||||||
quota_workers.push(tokio::spawn(async move {
|
|
||||||
gate.wait().await;
|
|
||||||
let user_lock = quota_user_lock(&user);
|
|
||||||
let _quota_guard = user_lock.lock().await;
|
|
||||||
|
|
||||||
if quota_would_be_exceeded_for_user(&stats, &user, Some(1), 1) {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
stats.add_user_octets_to(&user, 1);
|
|
||||||
true
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut ok_count = 0usize;
|
|
||||||
let mut denied_count = 0usize;
|
|
||||||
for worker in quota_workers {
|
|
||||||
let result = timeout(Duration::from_secs(2), worker)
|
|
||||||
.await
|
|
||||||
.expect("quota worker must finish")
|
|
||||||
.expect("quota worker must not panic");
|
|
||||||
if result {
|
|
||||||
ok_count += 1;
|
|
||||||
} else {
|
|
||||||
denied_count += 1;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
for worker in pressure_workers {
|
|
||||||
timeout(Duration::from_secs(2), worker)
|
|
||||||
.await
|
|
||||||
.expect("pressure worker must finish")
|
|
||||||
.expect("pressure worker must not panic");
|
|
||||||
}
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
stats.get_user_total_octets(&user),
|
|
||||||
1,
|
|
||||||
"black-hat campaign must not overshoot same-user quota under saturation"
|
|
||||||
);
|
|
||||||
assert!(ok_count <= 1, "at most one quota contender may succeed");
|
|
||||||
assert!(
|
|
||||||
denied_count >= 15,
|
|
||||||
"all remaining contenders must be quota-denied"
|
|
||||||
);
|
|
||||||
|
|
||||||
let pressure_seq_after = relay_pressure_event_seq();
|
|
||||||
assert!(
|
|
||||||
pressure_seq_after > pressure_seq_before,
|
|
||||||
"queue pressure leg must trigger pressure accounting"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
pressure_errors.load(Ordering::Relaxed) >= 1,
|
|
||||||
"at least one pressure worker should fail from persistent backpressure"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(retained);
|
|
||||||
}
|
|
||||||
|
|
@ -1,777 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::crypto::AesCtr;
|
|
||||||
use crate::crypto::SecureRandom;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use crate::stream::{BufferPool, PooledBuffer};
|
|
||||||
use std::sync::Arc;
|
|
||||||
use tokio::io::AsyncReadExt;
|
|
||||||
use tokio::io::duplex;
|
|
||||||
use tokio::sync::mpsc;
|
|
||||||
use tokio::time::{Duration as TokioDuration, timeout};
|
|
||||||
|
|
||||||
fn make_pooled_payload(data: &[u8]) -> PooledBuffer {
|
|
||||||
let pool = Arc::new(BufferPool::with_config(data.len().max(1), 4));
|
|
||||||
let mut payload = pool.get();
|
|
||||||
payload.resize(data.len(), 0);
|
|
||||||
payload[..data.len()].copy_from_slice(data);
|
|
||||||
payload
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn write_client_payload_abridged_short_quickack_sets_flag_and_preserves_payload() {
|
|
||||||
let (mut read_side, write_side) = duplex(4096);
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
|
|
||||||
let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024);
|
|
||||||
let mut decryptor = AesCtr::new(&key, iv);
|
|
||||||
let rng = SecureRandom::new();
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let payload = vec![0xA1, 0xB2, 0xC3, 0xD4, 0x10, 0x20, 0x30, 0x40];
|
|
||||||
|
|
||||||
write_client_payload(
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Abridged,
|
|
||||||
RPC_FLAG_QUICKACK,
|
|
||||||
&payload,
|
|
||||||
&rng,
|
|
||||||
&mut frame_buf,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("abridged quickack payload should serialize");
|
|
||||||
writer.flush().await.expect("flush must succeed");
|
|
||||||
|
|
||||||
let mut encrypted = vec![0u8; 1 + payload.len()];
|
|
||||||
read_side
|
|
||||||
.read_exact(&mut encrypted)
|
|
||||||
.await
|
|
||||||
.expect("must read serialized abridged frame");
|
|
||||||
let plaintext = decryptor.decrypt(&encrypted);
|
|
||||||
|
|
||||||
assert_eq!(plaintext[0], 0x80 | ((payload.len() / 4) as u8));
|
|
||||||
assert_eq!(&plaintext[1..], payload.as_slice());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn write_client_payload_abridged_extended_header_is_encoded_correctly() {
|
|
||||||
let (mut read_side, write_side) = duplex(16 * 1024);
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
|
|
||||||
let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024);
|
|
||||||
let mut decryptor = AesCtr::new(&key, iv);
|
|
||||||
let rng = SecureRandom::new();
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
|
|
||||||
// Boundary where abridged switches to extended length encoding.
|
|
||||||
let payload = vec![0x5Au8; 0x7f * 4];
|
|
||||||
|
|
||||||
write_client_payload(
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Abridged,
|
|
||||||
RPC_FLAG_QUICKACK,
|
|
||||||
&payload,
|
|
||||||
&rng,
|
|
||||||
&mut frame_buf,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("extended abridged payload should serialize");
|
|
||||||
writer.flush().await.expect("flush must succeed");
|
|
||||||
|
|
||||||
let mut encrypted = vec![0u8; 4 + payload.len()];
|
|
||||||
read_side
|
|
||||||
.read_exact(&mut encrypted)
|
|
||||||
.await
|
|
||||||
.expect("must read serialized extended abridged frame");
|
|
||||||
let plaintext = decryptor.decrypt(&encrypted);
|
|
||||||
|
|
||||||
assert_eq!(plaintext[0], 0xff, "0x7f with quickack bit must be set");
|
|
||||||
assert_eq!(&plaintext[1..4], &[0x7f, 0x00, 0x00]);
|
|
||||||
assert_eq!(&plaintext[4..], payload.as_slice());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn write_client_payload_abridged_misaligned_is_rejected_fail_closed() {
|
|
||||||
let (_read_side, write_side) = duplex(1024);
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
|
|
||||||
let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024);
|
|
||||||
let rng = SecureRandom::new();
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
|
|
||||||
let err = write_client_payload(
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Abridged,
|
|
||||||
0,
|
|
||||||
&[1, 2, 3],
|
|
||||||
&rng,
|
|
||||||
&mut frame_buf,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect_err("misaligned abridged payload must be rejected");
|
|
||||||
|
|
||||||
let msg = format!("{err}");
|
|
||||||
assert!(
|
|
||||||
msg.contains("4-byte aligned"),
|
|
||||||
"error should explain alignment contract, got: {msg}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn write_client_payload_secure_misaligned_is_rejected_fail_closed() {
|
|
||||||
let (_read_side, write_side) = duplex(1024);
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
|
|
||||||
let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024);
|
|
||||||
let rng = SecureRandom::new();
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
|
|
||||||
let err = write_client_payload(
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Secure,
|
|
||||||
0,
|
|
||||||
&[9, 8, 7, 6, 5],
|
|
||||||
&rng,
|
|
||||||
&mut frame_buf,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect_err("misaligned secure payload must be rejected");
|
|
||||||
|
|
||||||
let msg = format!("{err}");
|
|
||||||
assert!(
|
|
||||||
msg.contains("Secure payload must be 4-byte aligned"),
|
|
||||||
"error should be explicit for fail-closed triage, got: {msg}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn write_client_payload_intermediate_quickack_sets_length_msb() {
|
|
||||||
let (mut read_side, write_side) = duplex(4096);
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
|
|
||||||
let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024);
|
|
||||||
let mut decryptor = AesCtr::new(&key, iv);
|
|
||||||
let rng = SecureRandom::new();
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let payload = b"hello-middle-relay";
|
|
||||||
|
|
||||||
write_client_payload(
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
RPC_FLAG_QUICKACK,
|
|
||||||
payload,
|
|
||||||
&rng,
|
|
||||||
&mut frame_buf,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("intermediate quickack payload should serialize");
|
|
||||||
writer.flush().await.expect("flush must succeed");
|
|
||||||
|
|
||||||
let mut encrypted = vec![0u8; 4 + payload.len()];
|
|
||||||
read_side
|
|
||||||
.read_exact(&mut encrypted)
|
|
||||||
.await
|
|
||||||
.expect("must read intermediate frame");
|
|
||||||
let plaintext = decryptor.decrypt(&encrypted);
|
|
||||||
|
|
||||||
let mut len_bytes = [0u8; 4];
|
|
||||||
len_bytes.copy_from_slice(&plaintext[..4]);
|
|
||||||
let len_with_flags = u32::from_le_bytes(len_bytes);
|
|
||||||
assert_ne!(len_with_flags & 0x8000_0000, 0, "quickack bit must be set");
|
|
||||||
assert_eq!((len_with_flags & 0x7fff_ffff) as usize, payload.len());
|
|
||||||
assert_eq!(&plaintext[4..], payload);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn write_client_payload_secure_quickack_prefix_and_padding_bounds_hold() {
|
|
||||||
let (mut read_side, write_side) = duplex(4096);
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
|
|
||||||
let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024);
|
|
||||||
let mut decryptor = AesCtr::new(&key, iv);
|
|
||||||
let rng = SecureRandom::new();
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let payload = vec![0x33u8; 100]; // 4-byte aligned as required by secure mode.
|
|
||||||
|
|
||||||
write_client_payload(
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Secure,
|
|
||||||
RPC_FLAG_QUICKACK,
|
|
||||||
&payload,
|
|
||||||
&rng,
|
|
||||||
&mut frame_buf,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("secure quickack payload should serialize");
|
|
||||||
writer.flush().await.expect("flush must succeed");
|
|
||||||
|
|
||||||
// Secure mode adds 1..=3 bytes of randomized tail padding.
|
|
||||||
let mut encrypted_header = [0u8; 4];
|
|
||||||
read_side
|
|
||||||
.read_exact(&mut encrypted_header)
|
|
||||||
.await
|
|
||||||
.expect("must read secure header");
|
|
||||||
let decrypted_header = decryptor.decrypt(&encrypted_header);
|
|
||||||
let header: [u8; 4] = decrypted_header
|
|
||||||
.try_into()
|
|
||||||
.expect("decrypted secure header must be 4 bytes");
|
|
||||||
let wire_len_raw = u32::from_le_bytes(header);
|
|
||||||
|
|
||||||
assert_ne!(
|
|
||||||
wire_len_raw & 0x8000_0000,
|
|
||||||
0,
|
|
||||||
"secure quickack bit must be set"
|
|
||||||
);
|
|
||||||
|
|
||||||
let wire_len = (wire_len_raw & 0x7fff_ffff) as usize;
|
|
||||||
assert!(wire_len >= payload.len());
|
|
||||||
let padding_len = wire_len - payload.len();
|
|
||||||
assert!(
|
|
||||||
(1..=3).contains(&padding_len),
|
|
||||||
"secure writer must add bounded random tail padding, got {padding_len}"
|
|
||||||
);
|
|
||||||
|
|
||||||
let mut encrypted_body = vec![0u8; wire_len];
|
|
||||||
read_side
|
|
||||||
.read_exact(&mut encrypted_body)
|
|
||||||
.await
|
|
||||||
.expect("must read secure body");
|
|
||||||
let decrypted_body = decryptor.decrypt(&encrypted_body);
|
|
||||||
assert_eq!(&decrypted_body[..payload.len()], payload.as_slice());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
#[ignore = "heavy: allocates >64MiB to validate abridged too-large fail-closed branch"]
|
|
||||||
async fn write_client_payload_abridged_too_large_is_rejected_fail_closed() {
|
|
||||||
let (_read_side, write_side) = duplex(1024);
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
|
|
||||||
let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024);
|
|
||||||
let rng = SecureRandom::new();
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
|
|
||||||
// Exactly one 4-byte word above the encodable 24-bit abridged length range.
|
|
||||||
let payload = vec![0x00u8; (1 << 24) * 4];
|
|
||||||
let err = write_client_payload(
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Abridged,
|
|
||||||
0,
|
|
||||||
&payload,
|
|
||||||
&rng,
|
|
||||||
&mut frame_buf,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect_err("oversized abridged payload must be rejected");
|
|
||||||
|
|
||||||
let msg = format!("{err}");
|
|
||||||
assert!(
|
|
||||||
msg.contains("Abridged frame too large"),
|
|
||||||
"error must clearly indicate oversize fail-close path, got: {msg}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn write_client_ack_intermediate_is_little_endian() {
|
|
||||||
let (mut read_side, write_side) = duplex(1024);
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024);
|
|
||||||
let mut decryptor = AesCtr::new(&key, iv);
|
|
||||||
|
|
||||||
write_client_ack(&mut writer, ProtoTag::Intermediate, 0x11_22_33_44)
|
|
||||||
.await
|
|
||||||
.expect("ack serialization should succeed");
|
|
||||||
writer.flush().await.expect("flush must succeed");
|
|
||||||
|
|
||||||
let mut encrypted = [0u8; 4];
|
|
||||||
read_side
|
|
||||||
.read_exact(&mut encrypted)
|
|
||||||
.await
|
|
||||||
.expect("must read ack bytes");
|
|
||||||
let plain = decryptor.decrypt(&encrypted);
|
|
||||||
assert_eq!(plain.as_slice(), &0x11_22_33_44u32.to_le_bytes());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn write_client_ack_abridged_is_big_endian() {
|
|
||||||
let (mut read_side, write_side) = duplex(1024);
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024);
|
|
||||||
let mut decryptor = AesCtr::new(&key, iv);
|
|
||||||
|
|
||||||
write_client_ack(&mut writer, ProtoTag::Abridged, 0xDE_AD_BE_EF)
|
|
||||||
.await
|
|
||||||
.expect("ack serialization should succeed");
|
|
||||||
writer.flush().await.expect("flush must succeed");
|
|
||||||
|
|
||||||
let mut encrypted = [0u8; 4];
|
|
||||||
read_side
|
|
||||||
.read_exact(&mut encrypted)
|
|
||||||
.await
|
|
||||||
.expect("must read ack bytes");
|
|
||||||
let plain = decryptor.decrypt(&encrypted);
|
|
||||||
assert_eq!(plain.as_slice(), &0xDE_AD_BE_EFu32.to_be_bytes());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn write_client_payload_abridged_short_boundary_0x7e_is_single_byte_header() {
|
|
||||||
let (mut read_side, write_side) = duplex(1024 * 1024);
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024);
|
|
||||||
let mut decryptor = AesCtr::new(&key, iv);
|
|
||||||
let rng = SecureRandom::new();
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let payload = vec![0xABu8; 0x7e * 4];
|
|
||||||
|
|
||||||
write_client_payload(
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Abridged,
|
|
||||||
0,
|
|
||||||
&payload,
|
|
||||||
&rng,
|
|
||||||
&mut frame_buf,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("boundary payload should serialize");
|
|
||||||
writer.flush().await.expect("flush must succeed");
|
|
||||||
|
|
||||||
let mut encrypted = vec![0u8; 1 + payload.len()];
|
|
||||||
read_side.read_exact(&mut encrypted).await.unwrap();
|
|
||||||
let plain = decryptor.decrypt(&encrypted);
|
|
||||||
assert_eq!(plain[0], 0x7e);
|
|
||||||
assert_eq!(&plain[1..], payload.as_slice());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn write_client_payload_abridged_extended_without_quickack_has_clean_prefix() {
|
|
||||||
let (mut read_side, write_side) = duplex(16 * 1024);
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024);
|
|
||||||
let mut decryptor = AesCtr::new(&key, iv);
|
|
||||||
let rng = SecureRandom::new();
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let payload = vec![0x42u8; 0x80 * 4];
|
|
||||||
|
|
||||||
write_client_payload(
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Abridged,
|
|
||||||
0,
|
|
||||||
&payload,
|
|
||||||
&rng,
|
|
||||||
&mut frame_buf,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("extended payload should serialize");
|
|
||||||
writer.flush().await.expect("flush must succeed");
|
|
||||||
|
|
||||||
let mut encrypted = vec![0u8; 4 + payload.len()];
|
|
||||||
read_side.read_exact(&mut encrypted).await.unwrap();
|
|
||||||
let plain = decryptor.decrypt(&encrypted);
|
|
||||||
assert_eq!(plain[0], 0x7f);
|
|
||||||
assert_eq!(&plain[1..4], &[0x80, 0x00, 0x00]);
|
|
||||||
assert_eq!(&plain[4..], payload.as_slice());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn write_client_payload_intermediate_zero_length_emits_header_only() {
|
|
||||||
let (mut read_side, write_side) = duplex(1024);
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024);
|
|
||||||
let mut decryptor = AesCtr::new(&key, iv);
|
|
||||||
let rng = SecureRandom::new();
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
|
|
||||||
write_client_payload(
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
0,
|
|
||||||
&[],
|
|
||||||
&rng,
|
|
||||||
&mut frame_buf,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("zero-length intermediate payload should serialize");
|
|
||||||
writer.flush().await.expect("flush must succeed");
|
|
||||||
|
|
||||||
let mut encrypted = [0u8; 4];
|
|
||||||
read_side.read_exact(&mut encrypted).await.unwrap();
|
|
||||||
let plain = decryptor.decrypt(&encrypted);
|
|
||||||
assert_eq!(plain.as_slice(), &[0, 0, 0, 0]);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn write_client_payload_intermediate_ignores_unrelated_flags() {
|
|
||||||
let (mut read_side, write_side) = duplex(1024);
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024);
|
|
||||||
let mut decryptor = AesCtr::new(&key, iv);
|
|
||||||
let rng = SecureRandom::new();
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let payload = [7u8; 12];
|
|
||||||
|
|
||||||
write_client_payload(
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
0x4000_0000,
|
|
||||||
&payload,
|
|
||||||
&rng,
|
|
||||||
&mut frame_buf,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("payload should serialize");
|
|
||||||
writer.flush().await.expect("flush must succeed");
|
|
||||||
|
|
||||||
let mut encrypted = [0u8; 16];
|
|
||||||
read_side.read_exact(&mut encrypted).await.unwrap();
|
|
||||||
let plain = decryptor.decrypt(&encrypted);
|
|
||||||
let len = u32::from_le_bytes(plain[0..4].try_into().unwrap());
|
|
||||||
assert_eq!(len, payload.len() as u32, "only quickack bit may affect header");
|
|
||||||
assert_eq!(&plain[4..], payload.as_slice());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn write_client_payload_secure_without_quickack_keeps_msb_clear() {
|
|
||||||
let (mut read_side, write_side) = duplex(4096);
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024);
|
|
||||||
let mut decryptor = AesCtr::new(&key, iv);
|
|
||||||
let rng = SecureRandom::new();
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let payload = [0x1Du8; 64];
|
|
||||||
|
|
||||||
write_client_payload(
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Secure,
|
|
||||||
0,
|
|
||||||
&payload,
|
|
||||||
&rng,
|
|
||||||
&mut frame_buf,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("payload should serialize");
|
|
||||||
writer.flush().await.expect("flush must succeed");
|
|
||||||
|
|
||||||
let mut encrypted_header = [0u8; 4];
|
|
||||||
read_side.read_exact(&mut encrypted_header).await.unwrap();
|
|
||||||
let plain_header = decryptor.decrypt(&encrypted_header);
|
|
||||||
let h: [u8; 4] = plain_header.as_slice().try_into().unwrap();
|
|
||||||
let wire_len_raw = u32::from_le_bytes(h);
|
|
||||||
assert_eq!(wire_len_raw & 0x8000_0000, 0, "quickack bit must stay clear");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn secure_padding_light_fuzz_distribution_has_multiple_outcomes() {
|
|
||||||
let (mut read_side, write_side) = duplex(256 * 1024);
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024);
|
|
||||||
let mut decryptor = AesCtr::new(&key, iv);
|
|
||||||
let rng = SecureRandom::new();
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let payload = [0x55u8; 100];
|
|
||||||
let mut seen = [false; 4];
|
|
||||||
|
|
||||||
for _ in 0..96 {
|
|
||||||
write_client_payload(
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Secure,
|
|
||||||
0,
|
|
||||||
&payload,
|
|
||||||
&rng,
|
|
||||||
&mut frame_buf,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("secure payload should serialize");
|
|
||||||
writer.flush().await.expect("flush must succeed");
|
|
||||||
|
|
||||||
let mut encrypted_header = [0u8; 4];
|
|
||||||
read_side.read_exact(&mut encrypted_header).await.unwrap();
|
|
||||||
let plain_header = decryptor.decrypt(&encrypted_header);
|
|
||||||
let h: [u8; 4] = plain_header.as_slice().try_into().unwrap();
|
|
||||||
let wire_len = (u32::from_le_bytes(h) & 0x7fff_ffff) as usize;
|
|
||||||
let padding_len = wire_len - payload.len();
|
|
||||||
assert!((1..=3).contains(&padding_len));
|
|
||||||
seen[padding_len] = true;
|
|
||||||
|
|
||||||
let mut encrypted_body = vec![0u8; wire_len];
|
|
||||||
read_side.read_exact(&mut encrypted_body).await.unwrap();
|
|
||||||
let _ = decryptor.decrypt(&encrypted_body);
|
|
||||||
}
|
|
||||||
|
|
||||||
let distinct = (1..=3).filter(|idx| seen[*idx]).count();
|
|
||||||
assert!(
|
|
||||||
distinct >= 2,
|
|
||||||
"padding generator should not collapse to a single outcome under campaign"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn write_client_payload_mixed_proto_sequence_preserves_stream_sync() {
|
|
||||||
let (mut read_side, write_side) = duplex(128 * 1024);
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024);
|
|
||||||
let mut decryptor = AesCtr::new(&key, iv);
|
|
||||||
let rng = SecureRandom::new();
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
|
|
||||||
let p1 = vec![1u8; 8];
|
|
||||||
let p2 = vec![2u8; 16];
|
|
||||||
let p3 = vec![3u8; 20];
|
|
||||||
|
|
||||||
write_client_payload(&mut writer, ProtoTag::Abridged, 0, &p1, &rng, &mut frame_buf)
|
|
||||||
.await
|
|
||||||
.unwrap();
|
|
||||||
write_client_payload(
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
RPC_FLAG_QUICKACK,
|
|
||||||
&p2,
|
|
||||||
&rng,
|
|
||||||
&mut frame_buf,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.unwrap();
|
|
||||||
write_client_payload(&mut writer, ProtoTag::Secure, 0, &p3, &rng, &mut frame_buf)
|
|
||||||
.await
|
|
||||||
.unwrap();
|
|
||||||
writer.flush().await.unwrap();
|
|
||||||
|
|
||||||
// Frame 1: abridged short.
|
|
||||||
let mut e1 = vec![0u8; 1 + p1.len()];
|
|
||||||
read_side.read_exact(&mut e1).await.unwrap();
|
|
||||||
let d1 = decryptor.decrypt(&e1);
|
|
||||||
assert_eq!(d1[0], (p1.len() / 4) as u8);
|
|
||||||
assert_eq!(&d1[1..], p1.as_slice());
|
|
||||||
|
|
||||||
// Frame 2: intermediate with quickack.
|
|
||||||
let mut e2 = vec![0u8; 4 + p2.len()];
|
|
||||||
read_side.read_exact(&mut e2).await.unwrap();
|
|
||||||
let d2 = decryptor.decrypt(&e2);
|
|
||||||
let l2 = u32::from_le_bytes(d2[0..4].try_into().unwrap());
|
|
||||||
assert_ne!(l2 & 0x8000_0000, 0);
|
|
||||||
assert_eq!((l2 & 0x7fff_ffff) as usize, p2.len());
|
|
||||||
assert_eq!(&d2[4..], p2.as_slice());
|
|
||||||
|
|
||||||
// Frame 3: secure with bounded tail.
|
|
||||||
let mut e3h = [0u8; 4];
|
|
||||||
read_side.read_exact(&mut e3h).await.unwrap();
|
|
||||||
let d3h = decryptor.decrypt(&e3h);
|
|
||||||
let l3 = (u32::from_le_bytes(d3h.as_slice().try_into().unwrap()) & 0x7fff_ffff) as usize;
|
|
||||||
assert!(l3 >= p3.len());
|
|
||||||
assert!((1..=3).contains(&(l3 - p3.len())));
|
|
||||||
let mut e3b = vec![0u8; l3];
|
|
||||||
read_side.read_exact(&mut e3b).await.unwrap();
|
|
||||||
let d3b = decryptor.decrypt(&e3b);
|
|
||||||
assert_eq!(&d3b[..p3.len()], p3.as_slice());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn should_yield_sender_boundary_matrix_blackhat() {
|
|
||||||
assert!(!should_yield_c2me_sender(0, false));
|
|
||||||
assert!(!should_yield_c2me_sender(0, true));
|
|
||||||
assert!(!should_yield_c2me_sender(C2ME_SENDER_FAIRNESS_BUDGET - 1, true));
|
|
||||||
assert!(!should_yield_c2me_sender(C2ME_SENDER_FAIRNESS_BUDGET, false));
|
|
||||||
assert!(should_yield_c2me_sender(C2ME_SENDER_FAIRNESS_BUDGET, true));
|
|
||||||
assert!(should_yield_c2me_sender(
|
|
||||||
C2ME_SENDER_FAIRNESS_BUDGET.saturating_add(1024),
|
|
||||||
true
|
|
||||||
));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn should_yield_sender_light_fuzz_matches_oracle() {
|
|
||||||
let mut s: u64 = 0xD00D_BAAD_F00D_CAFE;
|
|
||||||
for _ in 0..5000 {
|
|
||||||
s ^= s << 7;
|
|
||||||
s ^= s >> 9;
|
|
||||||
s ^= s << 8;
|
|
||||||
let sent = (s as usize) & 0x1fff;
|
|
||||||
let backlog = (s & 1) != 0;
|
|
||||||
|
|
||||||
let expected = backlog && sent >= C2ME_SENDER_FAIRNESS_BUDGET;
|
|
||||||
assert_eq!(should_yield_c2me_sender(sent, backlog), expected);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn quota_would_be_exceeded_exact_remaining_one_byte() {
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = "quota-edge";
|
|
||||||
let quota = 100u64;
|
|
||||||
stats.add_user_octets_to(user, 99);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
!quota_would_be_exceeded_for_user(&stats, user, Some(quota), 1),
|
|
||||||
"exactly remaining budget should be allowed"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
quota_would_be_exceeded_for_user(&stats, user, Some(quota), 2),
|
|
||||||
"one byte beyond remaining budget must be rejected"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn quota_would_be_exceeded_saturating_edge_remains_fail_closed() {
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = "quota-saturating-edge";
|
|
||||||
let quota = u64::MAX - 3;
|
|
||||||
stats.add_user_octets_to(user, u64::MAX - 4);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
quota_would_be_exceeded_for_user(&stats, user, Some(quota), 2),
|
|
||||||
"saturating arithmetic edge must stay fail-closed"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn quota_exceeded_boundary_is_inclusive() {
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = "quota-inclusive-boundary";
|
|
||||||
stats.add_user_octets_to(user, 50);
|
|
||||||
|
|
||||||
assert!(quota_exceeded_for_user(&stats, user, Some(50)));
|
|
||||||
assert!(!quota_exceeded_for_user(&stats, user, Some(51)));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn quota_soft_helper_matches_capped_generic_helper_matrix() {
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = "quota-soft-parity";
|
|
||||||
|
|
||||||
for used in [0u64, 1, 7, 63, 127, 255] {
|
|
||||||
stats.sub_user_octets_to(user, stats.get_user_total_octets(user));
|
|
||||||
stats.add_user_octets_to(user, used);
|
|
||||||
|
|
||||||
for quota in [8u64, 64, 128, 256] {
|
|
||||||
for overshoot in [0u64, 1, 5, 32] {
|
|
||||||
for bytes in [0u64, 1, 2, 7, 31, 64] {
|
|
||||||
let soft = quota_would_be_exceeded_for_user_soft(
|
|
||||||
&stats,
|
|
||||||
user,
|
|
||||||
Some(quota),
|
|
||||||
bytes,
|
|
||||||
overshoot,
|
|
||||||
);
|
|
||||||
let capped = quota_would_be_exceeded_for_user(
|
|
||||||
&stats,
|
|
||||||
user,
|
|
||||||
Some(quota_soft_cap(quota, overshoot)),
|
|
||||||
bytes,
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
soft, capped,
|
|
||||||
"soft helper parity mismatch: used={used} quota={quota} overshoot={overshoot} bytes={bytes}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn quota_soft_helper_none_limit_never_rejects() {
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = "quota-soft-none";
|
|
||||||
stats.add_user_octets_to(user, u64::MAX);
|
|
||||||
|
|
||||||
assert!(!quota_would_be_exceeded_for_user_soft(
|
|
||||||
&stats,
|
|
||||||
user,
|
|
||||||
None,
|
|
||||||
u64::MAX,
|
|
||||||
u64::MAX,
|
|
||||||
));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn quota_soft_cap_saturates_and_stays_fail_closed() {
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = "quota-soft-saturating";
|
|
||||||
let quota = u64::MAX - 2;
|
|
||||||
let overshoot = 100;
|
|
||||||
|
|
||||||
assert_eq!(quota_soft_cap(quota, overshoot), u64::MAX);
|
|
||||||
|
|
||||||
stats.add_user_octets_to(user, u64::MAX - 1);
|
|
||||||
assert!(quota_would_be_exceeded_for_user_soft(
|
|
||||||
&stats,
|
|
||||||
user,
|
|
||||||
Some(quota),
|
|
||||||
2,
|
|
||||||
overshoot,
|
|
||||||
));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn enqueue_c2me_close_fast_path_succeeds_without_backpressure() {
|
|
||||||
let (tx, mut rx) = mpsc::channel::<C2MeCommand>(4);
|
|
||||||
enqueue_c2me_command(&tx, C2MeCommand::Close)
|
|
||||||
.await
|
|
||||||
.expect("close should enqueue on fast path");
|
|
||||||
|
|
||||||
let recv = timeout(TokioDuration::from_millis(50), rx.recv())
|
|
||||||
.await
|
|
||||||
.expect("must receive close command")
|
|
||||||
.expect("close command should be present");
|
|
||||||
assert!(matches!(recv, C2MeCommand::Close));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn enqueue_c2me_data_full_then_drain_preserves_order() {
|
|
||||||
let (tx, mut rx) = mpsc::channel::<C2MeCommand>(1);
|
|
||||||
tx.send(C2MeCommand::Data {
|
|
||||||
payload: make_pooled_payload(&[1]),
|
|
||||||
flags: 10,
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.unwrap();
|
|
||||||
|
|
||||||
let tx2 = tx.clone();
|
|
||||||
let producer = tokio::spawn(async move {
|
|
||||||
enqueue_c2me_command(
|
|
||||||
&tx2,
|
|
||||||
C2MeCommand::Data {
|
|
||||||
payload: make_pooled_payload(&[2, 2]),
|
|
||||||
flags: 20,
|
|
||||||
},
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
});
|
|
||||||
|
|
||||||
tokio::time::sleep(TokioDuration::from_millis(10)).await;
|
|
||||||
|
|
||||||
let first = rx.recv().await.expect("first item should exist");
|
|
||||||
match first {
|
|
||||||
C2MeCommand::Data { payload, flags } => {
|
|
||||||
assert_eq!(payload.as_ref(), &[1]);
|
|
||||||
assert_eq!(flags, 10);
|
|
||||||
}
|
|
||||||
C2MeCommand::Close => panic!("unexpected close as first item"),
|
|
||||||
}
|
|
||||||
|
|
||||||
producer.await.unwrap().expect("producer should complete");
|
|
||||||
|
|
||||||
let second = timeout(TokioDuration::from_millis(100), rx.recv())
|
|
||||||
.await
|
|
||||||
.unwrap()
|
|
||||||
.expect("second item should exist");
|
|
||||||
match second {
|
|
||||||
C2MeCommand::Data { payload, flags } => {
|
|
||||||
assert_eq!(payload.as_ref(), &[2, 2]);
|
|
||||||
assert_eq!(flags, 20);
|
|
||||||
}
|
|
||||||
C2MeCommand::Close => panic!("unexpected close as second item"),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
@ -1,295 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::crypto::{AesCtr, SecureRandom};
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use crate::stream::CryptoWriter;
|
|
||||||
use bytes::Bytes;
|
|
||||||
use std::pin::Pin;
|
|
||||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
|
||||||
use std::sync::{Arc, Mutex};
|
|
||||||
use std::task::{Context, Poll, Waker};
|
|
||||||
use tokio::io::AsyncWrite;
|
|
||||||
use tokio::sync::Notify;
|
|
||||||
use tokio::task::JoinSet;
|
|
||||||
use tokio::time::{Duration, timeout};
|
|
||||||
|
|
||||||
fn make_crypto_writer<W>(writer: W) -> CryptoWriter<W>
|
|
||||||
where
|
|
||||||
W: tokio::io::AsyncWrite + Unpin,
|
|
||||||
{
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
CryptoWriter::new(writer, AesCtr::new(&key, iv), 8 * 1024)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Default)]
|
|
||||||
struct BlockingWriteState {
|
|
||||||
write_entered: AtomicBool,
|
|
||||||
released: AtomicBool,
|
|
||||||
write_waker: Mutex<Option<Waker>>,
|
|
||||||
write_entered_notify: Notify,
|
|
||||||
}
|
|
||||||
|
|
||||||
struct BlockingWrite {
|
|
||||||
state: Arc<BlockingWriteState>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl BlockingWrite {
|
|
||||||
fn new(state: Arc<BlockingWriteState>) -> Self {
|
|
||||||
Self { state }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl AsyncWrite for BlockingWrite {
|
|
||||||
fn poll_write(
|
|
||||||
self: Pin<&mut Self>,
|
|
||||||
cx: &mut Context<'_>,
|
|
||||||
buf: &[u8],
|
|
||||||
) -> Poll<std::io::Result<usize>> {
|
|
||||||
self.state.write_entered.store(true, Ordering::Release);
|
|
||||||
self.state.write_entered_notify.notify_waiters();
|
|
||||||
|
|
||||||
if self.state.released.load(Ordering::Acquire) {
|
|
||||||
return Poll::Ready(Ok(buf.len()));
|
|
||||||
}
|
|
||||||
|
|
||||||
if let Ok(mut slot) = self.state.write_waker.lock() {
|
|
||||||
*slot = Some(cx.waker().clone());
|
|
||||||
}
|
|
||||||
|
|
||||||
Poll::Pending
|
|
||||||
}
|
|
||||||
|
|
||||||
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
|
|
||||||
Poll::Ready(Ok(()))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
|
|
||||||
Poll::Ready(Ok(()))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn wait_until_blocking_write_entered(state: &Arc<BlockingWriteState>) {
|
|
||||||
for _ in 0..8 {
|
|
||||||
if state.write_entered.load(Ordering::Acquire) {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let _ = timeout(Duration::from_millis(25), state.write_entered_notify.notified()).await;
|
|
||||||
}
|
|
||||||
|
|
||||||
panic!("blocking writer did not enter poll_write in bounded time");
|
|
||||||
}
|
|
||||||
|
|
||||||
fn release_blocking_write(state: &Arc<BlockingWriteState>) {
|
|
||||||
state.released.store(true, Ordering::Release);
|
|
||||||
if let Ok(mut slot) = state.write_waker.lock()
|
|
||||||
&& let Some(waker) = slot.take()
|
|
||||||
{
|
|
||||||
waker.wake();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
|
||||||
async fn adversarial_blocked_write_releases_cross_mode_lock_and_preserves_fail_closed_quota() {
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let user = format!("middle-cross-release-regression-{}", std::process::id());
|
|
||||||
let cross_mode_lock = Arc::new(cross_mode_quota_user_lock_for_tests(&user));
|
|
||||||
let bytes_me2c = Arc::new(AtomicU64::new(0));
|
|
||||||
let writer_state = Arc::new(BlockingWriteState::default());
|
|
||||||
|
|
||||||
let first = {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let user = user.clone();
|
|
||||||
let cross_mode_lock = Arc::clone(&cross_mode_lock);
|
|
||||||
let bytes_me2c = Arc::clone(&bytes_me2c);
|
|
||||||
let writer_state = Arc::clone(&writer_state);
|
|
||||||
tokio::spawn(async move {
|
|
||||||
let mut writer = make_crypto_writer(BlockingWrite::new(writer_state));
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0xAA, 0xBB, 0xCC, 0xDD]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
stats.as_ref(),
|
|
||||||
&user,
|
|
||||||
Some(4),
|
|
||||||
0,
|
|
||||||
Some(&cross_mode_lock),
|
|
||||||
bytes_me2c.as_ref(),
|
|
||||||
41_000,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
})
|
|
||||||
};
|
|
||||||
|
|
||||||
wait_until_blocking_write_entered(&writer_state).await;
|
|
||||||
|
|
||||||
let guard = timeout(Duration::from_millis(40), cross_mode_lock.lock())
|
|
||||||
.await
|
|
||||||
.expect("cross-mode lock must be released while first write is pending");
|
|
||||||
drop(guard);
|
|
||||||
|
|
||||||
let second = {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let user = user.clone();
|
|
||||||
let cross_mode_lock = Arc::clone(&cross_mode_lock);
|
|
||||||
let bytes_me2c = Arc::clone(&bytes_me2c);
|
|
||||||
tokio::spawn(async move {
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
timeout(
|
|
||||||
Duration::from_millis(150),
|
|
||||||
process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0xEE]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
stats.as_ref(),
|
|
||||||
&user,
|
|
||||||
Some(4),
|
|
||||||
0,
|
|
||||||
Some(&cross_mode_lock),
|
|
||||||
bytes_me2c.as_ref(),
|
|
||||||
41_001,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
})
|
|
||||||
};
|
|
||||||
|
|
||||||
let second_result = second
|
|
||||||
.await
|
|
||||||
.expect("second task must not panic")
|
|
||||||
.expect("second write must not block on cross-mode lock");
|
|
||||||
assert!(
|
|
||||||
matches!(second_result, Err(ProxyError::DataQuotaExceeded { .. })),
|
|
||||||
"second write must fail closed due to first write reservation"
|
|
||||||
);
|
|
||||||
|
|
||||||
release_blocking_write(&writer_state);
|
|
||||||
|
|
||||||
let first_result = timeout(Duration::from_millis(300), first)
|
|
||||||
.await
|
|
||||||
.expect("first task timed out")
|
|
||||||
.expect("first task must not panic");
|
|
||||||
assert!(first_result.is_ok());
|
|
||||||
|
|
||||||
assert_eq!(stats.get_user_total_octets(&user), 4);
|
|
||||||
assert_eq!(bytes_me2c.load(Ordering::Relaxed), 4);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn stress_pending_write_does_not_starve_same_user_waiters_after_quota_boundary() {
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let user = format!("middle-cross-release-stress-{}", std::process::id());
|
|
||||||
let cross_mode_lock = Arc::new(cross_mode_quota_user_lock_for_tests(&user));
|
|
||||||
let bytes_me2c = Arc::new(AtomicU64::new(0));
|
|
||||||
let writer_state = Arc::new(BlockingWriteState::default());
|
|
||||||
|
|
||||||
let first = {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let user = user.clone();
|
|
||||||
let cross_mode_lock = Arc::clone(&cross_mode_lock);
|
|
||||||
let bytes_me2c = Arc::clone(&bytes_me2c);
|
|
||||||
let writer_state = Arc::clone(&writer_state);
|
|
||||||
tokio::spawn(async move {
|
|
||||||
let mut writer = make_crypto_writer(BlockingWrite::new(writer_state));
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0x01, 0x02]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
stats.as_ref(),
|
|
||||||
&user,
|
|
||||||
Some(3),
|
|
||||||
0,
|
|
||||||
Some(&cross_mode_lock),
|
|
||||||
bytes_me2c.as_ref(),
|
|
||||||
41_100,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
})
|
|
||||||
};
|
|
||||||
|
|
||||||
wait_until_blocking_write_entered(&writer_state).await;
|
|
||||||
|
|
||||||
let mut set = JoinSet::new();
|
|
||||||
for idx in 0..48u64 {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let user = user.clone();
|
|
||||||
let cross_mode_lock = Arc::clone(&cross_mode_lock);
|
|
||||||
let bytes_me2c = Arc::clone(&bytes_me2c);
|
|
||||||
set.spawn(async move {
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
timeout(
|
|
||||||
Duration::from_millis(200),
|
|
||||||
process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0x10]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
stats.as_ref(),
|
|
||||||
&user,
|
|
||||||
Some(3),
|
|
||||||
0,
|
|
||||||
Some(&cross_mode_lock),
|
|
||||||
bytes_me2c.as_ref(),
|
|
||||||
41_200 + idx,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut ok = 0usize;
|
|
||||||
let mut quota_exceeded = 0usize;
|
|
||||||
while let Some(done) = set.join_next().await {
|
|
||||||
let timed = done.expect("waiter task must not panic");
|
|
||||||
let result = timed.expect("waiter must not block behind pending first write");
|
|
||||||
match result {
|
|
||||||
Ok(_) => ok += 1,
|
|
||||||
Err(ProxyError::DataQuotaExceeded { .. }) => quota_exceeded += 1,
|
|
||||||
Err(other) => panic!("unexpected error in waiter: {other:?}"),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
assert_eq!(ok, 1, "exactly one waiter should consume remaining one-byte quota");
|
|
||||||
assert_eq!(quota_exceeded, 47);
|
|
||||||
|
|
||||||
release_blocking_write(&writer_state);
|
|
||||||
|
|
||||||
let first_result = timeout(Duration::from_millis(300), first)
|
|
||||||
.await
|
|
||||||
.expect("first task timed out")
|
|
||||||
.expect("first task must not panic");
|
|
||||||
assert!(first_result.is_ok());
|
|
||||||
|
|
||||||
assert_eq!(stats.get_user_total_octets(&user), 3);
|
|
||||||
assert_eq!(bytes_me2c.load(Ordering::Relaxed), 3);
|
|
||||||
}
|
|
||||||
|
|
@ -1,116 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::crypto::{AesCtr, SecureRandom};
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use crate::stream::CryptoWriter;
|
|
||||||
use bytes::Bytes;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::{AtomicU64, Ordering};
|
|
||||||
use std::sync::{Mutex, OnceLock};
|
|
||||||
|
|
||||||
fn make_crypto_writer<W>(writer: W) -> CryptoWriter<W>
|
|
||||||
where
|
|
||||||
W: tokio::io::AsyncWrite + Unpin,
|
|
||||||
{
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
CryptoWriter::new(writer, AesCtr::new(&key, iv), 8 * 1024)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn lookup_counter_test_lock() -> &'static Mutex<()> {
|
|
||||||
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
|
|
||||||
LOCK.get_or_init(|| Mutex::new(()))
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn tdd_prefetched_cross_mode_lock_avoids_per_frame_registry_lookup_in_me_to_client_writer() {
|
|
||||||
let _guard = lookup_counter_test_lock()
|
|
||||||
.lock()
|
|
||||||
.unwrap_or_else(|poison| poison.into_inner());
|
|
||||||
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("middle-cross-mode-lookup-{}", std::process::id());
|
|
||||||
let cross_mode_lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
|
|
||||||
crate::proxy::quota_lock_registry::reset_cross_mode_quota_user_lock_lookup_count_for_tests();
|
|
||||||
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
|
|
||||||
for idx in 0..8u64 {
|
|
||||||
let outcome = process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0xAB]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(1024),
|
|
||||||
0,
|
|
||||||
Some(&cross_mode_lock),
|
|
||||||
&bytes_me2c,
|
|
||||||
20_000 + idx,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
assert!(outcome.is_ok());
|
|
||||||
}
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
crate::proxy::quota_lock_registry::cross_mode_quota_user_lock_lookup_count_for_user_for_tests(&user),
|
|
||||||
0,
|
|
||||||
"prefetched lock path must not re-query lock registry per frame"
|
|
||||||
);
|
|
||||||
assert_eq!(stats.get_user_total_octets(&user), 8);
|
|
||||||
assert_eq!(bytes_me2c.load(Ordering::Relaxed), 8);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn control_without_prefetched_lock_still_uses_registry_lookup_path() {
|
|
||||||
let _guard = lookup_counter_test_lock()
|
|
||||||
.lock()
|
|
||||||
.unwrap_or_else(|poison| poison.into_inner());
|
|
||||||
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("middle-cross-mode-lookup-control-{}", std::process::id());
|
|
||||||
|
|
||||||
crate::proxy::quota_lock_registry::reset_cross_mode_quota_user_lock_lookup_count_for_tests();
|
|
||||||
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
|
|
||||||
let outcome = process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0xCD]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(1024),
|
|
||||||
0,
|
|
||||||
None,
|
|
||||||
&bytes_me2c,
|
|
||||||
20_100,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
assert!(outcome.is_ok());
|
|
||||||
assert_eq!(
|
|
||||||
crate::proxy::quota_lock_registry::cross_mode_quota_user_lock_lookup_count_for_user_for_tests(&user),
|
|
||||||
1,
|
|
||||||
"fallback path without prefetched lock should perform a registry lookup"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
@ -1,376 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::crypto::{AesCtr, SecureRandom};
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use crate::stream::CryptoWriter;
|
|
||||||
use bytes::Bytes;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::{AtomicU64, Ordering};
|
|
||||||
use tokio::time::{Duration, timeout};
|
|
||||||
|
|
||||||
fn make_crypto_writer<W>(writer: W) -> CryptoWriter<W>
|
|
||||||
where
|
|
||||||
W: tokio::io::AsyncWrite + Unpin,
|
|
||||||
{
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
CryptoWriter::new(writer, AesCtr::new(&key, iv), 8 * 1024)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn positive_quota_limited_me_to_client_write_updates_counters_exactly_once() {
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("middle-cross-matrix-positive-{}", std::process::id());
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
|
|
||||||
let result = process_me_writer_response(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[1, 2, 3, 4]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(128),
|
|
||||||
0,
|
|
||||||
&bytes_me2c,
|
|
||||||
10_001,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
assert!(result.is_ok());
|
|
||||||
assert_eq!(stats.get_user_total_octets(&user), 4);
|
|
||||||
assert_eq!(bytes_me2c.load(Ordering::Relaxed), 4);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn negative_held_cross_mode_lock_blocks_quota_limited_me_to_client_path() {
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("middle-cross-matrix-negative-{}", std::process::id());
|
|
||||||
let held = cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold lock before ME->C call");
|
|
||||||
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
|
|
||||||
let blocked = timeout(
|
|
||||||
Duration::from_millis(25),
|
|
||||||
process_me_writer_response(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0x41]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(256),
|
|
||||||
0,
|
|
||||||
&bytes_me2c,
|
|
||||||
10_002,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
assert!(blocked.is_err());
|
|
||||||
drop(held_guard);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn edge_quota_none_bypasses_cross_mode_lock_guard_in_me_to_client_path() {
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("middle-cross-matrix-edge-none-{}", std::process::id());
|
|
||||||
let held = cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold lock while quota is disabled");
|
|
||||||
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
|
|
||||||
let outcome = timeout(
|
|
||||||
Duration::from_millis(80),
|
|
||||||
process_me_writer_response(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0x11, 0x22]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
None,
|
|
||||||
0,
|
|
||||||
&bytes_me2c,
|
|
||||||
10_003,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("quota-none path must not wait on cross-mode lock");
|
|
||||||
|
|
||||||
assert!(outcome.is_ok());
|
|
||||||
drop(held_guard);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn adversarial_same_user_parallel_quota_limited_writes_stay_hard_capped() {
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let user = format!("middle-cross-matrix-adversarial-{}", std::process::id());
|
|
||||||
let limit = 64u64;
|
|
||||||
let bytes_me2c = Arc::new(AtomicU64::new(0));
|
|
||||||
let mut tasks = Vec::new();
|
|
||||||
|
|
||||||
for idx in 0..256u64 {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let bytes_me2c = Arc::clone(&bytes_me2c);
|
|
||||||
let user = user.clone();
|
|
||||||
tasks.push(tokio::spawn(async move {
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
process_me_writer_response(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0xEE]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
stats.as_ref(),
|
|
||||||
&user,
|
|
||||||
Some(limit),
|
|
||||||
0,
|
|
||||||
bytes_me2c.as_ref(),
|
|
||||||
11_000 + idx,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut ok = 0usize;
|
|
||||||
for task in tasks {
|
|
||||||
match task.await.expect("task must not panic") {
|
|
||||||
Ok(_) => ok += 1,
|
|
||||||
Err(ProxyError::DataQuotaExceeded { .. }) => {}
|
|
||||||
Err(other) => panic!("unexpected error in adversarial parallel case: {other:?}"),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
assert_eq!(ok, limit as usize);
|
|
||||||
assert_eq!(stats.get_user_total_octets(&user), limit);
|
|
||||||
assert_eq!(bytes_me2c.load(Ordering::Relaxed), limit);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn integration_shared_lock_blocks_direct_relay_and_middle_relay_for_same_user() {
|
|
||||||
let user = format!("middle-cross-matrix-integration-{}", std::process::id());
|
|
||||||
let relay_lock = crate::proxy::relay::cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
let middle_lock = cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
assert!(
|
|
||||||
Arc::ptr_eq(&relay_lock, &middle_lock),
|
|
||||||
"relay and middle-relay must share the same cross-mode lock identity"
|
|
||||||
);
|
|
||||||
|
|
||||||
let held_guard = relay_lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold shared cross-mode lock");
|
|
||||||
|
|
||||||
let stats = Stats::new();
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
|
|
||||||
let middle_blocked = timeout(
|
|
||||||
Duration::from_millis(25),
|
|
||||||
process_me_writer_response(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0x92]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(1024),
|
|
||||||
0,
|
|
||||||
&bytes_me2c,
|
|
||||||
12_001,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
assert!(middle_blocked.is_err());
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
let middle_ready = timeout(
|
|
||||||
Duration::from_millis(250),
|
|
||||||
process_me_writer_response(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0x94]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(1024),
|
|
||||||
0,
|
|
||||||
&bytes_me2c,
|
|
||||||
12_002,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("middle path must complete after release");
|
|
||||||
|
|
||||||
assert!(middle_ready.is_ok());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn light_fuzz_mixed_payload_sizes_with_periodic_lock_holds_keeps_accounting_consistent() {
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("middle-cross-matrix-fuzz-{}", std::process::id());
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
let mut seed = 0xC0DE_1234_55AA_9988u64;
|
|
||||||
|
|
||||||
for case in 0..96u32 {
|
|
||||||
seed ^= seed << 7;
|
|
||||||
seed ^= seed >> 9;
|
|
||||||
seed ^= seed << 8;
|
|
||||||
|
|
||||||
let hold = (seed & 0x03) == 0;
|
|
||||||
let mut held_lock = None;
|
|
||||||
let maybe_guard = if hold {
|
|
||||||
held_lock = Some(cross_mode_quota_user_lock_for_tests(&user));
|
|
||||||
Some(
|
|
||||||
held_lock
|
|
||||||
.as_ref()
|
|
||||||
.expect("held lock should be present")
|
|
||||||
.try_lock()
|
|
||||||
.expect("cross-mode lock should be acquirable in fuzz round"),
|
|
||||||
)
|
|
||||||
} else {
|
|
||||||
None
|
|
||||||
};
|
|
||||||
|
|
||||||
let payload_len = ((seed >> 8) as usize % 8) + 1;
|
|
||||||
let payload = vec![(seed & 0xff) as u8; payload_len];
|
|
||||||
let before = stats.get_user_total_octets(&user);
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
|
|
||||||
let timed = timeout(
|
|
||||||
Duration::from_millis(20),
|
|
||||||
process_me_writer_response(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from(payload),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(1024),
|
|
||||||
0,
|
|
||||||
&bytes_me2c,
|
|
||||||
13_000 + case as u64,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
if hold {
|
|
||||||
assert!(timed.is_err(), "held-lock fuzz round must block within timeout");
|
|
||||||
assert_eq!(stats.get_user_total_octets(&user), before);
|
|
||||||
} else {
|
|
||||||
let done = timed.expect("unheld fuzz round must complete in time");
|
|
||||||
assert!(done.is_ok());
|
|
||||||
}
|
|
||||||
|
|
||||||
drop(maybe_guard);
|
|
||||||
drop(held_lock);
|
|
||||||
assert_eq!(bytes_me2c.load(Ordering::Relaxed), stats.get_user_total_octets(&user));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn stress_held_user_lock_does_not_block_other_users_me_to_client_writes() {
|
|
||||||
let held_user = format!("middle-cross-matrix-stress-held-{}", std::process::id());
|
|
||||||
let free_user = format!("middle-cross-matrix-stress-free-{}", std::process::id());
|
|
||||||
|
|
||||||
let held = cross_mode_quota_user_lock_for_tests(&held_user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold lock for blocked user");
|
|
||||||
|
|
||||||
let mut tasks = Vec::new();
|
|
||||||
for idx in 0..64u64 {
|
|
||||||
let user = free_user.clone();
|
|
||||||
tasks.push(tokio::spawn(async move {
|
|
||||||
let stats = Stats::new();
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
process_me_writer_response(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0xA0]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(1),
|
|
||||||
0,
|
|
||||||
&bytes_me2c,
|
|
||||||
14_000 + idx,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
timeout(Duration::from_secs(2), async {
|
|
||||||
for task in tasks {
|
|
||||||
let done = task.await.expect("free-user task must not panic");
|
|
||||||
assert!(done.is_ok());
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("free-user tasks should complete without waiting for held user's lock");
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
}
|
|
||||||
|
|
@ -1,254 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::crypto::{AesCtr, SecureRandom};
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use crate::stream::CryptoWriter;
|
|
||||||
use bytes::Bytes;
|
|
||||||
use std::pin::Pin;
|
|
||||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
|
||||||
use std::sync::{Arc, Mutex};
|
|
||||||
use std::task::{Context, Poll, Waker};
|
|
||||||
use tokio::io::AsyncWrite;
|
|
||||||
use tokio::sync::Notify;
|
|
||||||
use tokio::time::{Duration, timeout};
|
|
||||||
|
|
||||||
fn make_crypto_writer<W>(writer: W) -> CryptoWriter<W>
|
|
||||||
where
|
|
||||||
W: tokio::io::AsyncWrite + Unpin,
|
|
||||||
{
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
CryptoWriter::new(writer, AesCtr::new(&key, iv), 8 * 1024)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Default)]
|
|
||||||
struct BlockingWriteState {
|
|
||||||
write_entered: AtomicBool,
|
|
||||||
released: AtomicBool,
|
|
||||||
write_waker: Mutex<Option<Waker>>,
|
|
||||||
write_entered_notify: Notify,
|
|
||||||
}
|
|
||||||
|
|
||||||
struct BlockingWrite {
|
|
||||||
state: Arc<BlockingWriteState>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl BlockingWrite {
|
|
||||||
fn new(state: Arc<BlockingWriteState>) -> Self {
|
|
||||||
Self { state }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl AsyncWrite for BlockingWrite {
|
|
||||||
fn poll_write(
|
|
||||||
self: Pin<&mut Self>,
|
|
||||||
cx: &mut Context<'_>,
|
|
||||||
buf: &[u8],
|
|
||||||
) -> Poll<std::io::Result<usize>> {
|
|
||||||
self.state.write_entered.store(true, Ordering::Release);
|
|
||||||
self.state.write_entered_notify.notify_waiters();
|
|
||||||
|
|
||||||
if self.state.released.load(Ordering::Acquire) {
|
|
||||||
return Poll::Ready(Ok(buf.len()));
|
|
||||||
}
|
|
||||||
|
|
||||||
if let Ok(mut slot) = self.state.write_waker.lock() {
|
|
||||||
*slot = Some(cx.waker().clone());
|
|
||||||
}
|
|
||||||
Poll::Pending
|
|
||||||
}
|
|
||||||
|
|
||||||
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
|
|
||||||
Poll::Ready(Ok(()))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
|
|
||||||
Poll::Ready(Ok(()))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn wait_until_blocking_write_entered(state: &Arc<BlockingWriteState>) {
|
|
||||||
for _ in 0..8 {
|
|
||||||
if state.write_entered.load(Ordering::Acquire) {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
let _ = timeout(Duration::from_millis(25), state.write_entered_notify.notified()).await;
|
|
||||||
}
|
|
||||||
|
|
||||||
panic!("blocking writer did not enter poll_write in bounded time");
|
|
||||||
}
|
|
||||||
|
|
||||||
fn release_blocking_write(state: &Arc<BlockingWriteState>) {
|
|
||||||
state.released.store(true, Ordering::Release);
|
|
||||||
if let Ok(mut slot) = state.write_waker.lock()
|
|
||||||
&& let Some(waker) = slot.take()
|
|
||||||
{
|
|
||||||
waker.wake();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn adversarial_held_cross_mode_lock_blocks_me_to_client_quota_reservation_path() {
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("middle-me2c-cross-mode-held-{}", std::process::id());
|
|
||||||
let held = cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold shared cross-mode lock before ME->C write path");
|
|
||||||
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
|
|
||||||
let blocked = timeout(
|
|
||||||
Duration::from_millis(25),
|
|
||||||
process_me_writer_response(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0x41]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(1024),
|
|
||||||
0,
|
|
||||||
&bytes_me2c,
|
|
||||||
9901,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
blocked.is_err(),
|
|
||||||
"ME->C quota reservation path must be serialized by held shared cross-mode lock"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
let released = timeout(
|
|
||||||
Duration::from_millis(250),
|
|
||||||
process_me_writer_response(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0x42]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(1024),
|
|
||||||
0,
|
|
||||||
&bytes_me2c,
|
|
||||||
9902,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("ME->C write must complete after cross-mode lock release");
|
|
||||||
|
|
||||||
assert!(released.is_ok());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn business_uncontended_cross_mode_lock_allows_me_to_client_quota_reservation() {
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("middle-me2c-cross-mode-free-{}", std::process::id());
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
|
|
||||||
let outcome = timeout(
|
|
||||||
Duration::from_millis(250),
|
|
||||||
process_me_writer_response(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0x55, 0x66]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(1024),
|
|
||||||
0,
|
|
||||||
&bytes_me2c,
|
|
||||||
9903,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("uncontended ME->C path should not stall");
|
|
||||||
|
|
||||||
assert!(outcome.is_ok());
|
|
||||||
assert_eq!(stats.get_user_total_octets(&user), 2);
|
|
||||||
assert_eq!(bytes_me2c.load(std::sync::atomic::Ordering::Relaxed), 2);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
|
||||||
async fn adversarial_cross_mode_lock_is_released_before_me_to_client_write_await() {
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let user = format!("middle-me2c-lock-drop-before-write-{}", std::process::id());
|
|
||||||
let cross_mode_lock = cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
let bytes_me2c = Arc::new(AtomicU64::new(0));
|
|
||||||
let writer_state = Arc::new(BlockingWriteState::default());
|
|
||||||
|
|
||||||
let worker = {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let user = user.clone();
|
|
||||||
let cross_mode_lock = Arc::clone(&cross_mode_lock);
|
|
||||||
let bytes_me2c = Arc::clone(&bytes_me2c);
|
|
||||||
let writer_state = Arc::clone(&writer_state);
|
|
||||||
tokio::spawn(async move {
|
|
||||||
let mut writer = make_crypto_writer(BlockingWrite::new(writer_state));
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let rng = SecureRandom::new();
|
|
||||||
process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0xDE, 0xAD, 0xBE, 0xEF]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&rng,
|
|
||||||
&mut frame_buf,
|
|
||||||
stats.as_ref(),
|
|
||||||
&user,
|
|
||||||
Some(1024),
|
|
||||||
0,
|
|
||||||
Some(&cross_mode_lock),
|
|
||||||
bytes_me2c.as_ref(),
|
|
||||||
9910,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
})
|
|
||||||
};
|
|
||||||
|
|
||||||
wait_until_blocking_write_entered(&writer_state).await;
|
|
||||||
|
|
||||||
let acquired_guard = timeout(Duration::from_millis(40), cross_mode_lock.lock())
|
|
||||||
.await
|
|
||||||
.expect("cross-mode lock must be free while ME->C write is pending");
|
|
||||||
drop(acquired_guard);
|
|
||||||
|
|
||||||
release_blocking_write(&writer_state);
|
|
||||||
|
|
||||||
let result = timeout(Duration::from_millis(300), worker)
|
|
||||||
.await
|
|
||||||
.expect("ME->C worker timed out after releasing blocking writer")
|
|
||||||
.expect("ME->C worker must not panic");
|
|
||||||
|
|
||||||
assert!(result.is_ok());
|
|
||||||
assert_eq!(stats.get_user_total_octets(&user), 4);
|
|
||||||
assert_eq!(bytes_me2c.load(Ordering::Relaxed), 4);
|
|
||||||
}
|
|
||||||
|
|
@ -1,232 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::crypto::{AesCtr, SecureRandom};
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use crate::stream::CryptoWriter;
|
|
||||||
use bytes::Bytes;
|
|
||||||
use std::io;
|
|
||||||
use std::pin::Pin;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
|
||||||
use std::task::{Context, Poll, Waker};
|
|
||||||
use tokio::io::AsyncWrite;
|
|
||||||
use tokio::time::{Duration, timeout};
|
|
||||||
|
|
||||||
fn make_crypto_writer<W>(writer: W) -> CryptoWriter<W>
|
|
||||||
where
|
|
||||||
W: tokio::io::AsyncWrite + Unpin,
|
|
||||||
{
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
CryptoWriter::new(writer, AesCtr::new(&key, iv), 8 * 1024)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Default)]
|
|
||||||
struct GateState {
|
|
||||||
open: AtomicBool,
|
|
||||||
parked_waker: std::sync::Mutex<Option<Waker>>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl GateState {
|
|
||||||
fn open(&self) {
|
|
||||||
self.open.store(true, Ordering::Relaxed);
|
|
||||||
if let Ok(mut guard) = self.parked_waker.lock()
|
|
||||||
&& let Some(w) = guard.take()
|
|
||||||
{
|
|
||||||
w.wake();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn has_waiter(&self) -> bool {
|
|
||||||
self.parked_waker
|
|
||||||
.lock()
|
|
||||||
.map(|guard| guard.is_some())
|
|
||||||
.unwrap_or(false)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[derive(Default)]
|
|
||||||
struct GateWriter {
|
|
||||||
gate: Arc<GateState>,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl GateWriter {
|
|
||||||
fn new(gate: Arc<GateState>) -> Self {
|
|
||||||
Self { gate }
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
impl AsyncWrite for GateWriter {
|
|
||||||
fn poll_write(
|
|
||||||
self: Pin<&mut Self>,
|
|
||||||
cx: &mut Context<'_>,
|
|
||||||
buf: &[u8],
|
|
||||||
) -> Poll<io::Result<usize>> {
|
|
||||||
if self.gate.open.load(Ordering::Relaxed) {
|
|
||||||
return Poll::Ready(Ok(buf.len()));
|
|
||||||
}
|
|
||||||
|
|
||||||
if let Ok(mut guard) = self.gate.parked_waker.lock() {
|
|
||||||
*guard = Some(cx.waker().clone());
|
|
||||||
}
|
|
||||||
Poll::Pending
|
|
||||||
}
|
|
||||||
|
|
||||||
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
|
||||||
Poll::Ready(Ok(()))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
|
||||||
Poll::Ready(Ok(()))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
struct FailingWriter;
|
|
||||||
|
|
||||||
impl AsyncWrite for FailingWriter {
|
|
||||||
fn poll_write(
|
|
||||||
self: Pin<&mut Self>,
|
|
||||||
_cx: &mut Context<'_>,
|
|
||||||
_buf: &[u8],
|
|
||||||
) -> Poll<io::Result<usize>> {
|
|
||||||
Poll::Ready(Err(io::Error::new(
|
|
||||||
io::ErrorKind::BrokenPipe,
|
|
||||||
"injected writer failure",
|
|
||||||
)))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
|
||||||
Poll::Ready(Ok(()))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
|
|
||||||
Poll::Ready(Ok(()))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
|
||||||
async fn adversarial_same_user_slow_writer_must_not_hol_block_peer_connection() {
|
|
||||||
let stats = Stats::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
let rng = SecureRandom::new();
|
|
||||||
let quota_limit = Some(1024);
|
|
||||||
let user = "hol-quota-user";
|
|
||||||
|
|
||||||
let gate = Arc::new(GateState::default());
|
|
||||||
|
|
||||||
let mut blocked_writer = make_crypto_writer(GateWriter::new(Arc::clone(&gate)));
|
|
||||||
let slow_task = tokio::spawn(async move {
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
process_me_writer_response(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0x10, 0x20, 0x30, 0x40]),
|
|
||||||
},
|
|
||||||
&mut blocked_writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&rng,
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
user,
|
|
||||||
quota_limit,
|
|
||||||
0,
|
|
||||||
&bytes_me2c,
|
|
||||||
7001,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
});
|
|
||||||
|
|
||||||
timeout(Duration::from_millis(100), async {
|
|
||||||
loop {
|
|
||||||
if gate.has_waiter() {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
tokio::task::yield_now().await;
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("first writer must reach backpressure and park");
|
|
||||||
|
|
||||||
let stats_fast = Stats::new();
|
|
||||||
let bytes_fast = AtomicU64::new(0);
|
|
||||||
let rng_fast = SecureRandom::new();
|
|
||||||
let mut fast_writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf_fast = Vec::new();
|
|
||||||
|
|
||||||
timeout(
|
|
||||||
Duration::from_millis(50),
|
|
||||||
process_me_writer_response(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0x41]),
|
|
||||||
},
|
|
||||||
&mut fast_writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&rng_fast,
|
|
||||||
&mut frame_buf_fast,
|
|
||||||
&stats_fast,
|
|
||||||
user,
|
|
||||||
quota_limit,
|
|
||||||
0,
|
|
||||||
&bytes_fast,
|
|
||||||
7002,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("peer connection must not be blocked by same-user stalled write")
|
|
||||||
.expect("fast peer write must succeed");
|
|
||||||
|
|
||||||
gate.open();
|
|
||||||
let slow_result = timeout(Duration::from_secs(1), slow_task)
|
|
||||||
.await
|
|
||||||
.expect("stalled task must complete once gate opens")
|
|
||||||
.expect("stalled task must not panic");
|
|
||||||
assert!(slow_result.is_ok());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn negative_write_failure_rolls_back_pre_accounted_quota_and_forensics_bytes() {
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = "rollback-user";
|
|
||||||
stats.add_user_octets_from(user, 7);
|
|
||||||
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
let rng = SecureRandom::new();
|
|
||||||
let mut writer = make_crypto_writer(FailingWriter);
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
|
|
||||||
let result = process_me_writer_response(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[1, 2, 3, 4]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&rng,
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
user,
|
|
||||||
Some(64),
|
|
||||||
0,
|
|
||||||
&bytes_me2c,
|
|
||||||
7003,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
assert!(matches!(result, Err(ProxyError::Io(_))));
|
|
||||||
assert_eq!(
|
|
||||||
stats.get_user_total_octets(user),
|
|
||||||
7,
|
|
||||||
"failed client write must not overcharge user quota accounting"
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
bytes_me2c.load(Ordering::Relaxed),
|
|
||||||
0,
|
|
||||||
"failed client write must not inflate ME->C forensic byte counter"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
@ -1,372 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::crypto::{AesCtr, SecureRandom};
|
|
||||||
use crate::error::ProxyError;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use crate::stream::CryptoWriter;
|
|
||||||
use bytes::Bytes;
|
|
||||||
use std::sync::atomic::{AtomicU64, Ordering};
|
|
||||||
use std::sync::{Arc, OnceLock, Mutex};
|
|
||||||
use tokio::sync::Mutex as AsyncMutex;
|
|
||||||
use tokio::task::JoinSet;
|
|
||||||
use tokio::time::{Duration, timeout};
|
|
||||||
|
|
||||||
fn make_crypto_writer<W>(writer: W) -> CryptoWriter<W>
|
|
||||||
where
|
|
||||||
W: tokio::io::AsyncWrite + Unpin,
|
|
||||||
{
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
CryptoWriter::new(writer, AesCtr::new(&key, iv), 8 * 1024)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn lookup_test_lock() -> &'static Mutex<()> {
|
|
||||||
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
|
|
||||||
LOCK.get_or_init(|| Mutex::new(()))
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn positive_me2c_quota_counts_bytes_exactly_once() {
|
|
||||||
let _guard = lookup_test_lock().lock().unwrap();
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("quota-middle-ext-positive-{}", std::process::id());
|
|
||||||
let lock = Arc::new(AsyncMutex::new(()));
|
|
||||||
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
|
|
||||||
let result = process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[1, 2, 3, 4, 5]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(64),
|
|
||||||
0,
|
|
||||||
Some(&lock),
|
|
||||||
&bytes_me2c,
|
|
||||||
70_001,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
assert!(result.is_ok());
|
|
||||||
assert_eq!(stats.get_user_total_octets(&user), 5);
|
|
||||||
assert_eq!(bytes_me2c.load(Ordering::Relaxed), 5);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn negative_held_crossmode_lock_blocks_me2c_write() {
|
|
||||||
let _guard = lookup_test_lock().lock().unwrap();
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("quota-middle-ext-negative-{}", std::process::id());
|
|
||||||
|
|
||||||
let lock = Arc::new(AsyncMutex::new(()));
|
|
||||||
let _held = lock.try_lock().expect("lock must be held");
|
|
||||||
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
|
|
||||||
let blocked = timeout(
|
|
||||||
Duration::from_millis(25),
|
|
||||||
process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0xFE]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(16),
|
|
||||||
0,
|
|
||||||
Some(&lock),
|
|
||||||
&bytes_me2c,
|
|
||||||
70_101,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
assert!(blocked.is_err());
|
|
||||||
assert_eq!(stats.get_user_total_octets(&user), 0);
|
|
||||||
assert_eq!(bytes_me2c.load(Ordering::Relaxed), 0);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn edge_zero_quota_zero_payload_is_fail_closed() {
|
|
||||||
let _guard = lookup_test_lock().lock().unwrap();
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("quota-middle-ext-edge-{}", std::process::id());
|
|
||||||
|
|
||||||
let lock = Arc::new(AsyncMutex::new(()));
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
|
|
||||||
let result = process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::new(),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(0),
|
|
||||||
0,
|
|
||||||
Some(&lock),
|
|
||||||
&bytes_me2c,
|
|
||||||
70_201,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
assert!(matches!(result, Err(ProxyError::DataQuotaExceeded { .. })));
|
|
||||||
assert_eq!(stats.get_user_total_octets(&user), 0);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn adversarial_parallel_me2c_race_falls_back_to_quota_error() {
|
|
||||||
let _guard = lookup_test_lock().lock().unwrap();
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let user = format!("quota-middle-ext-blackhat-{}", std::process::id());
|
|
||||||
let quota = 64u64;
|
|
||||||
let lock = Arc::new(AsyncMutex::new(()));
|
|
||||||
let bytes_me2c = Arc::new(AtomicU64::new(0));
|
|
||||||
|
|
||||||
let mut set = JoinSet::new();
|
|
||||||
for i in 0..256u64 {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let user = user.clone();
|
|
||||||
let lock = Arc::clone(&lock);
|
|
||||||
let bytes_me2c = Arc::clone(&bytes_me2c);
|
|
||||||
|
|
||||||
set.spawn(async move {
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let payload = vec![((i & 0xFF) as u8); (i % 4 + 1) as usize];
|
|
||||||
|
|
||||||
process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from(payload),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
stats.as_ref(),
|
|
||||||
&user,
|
|
||||||
Some(quota),
|
|
||||||
0,
|
|
||||||
Some(&lock),
|
|
||||||
bytes_me2c.as_ref(),
|
|
||||||
70_301 + i,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut succeeded = 0usize;
|
|
||||||
while let Some(done) = set.join_next().await {
|
|
||||||
match done.expect("task must not panic") {
|
|
||||||
Ok(_) => succeeded += 1,
|
|
||||||
Err(ProxyError::DataQuotaExceeded { .. }) => {}
|
|
||||||
Err(other) => panic!("unexpected error {other:?}"),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
assert_eq!(stats.get_user_total_octets(&user), bytes_me2c.load(Ordering::Relaxed));
|
|
||||||
assert!(stats.get_user_total_octets(&user) <= quota);
|
|
||||||
assert!(succeeded <= quota as usize);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn integration_shared_prefetched_lock_blocks_then_releases_writer() {
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("quota-middle-ext-integration-{}", std::process::id());
|
|
||||||
let lock = Arc::new(AsyncMutex::new(()));
|
|
||||||
let held = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("integration test must hold prefetched lock first");
|
|
||||||
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
|
|
||||||
let blocked = timeout(
|
|
||||||
Duration::from_millis(25),
|
|
||||||
process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0xA1]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(8),
|
|
||||||
0,
|
|
||||||
Some(&lock),
|
|
||||||
&bytes_me2c,
|
|
||||||
70_360,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
assert!(blocked.is_err());
|
|
||||||
|
|
||||||
drop(held);
|
|
||||||
|
|
||||||
let after_release = timeout(
|
|
||||||
Duration::from_millis(150),
|
|
||||||
process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0xA2]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(8),
|
|
||||||
0,
|
|
||||||
Some(&lock),
|
|
||||||
&bytes_me2c,
|
|
||||||
70_361,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("writer should progress once the shared lock is released");
|
|
||||||
|
|
||||||
assert!(after_release.is_ok());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn light_fuzz_small_payloads_toggle_lock_state_stays_consistent() {
|
|
||||||
let _guard = lookup_test_lock().lock().unwrap();
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("quota-middle-ext-fuzz-{}", std::process::id());
|
|
||||||
let mut seed = 0xCAFE_BABE_1234u64;
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
|
|
||||||
for case in 0..48u32 {
|
|
||||||
seed ^= seed << 5;
|
|
||||||
seed ^= seed >> 12;
|
|
||||||
seed ^= seed << 13;
|
|
||||||
let hold = (seed & 0x1) == 0;
|
|
||||||
|
|
||||||
let lock = Arc::new(AsyncMutex::new(()));
|
|
||||||
let maybe_guard = if hold {
|
|
||||||
Some(lock.try_lock().unwrap())
|
|
||||||
} else {
|
|
||||||
None
|
|
||||||
};
|
|
||||||
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
|
|
||||||
let result = timeout(
|
|
||||||
Duration::from_millis(30),
|
|
||||||
process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from(vec![(seed & 0xFF) as u8; ((seed as usize % 5) + 1)]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(128),
|
|
||||||
0,
|
|
||||||
Some(&lock),
|
|
||||||
&bytes_me2c,
|
|
||||||
70_401 + case as u64,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
if hold {
|
|
||||||
assert!(result.is_err());
|
|
||||||
} else {
|
|
||||||
assert!(result.unwrap().is_ok());
|
|
||||||
}
|
|
||||||
|
|
||||||
drop(maybe_guard);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn stress_parallel_free_users_during_held_user_lock_maintains_liveness() {
|
|
||||||
let _guard = lookup_test_lock().lock().unwrap();
|
|
||||||
let held = Arc::new(AsyncMutex::new(()));
|
|
||||||
let _held_guard = held.try_lock().unwrap();
|
|
||||||
|
|
||||||
let mut set = JoinSet::new();
|
|
||||||
for i in 0..48u64 {
|
|
||||||
set.spawn(async move {
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("quota-middle-ext-stress-free-{i}");
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
let free_lock = Arc::new(AsyncMutex::new(()));
|
|
||||||
|
|
||||||
process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0xEE]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(1),
|
|
||||||
0,
|
|
||||||
Some(&free_lock),
|
|
||||||
&bytes_me2c,
|
|
||||||
70_500 + i,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
timeout(Duration::from_secs(2), async {
|
|
||||||
while let Some(task) = set.join_next().await {
|
|
||||||
task.unwrap().unwrap();
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.unwrap();
|
|
||||||
}
|
|
||||||
|
|
@ -1,131 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use dashmap::DashMap;
|
|
||||||
use std::sync::Arc;
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn saturation_uses_stable_overflow_lock_without_cache_growth() {
|
|
||||||
let _guard = super::quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let prefix = format!("middle-quota-held-{}", std::process::id());
|
|
||||||
let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX);
|
|
||||||
for idx in 0..QUOTA_USER_LOCKS_MAX {
|
|
||||||
retained.push(quota_user_lock(&format!("{prefix}-{idx}")));
|
|
||||||
}
|
|
||||||
|
|
||||||
assert_eq!(map.len(), QUOTA_USER_LOCKS_MAX);
|
|
||||||
|
|
||||||
let user = format!("middle-quota-overflow-{}", std::process::id());
|
|
||||||
let first = quota_user_lock(&user);
|
|
||||||
let second = quota_user_lock(&user);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
Arc::ptr_eq(&first, &second),
|
|
||||||
"overflow user must get deterministic same lock while cache is saturated"
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
map.len(),
|
|
||||||
QUOTA_USER_LOCKS_MAX,
|
|
||||||
"overflow path must not grow bounded lock map"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
map.get(&user).is_none(),
|
|
||||||
"overflow user should stay outside bounded lock map under saturation"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(retained);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn overflow_striping_keeps_different_users_distributed() {
|
|
||||||
let _guard = super::quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let prefix = format!("middle-quota-dist-held-{}", std::process::id());
|
|
||||||
let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX);
|
|
||||||
for idx in 0..QUOTA_USER_LOCKS_MAX {
|
|
||||||
retained.push(quota_user_lock(&format!("{prefix}-{idx}")));
|
|
||||||
}
|
|
||||||
|
|
||||||
let a = quota_user_lock("middle-overflow-user-a");
|
|
||||||
let b = quota_user_lock("middle-overflow-user-b");
|
|
||||||
let c = quota_user_lock("middle-overflow-user-c");
|
|
||||||
|
|
||||||
let distinct = [
|
|
||||||
Arc::as_ptr(&a) as usize,
|
|
||||||
Arc::as_ptr(&b) as usize,
|
|
||||||
Arc::as_ptr(&c) as usize,
|
|
||||||
]
|
|
||||||
.iter()
|
|
||||||
.copied()
|
|
||||||
.collect::<std::collections::HashSet<_>>()
|
|
||||||
.len();
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
distinct >= 2,
|
|
||||||
"striped overflow lock set should avoid collapsing all users to one lock"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(retained);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn reclaim_path_caches_new_user_after_stale_entries_drop() {
|
|
||||||
let _guard = super::quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let prefix = format!("middle-quota-reclaim-held-{}", std::process::id());
|
|
||||||
let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX);
|
|
||||||
for idx in 0..QUOTA_USER_LOCKS_MAX {
|
|
||||||
retained.push(quota_user_lock(&format!("{prefix}-{idx}")));
|
|
||||||
}
|
|
||||||
|
|
||||||
drop(retained);
|
|
||||||
|
|
||||||
let user = format!("middle-quota-reclaim-user-{}", std::process::id());
|
|
||||||
let got = quota_user_lock(&user);
|
|
||||||
assert!(map.get(&user).is_some());
|
|
||||||
assert!(
|
|
||||||
Arc::strong_count(&got) >= 2,
|
|
||||||
"after reclaim, lock should be held both by caller and map"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn overflow_path_same_user_is_stable_across_parallel_threads() {
|
|
||||||
let _guard = super::quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX);
|
|
||||||
for idx in 0..QUOTA_USER_LOCKS_MAX {
|
|
||||||
retained.push(quota_user_lock(&format!(
|
|
||||||
"middle-quota-thread-held-{}-{idx}",
|
|
||||||
std::process::id()
|
|
||||||
)));
|
|
||||||
}
|
|
||||||
|
|
||||||
let user = format!("middle-quota-overflow-thread-user-{}", std::process::id());
|
|
||||||
let mut workers = Vec::new();
|
|
||||||
for _ in 0..32 {
|
|
||||||
let user = user.clone();
|
|
||||||
workers.push(std::thread::spawn(move || quota_user_lock(&user)));
|
|
||||||
}
|
|
||||||
|
|
||||||
let first = workers
|
|
||||||
.remove(0)
|
|
||||||
.join()
|
|
||||||
.expect("thread must return lock handle");
|
|
||||||
for worker in workers {
|
|
||||||
let got = worker.join().expect("thread must return lock handle");
|
|
||||||
assert!(
|
|
||||||
Arc::ptr_eq(&first, &got),
|
|
||||||
"same overflow user should resolve to one striped lock even under contention"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
drop(retained);
|
|
||||||
}
|
|
||||||
File diff suppressed because it is too large
Load Diff
|
|
@ -1,399 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::crypto::{AesCtr, SecureRandom};
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use crate::stream::CryptoWriter;
|
|
||||||
use bytes::Bytes;
|
|
||||||
use std::sync::atomic::{AtomicU64, Ordering};
|
|
||||||
use std::sync::{Arc, Mutex, OnceLock};
|
|
||||||
use tokio::sync::Mutex as AsyncMutex;
|
|
||||||
use tokio::task::JoinSet;
|
|
||||||
use tokio::time::{Duration, timeout};
|
|
||||||
|
|
||||||
fn make_crypto_writer<W>(writer: W) -> CryptoWriter<W>
|
|
||||||
where
|
|
||||||
W: tokio::io::AsyncWrite + Unpin,
|
|
||||||
{
|
|
||||||
let key = [0u8; 32];
|
|
||||||
let iv = 0u128;
|
|
||||||
CryptoWriter::new(writer, AesCtr::new(&key, iv), 8 * 1024)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn lookup_counter_test_lock() -> &'static Mutex<()> {
|
|
||||||
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
|
|
||||||
LOCK.get_or_init(|| Mutex::new(()))
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn positive_prefetched_cross_mode_lock_multi_frame_accounting_is_exact() {
|
|
||||||
let _guard = lookup_counter_test_lock()
|
|
||||||
.lock()
|
|
||||||
.unwrap_or_else(|poison| poison.into_inner());
|
|
||||||
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("quota-extreme-positive-{}", std::process::id());
|
|
||||||
let lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
crate::proxy::quota_lock_registry::reset_cross_mode_quota_user_lock_lookup_count_for_tests();
|
|
||||||
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
|
|
||||||
for idx in 0..12u64 {
|
|
||||||
let payload = vec![0x5A; ((idx % 4) + 1) as usize];
|
|
||||||
let result = process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from(payload),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(512),
|
|
||||||
0,
|
|
||||||
Some(&lock),
|
|
||||||
&bytes_me2c,
|
|
||||||
31_000 + idx,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
assert!(result.is_ok());
|
|
||||||
}
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
crate::proxy::quota_lock_registry::cross_mode_quota_user_lock_lookup_count_for_user_for_tests(&user),
|
|
||||||
0,
|
|
||||||
"prefetched lock path must avoid hot-path registry lookups"
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
stats.get_user_total_octets(&user),
|
|
||||||
bytes_me2c.load(Ordering::Relaxed),
|
|
||||||
"forensics and quota accounting must remain synchronized"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn negative_held_prefetched_lock_blocks_writer_without_accounting_mutation() {
|
|
||||||
let _guard = lookup_counter_test_lock()
|
|
||||||
.lock()
|
|
||||||
.unwrap_or_else(|poison| poison.into_inner());
|
|
||||||
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("quota-extreme-negative-{}", std::process::id());
|
|
||||||
let lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold lock before calling ME->C writer");
|
|
||||||
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
|
|
||||||
let blocked = timeout(
|
|
||||||
Duration::from_millis(25),
|
|
||||||
process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[1, 2, 3]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(64),
|
|
||||||
0,
|
|
||||||
Some(&lock),
|
|
||||||
&bytes_me2c,
|
|
||||||
31_100,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
assert!(blocked.is_err());
|
|
||||||
assert_eq!(stats.get_user_total_octets(&user), 0);
|
|
||||||
assert_eq!(bytes_me2c.load(Ordering::Relaxed), 0);
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn edge_zero_quota_and_zero_payload_is_fail_closed() {
|
|
||||||
let _guard = lookup_counter_test_lock()
|
|
||||||
.lock()
|
|
||||||
.unwrap_or_else(|poison| poison.into_inner());
|
|
||||||
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("quota-extreme-edge-{}", std::process::id());
|
|
||||||
let lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
|
|
||||||
let result = process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::new(),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(0),
|
|
||||||
0,
|
|
||||||
Some(&lock),
|
|
||||||
&bytes_me2c,
|
|
||||||
31_200,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
assert!(matches!(result, Err(ProxyError::DataQuotaExceeded { .. })));
|
|
||||||
assert_eq!(stats.get_user_total_octets(&user), 0);
|
|
||||||
assert_eq!(bytes_me2c.load(Ordering::Relaxed), 0);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn adversarial_blackhat_parallel_quota_race_never_overshoots_soft_cap() {
|
|
||||||
let _guard = lookup_counter_test_lock()
|
|
||||||
.lock()
|
|
||||||
.unwrap_or_else(|poison| poison.into_inner());
|
|
||||||
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let user = format!("quota-extreme-blackhat-{}", std::process::id());
|
|
||||||
let quota = 80u64;
|
|
||||||
let overshoot = 7u64;
|
|
||||||
let soft_limit = quota + overshoot;
|
|
||||||
let lock = Arc::new(crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user));
|
|
||||||
let bytes_me2c = Arc::new(AtomicU64::new(0));
|
|
||||||
|
|
||||||
let mut set = JoinSet::new();
|
|
||||||
for idx in 0..256u64 {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let user = user.clone();
|
|
||||||
let lock = Arc::clone(&lock);
|
|
||||||
let bytes_me2c = Arc::clone(&bytes_me2c);
|
|
||||||
|
|
||||||
set.spawn(async move {
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let len = ((idx % 5) + 1) as usize;
|
|
||||||
let payload = vec![0xAA; len];
|
|
||||||
|
|
||||||
process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from(payload),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
stats.as_ref(),
|
|
||||||
&user,
|
|
||||||
Some(quota),
|
|
||||||
overshoot,
|
|
||||||
Some(&lock),
|
|
||||||
bytes_me2c.as_ref(),
|
|
||||||
31_300 + idx,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
while let Some(done) = set.join_next().await {
|
|
||||||
match done.expect("task must not panic") {
|
|
||||||
Ok(_) | Err(ProxyError::DataQuotaExceeded { .. }) => {}
|
|
||||||
Err(other) => panic!("unexpected error variant under black-hat race: {other:?}"),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let total = stats.get_user_total_octets(&user);
|
|
||||||
assert!(
|
|
||||||
total <= soft_limit,
|
|
||||||
"parallel adversarial race must stay under soft cap"
|
|
||||||
);
|
|
||||||
assert_eq!(bytes_me2c.load(Ordering::Relaxed), total);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn integration_without_prefetched_lock_uses_registry_lookup_path() {
|
|
||||||
let _guard = lookup_counter_test_lock()
|
|
||||||
.lock()
|
|
||||||
.unwrap_or_else(|poison| poison.into_inner());
|
|
||||||
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("quota-extreme-integration-{}", std::process::id());
|
|
||||||
crate::proxy::quota_lock_registry::reset_cross_mode_quota_user_lock_lookup_count_for_tests();
|
|
||||||
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
|
|
||||||
for idx in 0..3u64 {
|
|
||||||
let result = process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0x41]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(16),
|
|
||||||
0,
|
|
||||||
None,
|
|
||||||
&bytes_me2c,
|
|
||||||
31_400 + idx,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
assert!(result.is_ok());
|
|
||||||
}
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
crate::proxy::quota_lock_registry::cross_mode_quota_user_lock_lookup_count_for_user_for_tests(&user),
|
|
||||||
3,
|
|
||||||
"control path should perform one lock-registry lookup per call"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn light_fuzz_quota_matrix_preserves_fail_closed_accounting() {
|
|
||||||
let _guard = lookup_counter_test_lock()
|
|
||||||
.lock()
|
|
||||||
.unwrap_or_else(|poison| poison.into_inner());
|
|
||||||
|
|
||||||
let stats = Stats::new();
|
|
||||||
let user = format!("quota-extreme-fuzz-{}", std::process::id());
|
|
||||||
let lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
let bytes_me2c = AtomicU64::new(0);
|
|
||||||
let mut seed = 0xA11C_55EE_2026_0323u64;
|
|
||||||
|
|
||||||
for idx in 0..512u64 {
|
|
||||||
seed ^= seed << 7;
|
|
||||||
seed ^= seed >> 9;
|
|
||||||
seed ^= seed << 8;
|
|
||||||
|
|
||||||
let quota = 24 + (seed & 0x3f);
|
|
||||||
let overshoot = (seed >> 13) & 0x0f;
|
|
||||||
let len = ((seed >> 19) & 0x07) + 1;
|
|
||||||
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
let before = stats.get_user_total_octets(&user);
|
|
||||||
|
|
||||||
let result = process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from(vec![0x11; len as usize]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
&stats,
|
|
||||||
&user,
|
|
||||||
Some(quota),
|
|
||||||
overshoot,
|
|
||||||
Some(&lock),
|
|
||||||
&bytes_me2c,
|
|
||||||
31_500 + idx,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
let after = stats.get_user_total_octets(&user);
|
|
||||||
if result.is_ok() {
|
|
||||||
assert!(after >= before);
|
|
||||||
} else {
|
|
||||||
assert!(matches!(result, Err(ProxyError::DataQuotaExceeded { .. })));
|
|
||||||
assert_eq!(after, before);
|
|
||||||
}
|
|
||||||
assert_eq!(bytes_me2c.load(Ordering::Relaxed), after);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn stress_prefetched_lock_high_fanout_exact_quota_success_count() {
|
|
||||||
let _guard = lookup_counter_test_lock()
|
|
||||||
.lock()
|
|
||||||
.unwrap_or_else(|poison| poison.into_inner());
|
|
||||||
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let user = format!("quota-extreme-stress-{}", std::process::id());
|
|
||||||
let quota = 96u64;
|
|
||||||
let lock: Arc<AsyncMutex<()>> = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
let bytes_me2c = Arc::new(AtomicU64::new(0));
|
|
||||||
|
|
||||||
crate::proxy::quota_lock_registry::reset_cross_mode_quota_user_lock_lookup_count_for_tests();
|
|
||||||
|
|
||||||
let mut set = JoinSet::new();
|
|
||||||
for idx in 0..384u64 {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let user = user.clone();
|
|
||||||
let lock = Arc::clone(&lock);
|
|
||||||
let bytes_me2c = Arc::clone(&bytes_me2c);
|
|
||||||
|
|
||||||
set.spawn(async move {
|
|
||||||
let mut writer = make_crypto_writer(tokio::io::sink());
|
|
||||||
let mut frame_buf = Vec::new();
|
|
||||||
process_me_writer_response_with_cross_mode_lock(
|
|
||||||
MeResponse::Data {
|
|
||||||
flags: 0,
|
|
||||||
data: Bytes::from_static(&[0xFF]),
|
|
||||||
},
|
|
||||||
&mut writer,
|
|
||||||
ProtoTag::Intermediate,
|
|
||||||
&SecureRandom::new(),
|
|
||||||
&mut frame_buf,
|
|
||||||
stats.as_ref(),
|
|
||||||
&user,
|
|
||||||
Some(quota),
|
|
||||||
0,
|
|
||||||
Some(&lock),
|
|
||||||
bytes_me2c.as_ref(),
|
|
||||||
31_600 + idx,
|
|
||||||
false,
|
|
||||||
false,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut success = 0usize;
|
|
||||||
while let Some(done) = set.join_next().await {
|
|
||||||
match done.expect("task must not panic") {
|
|
||||||
Ok(_) => success += 1,
|
|
||||||
Err(ProxyError::DataQuotaExceeded { .. }) => {}
|
|
||||||
Err(other) => panic!("unexpected error variant in stress fanout: {other:?}"),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
assert_eq!(success, quota as usize);
|
|
||||||
assert_eq!(stats.get_user_total_octets(&user), quota);
|
|
||||||
assert_eq!(bytes_me2c.load(Ordering::Relaxed), quota);
|
|
||||||
assert_eq!(
|
|
||||||
crate::proxy::quota_lock_registry::cross_mode_quota_user_lock_lookup_count_for_user_for_tests(&user),
|
|
||||||
0,
|
|
||||||
"stress prefetched path must not use lock registry lookups"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
File diff suppressed because it is too large
Load Diff
|
|
@ -1,108 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::{Mutex, OnceLock};
|
|
||||||
|
|
||||||
fn cross_mode_lock_test_guard() -> std::sync::MutexGuard<'static, ()> {
|
|
||||||
static TEST_LOCK: OnceLock<Mutex<()>> = OnceLock::new();
|
|
||||||
TEST_LOCK
|
|
||||||
.get_or_init(|| Mutex::new(()))
|
|
||||||
.lock()
|
|
||||||
.unwrap_or_else(|poisoned| poisoned.into_inner())
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn same_user_returns_same_lock_identity() {
|
|
||||||
let _guard = cross_mode_lock_test_guard();
|
|
||||||
let locks = CROSS_MODE_QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
locks.clear();
|
|
||||||
|
|
||||||
let a = cross_mode_quota_user_lock("cross-mode-same-user");
|
|
||||||
let b = cross_mode_quota_user_lock("cross-mode-same-user");
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
Arc::ptr_eq(&a, &b),
|
|
||||||
"same user must reuse a stable lock identity"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn saturation_overflow_path_returns_stable_striped_lock_without_cache_growth() {
|
|
||||||
let _guard = cross_mode_lock_test_guard();
|
|
||||||
let locks = CROSS_MODE_QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
locks.clear();
|
|
||||||
|
|
||||||
let prefix = format!("cross-mode-saturated-{}", std::process::id());
|
|
||||||
let mut retained = Vec::with_capacity(CROSS_MODE_QUOTA_USER_LOCKS_MAX);
|
|
||||||
for idx in 0..CROSS_MODE_QUOTA_USER_LOCKS_MAX {
|
|
||||||
retained.push(cross_mode_quota_user_lock(&format!("{prefix}-{idx}")));
|
|
||||||
}
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
locks.len(),
|
|
||||||
CROSS_MODE_QUOTA_USER_LOCKS_MAX,
|
|
||||||
"lock cache must be saturated for overflow check"
|
|
||||||
);
|
|
||||||
|
|
||||||
let overflow_user = format!("cross-mode-overflow-{}", std::process::id());
|
|
||||||
let overflow_a = cross_mode_quota_user_lock(&overflow_user);
|
|
||||||
let overflow_b = cross_mode_quota_user_lock(&overflow_user);
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
locks.len(),
|
|
||||||
CROSS_MODE_QUOTA_USER_LOCKS_MAX,
|
|
||||||
"overflow path must not grow bounded lock cache"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
locks.get(&overflow_user).is_none(),
|
|
||||||
"overflow user must stay on striped fallback while cache is saturated"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
Arc::ptr_eq(&overflow_a, &overflow_b),
|
|
||||||
"overflow user must receive a stable striped lock across repeated lookups"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(retained);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn reclaim_drops_stale_entries_but_preserves_active_user_lock_identity() {
|
|
||||||
let _guard = cross_mode_lock_test_guard();
|
|
||||||
let locks = CROSS_MODE_QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
locks.clear();
|
|
||||||
|
|
||||||
let prefix = format!("cross-mode-reclaim-{}", std::process::id());
|
|
||||||
let protected_user = format!("{prefix}-protected");
|
|
||||||
|
|
||||||
let protected_lock = cross_mode_quota_user_lock(&protected_user);
|
|
||||||
let mut retained = Vec::with_capacity(CROSS_MODE_QUOTA_USER_LOCKS_MAX.saturating_sub(1));
|
|
||||||
for idx in 0..(CROSS_MODE_QUOTA_USER_LOCKS_MAX.saturating_sub(1)) {
|
|
||||||
retained.push(cross_mode_quota_user_lock(&format!("{prefix}-{idx}")));
|
|
||||||
}
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
locks.len(),
|
|
||||||
CROSS_MODE_QUOTA_USER_LOCKS_MAX,
|
|
||||||
"fixture must saturate lock cache before reclaim path is exercised"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(retained);
|
|
||||||
|
|
||||||
let newcomer_user = format!("{prefix}-newcomer");
|
|
||||||
let _newcomer = cross_mode_quota_user_lock(&newcomer_user);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
locks.get(&protected_user).is_some(),
|
|
||||||
"active protected user must remain cache-resident after reclaim"
|
|
||||||
);
|
|
||||||
let locked = locks
|
|
||||||
.get(&protected_user)
|
|
||||||
.expect("protected user must remain in map after reclaim");
|
|
||||||
assert!(
|
|
||||||
Arc::ptr_eq(locked.value(), &protected_lock),
|
|
||||||
"reclaim must not swap active user lock identity"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
locks.get(&newcomer_user).is_some(),
|
|
||||||
"newcomer should become cacheable after stale entries are reclaimed"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
@ -1,267 +0,0 @@
|
||||||
use super::relay_bidirectional;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use crate::stream::BufferPool;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use tokio::io::{AsyncReadExt, AsyncWriteExt, duplex};
|
|
||||||
use tokio::time::{Duration, timeout};
|
|
||||||
|
|
||||||
fn quota_test_guard() -> impl Drop {
|
|
||||||
super::quota_user_lock_test_scope()
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn negative_same_user_pipeline_stalls_while_middle_lock_is_held() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("relay-pipeline-stall-{}", std::process::id());
|
|
||||||
let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold shared cross-mode lock");
|
|
||||||
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let (mut client_peer, relay_client) = duplex(1024);
|
|
||||||
let (relay_server, mut server_peer) = duplex(1024);
|
|
||||||
let (client_reader, client_writer) = tokio::io::split(relay_client);
|
|
||||||
let (server_reader, server_writer) = tokio::io::split(relay_server);
|
|
||||||
|
|
||||||
let relay_user = user.clone();
|
|
||||||
let relay_stats = Arc::clone(&stats);
|
|
||||||
let relay_task = tokio::spawn(async move {
|
|
||||||
relay_bidirectional(
|
|
||||||
client_reader,
|
|
||||||
client_writer,
|
|
||||||
server_reader,
|
|
||||||
server_writer,
|
|
||||||
256,
|
|
||||||
256,
|
|
||||||
&relay_user,
|
|
||||||
relay_stats,
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(BufferPool::new()),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
});
|
|
||||||
|
|
||||||
server_peer
|
|
||||||
.write_all(&[0xA1])
|
|
||||||
.await
|
|
||||||
.expect("server write should enqueue while relay is stalled");
|
|
||||||
|
|
||||||
let mut one = [0u8; 1];
|
|
||||||
let blocked_read = timeout(Duration::from_millis(40), client_peer.read_exact(&mut one)).await;
|
|
||||||
assert!(
|
|
||||||
blocked_read.is_err(),
|
|
||||||
"same-user relay must remain blocked while cross-mode lock is held"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
timeout(Duration::from_millis(400), client_peer.read_exact(&mut one))
|
|
||||||
.await
|
|
||||||
.expect("blocked relay must resume after cross-mode lock release")
|
|
||||||
.expect("resumed relay must deliver queued byte");
|
|
||||||
assert_eq!(one, [0xA1]);
|
|
||||||
|
|
||||||
drop(client_peer);
|
|
||||||
drop(server_peer);
|
|
||||||
|
|
||||||
let relay_result = timeout(Duration::from_secs(1), relay_task)
|
|
||||||
.await
|
|
||||||
.expect("relay task must complete")
|
|
||||||
.expect("relay task must not panic");
|
|
||||||
assert!(relay_result.is_ok());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn integration_other_user_pipeline_progresses_while_blocked_user_is_stalled() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let blocked_user = format!("relay-pipeline-blocked-{}", std::process::id());
|
|
||||||
let free_user = format!("relay-pipeline-free-{}", std::process::id());
|
|
||||||
let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&blocked_user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold blocked user's shared cross-mode lock");
|
|
||||||
|
|
||||||
let stats_blocked = Arc::new(Stats::new());
|
|
||||||
let stats_free = Arc::new(Stats::new());
|
|
||||||
|
|
||||||
let (mut blocked_client, blocked_relay_client) = duplex(1024);
|
|
||||||
let (blocked_relay_server, mut blocked_server) = duplex(1024);
|
|
||||||
let (blocked_client_reader, blocked_client_writer) = tokio::io::split(blocked_relay_client);
|
|
||||||
let (blocked_server_reader, blocked_server_writer) = tokio::io::split(blocked_relay_server);
|
|
||||||
|
|
||||||
let (mut free_client, free_relay_client) = duplex(1024);
|
|
||||||
let (free_relay_server, mut free_server) = duplex(1024);
|
|
||||||
let (free_client_reader, free_client_writer) = tokio::io::split(free_relay_client);
|
|
||||||
let (free_server_reader, free_server_writer) = tokio::io::split(free_relay_server);
|
|
||||||
|
|
||||||
let blocked_task = {
|
|
||||||
let user = blocked_user.clone();
|
|
||||||
let stats = Arc::clone(&stats_blocked);
|
|
||||||
tokio::spawn(async move {
|
|
||||||
relay_bidirectional(
|
|
||||||
blocked_client_reader,
|
|
||||||
blocked_client_writer,
|
|
||||||
blocked_server_reader,
|
|
||||||
blocked_server_writer,
|
|
||||||
256,
|
|
||||||
256,
|
|
||||||
&user,
|
|
||||||
stats,
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(BufferPool::new()),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
})
|
|
||||||
};
|
|
||||||
|
|
||||||
let free_task = {
|
|
||||||
let user = free_user.clone();
|
|
||||||
let stats = Arc::clone(&stats_free);
|
|
||||||
tokio::spawn(async move {
|
|
||||||
relay_bidirectional(
|
|
||||||
free_client_reader,
|
|
||||||
free_client_writer,
|
|
||||||
free_server_reader,
|
|
||||||
free_server_writer,
|
|
||||||
256,
|
|
||||||
256,
|
|
||||||
&user,
|
|
||||||
stats,
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(BufferPool::new()),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
})
|
|
||||||
};
|
|
||||||
|
|
||||||
blocked_server
|
|
||||||
.write_all(&[0xB1])
|
|
||||||
.await
|
|
||||||
.expect("blocked user server write should queue");
|
|
||||||
free_server
|
|
||||||
.write_all(&[0xC1])
|
|
||||||
.await
|
|
||||||
.expect("free user server write should queue");
|
|
||||||
|
|
||||||
let mut blocked_buf = [0u8; 1];
|
|
||||||
let mut free_buf = [0u8; 1];
|
|
||||||
|
|
||||||
let blocked_stalled = timeout(
|
|
||||||
Duration::from_millis(40),
|
|
||||||
blocked_client.read_exact(&mut blocked_buf),
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
assert!(
|
|
||||||
blocked_stalled.is_err(),
|
|
||||||
"blocked user must remain stalled while its lock is held"
|
|
||||||
);
|
|
||||||
|
|
||||||
timeout(Duration::from_millis(250), free_client.read_exact(&mut free_buf))
|
|
||||||
.await
|
|
||||||
.expect("free user must make progress while other user is blocked")
|
|
||||||
.expect("free user read must succeed");
|
|
||||||
assert_eq!(free_buf, [0xC1]);
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
timeout(Duration::from_millis(400), blocked_client.read_exact(&mut blocked_buf))
|
|
||||||
.await
|
|
||||||
.expect("blocked user must resume after release")
|
|
||||||
.expect("blocked user resumed read must succeed");
|
|
||||||
assert_eq!(blocked_buf, [0xB1]);
|
|
||||||
|
|
||||||
drop(blocked_client);
|
|
||||||
drop(blocked_server);
|
|
||||||
drop(free_client);
|
|
||||||
drop(free_server);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
timeout(Duration::from_secs(1), blocked_task)
|
|
||||||
.await
|
|
||||||
.expect("blocked relay task must complete")
|
|
||||||
.expect("blocked relay task must not panic")
|
|
||||||
.is_ok()
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
timeout(Duration::from_secs(1), free_task)
|
|
||||||
.await
|
|
||||||
.expect("free relay task must complete")
|
|
||||||
.expect("free relay task must not panic")
|
|
||||||
.is_ok()
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn stress_jittered_hold_release_cycles_preserve_pipeline_liveness() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let mut seed = 0x5EED_C0DE_2026_0323u64;
|
|
||||||
for round in 0..24u32 {
|
|
||||||
seed ^= seed << 7;
|
|
||||||
seed ^= seed >> 9;
|
|
||||||
seed ^= seed << 8;
|
|
||||||
|
|
||||||
let hold_ms = 2 + (seed % 10);
|
|
||||||
let user = format!("relay-pipeline-fuzz-{}-{round}", std::process::id());
|
|
||||||
let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold lock during fuzz round");
|
|
||||||
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let (mut client_peer, relay_client) = duplex(1024);
|
|
||||||
let (relay_server, mut server_peer) = duplex(1024);
|
|
||||||
let (client_reader, client_writer) = tokio::io::split(relay_client);
|
|
||||||
let (server_reader, server_writer) = tokio::io::split(relay_server);
|
|
||||||
|
|
||||||
let relay_user = user.clone();
|
|
||||||
let relay_stats = Arc::clone(&stats);
|
|
||||||
let relay_task = tokio::spawn(async move {
|
|
||||||
relay_bidirectional(
|
|
||||||
client_reader,
|
|
||||||
client_writer,
|
|
||||||
server_reader,
|
|
||||||
server_writer,
|
|
||||||
256,
|
|
||||||
256,
|
|
||||||
&relay_user,
|
|
||||||
relay_stats,
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(BufferPool::new()),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
});
|
|
||||||
|
|
||||||
server_peer
|
|
||||||
.write_all(&[0xD1])
|
|
||||||
.await
|
|
||||||
.expect("server write should queue in fuzz round");
|
|
||||||
|
|
||||||
let mut one = [0u8; 1];
|
|
||||||
let stalled = timeout(Duration::from_millis(30), client_peer.read_exact(&mut one)).await;
|
|
||||||
assert!(stalled.is_err(), "held phase must stall same-user relay");
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(hold_ms)).await;
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
timeout(Duration::from_millis(400), client_peer.read_exact(&mut one))
|
|
||||||
.await
|
|
||||||
.expect("released phase must resume same-user relay")
|
|
||||||
.expect("released phase read must succeed");
|
|
||||||
assert_eq!(one, [0xD1]);
|
|
||||||
|
|
||||||
drop(client_peer);
|
|
||||||
drop(server_peer);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
timeout(Duration::from_secs(1), relay_task)
|
|
||||||
.await
|
|
||||||
.expect("fuzz relay task must complete")
|
|
||||||
.expect("fuzz relay task must not panic")
|
|
||||||
.is_ok()
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
@ -1,213 +0,0 @@
|
||||||
use super::relay_bidirectional;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use crate::stream::BufferPool;
|
|
||||||
use std::sync::{Arc, Mutex};
|
|
||||||
use tokio::io::{AsyncReadExt, AsyncWriteExt, duplex};
|
|
||||||
use tokio::sync::{Barrier, watch};
|
|
||||||
use tokio::time::{Duration, Instant, timeout};
|
|
||||||
|
|
||||||
fn quota_test_guard() -> impl Drop {
|
|
||||||
super::quota_user_lock_test_scope()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn percentile_index(len: usize, percentile: usize) -> usize {
|
|
||||||
((len * percentile) / 100).min(len.saturating_sub(1))
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn micro_benchmark_pipeline_release_to_delivery_latency_stays_bounded() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let rounds = 64usize;
|
|
||||||
let user = format!("relay-pipeline-latency-single-{}", std::process::id());
|
|
||||||
let mut samples_ms = Vec::with_capacity(rounds);
|
|
||||||
|
|
||||||
for round in 0..rounds {
|
|
||||||
let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold shared cross-mode lock before round");
|
|
||||||
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let (mut client_peer, relay_client) = duplex(1024);
|
|
||||||
let (relay_server, mut server_peer) = duplex(1024);
|
|
||||||
let (client_reader, client_writer) = tokio::io::split(relay_client);
|
|
||||||
let (server_reader, server_writer) = tokio::io::split(relay_server);
|
|
||||||
|
|
||||||
let relay_user = user.clone();
|
|
||||||
let relay_stats = Arc::clone(&stats);
|
|
||||||
let relay_task = tokio::spawn(async move {
|
|
||||||
relay_bidirectional(
|
|
||||||
client_reader,
|
|
||||||
client_writer,
|
|
||||||
server_reader,
|
|
||||||
server_writer,
|
|
||||||
256,
|
|
||||||
256,
|
|
||||||
&relay_user,
|
|
||||||
relay_stats,
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(BufferPool::new()),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
});
|
|
||||||
|
|
||||||
server_peer
|
|
||||||
.write_all(&[(round as u8) ^ 0xA5])
|
|
||||||
.await
|
|
||||||
.expect("server write should queue before release");
|
|
||||||
|
|
||||||
let release_at = Instant::now();
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
let mut one = [0u8; 1];
|
|
||||||
timeout(Duration::from_millis(450), client_peer.read_exact(&mut one))
|
|
||||||
.await
|
|
||||||
.expect("client must receive queued byte after release")
|
|
||||||
.expect("queued byte read must succeed");
|
|
||||||
samples_ms.push(release_at.elapsed().as_millis() as u64);
|
|
||||||
|
|
||||||
drop(client_peer);
|
|
||||||
drop(server_peer);
|
|
||||||
|
|
||||||
let relay_result = timeout(Duration::from_secs(1), relay_task)
|
|
||||||
.await
|
|
||||||
.expect("relay task must complete")
|
|
||||||
.expect("relay task must not panic");
|
|
||||||
assert!(relay_result.is_ok());
|
|
||||||
}
|
|
||||||
|
|
||||||
samples_ms.sort_unstable();
|
|
||||||
let p50_ms = samples_ms[percentile_index(samples_ms.len(), 50)];
|
|
||||||
let p95_ms = samples_ms[percentile_index(samples_ms.len(), 95)];
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
p50_ms <= 45,
|
|
||||||
"single-flow release latency p50 must stay bounded; p50_ms={p50_ms}, samples={samples_ms:?}"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
p95_ms <= 130,
|
|
||||||
"single-flow release latency p95 must stay bounded; p95_ms={p95_ms}, samples={samples_ms:?}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn stress_128_waiter_pipeline_release_latency_p95_stays_bounded() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let waiters = 128usize;
|
|
||||||
let user = format!("relay-pipeline-latency-fanout-{}", std::process::id());
|
|
||||||
let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold shared lock before fanout release benchmark");
|
|
||||||
|
|
||||||
let ready_barrier = Arc::new(Barrier::new(waiters + 1));
|
|
||||||
let release_at = Arc::new(Mutex::new(None::<Instant>));
|
|
||||||
let (release_tx, release_rx) = watch::channel(false);
|
|
||||||
let mut tasks = Vec::with_capacity(waiters);
|
|
||||||
|
|
||||||
for idx in 0..waiters {
|
|
||||||
let user = user.clone();
|
|
||||||
let barrier = Arc::clone(&ready_barrier);
|
|
||||||
let release_at = Arc::clone(&release_at);
|
|
||||||
let mut release_rx = release_rx.clone();
|
|
||||||
|
|
||||||
tasks.push(tokio::spawn(async move {
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let (mut client_peer, relay_client) = duplex(512);
|
|
||||||
let (relay_server, mut server_peer) = duplex(512);
|
|
||||||
let (client_reader, client_writer) = tokio::io::split(relay_client);
|
|
||||||
let (server_reader, server_writer) = tokio::io::split(relay_server);
|
|
||||||
|
|
||||||
let relay_user = user;
|
|
||||||
let relay_stats = Arc::clone(&stats);
|
|
||||||
let relay_task = tokio::spawn(async move {
|
|
||||||
relay_bidirectional(
|
|
||||||
client_reader,
|
|
||||||
client_writer,
|
|
||||||
server_reader,
|
|
||||||
server_writer,
|
|
||||||
256,
|
|
||||||
256,
|
|
||||||
&relay_user,
|
|
||||||
relay_stats,
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(BufferPool::new()),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
});
|
|
||||||
|
|
||||||
server_peer
|
|
||||||
.write_all(&[(idx as u8) ^ 0x5A])
|
|
||||||
.await
|
|
||||||
.expect("fanout server write should queue before release");
|
|
||||||
|
|
||||||
barrier.wait().await;
|
|
||||||
release_rx
|
|
||||||
.changed()
|
|
||||||
.await
|
|
||||||
.expect("release signal should remain available");
|
|
||||||
|
|
||||||
let started = {
|
|
||||||
let guard = release_at.lock().unwrap_or_else(|poison| poison.into_inner());
|
|
||||||
guard.expect("release timestamp must be populated before signal")
|
|
||||||
};
|
|
||||||
|
|
||||||
let mut one = [0u8; 1];
|
|
||||||
timeout(Duration::from_millis(900), client_peer.read_exact(&mut one))
|
|
||||||
.await
|
|
||||||
.expect("fanout waiter must receive queued byte after release")
|
|
||||||
.expect("fanout waiter read must succeed");
|
|
||||||
|
|
||||||
drop(client_peer);
|
|
||||||
drop(server_peer);
|
|
||||||
|
|
||||||
let relay_result = timeout(Duration::from_secs(2), relay_task)
|
|
||||||
.await
|
|
||||||
.expect("fanout relay task must complete")
|
|
||||||
.expect("fanout relay task must not panic");
|
|
||||||
assert!(relay_result.is_ok());
|
|
||||||
|
|
||||||
started.elapsed().as_millis() as u64
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
ready_barrier.wait().await;
|
|
||||||
{
|
|
||||||
let mut guard = release_at.lock().unwrap_or_else(|poison| poison.into_inner());
|
|
||||||
*guard = Some(Instant::now());
|
|
||||||
}
|
|
||||||
drop(held_guard);
|
|
||||||
release_tx
|
|
||||||
.send(true)
|
|
||||||
.expect("release broadcast must succeed");
|
|
||||||
|
|
||||||
let mut samples_ms = Vec::with_capacity(waiters);
|
|
||||||
timeout(Duration::from_secs(8), async {
|
|
||||||
for task in tasks {
|
|
||||||
let elapsed = task.await.expect("fanout waiter must not panic");
|
|
||||||
samples_ms.push(elapsed);
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("fanout benchmark must complete in bounded time");
|
|
||||||
|
|
||||||
samples_ms.sort_unstable();
|
|
||||||
let p50_ms = samples_ms[percentile_index(samples_ms.len(), 50)];
|
|
||||||
let p95_ms = samples_ms[percentile_index(samples_ms.len(), 95)];
|
|
||||||
let max_ms = *samples_ms.last().unwrap_or(&0);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
p50_ms <= 120,
|
|
||||||
"fanout release latency p50 must stay bounded; p50_ms={p50_ms}, p95_ms={p95_ms}, max_ms={max_ms}"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
p95_ms <= 260,
|
|
||||||
"fanout release latency p95 must stay bounded; p50_ms={p50_ms}, p95_ms={p95_ms}, max_ms={max_ms}"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
max_ms <= 700,
|
|
||||||
"fanout release latency max must stay bounded; p50_ms={p50_ms}, p95_ms={p95_ms}, max_ms={max_ms}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
@ -1,604 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use std::pin::Pin;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
|
||||||
use std::task::{Context, Poll, Waker};
|
|
||||||
use tokio::io::{AsyncReadExt, AsyncWriteExt, ReadBuf};
|
|
||||||
use tokio::sync::Barrier;
|
|
||||||
use tokio::time::{Duration, timeout};
|
|
||||||
|
|
||||||
#[derive(Default)]
|
|
||||||
struct WakeCounter {
|
|
||||||
wakes: AtomicUsize,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl std::task::Wake for WakeCounter {
|
|
||||||
fn wake(self: Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn wake_by_ref(self: &Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn quota_test_guard() -> impl Drop {
|
|
||||||
super::quota_user_lock_test_scope()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn build_context() -> (Arc<WakeCounter>, Context<'static>) {
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let leaked_waker: &'static Waker = Box::leak(Box::new(waker));
|
|
||||||
(wake_counter, Context::from_waker(leaked_waker))
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn positive_cross_mode_uncontended_writer_progresses() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
"cross-mode-tdd-uncontended".to_string(),
|
|
||||||
Some(4096),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let result = io.write_all(&[0x11, 0x22]).await;
|
|
||||||
assert!(result.is_ok(), "uncontended writer must progress");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn adversarial_held_cross_mode_lock_blocks_writer_even_if_local_lock_free() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("cross-mode-tdd-held-{}", std::process::id());
|
|
||||||
let held = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
let _held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold cross-mode lock before polling writer");
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(4096),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
|
|
||||||
let poll = Pin::new(&mut io).poll_write(&mut cx, &[0xAA]);
|
|
||||||
assert!(poll.is_pending(), "writer must not bypass held cross-mode lock");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn integration_parallel_waiters_resume_after_cross_mode_release() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("cross-mode-tdd-resume-{}", std::process::id());
|
|
||||||
let held = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold cross-mode lock before launching waiters");
|
|
||||||
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let mut waiters = Vec::new();
|
|
||||||
for _ in 0..16 {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let user = user.clone();
|
|
||||||
waiters.push(tokio::spawn(async move {
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
stats,
|
|
||||||
user,
|
|
||||||
Some(4096),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
io.write_all(&[0x7F]).await
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(5)).await;
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
timeout(Duration::from_secs(1), async {
|
|
||||||
for waiter in waiters {
|
|
||||||
let result = waiter.await.expect("waiter task must not panic");
|
|
||||||
assert!(result.is_ok(), "waiter must complete after cross-mode release");
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("all waiters must complete in bounded time");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn adversarial_cross_mode_contention_wake_budget_stays_bounded() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("cross-mode-tdd-wakes-{}", std::process::id());
|
|
||||||
let held = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
let _held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold cross-mode lock before polling");
|
|
||||||
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let mut ios = Vec::new();
|
|
||||||
let mut counters = Vec::new();
|
|
||||||
for _ in 0..20 {
|
|
||||||
ios.push(StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.clone(),
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
));
|
|
||||||
}
|
|
||||||
|
|
||||||
for io in &mut ios {
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
let poll = Pin::new(io).poll_write(&mut cx, &[0x33]);
|
|
||||||
assert!(poll.is_pending());
|
|
||||||
counters.push(wake_counter);
|
|
||||||
}
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(25)).await;
|
|
||||||
let total_wakes: usize = counters
|
|
||||||
.iter()
|
|
||||||
.map(|counter| counter.wakes.load(Ordering::Relaxed))
|
|
||||||
.sum();
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
total_wakes <= 20 * 4,
|
|
||||||
"cross-mode contention should not create wake storms; wakes={total_wakes}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn light_fuzz_cross_mode_release_timing_preserves_read_write_liveness() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let mut seed = 0xC0DE_BAAD_2026_0322u64;
|
|
||||||
for round in 0..16u32 {
|
|
||||||
seed ^= seed << 7;
|
|
||||||
seed ^= seed >> 9;
|
|
||||||
seed ^= seed << 8;
|
|
||||||
|
|
||||||
let sleep_ms = 2 + (seed as u64 % 8);
|
|
||||||
let user = format!("cross-mode-tdd-fuzz-{}-{round}", std::process::id());
|
|
||||||
let held = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold cross-mode lock in fuzz round");
|
|
||||||
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let user_reader = user.clone();
|
|
||||||
let reader_task = tokio::spawn(async move {
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::empty(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user_reader,
|
|
||||||
Some(4096),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
let mut one = [0u8; 1];
|
|
||||||
io.read(&mut one).await
|
|
||||||
});
|
|
||||||
|
|
||||||
let user_writer = user.clone();
|
|
||||||
let writer_task = tokio::spawn(async move {
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user_writer,
|
|
||||||
Some(4096),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
io.write_all(&[0x44]).await
|
|
||||||
});
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(sleep_ms)).await;
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
let read_done = timeout(Duration::from_millis(350), reader_task)
|
|
||||||
.await
|
|
||||||
.expect("reader task must complete after release")
|
|
||||||
.expect("reader task must not panic");
|
|
||||||
assert!(read_done.is_ok());
|
|
||||||
|
|
||||||
let write_done = timeout(Duration::from_millis(350), writer_task)
|
|
||||||
.await
|
|
||||||
.expect("writer task must complete after release")
|
|
||||||
.expect("writer task must not panic");
|
|
||||||
assert!(write_done.is_ok());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn integration_middle_lock_blocks_relay_reader_for_same_user() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("cross-mode-middle-reader-block-{}", std::process::id());
|
|
||||||
let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
let _held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold middle-relay shared lock");
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::empty(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let (_wake_counter, mut cx) = build_context();
|
|
||||||
let mut one = [0u8; 1];
|
|
||||||
let mut buf = ReadBuf::new(&mut one);
|
|
||||||
let poll = Pin::new(&mut io).poll_read(&mut cx, &mut buf);
|
|
||||||
assert!(poll.is_pending());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
|
||||||
async fn integration_middle_lock_release_unblocks_relay_reader() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("cross-mode-middle-reader-release-{}", std::process::id());
|
|
||||||
let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold middle-relay shared lock");
|
|
||||||
|
|
||||||
let task = tokio::spawn({
|
|
||||||
let user = user.clone();
|
|
||||||
async move {
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::empty(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
let mut one = [0u8; 1];
|
|
||||||
io.read(&mut one).await
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(5)).await;
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
let done = timeout(Duration::from_millis(300), task)
|
|
||||||
.await
|
|
||||||
.expect("reader task must complete after release")
|
|
||||||
.expect("reader task must not panic");
|
|
||||||
assert!(done.is_ok());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn business_different_user_middle_lock_does_not_block_relay_writer() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let held_user = format!("cross-mode-middle-held-{}", std::process::id());
|
|
||||||
let active_user = format!("cross-mode-middle-active-{}", std::process::id());
|
|
||||||
let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&held_user);
|
|
||||||
let _held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold middle-relay lock for other user");
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
active_user,
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let (_wake_counter, mut cx) = build_context();
|
|
||||||
let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x61]);
|
|
||||||
assert!(matches!(poll, Poll::Ready(Ok(1))));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn edge_quota_none_bypasses_cross_mode_lock_even_when_held() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("cross-mode-none-limit-{}", std::process::id());
|
|
||||||
let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
let _held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold lock while quota is disabled");
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
None,
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let (_wake_counter, mut cx) = build_context();
|
|
||||||
let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x62, 0x63]);
|
|
||||||
assert!(matches!(poll, Poll::Ready(Ok(2))));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn edge_quota_exceeded_flag_short_circuits_before_lock_path() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("cross-mode-pre-exceeded-{}", std::process::id());
|
|
||||||
let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
let _held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold shared lock before poll");
|
|
||||||
|
|
||||||
let quota_exceeded = Arc::new(AtomicBool::new(true));
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(1024),
|
|
||||||
Arc::clone("a_exceeded),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let (_wake_counter, mut cx) = build_context();
|
|
||||||
let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x64]);
|
|
||||||
assert!(matches!(poll, Poll::Ready(Err(ref e)) if is_quota_io_error(e)));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn adversarial_repoll_while_middle_lock_held_keeps_pending_without_usage_leak() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("cross-mode-repoll-held-{}", std::process::id());
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
let _held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold lock for repoll sequence");
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.clone(),
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let (_wake_counter, mut cx) = build_context();
|
|
||||||
for _ in 0..8 {
|
|
||||||
let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x65]);
|
|
||||||
assert!(poll.is_pending());
|
|
||||||
}
|
|
||||||
|
|
||||||
assert_eq!(stats.get_user_total_octets(&user), 0);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn integration_same_user_mixed_read_write_waiters_resume_after_release() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("cross-mode-mixed-resume-{}", std::process::id());
|
|
||||||
let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold lock before spawning mixed waiters");
|
|
||||||
|
|
||||||
let mut tasks = Vec::new();
|
|
||||||
for i in 0..12usize {
|
|
||||||
let user = user.clone();
|
|
||||||
tasks.push(tokio::spawn(async move {
|
|
||||||
if i % 2 == 0 {
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::empty(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
let mut b = [0u8; 1];
|
|
||||||
io.read(&mut b).await.map(|_| ())
|
|
||||||
} else {
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
io.write_all(&[0x66]).await
|
|
||||||
}
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(8)).await;
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
timeout(Duration::from_secs(1), async {
|
|
||||||
for task in tasks {
|
|
||||||
let result = task.await.expect("mixed waiter task must not panic");
|
|
||||||
assert!(result.is_ok());
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("all mixed waiters must finish after release");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn integration_one_user_blocked_other_user_progresses_under_middle_lock() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let blocked_user = format!("cross-mode-blocked-{}", std::process::id());
|
|
||||||
let free_user = format!("cross-mode-free-{}", std::process::id());
|
|
||||||
let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&blocked_user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold blocked user lock");
|
|
||||||
|
|
||||||
let blocked_task = tokio::spawn({
|
|
||||||
let blocked_user = blocked_user.clone();
|
|
||||||
async move {
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
blocked_user,
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
io.write_all(&[0x77]).await
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
let free_task = tokio::spawn({
|
|
||||||
let free_user = free_user.clone();
|
|
||||||
async move {
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
free_user,
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
io.write_all(&[0x78]).await
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
let free_done = timeout(Duration::from_millis(250), free_task)
|
|
||||||
.await
|
|
||||||
.expect("free user must not be blocked")
|
|
||||||
.expect("free user task must not panic");
|
|
||||||
assert!(free_done.is_ok());
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
let blocked_done = timeout(Duration::from_secs(1), blocked_task)
|
|
||||||
.await
|
|
||||||
.expect("blocked user must resume after release")
|
|
||||||
.expect("blocked user task must not panic");
|
|
||||||
assert!(blocked_done.is_ok());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn stress_middle_lock_release_allows_high_waiter_fanout_completion() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("cross-mode-fanout-{}", std::process::id());
|
|
||||||
let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold lock before fanout");
|
|
||||||
|
|
||||||
let waiters = 48usize;
|
|
||||||
let gate = Arc::new(Barrier::new(waiters + 1));
|
|
||||||
let mut tasks = Vec::new();
|
|
||||||
for _ in 0..waiters {
|
|
||||||
let user = user.clone();
|
|
||||||
let gate = Arc::clone(&gate);
|
|
||||||
tasks.push(tokio::spawn(async move {
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
gate.wait().await;
|
|
||||||
io.write_all(&[0x79]).await
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
gate.wait().await;
|
|
||||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
timeout(Duration::from_secs(2), async {
|
|
||||||
for task in tasks {
|
|
||||||
let result = task.await.expect("fanout task must not panic");
|
|
||||||
assert!(result.is_ok());
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("fanout waiters must complete after release");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn light_fuzz_middle_lock_hold_release_cycles_preserve_same_user_liveness() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let mut seed = 0xA11C_EE55_2026_0323u64;
|
|
||||||
for round in 0..20u32 {
|
|
||||||
seed ^= seed << 7;
|
|
||||||
seed ^= seed >> 9;
|
|
||||||
seed ^= seed << 8;
|
|
||||||
|
|
||||||
let hold_ms = 2 + (seed % 10);
|
|
||||||
let user = format!("cross-mode-middle-fuzz-{}-{round}", std::process::id());
|
|
||||||
let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold lock in fuzz round");
|
|
||||||
|
|
||||||
let writer = tokio::spawn({
|
|
||||||
let user = user.clone();
|
|
||||||
async move {
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
io.write_all(&[0x7A]).await
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(hold_ms)).await;
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
let done = timeout(Duration::from_millis(400), writer)
|
|
||||||
.await
|
|
||||||
.expect("writer must complete after lock release")
|
|
||||||
.expect("writer task must not panic");
|
|
||||||
assert!(done.is_ok());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
@ -1,81 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use std::pin::Pin;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
|
||||||
use std::task::Waker;
|
|
||||||
use std::task::{Context, Poll};
|
|
||||||
|
|
||||||
#[derive(Default)]
|
|
||||||
struct WakeCounter {
|
|
||||||
wakes: AtomicUsize,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl std::task::Wake for WakeCounter {
|
|
||||||
fn wake(self: Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn wake_by_ref(self: &Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn build_context() -> (Arc<WakeCounter>, Context<'static>) {
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let leaked_waker: &'static Waker = Box::leak(Box::new(waker));
|
|
||||||
(wake_counter, Context::from_waker(leaked_waker))
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn adversarial_middle_held_cross_mode_lock_blocks_relay_writer() {
|
|
||||||
let _guard = quota_user_lock_test_scope();
|
|
||||||
|
|
||||||
let user = "cross-mode-lock-shared-user";
|
|
||||||
let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(user);
|
|
||||||
let _held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold shared cross-mode lock before relay poll");
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(crate::stats::Stats::new()),
|
|
||||||
user.to_string(),
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let (_wake_counter, mut cx) = build_context();
|
|
||||||
let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x41, 0x42, 0x43]);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
matches!(poll, Poll::Pending),
|
|
||||||
"relay writer must not bypass cross-mode lock held by middle-relay path"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn business_cross_mode_lock_uncontended_allows_relay_writer_progress() {
|
|
||||||
let _guard = quota_user_lock_test_scope();
|
|
||||||
|
|
||||||
let user = "cross-mode-lock-progress-user";
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(crate::stats::Stats::new()),
|
|
||||||
user.to_string(),
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let (_wake_counter, mut cx) = build_context();
|
|
||||||
let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x51, 0x52]);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
matches!(poll, Poll::Ready(Ok(2))),
|
|
||||||
"relay writer should progress when shared cross-mode lock is uncontended"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
@ -1,340 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use std::pin::Pin;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
|
||||||
use std::task::{Context, Waker};
|
|
||||||
use tokio::io::AsyncWriteExt;
|
|
||||||
use tokio::time::{Duration, Instant, timeout};
|
|
||||||
|
|
||||||
#[derive(Default)]
|
|
||||||
struct WakeCounter {
|
|
||||||
wakes: AtomicUsize,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl std::task::Wake for WakeCounter {
|
|
||||||
fn wake(self: Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn wake_by_ref(self: &Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn quota_test_guard() -> impl Drop {
|
|
||||||
super::quota_user_lock_test_scope()
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn positive_uncontended_dual_lock_writer_has_zero_retry_attempt() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
format!("dual-lock-alt-positive-{}", std::process::id()),
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let write = io.write_all(&[0xAA, 0xBB]).await;
|
|
||||||
assert!(write.is_ok(), "uncontended write must complete");
|
|
||||||
assert_eq!(
|
|
||||||
io.quota_write_retry_attempt, 0,
|
|
||||||
"uncontended write must not advance retry backoff"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn adversarial_alternating_local_and_cross_mode_contention_preserves_backoff_growth() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("dual-lock-alt-adversarial-{}", std::process::id());
|
|
||||||
let local_lock = quota_user_lock(&user);
|
|
||||||
let cross_mode_lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
|
|
||||||
let mut local_guard = Some(
|
|
||||||
local_lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold local quota lock initially"),
|
|
||||||
);
|
|
||||||
let mut cross_guard = None;
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
|
|
||||||
let first = Pin::new(&mut io).poll_write(&mut cx, &[0x11]);
|
|
||||||
assert!(first.is_pending(), "held local lock must block first poll");
|
|
||||||
|
|
||||||
let mut observed_wakes = 0usize;
|
|
||||||
for idx in 0..18usize {
|
|
||||||
tokio::time::sleep(Duration::from_millis(6)).await;
|
|
||||||
|
|
||||||
if idx % 2 == 0 {
|
|
||||||
drop(local_guard.take());
|
|
||||||
cross_guard = Some(
|
|
||||||
cross_mode_lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("cross-mode lock should be acquirable while local lock released"),
|
|
||||||
);
|
|
||||||
} else {
|
|
||||||
drop(cross_guard.take());
|
|
||||||
local_guard = Some(
|
|
||||||
local_lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("local lock should be acquirable while cross lock released"),
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
let wakes = wake_counter.wakes.load(Ordering::Relaxed);
|
|
||||||
if wakes > observed_wakes {
|
|
||||||
observed_wakes = wakes;
|
|
||||||
let pending = Pin::new(&mut io).poll_write(&mut cx, &[0x12]);
|
|
||||||
assert!(
|
|
||||||
pending.is_pending(),
|
|
||||||
"alternating contention must keep write pending while one lock is held"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
io.quota_write_retry_attempt >= 2,
|
|
||||||
"alternating contention must still ramp retry backoff; got {}",
|
|
||||||
io.quota_write_retry_attempt
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
wake_counter.wakes.load(Ordering::Relaxed) <= 32,
|
|
||||||
"alternating contention must stay wake-rate-limited"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(local_guard);
|
|
||||||
drop(cross_guard);
|
|
||||||
let ready = Pin::new(&mut io).poll_write(&mut cx, &[0x13]);
|
|
||||||
assert!(ready.is_ready(), "writer must resume after both locks released");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn edge_retry_scheduler_resets_after_alternating_contention_clears() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("dual-lock-alt-edge-reset-{}", std::process::id());
|
|
||||||
let local_lock = quota_user_lock(&user);
|
|
||||||
let local_guard = local_lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold local lock for edge scenario");
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
|
|
||||||
let first = Pin::new(&mut io).poll_write(&mut cx, &[0x21]);
|
|
||||||
assert!(first.is_pending());
|
|
||||||
tokio::time::sleep(Duration::from_millis(15)).await;
|
|
||||||
if wake_counter.wakes.load(Ordering::Relaxed) > 0 {
|
|
||||||
let next = Pin::new(&mut io).poll_write(&mut cx, &[0x22]);
|
|
||||||
assert!(next.is_pending());
|
|
||||||
}
|
|
||||||
|
|
||||||
drop(local_guard);
|
|
||||||
|
|
||||||
let ready = Pin::new(&mut io).poll_write(&mut cx, &[0x23]);
|
|
||||||
assert!(ready.is_ready());
|
|
||||||
assert_eq!(
|
|
||||||
io.quota_write_retry_attempt, 0,
|
|
||||||
"successful dual-lock acquisition must reset retry scheduler"
|
|
||||||
);
|
|
||||||
assert!(!io.quota_write_wake_scheduled);
|
|
||||||
assert!(io.quota_write_retry_sleep.is_none());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn integration_cross_mode_waiters_remain_live_under_alternating_contention_then_resume() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("dual-lock-alt-integration-{}", std::process::id());
|
|
||||||
let local_lock = quota_user_lock(&user);
|
|
||||||
let cross_mode_lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
|
|
||||||
let mut waiters = Vec::new();
|
|
||||||
for _ in 0..16usize {
|
|
||||||
let user = user.clone();
|
|
||||||
waiters.push(tokio::spawn(async move {
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
timeout(Duration::from_secs(2), io.write_all(&[0x31])).await
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut local_guard = Some(
|
|
||||||
local_lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("integration toggle must acquire local lock first"),
|
|
||||||
);
|
|
||||||
let mut cross_guard = None;
|
|
||||||
|
|
||||||
for idx in 0..24usize {
|
|
||||||
tokio::time::sleep(Duration::from_millis(4)).await;
|
|
||||||
if idx % 2 == 0 {
|
|
||||||
drop(local_guard.take());
|
|
||||||
cross_guard = cross_mode_lock.try_lock().ok();
|
|
||||||
} else {
|
|
||||||
drop(cross_guard.take());
|
|
||||||
local_guard = local_lock.try_lock().ok();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
drop(local_guard);
|
|
||||||
drop(cross_guard);
|
|
||||||
|
|
||||||
for waiter in waiters {
|
|
||||||
let done = waiter.await.expect("waiter task must not panic");
|
|
||||||
assert!(
|
|
||||||
done.is_ok(),
|
|
||||||
"waiter must finish once alternating contention window ends"
|
|
||||||
);
|
|
||||||
assert!(done.expect("waiter timeout must not fire").is_ok());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn light_fuzz_alternating_contention_matrix_preserves_lock_gating() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("dual-lock-alt-fuzz-{}", std::process::id());
|
|
||||||
let local_lock = quota_user_lock(&user);
|
|
||||||
let cross_mode_lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
let mut seed = 0xD00D_BAAD_F00D_2026u64;
|
|
||||||
|
|
||||||
for _round in 0..64u32 {
|
|
||||||
seed ^= seed << 7;
|
|
||||||
seed ^= seed >> 9;
|
|
||||||
seed ^= seed << 8;
|
|
||||||
|
|
||||||
let hold_mode = (seed % 3) as u8;
|
|
||||||
let local_guard = if hold_mode == 0 {
|
|
||||||
Some(
|
|
||||||
local_lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("fuzz local lock should be acquirable"),
|
|
||||||
)
|
|
||||||
} else {
|
|
||||||
None
|
|
||||||
};
|
|
||||||
let cross_guard = if hold_mode == 1 {
|
|
||||||
Some(
|
|
||||||
cross_mode_lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("fuzz cross lock should be acquirable"),
|
|
||||||
)
|
|
||||||
} else {
|
|
||||||
None
|
|
||||||
};
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user.clone(),
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let write = timeout(Duration::from_millis(35), io.write_all(&[0x51])).await;
|
|
||||||
if hold_mode == 2 {
|
|
||||||
assert!(write.is_ok(), "unheld fuzz round must make progress");
|
|
||||||
assert!(write.expect("unheld round timeout").is_ok());
|
|
||||||
} else {
|
|
||||||
assert!(
|
|
||||||
write.is_err(),
|
|
||||||
"held-lock fuzz round must remain pending inside bounded window"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
drop(local_guard);
|
|
||||||
drop(cross_guard);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn stress_fanout_alternating_contention_recovers_without_hanging() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("dual-lock-alt-stress-{}", std::process::id());
|
|
||||||
let local_lock = quota_user_lock(&user);
|
|
||||||
let cross_mode_lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
|
|
||||||
let mut waiters = Vec::new();
|
|
||||||
for _ in 0..48usize {
|
|
||||||
let user = user.clone();
|
|
||||||
waiters.push(tokio::spawn(async move {
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(4096),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
timeout(Duration::from_secs(3), io.write_all(&[0xA0, 0xA1])).await
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut local_guard = Some(
|
|
||||||
local_lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("stress toggle must acquire local lock first"),
|
|
||||||
);
|
|
||||||
let mut cross_guard = None;
|
|
||||||
for idx in 0..40usize {
|
|
||||||
tokio::time::sleep(Duration::from_millis(3)).await;
|
|
||||||
if idx % 2 == 0 {
|
|
||||||
drop(local_guard.take());
|
|
||||||
cross_guard = cross_mode_lock.try_lock().ok();
|
|
||||||
} else {
|
|
||||||
drop(cross_guard.take());
|
|
||||||
local_guard = local_lock.try_lock().ok();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
drop(local_guard);
|
|
||||||
drop(cross_guard);
|
|
||||||
|
|
||||||
for waiter in waiters {
|
|
||||||
let done = waiter.await.expect("stress waiter task must not panic");
|
|
||||||
assert!(done.is_ok(), "stress waiter timed out under alternating contention");
|
|
||||||
assert!(done.expect("stress waiter timeout should not fire").is_ok());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
@ -1,74 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use std::pin::Pin;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
|
||||||
use std::task::{Context, Waker};
|
|
||||||
use tokio::time::{Duration, Instant};
|
|
||||||
|
|
||||||
#[derive(Default)]
|
|
||||||
struct WakeCounter {
|
|
||||||
wakes: AtomicUsize,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl std::task::Wake for WakeCounter {
|
|
||||||
fn wake(self: Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn wake_by_ref(self: &Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn quota_test_guard() -> impl Drop {
|
|
||||||
super::quota_user_lock_test_scope()
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn adversarial_cross_mode_only_contention_backoff_attempt_must_ramp() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("dual-lock-backoff-{}", std::process::id());
|
|
||||||
let cross_mode_lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
let held_cross_mode_guard = cross_mode_lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold cross-mode lock before polling");
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
|
|
||||||
let first = Pin::new(&mut io).poll_write(&mut cx, &[0xAA]);
|
|
||||||
assert!(first.is_pending(), "held cross-mode lock must block writer");
|
|
||||||
|
|
||||||
let started = Instant::now();
|
|
||||||
let mut last_wakes = 0usize;
|
|
||||||
while started.elapsed() < Duration::from_millis(120) {
|
|
||||||
let wakes = wake_counter.wakes.load(Ordering::Relaxed);
|
|
||||||
if wakes > last_wakes {
|
|
||||||
last_wakes = wakes;
|
|
||||||
let next = Pin::new(&mut io).poll_write(&mut cx, &[0xAB]);
|
|
||||||
assert!(next.is_pending(), "writer must remain blocked while lock is held");
|
|
||||||
}
|
|
||||||
tokio::time::sleep(Duration::from_millis(1)).await;
|
|
||||||
}
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
io.quota_write_retry_attempt >= 2,
|
|
||||||
"retry attempt must ramp under sustained second-lock contention; got {}",
|
|
||||||
io.quota_write_retry_attempt
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(held_cross_mode_guard);
|
|
||||||
}
|
|
||||||
|
|
@ -1,325 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use std::pin::Pin;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
|
||||||
use std::task::{Context, Waker};
|
|
||||||
use tokio::io::{AsyncReadExt, AsyncWriteExt, ReadBuf};
|
|
||||||
use tokio::time::{Duration, Instant, timeout};
|
|
||||||
|
|
||||||
#[derive(Default)]
|
|
||||||
struct WakeCounter {
|
|
||||||
wakes: AtomicUsize,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl std::task::Wake for WakeCounter {
|
|
||||||
fn wake(self: Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn wake_by_ref(self: &Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn quota_test_guard() -> impl Drop {
|
|
||||||
super::quota_user_lock_test_scope()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn build_context() -> (Arc<WakeCounter>, Context<'static>) {
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let leaked_waker: &'static Waker = Box::leak(Box::new(waker));
|
|
||||||
(wake_counter, Context::from_waker(leaked_waker))
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn positive_uncontended_dual_locks_writer_completes_without_retry_state() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
format!("dual-lock-positive-{}", std::process::id()),
|
|
||||||
Some(4096),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let (_wake_counter, mut cx) = build_context();
|
|
||||||
let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x01, 0x02, 0x03]);
|
|
||||||
assert!(poll.is_ready());
|
|
||||||
assert_eq!(io.quota_write_retry_attempt, 0);
|
|
||||||
assert!(!io.quota_write_wake_scheduled);
|
|
||||||
assert!(io.quota_write_retry_sleep.is_none());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn negative_local_lock_contention_read_retry_attempt_ramps() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("dual-lock-local-contention-{}", std::process::id());
|
|
||||||
let held = quota_user_lock(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold local quota lock before polling");
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::empty(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let (wake_counter, mut cx) = build_context();
|
|
||||||
let mut one = [0u8; 1];
|
|
||||||
let mut buf = ReadBuf::new(&mut one);
|
|
||||||
let first = Pin::new(&mut io).poll_read(&mut cx, &mut buf);
|
|
||||||
assert!(first.is_pending());
|
|
||||||
|
|
||||||
let started = Instant::now();
|
|
||||||
let mut observed = 0usize;
|
|
||||||
while started.elapsed() < Duration::from_millis(120) {
|
|
||||||
let wakes = wake_counter.wakes.load(Ordering::Relaxed);
|
|
||||||
if wakes > observed {
|
|
||||||
observed = wakes;
|
|
||||||
let mut step_buf = ReadBuf::new(&mut one);
|
|
||||||
let next = Pin::new(&mut io).poll_read(&mut cx, &mut step_buf);
|
|
||||||
assert!(next.is_pending());
|
|
||||||
}
|
|
||||||
tokio::time::sleep(Duration::from_millis(1)).await;
|
|
||||||
}
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
io.quota_read_retry_attempt >= 2,
|
|
||||||
"retry attempt must ramp under sustained local-lock contention; got {}",
|
|
||||||
io.quota_read_retry_attempt
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn edge_cross_mode_contention_release_resets_retry_scheduler_on_success() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("dual-lock-reset-{}", std::process::id());
|
|
||||||
let cross_mode = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
let held_guard = cross_mode
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold cross-mode lock before polling");
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let (wake_counter, mut cx) = build_context();
|
|
||||||
let first = Pin::new(&mut io).poll_write(&mut cx, &[0x10]);
|
|
||||||
assert!(first.is_pending());
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
|
||||||
if wake_counter.wakes.load(Ordering::Relaxed) > 0 {
|
|
||||||
let next = Pin::new(&mut io).poll_write(&mut cx, &[0x11]);
|
|
||||||
assert!(next.is_pending());
|
|
||||||
}
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
let ready = Pin::new(&mut io).poll_write(&mut cx, &[0x12]);
|
|
||||||
assert!(ready.is_ready());
|
|
||||||
assert_eq!(io.quota_write_retry_attempt, 0);
|
|
||||||
assert!(!io.quota_write_wake_scheduled);
|
|
||||||
assert!(io.quota_write_retry_sleep.is_none());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn adversarial_cross_mode_hold_blocks_many_waiters_without_usage_leak() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("dual-lock-adversarial-{}", std::process::id());
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let held = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold cross-mode lock before launching waiters");
|
|
||||||
|
|
||||||
let mut tasks = Vec::new();
|
|
||||||
for _ in 0..24usize {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let user = user.clone();
|
|
||||||
tasks.push(tokio::spawn(async move {
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
stats,
|
|
||||||
user,
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
timeout(Duration::from_millis(40), io.write_all(&[0x33])).await
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
for task in tasks {
|
|
||||||
let timed = task.await.expect("waiter task must not panic");
|
|
||||||
assert!(timed.is_err(), "held cross-mode lock must keep waiter pending");
|
|
||||||
}
|
|
||||||
|
|
||||||
assert_eq!(stats.get_user_total_octets(&user), 0);
|
|
||||||
drop(held_guard);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn integration_waiters_resume_after_cross_mode_release() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("dual-lock-integration-{}", std::process::id());
|
|
||||||
let held = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold cross-mode lock before starting waiter");
|
|
||||||
|
|
||||||
let task = tokio::spawn({
|
|
||||||
let user = user.clone();
|
|
||||||
async move {
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
io.write_all(&[0x44]).await
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
let done = timeout(Duration::from_secs(1), task)
|
|
||||||
.await
|
|
||||||
.expect("waiter task must complete after release")
|
|
||||||
.expect("waiter task must not panic");
|
|
||||||
assert!(done.is_ok());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn light_fuzz_randomized_lock_holds_preserve_liveness_and_quota_bounds() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("dual-lock-fuzz-{}", std::process::id());
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let mut seed = 0xA55A_55AA_C3D2_E1F0u64;
|
|
||||||
|
|
||||||
for _round in 0..48u32 {
|
|
||||||
seed ^= seed << 7;
|
|
||||||
seed ^= seed >> 9;
|
|
||||||
seed ^= seed << 8;
|
|
||||||
|
|
||||||
let hold_mode = (seed % 3) as u8;
|
|
||||||
let mut local_lock = None;
|
|
||||||
let mut cross_lock = None;
|
|
||||||
let mut local_guard = None;
|
|
||||||
let mut cross_guard = None;
|
|
||||||
|
|
||||||
if hold_mode == 0 {
|
|
||||||
local_lock = Some(quota_user_lock(&user));
|
|
||||||
local_guard = Some(
|
|
||||||
local_lock
|
|
||||||
.as_ref()
|
|
||||||
.expect("local lock should be present")
|
|
||||||
.try_lock()
|
|
||||||
.expect("local lock should be acquirable in fuzz round"),
|
|
||||||
);
|
|
||||||
} else if hold_mode == 1 {
|
|
||||||
cross_lock = Some(crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(
|
|
||||||
&user,
|
|
||||||
));
|
|
||||||
cross_guard = Some(
|
|
||||||
cross_lock
|
|
||||||
.as_ref()
|
|
||||||
.expect("cross lock should be present")
|
|
||||||
.try_lock()
|
|
||||||
.expect("cross lock should be acquirable in fuzz round"),
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.clone(),
|
|
||||||
Some(4096),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let write = timeout(Duration::from_millis(25), io.write_all(&[0x7A])).await;
|
|
||||||
if hold_mode == 2 {
|
|
||||||
assert!(write.is_ok(), "unheld round must make progress");
|
|
||||||
} else {
|
|
||||||
assert!(write.is_err(), "held-lock round must stay blocked within timeout");
|
|
||||||
}
|
|
||||||
|
|
||||||
drop(local_guard);
|
|
||||||
drop(cross_guard);
|
|
||||||
drop(local_lock);
|
|
||||||
drop(cross_lock);
|
|
||||||
}
|
|
||||||
|
|
||||||
assert!(stats.get_user_total_octets(&user) <= 4096);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn stress_fanout_waiters_complete_after_release_without_panics() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("dual-lock-stress-{}", std::process::id());
|
|
||||||
let held = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold cross-mode lock before stress fanout");
|
|
||||||
|
|
||||||
let waiters = 64usize;
|
|
||||||
let mut tasks = Vec::new();
|
|
||||||
for _ in 0..waiters {
|
|
||||||
let user = user.clone();
|
|
||||||
tasks.push(tokio::spawn(async move {
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::empty(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
let mut one = [0u8; 1];
|
|
||||||
io.read(&mut one).await
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(12)).await;
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
timeout(Duration::from_secs(2), async {
|
|
||||||
for task in tasks {
|
|
||||||
let result = task.await.expect("stress waiter task must not panic");
|
|
||||||
assert!(result.is_ok());
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("all stress waiters must complete after release");
|
|
||||||
}
|
|
||||||
|
|
@ -1,128 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::AtomicBool;
|
|
||||||
use tokio::io::AsyncWriteExt;
|
|
||||||
use tokio::time::{Duration, timeout};
|
|
||||||
|
|
||||||
fn quota_test_guard() -> impl Drop {
|
|
||||||
super::quota_user_lock_test_scope()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn make_stats_io(user: String) -> StatsIo<tokio::io::Sink> {
|
|
||||||
StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(4096),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn light_fuzz_1024_round_hold_release_cycles_preserve_same_user_liveness() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("dual-lock-race-fuzz-{}", std::process::id());
|
|
||||||
let mut seed = 0xD1CE_BAAD_5EED_1234u64;
|
|
||||||
|
|
||||||
for round in 0..1024u32 {
|
|
||||||
seed ^= seed << 7;
|
|
||||||
seed ^= seed >> 9;
|
|
||||||
seed ^= seed << 8;
|
|
||||||
|
|
||||||
let hold = (seed & 1) == 0;
|
|
||||||
let hold_ms = (seed % 3) as u64;
|
|
||||||
|
|
||||||
let maybe_lock = if hold {
|
|
||||||
Some(crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(
|
|
||||||
&user,
|
|
||||||
))
|
|
||||||
} else {
|
|
||||||
None
|
|
||||||
};
|
|
||||||
|
|
||||||
let maybe_guard = maybe_lock.as_ref().map(|lock| {
|
|
||||||
lock.try_lock()
|
|
||||||
.expect("cross-mode lock must be acquirable in fuzz round")
|
|
||||||
});
|
|
||||||
|
|
||||||
if hold {
|
|
||||||
let mut blocked_io = make_stats_io(user.clone());
|
|
||||||
let blocked = timeout(Duration::from_millis(5), blocked_io.write_all(&[0xA5])).await;
|
|
||||||
assert!(
|
|
||||||
blocked.is_err(),
|
|
||||||
"held round must block waiter before lock release (round={round})"
|
|
||||||
);
|
|
||||||
|
|
||||||
if hold_ms > 0 {
|
|
||||||
tokio::time::sleep(Duration::from_millis(hold_ms)).await;
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
let mut free_io = make_stats_io(user.clone());
|
|
||||||
let free = timeout(Duration::from_millis(120), free_io.write_all(&[0xA5])).await;
|
|
||||||
assert!(
|
|
||||||
free.is_ok(),
|
|
||||||
"unheld round must complete promptly (round={round})"
|
|
||||||
);
|
|
||||||
assert!(free.expect("unheld round should complete").is_ok());
|
|
||||||
}
|
|
||||||
|
|
||||||
drop(maybe_guard);
|
|
||||||
|
|
||||||
let done = timeout(Duration::from_millis(350), async {
|
|
||||||
let user = user.clone();
|
|
||||||
let mut io = make_stats_io(user);
|
|
||||||
io.write_all(&[0xA6]).await
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("post-release write must complete in bounded time");
|
|
||||||
assert!(done.is_ok());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn stress_jittered_three_waiter_rounds_do_not_starve_after_release() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("dual-lock-race-stress-{}", std::process::id());
|
|
||||||
let mut seed = 0xC0FF_EE77_4444_9999u64;
|
|
||||||
|
|
||||||
for round in 0..256u32 {
|
|
||||||
seed ^= seed << 7;
|
|
||||||
seed ^= seed >> 9;
|
|
||||||
seed ^= seed << 8;
|
|
||||||
|
|
||||||
let hold_ms = (seed % 4) as u64;
|
|
||||||
let lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user);
|
|
||||||
let guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("cross-mode lock must be acquirable at round start");
|
|
||||||
|
|
||||||
let mut waiters = Vec::new();
|
|
||||||
for _ in 0..3usize {
|
|
||||||
let user = user.clone();
|
|
||||||
waiters.push(tokio::spawn(async move {
|
|
||||||
let mut io = make_stats_io(user);
|
|
||||||
io.write_all(&[0x55]).await
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(hold_ms)).await;
|
|
||||||
drop(guard);
|
|
||||||
|
|
||||||
timeout(Duration::from_secs(1), async {
|
|
||||||
for waiter in waiters {
|
|
||||||
let done = waiter.await.expect("waiter task must not panic");
|
|
||||||
assert!(
|
|
||||||
done.is_ok(),
|
|
||||||
"waiter must complete after release (round={round})"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("all waiters must complete in bounded time after release");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
@ -1,79 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use dashmap::DashMap;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use tokio::time::{Duration, timeout};
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn tdd_explicit_quota_lock_evict_reclaims_only_unheld_entries() {
|
|
||||||
let _guard = quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let held_user = format!("quota-evict-held-{}", std::process::id());
|
|
||||||
let stale_a_user = format!("quota-evict-stale-a-{}", std::process::id());
|
|
||||||
let stale_b_user = format!("quota-evict-stale-b-{}", std::process::id());
|
|
||||||
|
|
||||||
let held = quota_user_lock(&held_user);
|
|
||||||
let stale_a = quota_user_lock(&stale_a_user);
|
|
||||||
let stale_b = quota_user_lock(&stale_b_user);
|
|
||||||
|
|
||||||
assert!(map.get(&held_user).is_some());
|
|
||||||
assert!(map.get(&stale_a_user).is_some());
|
|
||||||
assert!(map.get(&stale_b_user).is_some());
|
|
||||||
|
|
||||||
drop(stale_a);
|
|
||||||
drop(stale_b);
|
|
||||||
|
|
||||||
quota_user_lock_evict();
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
map.get(&held_user).is_some(),
|
|
||||||
"held entry must survive eviction"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
map.get(&stale_a_user).is_none(),
|
|
||||||
"unheld stale entry must be reclaimed"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
map.get(&stale_b_user).is_none(),
|
|
||||||
"unheld stale entry must be reclaimed"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(held);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
|
||||||
async fn tdd_periodic_quota_lock_evictor_reclaims_stale_entries_off_hot_path() {
|
|
||||||
let _guard = quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let held_user = format!("quota-evict-loop-held-{}", std::process::id());
|
|
||||||
let stale_user = format!("quota-evict-loop-stale-{}", std::process::id());
|
|
||||||
|
|
||||||
let held = quota_user_lock(&held_user);
|
|
||||||
let stale = quota_user_lock(&stale_user);
|
|
||||||
|
|
||||||
assert_eq!(map.len(), 2);
|
|
||||||
drop(stale);
|
|
||||||
|
|
||||||
let evictor = spawn_quota_user_lock_evictor_for_tests(Duration::from_millis(5));
|
|
||||||
|
|
||||||
timeout(Duration::from_millis(200), async {
|
|
||||||
loop {
|
|
||||||
if map.get(&stale_user).is_none() {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
tokio::time::sleep(Duration::from_millis(5)).await;
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("periodic quota lock evictor must reclaim stale entry");
|
|
||||||
|
|
||||||
evictor.abort();
|
|
||||||
|
|
||||||
assert!(map.get(&held_user).is_some());
|
|
||||||
assert!(map.get(&stale_user).is_none());
|
|
||||||
|
|
||||||
drop(held);
|
|
||||||
}
|
|
||||||
|
|
@ -1,153 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use dashmap::DashMap;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use tokio::task::JoinSet;
|
|
||||||
use tokio::time::{Duration, timeout};
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn stress_background_evictor_with_high_churn_keeps_cache_bounded_and_live() {
|
|
||||||
let _guard = quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let evictor = spawn_quota_user_lock_evictor_for_tests(Duration::from_millis(5));
|
|
||||||
|
|
||||||
let mut tasks = JoinSet::new();
|
|
||||||
for worker in 0..24u32 {
|
|
||||||
tasks.spawn(async move {
|
|
||||||
for round in 0..320u32 {
|
|
||||||
let user = format!(
|
|
||||||
"quota-evict-stress-user-{}-{}-{}",
|
|
||||||
std::process::id(),
|
|
||||||
worker,
|
|
||||||
round
|
|
||||||
);
|
|
||||||
let lock = quota_user_lock(&user);
|
|
||||||
if round % 19 == 0 {
|
|
||||||
tokio::task::yield_now().await;
|
|
||||||
}
|
|
||||||
drop(lock);
|
|
||||||
}
|
|
||||||
});
|
|
||||||
}
|
|
||||||
|
|
||||||
while let Some(done) = tasks.join_next().await {
|
|
||||||
done.expect("stress worker must not panic");
|
|
||||||
}
|
|
||||||
|
|
||||||
quota_user_lock_evict();
|
|
||||||
tokio::time::sleep(Duration::from_millis(20)).await;
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
map.len() <= QUOTA_USER_LOCKS_MAX,
|
|
||||||
"quota lock map must remain bounded after churn + eviction"
|
|
||||||
);
|
|
||||||
|
|
||||||
let sanity_user = format!("quota-evict-stress-sanity-{}", std::process::id());
|
|
||||||
let sanity_lock = quota_user_lock(&sanity_user);
|
|
||||||
assert!(
|
|
||||||
map.get(&sanity_user).is_some(),
|
|
||||||
"sanity user should be cacheable after eviction reclaimed stale entries"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(sanity_lock);
|
|
||||||
evictor.abort();
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn adversarial_held_lock_survives_repeated_eviction_then_reclaims_after_release() {
|
|
||||||
let _guard = quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let held_user = format!("quota-evict-held-survive-{}", std::process::id());
|
|
||||||
let held = quota_user_lock(&held_user);
|
|
||||||
|
|
||||||
let evictor = spawn_quota_user_lock_evictor_for_tests(Duration::from_millis(3));
|
|
||||||
|
|
||||||
for idx in 0..512u32 {
|
|
||||||
let user = format!("quota-evict-held-churn-{}-{}", std::process::id(), idx);
|
|
||||||
let temp = quota_user_lock(&user);
|
|
||||||
drop(temp);
|
|
||||||
if idx % 32 == 0 {
|
|
||||||
tokio::task::yield_now().await;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let reacquired = quota_user_lock(&held_user);
|
|
||||||
assert!(
|
|
||||||
Arc::ptr_eq(&held, &reacquired),
|
|
||||||
"held user lock identity must remain stable across repeated evictions"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
map.get(&held_user).is_some(),
|
|
||||||
"held user entry must not be reclaimed while externally referenced"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(reacquired);
|
|
||||||
drop(held);
|
|
||||||
|
|
||||||
timeout(Duration::from_millis(300), async {
|
|
||||||
loop {
|
|
||||||
if map.get(&held_user).is_none() {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
tokio::time::sleep(Duration::from_millis(5)).await;
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("released held lock must be reclaimed by periodic evictor");
|
|
||||||
|
|
||||||
evictor.abort();
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn stress_saturation_then_periodic_eviction_recovers_cacheability_without_inline_retain() {
|
|
||||||
let _guard = quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX);
|
|
||||||
let prefix = format!("quota-evict-saturated-{}", std::process::id());
|
|
||||||
for idx in 0..QUOTA_USER_LOCKS_MAX {
|
|
||||||
retained.push(quota_user_lock(&format!("{prefix}-{idx}")));
|
|
||||||
}
|
|
||||||
|
|
||||||
assert_eq!(map.len(), QUOTA_USER_LOCKS_MAX);
|
|
||||||
|
|
||||||
let overflow_user = format!("quota-evict-overflow-user-{}", std::process::id());
|
|
||||||
let overflow_before = quota_user_lock(&overflow_user);
|
|
||||||
assert!(
|
|
||||||
map.get(&overflow_user).is_none(),
|
|
||||||
"saturated map must initially route new user to overflow stripe"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(retained);
|
|
||||||
|
|
||||||
let evictor = spawn_quota_user_lock_evictor_for_tests(Duration::from_millis(4));
|
|
||||||
|
|
||||||
timeout(Duration::from_millis(400), async {
|
|
||||||
loop {
|
|
||||||
if map.len() < QUOTA_USER_LOCKS_MAX {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
tokio::time::sleep(Duration::from_millis(5)).await;
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("periodic evictor must reclaim stale saturated entries");
|
|
||||||
|
|
||||||
let overflow_after = quota_user_lock(&overflow_user);
|
|
||||||
assert!(
|
|
||||||
map.get(&overflow_user).is_some(),
|
|
||||||
"after eviction, overflow user should become cacheable again"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
Arc::strong_count(&overflow_after) >= 2,
|
|
||||||
"cacheable lock should be held by map and caller"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(overflow_before);
|
|
||||||
drop(overflow_after);
|
|
||||||
evictor.abort();
|
|
||||||
}
|
|
||||||
|
|
@ -1,135 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use std::pin::Pin;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
|
||||||
use std::task::Waker;
|
|
||||||
use std::task::{Context, Poll};
|
|
||||||
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
|
|
||||||
|
|
||||||
#[derive(Default)]
|
|
||||||
struct WakeCounter {
|
|
||||||
wakes: AtomicUsize,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl std::task::Wake for WakeCounter {
|
|
||||||
fn wake(self: Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn wake_by_ref(self: &Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn build_context() -> (Arc<WakeCounter>, Context<'static>) {
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
// Context stores a reference; leak one Waker for deterministic test scope.
|
|
||||||
let leaked_waker: &'static Waker = Box::leak(Box::new(waker));
|
|
||||||
(wake_counter, Context::from_waker(leaked_waker))
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn adversarial_map_churn_cannot_bypass_held_writer_lock() {
|
|
||||||
let _guard = quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let user = "quota-identity-writer-user";
|
|
||||||
let held_lock = quota_user_lock(user);
|
|
||||||
let _held_guard = held_lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold initial user lock before StatsIo poll");
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user.to_string(),
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
map.clear();
|
|
||||||
let churned_lock = quota_user_lock(user);
|
|
||||||
assert!(
|
|
||||||
!Arc::ptr_eq(&held_lock, &churned_lock),
|
|
||||||
"precondition: map churn should produce a distinct lock identity"
|
|
||||||
);
|
|
||||||
|
|
||||||
let (_wake_counter, mut cx) = build_context();
|
|
||||||
let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x11, 0x22, 0x33, 0x44]);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
matches!(poll, Poll::Pending),
|
|
||||||
"writer must remain pending on the originally-held lock identity"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn adversarial_map_churn_cannot_bypass_held_reader_lock() {
|
|
||||||
let _guard = quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let user = "quota-identity-reader-user";
|
|
||||||
let held_lock = quota_user_lock(user);
|
|
||||||
let _held_guard = held_lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold initial user lock before StatsIo poll");
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::empty(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user.to_string(),
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
map.clear();
|
|
||||||
let churned_lock = quota_user_lock(user);
|
|
||||||
assert!(
|
|
||||||
!Arc::ptr_eq(&held_lock, &churned_lock),
|
|
||||||
"precondition: map churn should produce a distinct lock identity"
|
|
||||||
);
|
|
||||||
|
|
||||||
let (_wake_counter, mut cx) = build_context();
|
|
||||||
let mut storage = [0u8; 8];
|
|
||||||
let mut read_buf = ReadBuf::new(&mut storage);
|
|
||||||
let poll = Pin::new(&mut io).poll_read(&mut cx, &mut read_buf);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
matches!(poll, Poll::Pending),
|
|
||||||
"reader must remain pending on the originally-held lock identity"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn business_no_lock_contention_keeps_writer_progress() {
|
|
||||||
let _guard = quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let user = "quota-identity-progress-user";
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user.to_string(),
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let (_wake_counter, mut cx) = build_context();
|
|
||||||
let poll = Pin::new(&mut io).poll_write(&mut cx, &[0xAA, 0xBB]);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
matches!(poll, Poll::Ready(Ok(2))),
|
|
||||||
"writer should progress immediately without contention"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
@ -1,440 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::error::ProxyError;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use crate::stream::BufferPool;
|
|
||||||
use dashmap::DashMap;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::AtomicBool;
|
|
||||||
use std::time::Duration;
|
|
||||||
use tokio::io::{AsyncReadExt, AsyncWriteExt, duplex};
|
|
||||||
use tokio::sync::Barrier;
|
|
||||||
use tokio::time::Instant;
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn quota_lock_same_user_returns_same_arc_instance() {
|
|
||||||
let _guard = super::quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let a = quota_user_lock("quota-lock-same-user");
|
|
||||||
let b = quota_user_lock("quota-lock-same-user");
|
|
||||||
assert!(Arc::ptr_eq(&a, &b));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn quota_lock_parallel_same_user_reuses_single_lock() {
|
|
||||||
let _guard = super::quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let user = "quota-lock-parallel-same";
|
|
||||||
let mut handles = Vec::new();
|
|
||||||
|
|
||||||
for _ in 0..64 {
|
|
||||||
handles.push(std::thread::spawn(move || quota_user_lock(user)));
|
|
||||||
}
|
|
||||||
|
|
||||||
let first = handles
|
|
||||||
.remove(0)
|
|
||||||
.join()
|
|
||||||
.expect("thread must return lock handle");
|
|
||||||
|
|
||||||
for handle in handles {
|
|
||||||
let got = handle.join().expect("thread must return lock handle");
|
|
||||||
assert!(Arc::ptr_eq(&first, &got));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn quota_lock_unique_users_materialize_distinct_entries() {
|
|
||||||
let _guard = super::quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let base = format!("quota-lock-distinct-{}", std::process::id());
|
|
||||||
let users: Vec<String> = (0..(QUOTA_USER_LOCKS_MAX / 2))
|
|
||||||
.map(|idx| format!("{base}-{idx}"))
|
|
||||||
.collect();
|
|
||||||
|
|
||||||
for user in &users {
|
|
||||||
let _ = quota_user_lock(user);
|
|
||||||
}
|
|
||||||
|
|
||||||
for user in &users {
|
|
||||||
assert!(
|
|
||||||
map.get(user).is_some(),
|
|
||||||
"lock cache must contain entry for {user}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn quota_lock_unique_churn_stress_keeps_all_inserted_keys_addressable() {
|
|
||||||
let _guard = super::quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let base = format!("quota-lock-churn-{}", std::process::id());
|
|
||||||
for idx in 0..(QUOTA_USER_LOCKS_MAX + 256) {
|
|
||||||
let _ = quota_user_lock(&format!("{base}-{idx}"));
|
|
||||||
}
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
map.len() <= QUOTA_USER_LOCKS_MAX,
|
|
||||||
"quota lock cache must stay bounded under unique-user churn"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn quota_lock_saturation_returns_stable_overflow_lock_without_cache_growth() {
|
|
||||||
let _guard = super::quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let prefix = format!("quota-held-{}", std::process::id());
|
|
||||||
let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX);
|
|
||||||
for idx in 0..QUOTA_USER_LOCKS_MAX {
|
|
||||||
retained.push(quota_user_lock(&format!("{prefix}-{idx}")));
|
|
||||||
}
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
map.len(),
|
|
||||||
QUOTA_USER_LOCKS_MAX,
|
|
||||||
"cache must be saturated for overflow check"
|
|
||||||
);
|
|
||||||
|
|
||||||
let overflow_user = format!("quota-overflow-{}", std::process::id());
|
|
||||||
let overflow_a = quota_user_lock(&overflow_user);
|
|
||||||
let overflow_b = quota_user_lock(&overflow_user);
|
|
||||||
|
|
||||||
assert_eq!(
|
|
||||||
map.len(),
|
|
||||||
QUOTA_USER_LOCKS_MAX,
|
|
||||||
"overflow path must not grow lock cache"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
map.get(&overflow_user).is_none(),
|
|
||||||
"overflow user lock must stay outside bounded cache under saturation"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
Arc::ptr_eq(&overflow_a, &overflow_b),
|
|
||||||
"overflow user must receive stable striped overflow lock while saturated"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(retained);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn quota_lock_reclaims_unreferenced_entries_after_explicit_eviction_pass() {
|
|
||||||
let _guard = super::quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
// Saturate with retained strong references first so parallel tests cannot
|
|
||||||
// reclaim our fixture entries before we validate the reclaim path.
|
|
||||||
let prefix = format!("quota-reclaim-drop-{}", std::process::id());
|
|
||||||
let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX);
|
|
||||||
for idx in 0..QUOTA_USER_LOCKS_MAX {
|
|
||||||
retained.push(quota_user_lock(&format!("{prefix}-{idx}")));
|
|
||||||
}
|
|
||||||
|
|
||||||
drop(retained);
|
|
||||||
|
|
||||||
quota_user_lock_evict();
|
|
||||||
|
|
||||||
let overflow_user = format!("quota-reclaim-overflow-{}", std::process::id());
|
|
||||||
let overflow = quota_user_lock(&overflow_user);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
map.get(&overflow_user).is_some(),
|
|
||||||
"after reclaiming stale entries, overflow user should become cacheable"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
Arc::strong_count(&overflow) >= 2,
|
|
||||||
"cacheable overflow lock should be held by both map and caller"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn quota_lock_saturated_same_user_must_not_return_distinct_locks() {
|
|
||||||
let _guard = super::quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX);
|
|
||||||
for idx in 0..QUOTA_USER_LOCKS_MAX {
|
|
||||||
retained.push(quota_user_lock(&format!(
|
|
||||||
"quota-saturated-held-{}-{idx}",
|
|
||||||
std::process::id()
|
|
||||||
)));
|
|
||||||
}
|
|
||||||
|
|
||||||
let overflow_user = format!("quota-saturated-same-user-{}", std::process::id());
|
|
||||||
let a = quota_user_lock(&overflow_user);
|
|
||||||
let b = quota_user_lock(&overflow_user);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
Arc::ptr_eq(&a, &b),
|
|
||||||
"same user must not receive distinct locks under saturation because that enables quota race bypass"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(retained);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn quota_lock_saturation_concurrent_same_user_never_overshoots_quota() {
|
|
||||||
let _guard = super::quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX);
|
|
||||||
for idx in 0..QUOTA_USER_LOCKS_MAX {
|
|
||||||
retained.push(quota_user_lock(&format!(
|
|
||||||
"quota-saturated-race-held-{}-{idx}",
|
|
||||||
std::process::id()
|
|
||||||
)));
|
|
||||||
}
|
|
||||||
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let user = format!("quota-saturated-race-user-{}", std::process::id());
|
|
||||||
let gate = Arc::new(Barrier::new(2));
|
|
||||||
|
|
||||||
let worker = |label: u8, stats: Arc<Stats>, user: String, gate: Arc<Barrier>| {
|
|
||||||
tokio::spawn(async move {
|
|
||||||
let counters = Arc::new(SharedCounters::new());
|
|
||||||
let quota_exceeded = Arc::new(AtomicBool::new(false));
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
counters,
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user,
|
|
||||||
Some(1),
|
|
||||||
quota_exceeded,
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
gate.wait().await;
|
|
||||||
io.write_all(&[label]).await
|
|
||||||
})
|
|
||||||
};
|
|
||||||
|
|
||||||
let one = worker(0x11, Arc::clone(&stats), user.clone(), Arc::clone(&gate));
|
|
||||||
let two = worker(0x22, Arc::clone(&stats), user.clone(), Arc::clone(&gate));
|
|
||||||
|
|
||||||
let _ = tokio::time::timeout(Duration::from_secs(2), async {
|
|
||||||
let _ = one.await.expect("task one must not panic");
|
|
||||||
let _ = two.await.expect("task two must not panic");
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("quota race workers must complete");
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
stats.get_user_total_octets(&user) <= 1,
|
|
||||||
"saturated lock path must never overshoot quota for same user"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(retained);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn quota_lock_saturation_stress_same_user_never_overshoots_quota() {
|
|
||||||
let _guard = super::quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX);
|
|
||||||
for idx in 0..QUOTA_USER_LOCKS_MAX {
|
|
||||||
retained.push(quota_user_lock(&format!(
|
|
||||||
"quota-saturated-stress-held-{}-{idx}",
|
|
||||||
std::process::id()
|
|
||||||
)));
|
|
||||||
}
|
|
||||||
|
|
||||||
for round in 0..128u32 {
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let user = format!("quota-saturated-stress-user-{}-{round}", std::process::id());
|
|
||||||
let gate = Arc::new(Barrier::new(2));
|
|
||||||
|
|
||||||
let one = {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let user = user.clone();
|
|
||||||
let gate = Arc::clone(&gate);
|
|
||||||
tokio::spawn(async move {
|
|
||||||
let counters = Arc::new(SharedCounters::new());
|
|
||||||
let quota_exceeded = Arc::new(AtomicBool::new(false));
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
counters,
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user,
|
|
||||||
Some(1),
|
|
||||||
quota_exceeded,
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
gate.wait().await;
|
|
||||||
io.write_all(&[0x31]).await
|
|
||||||
})
|
|
||||||
};
|
|
||||||
|
|
||||||
let two = {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let user = user.clone();
|
|
||||||
let gate = Arc::clone(&gate);
|
|
||||||
tokio::spawn(async move {
|
|
||||||
let counters = Arc::new(SharedCounters::new());
|
|
||||||
let quota_exceeded = Arc::new(AtomicBool::new(false));
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
counters,
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user,
|
|
||||||
Some(1),
|
|
||||||
quota_exceeded,
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
gate.wait().await;
|
|
||||||
io.write_all(&[0x32]).await
|
|
||||||
})
|
|
||||||
};
|
|
||||||
|
|
||||||
let _ = one.await.expect("stress task one must not panic");
|
|
||||||
let _ = two.await.expect("stress task two must not panic");
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
stats.get_user_total_octets(&user) <= 1,
|
|
||||||
"round {round}: saturated path must not overshoot quota"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
drop(retained);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn quota_error_classifier_accepts_internal_quota_sentinel_only() {
|
|
||||||
let err = quota_io_error();
|
|
||||||
assert!(is_quota_io_error(&err));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn quota_error_classifier_rejects_plain_permission_denied() {
|
|
||||||
let err = std::io::Error::new(std::io::ErrorKind::PermissionDenied, "permission denied");
|
|
||||||
assert!(!is_quota_io_error(&err));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn quota_lock_test_scope_recovers_after_guard_poison() {
|
|
||||||
let poison_result = std::thread::spawn(|| {
|
|
||||||
let _guard = super::quota_user_lock_test_scope();
|
|
||||||
panic!("intentional test-only guard poison");
|
|
||||||
})
|
|
||||||
.join();
|
|
||||||
assert!(poison_result.is_err(), "poison setup thread must panic");
|
|
||||||
|
|
||||||
let _guard = super::quota_user_lock_test_scope();
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let a = quota_user_lock("quota-lock-poison-recovery-user");
|
|
||||||
let b = quota_user_lock("quota-lock-poison-recovery-user");
|
|
||||||
assert!(Arc::ptr_eq(&a, &b));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn quota_lock_integration_zero_quota_cuts_off_without_forwarding() {
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let user = "quota-zero-user";
|
|
||||||
|
|
||||||
let (mut client_peer, relay_client) = duplex(2048);
|
|
||||||
let (relay_server, mut server_peer) = duplex(2048);
|
|
||||||
let (client_reader, client_writer) = tokio::io::split(relay_client);
|
|
||||||
let (server_reader, server_writer) = tokio::io::split(relay_server);
|
|
||||||
|
|
||||||
let relay = tokio::spawn(relay_bidirectional(
|
|
||||||
client_reader,
|
|
||||||
client_writer,
|
|
||||||
server_reader,
|
|
||||||
server_writer,
|
|
||||||
512,
|
|
||||||
512,
|
|
||||||
user,
|
|
||||||
Arc::clone(&stats),
|
|
||||||
Some(0),
|
|
||||||
Arc::new(BufferPool::new()),
|
|
||||||
));
|
|
||||||
|
|
||||||
client_peer
|
|
||||||
.write_all(b"x")
|
|
||||||
.await
|
|
||||||
.expect("client write must succeed");
|
|
||||||
|
|
||||||
let mut probe = [0u8; 1];
|
|
||||||
let forwarded =
|
|
||||||
tokio::time::timeout(Duration::from_millis(80), server_peer.read(&mut probe)).await;
|
|
||||||
if let Ok(Ok(n)) = forwarded {
|
|
||||||
assert_eq!(n, 0, "zero quota path must not forward payload bytes");
|
|
||||||
}
|
|
||||||
|
|
||||||
let result = tokio::time::timeout(Duration::from_secs(2), relay)
|
|
||||||
.await
|
|
||||||
.expect("relay must terminate under zero quota")
|
|
||||||
.expect("relay task must not panic");
|
|
||||||
assert!(matches!(result, Err(ProxyError::DataQuotaExceeded { .. })));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn quota_lock_integration_no_quota_relays_both_directions_under_burst() {
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
|
|
||||||
let (mut client_peer, relay_client) = duplex(8192);
|
|
||||||
let (relay_server, mut server_peer) = duplex(8192);
|
|
||||||
let (client_reader, client_writer) = tokio::io::split(relay_client);
|
|
||||||
let (server_reader, server_writer) = tokio::io::split(relay_server);
|
|
||||||
|
|
||||||
let relay = tokio::spawn(relay_bidirectional(
|
|
||||||
client_reader,
|
|
||||||
client_writer,
|
|
||||||
server_reader,
|
|
||||||
server_writer,
|
|
||||||
1024,
|
|
||||||
1024,
|
|
||||||
"quota-none-burst-user",
|
|
||||||
Arc::clone(&stats),
|
|
||||||
None,
|
|
||||||
Arc::new(BufferPool::new()),
|
|
||||||
));
|
|
||||||
|
|
||||||
let c2s = vec![0xA5; 2048];
|
|
||||||
let s2c = vec![0x5A; 1536];
|
|
||||||
|
|
||||||
client_peer
|
|
||||||
.write_all(&c2s)
|
|
||||||
.await
|
|
||||||
.expect("client burst write must succeed");
|
|
||||||
let mut got_c2s = vec![0u8; c2s.len()];
|
|
||||||
server_peer
|
|
||||||
.read_exact(&mut got_c2s)
|
|
||||||
.await
|
|
||||||
.expect("server must receive c2s burst");
|
|
||||||
assert_eq!(got_c2s, c2s);
|
|
||||||
|
|
||||||
server_peer
|
|
||||||
.write_all(&s2c)
|
|
||||||
.await
|
|
||||||
.expect("server burst write must succeed");
|
|
||||||
let mut got_s2c = vec![0u8; s2c.len()];
|
|
||||||
client_peer
|
|
||||||
.read_exact(&mut got_s2c)
|
|
||||||
.await
|
|
||||||
.expect("client must receive s2c burst");
|
|
||||||
assert_eq!(got_s2c, s2c);
|
|
||||||
|
|
||||||
drop(client_peer);
|
|
||||||
drop(server_peer);
|
|
||||||
|
|
||||||
let done = tokio::time::timeout(Duration::from_secs(2), relay)
|
|
||||||
.await
|
|
||||||
.expect("relay must terminate after peers close")
|
|
||||||
.expect("relay task must not panic");
|
|
||||||
assert!(done.is_ok());
|
|
||||||
}
|
|
||||||
|
|
@ -1,249 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use std::pin::Pin;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
|
||||||
use std::task::{Context, Waker};
|
|
||||||
use tokio::io::AsyncWriteExt;
|
|
||||||
use tokio::time::{Duration, Instant, timeout};
|
|
||||||
|
|
||||||
#[derive(Default)]
|
|
||||||
struct WakeCounter {
|
|
||||||
wakes: AtomicUsize,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl std::task::Wake for WakeCounter {
|
|
||||||
fn wake(self: Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn wake_by_ref(self: &Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn quota_test_guard() -> impl Drop {
|
|
||||||
super::quota_user_lock_test_scope()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn build_context() -> (Arc<WakeCounter>, Context<'static>) {
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let leaked_waker: &'static Waker = Box::leak(Box::new(waker));
|
|
||||||
(wake_counter, Context::from_waker(leaked_waker))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn sleep_slot_ptr(slot: &Option<Pin<Box<tokio::time::Sleep>>>) -> usize {
|
|
||||||
slot.as_ref()
|
|
||||||
.map(|sleep| (&**sleep) as *const tokio::time::Sleep as usize)
|
|
||||||
.unwrap_or(0)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn tdd_single_pending_timer_does_not_allocate_on_each_repoll() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("retry-alloc-single-pending-{}", std::process::id());
|
|
||||||
let lock = quota_user_lock(&user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold local lock to force retry scheduling");
|
|
||||||
|
|
||||||
reset_quota_retry_sleep_allocs_for_tests();
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let (_wake_counter, mut cx) = build_context();
|
|
||||||
|
|
||||||
let first = Pin::new(&mut io).poll_write(&mut cx, &[0xA1]);
|
|
||||||
assert!(first.is_pending());
|
|
||||||
let allocs_after_first = quota_retry_sleep_allocs_for_tests();
|
|
||||||
let ptr_after_first = sleep_slot_ptr(&io.quota_write_retry_sleep);
|
|
||||||
|
|
||||||
let second = Pin::new(&mut io).poll_write(&mut cx, &[0xA2]);
|
|
||||||
assert!(second.is_pending());
|
|
||||||
let allocs_after_second = quota_retry_sleep_allocs_for_tests();
|
|
||||||
let ptr_after_second = sleep_slot_ptr(&io.quota_write_retry_sleep);
|
|
||||||
|
|
||||||
assert_eq!(allocs_after_first, 1, "first pending poll must allocate one timer");
|
|
||||||
assert_eq!(
|
|
||||||
allocs_after_second, 1,
|
|
||||||
"repoll while the same timer is pending must not allocate again"
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
ptr_after_first, ptr_after_second,
|
|
||||||
"repoll while pending should retain the same timer allocation"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn tdd_retry_cycle_allocates_once_per_fired_timer_cycle_not_per_poll() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("retry-alloc-per-cycle-{}", std::process::id());
|
|
||||||
let lock = quota_user_lock(&user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold local lock to keep write path pending");
|
|
||||||
|
|
||||||
reset_quota_retry_sleep_allocs_for_tests();
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let (wake_counter, mut cx) = build_context();
|
|
||||||
|
|
||||||
let mut polls = 0u64;
|
|
||||||
let mut observed_wakes = 0usize;
|
|
||||||
let started = Instant::now();
|
|
||||||
while started.elapsed() < Duration::from_millis(70) {
|
|
||||||
let poll = Pin::new(&mut io).poll_write(&mut cx, &[0xB1]);
|
|
||||||
polls = polls.saturating_add(1);
|
|
||||||
assert!(poll.is_pending());
|
|
||||||
|
|
||||||
let wakes = wake_counter.wakes.load(Ordering::Relaxed);
|
|
||||||
if wakes > observed_wakes {
|
|
||||||
observed_wakes = wakes;
|
|
||||||
}
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(1)).await;
|
|
||||||
}
|
|
||||||
|
|
||||||
let allocs = quota_retry_sleep_allocs_for_tests();
|
|
||||||
assert!(allocs >= 2, "multiple fired cycles should allocate multiple timers");
|
|
||||||
assert!(
|
|
||||||
allocs < polls,
|
|
||||||
"timer allocations must be bounded by cycles, not by every repoll (allocs={allocs}, polls={polls})"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn adversarial_backoff_latency_envelope_stays_bounded_under_contention() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("retry-latency-envelope-{}", std::process::id());
|
|
||||||
let lock = quota_user_lock(&user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold local lock for sustained contention");
|
|
||||||
|
|
||||||
reset_quota_retry_sleep_allocs_for_tests();
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let (wake_counter, mut cx) = build_context();
|
|
||||||
|
|
||||||
let first = Pin::new(&mut io).poll_write(&mut cx, &[0xC1]);
|
|
||||||
assert!(first.is_pending());
|
|
||||||
|
|
||||||
let started = Instant::now();
|
|
||||||
let mut last_wakes = 0usize;
|
|
||||||
let mut wake_instants = Vec::new();
|
|
||||||
|
|
||||||
while started.elapsed() < Duration::from_millis(120) {
|
|
||||||
let wakes = wake_counter.wakes.load(Ordering::Relaxed);
|
|
||||||
if wakes > last_wakes {
|
|
||||||
last_wakes = wakes;
|
|
||||||
wake_instants.push(Instant::now());
|
|
||||||
let pending = Pin::new(&mut io).poll_write(&mut cx, &[0xC2]);
|
|
||||||
assert!(pending.is_pending());
|
|
||||||
}
|
|
||||||
tokio::time::sleep(Duration::from_millis(1)).await;
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut max_gap = Duration::from_millis(0);
|
|
||||||
for idx in 1..wake_instants.len() {
|
|
||||||
let gap = wake_instants[idx].saturating_duration_since(wake_instants[idx - 1]);
|
|
||||||
if gap > max_gap {
|
|
||||||
max_gap = gap;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
max_gap <= Duration::from_millis(35),
|
|
||||||
"retry wake gap must remain bounded in test profile; observed max gap={max_gap:?}"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
quota_retry_sleep_allocs_for_tests() <= 16,
|
|
||||||
"allocation cycles must remain bounded during a short contention window"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn micro_benchmark_release_to_completion_latency_stays_bounded() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let rounds = 96usize;
|
|
||||||
let mut samples_ms = Vec::with_capacity(rounds);
|
|
||||||
|
|
||||||
for round in 0..rounds {
|
|
||||||
let user = format!("retry-release-latency-{}-{round}", std::process::id());
|
|
||||||
let lock = quota_user_lock(&user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold local lock before spawning blocked writer");
|
|
||||||
|
|
||||||
let writer = tokio::spawn(async move {
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::new(Stats::new()),
|
|
||||||
user,
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
Instant::now(),
|
|
||||||
);
|
|
||||||
io.write_all(&[0xD1]).await
|
|
||||||
});
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(2)).await;
|
|
||||||
let release_at = Instant::now();
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
let done = timeout(Duration::from_millis(120), writer)
|
|
||||||
.await
|
|
||||||
.expect("blocked writer must complete after release")
|
|
||||||
.expect("writer task must not panic");
|
|
||||||
assert!(done.is_ok());
|
|
||||||
|
|
||||||
samples_ms.push(release_at.elapsed().as_millis() as u64);
|
|
||||||
}
|
|
||||||
|
|
||||||
samples_ms.sort_unstable();
|
|
||||||
let p95_idx = ((samples_ms.len() * 95) / 100).min(samples_ms.len().saturating_sub(1));
|
|
||||||
let p95_ms = samples_ms[p95_idx];
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
p95_ms <= 40,
|
|
||||||
"contention release->completion p95 must stay bounded; p95_ms={p95_ms}, samples={samples_ms:?}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
@ -1,241 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use dashmap::DashMap;
|
|
||||||
use std::pin::Pin;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
|
||||||
use std::task::{Context, Waker};
|
|
||||||
use tokio::io::ReadBuf;
|
|
||||||
use tokio::time::{Duration, Instant};
|
|
||||||
|
|
||||||
#[derive(Default)]
|
|
||||||
struct WakeCounter {
|
|
||||||
wakes: AtomicUsize,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl std::task::Wake for WakeCounter {
|
|
||||||
fn wake(self: Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn wake_by_ref(self: &Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn quota_test_guard() -> impl Drop {
|
|
||||||
super::quota_user_lock_test_scope()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn saturate_quota_user_locks() -> Vec<Arc<std::sync::Mutex<()>>> {
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX);
|
|
||||||
for idx in 0..QUOTA_USER_LOCKS_MAX {
|
|
||||||
retained.push(quota_user_lock(&format!("quota-retry-bench-saturate-{idx}")));
|
|
||||||
}
|
|
||||||
retained
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn stress_contention_wake_rate_decays_with_backoff_curve() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_quota_user_locks();
|
|
||||||
let user = format!("quota-backoff-bench-{}", std::process::id());
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
|
|
||||||
let lock = quota_user_lock(&user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold quota lock before benchmark run");
|
|
||||||
|
|
||||||
let waiters = 64usize;
|
|
||||||
let mut ios = Vec::with_capacity(waiters);
|
|
||||||
let mut wake_counters = Vec::with_capacity(waiters);
|
|
||||||
|
|
||||||
for _ in 0..waiters {
|
|
||||||
ios.push(StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.clone(),
|
|
||||||
Some(4096),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
));
|
|
||||||
}
|
|
||||||
|
|
||||||
for io in &mut ios {
|
|
||||||
let counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
let pending = Pin::new(io).poll_write(&mut cx, &[0x71]);
|
|
||||||
assert!(pending.is_pending());
|
|
||||||
wake_counters.push(counter);
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut observed = vec![0usize; waiters];
|
|
||||||
let start = Instant::now();
|
|
||||||
let mut wakes_at_40ms = 0usize;
|
|
||||||
let mut wakes_at_160ms = 0usize;
|
|
||||||
|
|
||||||
while start.elapsed() < Duration::from_millis(200) {
|
|
||||||
for (idx, counter) in wake_counters.iter().enumerate() {
|
|
||||||
let wakes = counter.wakes.load(Ordering::Relaxed);
|
|
||||||
if wakes > observed[idx] {
|
|
||||||
observed[idx] = wakes;
|
|
||||||
let waker = Waker::from(Arc::clone(counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
let pending = Pin::new(&mut ios[idx]).poll_write(&mut cx, &[0x72]);
|
|
||||||
assert!(pending.is_pending());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let elapsed = start.elapsed();
|
|
||||||
if elapsed >= Duration::from_millis(40) && wakes_at_40ms == 0 {
|
|
||||||
wakes_at_40ms = wake_counters
|
|
||||||
.iter()
|
|
||||||
.map(|counter| counter.wakes.load(Ordering::Relaxed))
|
|
||||||
.sum();
|
|
||||||
}
|
|
||||||
if elapsed >= Duration::from_millis(160) && wakes_at_160ms == 0 {
|
|
||||||
wakes_at_160ms = wake_counters
|
|
||||||
.iter()
|
|
||||||
.map(|counter| counter.wakes.load(Ordering::Relaxed))
|
|
||||||
.sum();
|
|
||||||
}
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(1)).await;
|
|
||||||
}
|
|
||||||
|
|
||||||
let total_wakes: usize = wake_counters
|
|
||||||
.iter()
|
|
||||||
.map(|counter| counter.wakes.load(Ordering::Relaxed))
|
|
||||||
.sum();
|
|
||||||
|
|
||||||
let wakes_at_200ms = total_wakes;
|
|
||||||
let early_window_wakes = wakes_at_40ms;
|
|
||||||
let late_window_wakes = wakes_at_200ms.saturating_sub(wakes_at_160ms);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
total_wakes <= waiters * 28,
|
|
||||||
"backoff benchmark exceeded wake budget; waiters={waiters}, wakes={total_wakes}"
|
|
||||||
);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
early_window_wakes > 0,
|
|
||||||
"benchmark failed to observe early contention wakes"
|
|
||||||
);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
late_window_wakes * 4 <= early_window_wakes * 3,
|
|
||||||
"wake-rate decay invariant violated; early_0_40ms={early_window_wakes}, late_160_200ms={late_window_wakes}, total={total_wakes}"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn stress_read_contention_wake_rate_decays_with_backoff_curve() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_quota_user_locks();
|
|
||||||
let user = format!("quota-backoff-read-bench-{}", std::process::id());
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
|
|
||||||
let lock = quota_user_lock(&user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold quota lock before read benchmark run");
|
|
||||||
|
|
||||||
let waiters = 64usize;
|
|
||||||
let mut ios = Vec::with_capacity(waiters);
|
|
||||||
let mut wake_counters = Vec::with_capacity(waiters);
|
|
||||||
|
|
||||||
for _ in 0..waiters {
|
|
||||||
ios.push(StatsIo::new(
|
|
||||||
tokio::io::empty(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.clone(),
|
|
||||||
Some(4096),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
));
|
|
||||||
}
|
|
||||||
|
|
||||||
for io in &mut ios {
|
|
||||||
let counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
let mut storage = [0u8; 1];
|
|
||||||
let mut buf = ReadBuf::new(&mut storage);
|
|
||||||
let pending = Pin::new(io).poll_read(&mut cx, &mut buf);
|
|
||||||
assert!(pending.is_pending());
|
|
||||||
wake_counters.push(counter);
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut observed = vec![0usize; waiters];
|
|
||||||
let start = Instant::now();
|
|
||||||
let mut wakes_at_40ms = 0usize;
|
|
||||||
let mut wakes_at_160ms = 0usize;
|
|
||||||
|
|
||||||
while start.elapsed() < Duration::from_millis(200) {
|
|
||||||
for (idx, counter) in wake_counters.iter().enumerate() {
|
|
||||||
let wakes = counter.wakes.load(Ordering::Relaxed);
|
|
||||||
if wakes > observed[idx] {
|
|
||||||
observed[idx] = wakes;
|
|
||||||
let waker = Waker::from(Arc::clone(counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
let mut storage = [0u8; 1];
|
|
||||||
let mut buf = ReadBuf::new(&mut storage);
|
|
||||||
let pending = Pin::new(&mut ios[idx]).poll_read(&mut cx, &mut buf);
|
|
||||||
assert!(pending.is_pending());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
let elapsed = start.elapsed();
|
|
||||||
if elapsed >= Duration::from_millis(40) && wakes_at_40ms == 0 {
|
|
||||||
wakes_at_40ms = wake_counters
|
|
||||||
.iter()
|
|
||||||
.map(|counter| counter.wakes.load(Ordering::Relaxed))
|
|
||||||
.sum();
|
|
||||||
}
|
|
||||||
if elapsed >= Duration::from_millis(160) && wakes_at_160ms == 0 {
|
|
||||||
wakes_at_160ms = wake_counters
|
|
||||||
.iter()
|
|
||||||
.map(|counter| counter.wakes.load(Ordering::Relaxed))
|
|
||||||
.sum();
|
|
||||||
}
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(1)).await;
|
|
||||||
}
|
|
||||||
|
|
||||||
let total_wakes: usize = wake_counters
|
|
||||||
.iter()
|
|
||||||
.map(|counter| counter.wakes.load(Ordering::Relaxed))
|
|
||||||
.sum();
|
|
||||||
|
|
||||||
let wakes_at_200ms = total_wakes;
|
|
||||||
let early_window_wakes = wakes_at_40ms;
|
|
||||||
let late_window_wakes = wakes_at_200ms.saturating_sub(wakes_at_160ms);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
total_wakes <= waiters * 28,
|
|
||||||
"read backoff benchmark exceeded wake budget; waiters={waiters}, wakes={total_wakes}"
|
|
||||||
);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
early_window_wakes > 0,
|
|
||||||
"read benchmark failed to observe early contention wakes"
|
|
||||||
);
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
late_window_wakes * 4 <= early_window_wakes * 3,
|
|
||||||
"read wake-rate decay invariant violated; early_0_40ms={early_window_wakes}, late_160_200ms={late_window_wakes}, total={total_wakes}"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
}
|
|
||||||
|
|
@ -1,339 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use dashmap::DashMap;
|
|
||||||
use std::pin::Pin;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
|
||||||
use std::task::{Context, Waker};
|
|
||||||
use tokio::io::ReadBuf;
|
|
||||||
use tokio::time::{Duration, Instant};
|
|
||||||
|
|
||||||
#[derive(Default)]
|
|
||||||
struct WakeCounter {
|
|
||||||
wakes: AtomicUsize,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl std::task::Wake for WakeCounter {
|
|
||||||
fn wake(self: Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn wake_by_ref(self: &Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn quota_test_guard() -> impl Drop {
|
|
||||||
super::quota_user_lock_test_scope()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn saturate_quota_user_locks() -> Vec<Arc<std::sync::Mutex<()>>> {
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX);
|
|
||||||
for idx in 0..QUOTA_USER_LOCKS_MAX {
|
|
||||||
retained.push(quota_user_lock(&format!("quota-retry-backoff-saturate-{idx}")));
|
|
||||||
}
|
|
||||||
retained
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn positive_uncontended_writer_keeps_retry_wakes_zero() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::clone(&stats),
|
|
||||||
"quota-backoff-positive".to_string(),
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
|
|
||||||
let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x41, 0x42]);
|
|
||||||
assert!(poll.is_ready(), "uncontended writer must complete immediately");
|
|
||||||
assert_eq!(
|
|
||||||
wake_counter.wakes.load(Ordering::Relaxed),
|
|
||||||
0,
|
|
||||||
"uncontended path must not schedule deferred contention wakes"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn adversarial_writer_sustained_contention_executor_repoll_is_rate_limited() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_quota_user_locks();
|
|
||||||
let user = "quota-backoff-adversarial-writer";
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
|
|
||||||
let lock = quota_user_lock(user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold quota lock before polling writer");
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.to_string(),
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
|
|
||||||
let first = Pin::new(&mut io).poll_write(&mut cx, &[0xAA]);
|
|
||||||
assert!(first.is_pending());
|
|
||||||
|
|
||||||
let start = Instant::now();
|
|
||||||
let mut observed = 0usize;
|
|
||||||
while start.elapsed() < Duration::from_millis(80) {
|
|
||||||
let wakes = wake_counter.wakes.load(Ordering::Relaxed);
|
|
||||||
if wakes > observed {
|
|
||||||
observed = wakes;
|
|
||||||
let pending = Pin::new(&mut io).poll_write(&mut cx, &[0xAB]);
|
|
||||||
assert!(pending.is_pending());
|
|
||||||
}
|
|
||||||
tokio::time::sleep(Duration::from_millis(1)).await;
|
|
||||||
}
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
wake_counter.wakes.load(Ordering::Relaxed) <= 16,
|
|
||||||
"sustained contention must be rate limited; observed wakes={} in 80ms",
|
|
||||||
wake_counter.wakes.load(Ordering::Relaxed)
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
let ready = Pin::new(&mut io).poll_write(&mut cx, &[0xAC]);
|
|
||||||
assert!(ready.is_ready());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn adversarial_reader_sustained_contention_executor_repoll_is_rate_limited() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_quota_user_locks();
|
|
||||||
let user = "quota-backoff-adversarial-reader";
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
|
|
||||||
let lock = quota_user_lock(user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold quota lock before polling reader");
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::empty(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.to_string(),
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
let mut storage = [0u8; 1];
|
|
||||||
|
|
||||||
let mut buf = ReadBuf::new(&mut storage);
|
|
||||||
let first = Pin::new(&mut io).poll_read(&mut cx, &mut buf);
|
|
||||||
assert!(first.is_pending());
|
|
||||||
|
|
||||||
let start = Instant::now();
|
|
||||||
let mut observed = 0usize;
|
|
||||||
while start.elapsed() < Duration::from_millis(80) {
|
|
||||||
let wakes = wake_counter.wakes.load(Ordering::Relaxed);
|
|
||||||
if wakes > observed {
|
|
||||||
observed = wakes;
|
|
||||||
let mut next = ReadBuf::new(&mut storage);
|
|
||||||
let pending = Pin::new(&mut io).poll_read(&mut cx, &mut next);
|
|
||||||
assert!(pending.is_pending());
|
|
||||||
}
|
|
||||||
tokio::time::sleep(Duration::from_millis(1)).await;
|
|
||||||
}
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
wake_counter.wakes.load(Ordering::Relaxed) <= 16,
|
|
||||||
"sustained contention must be rate limited; observed wakes={} in 80ms",
|
|
||||||
wake_counter.wakes.load(Ordering::Relaxed)
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
let mut done = ReadBuf::new(&mut storage);
|
|
||||||
let ready = Pin::new(&mut io).poll_read(&mut cx, &mut done);
|
|
||||||
assert!(ready.is_ready());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn edge_backoff_attempt_resets_after_contention_release() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_quota_user_locks();
|
|
||||||
let user = "quota-backoff-edge-reset";
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
|
|
||||||
let lock = quota_user_lock(user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold quota lock before polling writer");
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.to_string(),
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
|
|
||||||
let initial = Pin::new(&mut io).poll_write(&mut cx, &[0x31]);
|
|
||||||
assert!(initial.is_pending());
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(15)).await;
|
|
||||||
let wakes = wake_counter.wakes.load(Ordering::Relaxed);
|
|
||||||
if wakes > 0 {
|
|
||||||
let pending = Pin::new(&mut io).poll_write(&mut cx, &[0x32]);
|
|
||||||
assert!(pending.is_pending());
|
|
||||||
}
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
let ready = Pin::new(&mut io).poll_write(&mut cx, &[0x33]);
|
|
||||||
assert!(ready.is_ready());
|
|
||||||
assert!(
|
|
||||||
!io.quota_write_wake_scheduled,
|
|
||||||
"successful write must clear deferred wake scheduling flag"
|
|
||||||
);
|
|
||||||
assert!(
|
|
||||||
io.quota_write_retry_sleep.is_none(),
|
|
||||||
"successful write must clear deferred sleep slot"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn light_fuzz_writer_repoll_schedule_keeps_wake_budget_bounded() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_quota_user_locks();
|
|
||||||
let user = "quota-backoff-fuzz-writer";
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
|
|
||||||
let lock = quota_user_lock(user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold quota lock before fuzz loop");
|
|
||||||
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.to_string(),
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
|
|
||||||
let mut seed = 0x5EED_CAFE_7788_9900u64;
|
|
||||||
for _ in 0..64 {
|
|
||||||
let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x51]);
|
|
||||||
assert!(poll.is_pending());
|
|
||||||
|
|
||||||
seed ^= seed << 7;
|
|
||||||
seed ^= seed >> 9;
|
|
||||||
seed ^= seed << 8;
|
|
||||||
let sleep_ms = (seed % 4) as u64;
|
|
||||||
tokio::time::sleep(Duration::from_millis(sleep_ms)).await;
|
|
||||||
}
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
wake_counter.wakes.load(Ordering::Relaxed) <= 24,
|
|
||||||
"fuzzed repoll schedule must keep wake budget bounded; observed wakes={}",
|
|
||||||
wake_counter.wakes.load(Ordering::Relaxed)
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn stress_multi_waiter_contention_keeps_global_wake_budget_bounded() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_quota_user_locks();
|
|
||||||
let user = format!("quota-backoff-stress-{}", std::process::id());
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
|
|
||||||
let lock = quota_user_lock(&user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold quota lock before launching stress waiters");
|
|
||||||
|
|
||||||
let waiters = 48usize;
|
|
||||||
let mut ios = Vec::with_capacity(waiters);
|
|
||||||
let mut wake_counters = Vec::with_capacity(waiters);
|
|
||||||
|
|
||||||
for _ in 0..waiters {
|
|
||||||
ios.push(StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.clone(),
|
|
||||||
Some(4096),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
));
|
|
||||||
}
|
|
||||||
|
|
||||||
for io in &mut ios {
|
|
||||||
let counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
let pending = Pin::new(io).poll_write(&mut cx, &[0x61]);
|
|
||||||
assert!(pending.is_pending());
|
|
||||||
wake_counters.push(counter);
|
|
||||||
}
|
|
||||||
|
|
||||||
let start = Instant::now();
|
|
||||||
while start.elapsed() < Duration::from_millis(120) {
|
|
||||||
for (idx, counter) in wake_counters.iter().enumerate() {
|
|
||||||
if counter.wakes.load(Ordering::Relaxed) > 0 {
|
|
||||||
let waker = Waker::from(Arc::clone(counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
let pending = Pin::new(&mut ios[idx]).poll_write(&mut cx, &[0x62]);
|
|
||||||
assert!(pending.is_pending());
|
|
||||||
}
|
|
||||||
}
|
|
||||||
tokio::time::sleep(Duration::from_millis(1)).await;
|
|
||||||
}
|
|
||||||
|
|
||||||
let total_wakes: usize = wake_counters
|
|
||||||
.iter()
|
|
||||||
.map(|counter| counter.wakes.load(Ordering::Relaxed))
|
|
||||||
.sum();
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
total_wakes <= waiters * 20,
|
|
||||||
"stress contention must keep aggregate wake budget bounded; waiters={waiters}, wakes={total_wakes}"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
}
|
|
||||||
|
|
@ -1,246 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use std::pin::Pin;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
|
||||||
use std::task::{Context, Poll, Waker};
|
|
||||||
use tokio::io::{AsyncReadExt, AsyncWriteExt, ReadBuf};
|
|
||||||
use tokio::time::{Duration, timeout};
|
|
||||||
|
|
||||||
#[derive(Default)]
|
|
||||||
struct WakeCounter {
|
|
||||||
wakes: AtomicUsize,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl std::task::Wake for WakeCounter {
|
|
||||||
fn wake(self: Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn wake_by_ref(self: &Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn quota_test_guard() -> impl Drop {
|
|
||||||
super::quota_user_lock_test_scope()
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn positive_uncontended_quota_limited_writer_completes() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::clone(&stats),
|
|
||||||
"tdd-uncontended".to_string(),
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let result = io.write_all(&[0x41, 0x42, 0x43]).await;
|
|
||||||
assert!(result.is_ok(), "uncontended writer must complete");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn adversarial_contended_writers_without_repoll_must_not_wake_storm() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("tdd-writer-storm-{}", std::process::id());
|
|
||||||
let held = quota_user_lock(&user);
|
|
||||||
let _held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold quota lock before polling writers");
|
|
||||||
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let writers = 24usize;
|
|
||||||
let mut ios = Vec::with_capacity(writers);
|
|
||||||
let mut wake_counters = Vec::with_capacity(writers);
|
|
||||||
|
|
||||||
for _ in 0..writers {
|
|
||||||
ios.push(StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.clone(),
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
));
|
|
||||||
}
|
|
||||||
|
|
||||||
for io in &mut ios {
|
|
||||||
let counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
let poll = Pin::new(io).poll_write(&mut cx, &[0xAA]);
|
|
||||||
assert!(poll.is_pending(), "writer must be pending under held lock");
|
|
||||||
wake_counters.push(counter);
|
|
||||||
}
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(25)).await;
|
|
||||||
|
|
||||||
let total_wakes: usize = wake_counters
|
|
||||||
.iter()
|
|
||||||
.map(|counter| counter.wakes.load(Ordering::Relaxed))
|
|
||||||
.sum();
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
total_wakes <= writers * 4,
|
|
||||||
"retry scheduler must remain bounded without repoll; observed wakes={total_wakes}, writers={writers}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn adversarial_contended_readers_without_repoll_must_not_wake_storm() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("tdd-reader-storm-{}", std::process::id());
|
|
||||||
let held = quota_user_lock(&user);
|
|
||||||
let _held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold quota lock before polling readers");
|
|
||||||
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let readers = 24usize;
|
|
||||||
let mut ios = Vec::with_capacity(readers);
|
|
||||||
let mut wake_counters = Vec::with_capacity(readers);
|
|
||||||
|
|
||||||
for _ in 0..readers {
|
|
||||||
ios.push(StatsIo::new(
|
|
||||||
tokio::io::empty(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.clone(),
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
));
|
|
||||||
}
|
|
||||||
|
|
||||||
for io in &mut ios {
|
|
||||||
let counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
let mut storage = [0u8; 1];
|
|
||||||
let mut buf = ReadBuf::new(&mut storage);
|
|
||||||
let poll = Pin::new(io).poll_read(&mut cx, &mut buf);
|
|
||||||
assert!(poll.is_pending(), "reader must be pending under held lock");
|
|
||||||
wake_counters.push(counter);
|
|
||||||
}
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(25)).await;
|
|
||||||
|
|
||||||
let total_wakes: usize = wake_counters
|
|
||||||
.iter()
|
|
||||||
.map(|counter| counter.wakes.load(Ordering::Relaxed))
|
|
||||||
.sum();
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
total_wakes <= readers * 4,
|
|
||||||
"retry scheduler must remain bounded without repoll; observed wakes={total_wakes}, readers={readers}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn integration_contended_waiters_resume_after_lock_release() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let user = format!("tdd-resume-{}", std::process::id());
|
|
||||||
let held = quota_user_lock(&user);
|
|
||||||
let held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold quota lock before launching waiters");
|
|
||||||
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let mut waiters = Vec::new();
|
|
||||||
for _ in 0..12 {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let user = user.clone();
|
|
||||||
waiters.push(tokio::spawn(async move {
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
stats,
|
|
||||||
user,
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
io.write_all(&[0x5A]).await
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(5)).await;
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
timeout(Duration::from_secs(1), async {
|
|
||||||
for waiter in waiters {
|
|
||||||
let result = waiter.await.expect("waiter task must not panic");
|
|
||||||
assert!(result.is_ok(), "waiter must complete after release");
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("all waiters must complete in bounded time");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn light_fuzz_contention_rounds_keep_retry_wakes_bounded() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let mut seed = 0x9E37_79B9_AA55_1234u64;
|
|
||||||
for round in 0..20u32 {
|
|
||||||
seed ^= seed << 7;
|
|
||||||
seed ^= seed >> 9;
|
|
||||||
seed ^= seed << 8;
|
|
||||||
|
|
||||||
let writers = 8 + (seed as usize % 12);
|
|
||||||
let sleep_ms = 10 + (seed as u64 % 15);
|
|
||||||
let user = format!("tdd-fuzz-{}-{round}", std::process::id());
|
|
||||||
|
|
||||||
let held = quota_user_lock(&user);
|
|
||||||
let _held_guard = held
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold quota lock in fuzz round");
|
|
||||||
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let mut ios = Vec::with_capacity(writers);
|
|
||||||
let mut wake_counters = Vec::with_capacity(writers);
|
|
||||||
|
|
||||||
for _ in 0..writers {
|
|
||||||
ios.push(StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
Arc::new(SharedCounters::new()),
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.clone(),
|
|
||||||
Some(2048),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
));
|
|
||||||
}
|
|
||||||
|
|
||||||
for io in &mut ios {
|
|
||||||
let counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
let poll = Pin::new(io).poll_write(&mut cx, &[0x7A]);
|
|
||||||
assert!(matches!(poll, Poll::Pending));
|
|
||||||
wake_counters.push(counter);
|
|
||||||
}
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(sleep_ms)).await;
|
|
||||||
|
|
||||||
let total_wakes: usize = wake_counters
|
|
||||||
.iter()
|
|
||||||
.map(|counter| counter.wakes.load(Ordering::Relaxed))
|
|
||||||
.sum();
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
total_wakes <= writers * 4,
|
|
||||||
"fuzz round must keep wakes bounded; round={round}, writers={writers}, wakes={total_wakes}, sleep_ms={sleep_ms}"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
@ -1,294 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use dashmap::DashMap;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::AtomicBool;
|
|
||||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
||||||
use tokio::sync::Barrier;
|
|
||||||
use tokio::time::{Duration, timeout};
|
|
||||||
|
|
||||||
fn saturate_lock_cache() -> Vec<Arc<std::sync::Mutex<()>>> {
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX);
|
|
||||||
for idx in 0..QUOTA_USER_LOCKS_MAX {
|
|
||||||
retained.push(quota_user_lock(&format!("quota-liveness-saturated-{idx}")));
|
|
||||||
}
|
|
||||||
retained
|
|
||||||
}
|
|
||||||
|
|
||||||
fn quota_test_guard() -> impl Drop {
|
|
||||||
super::quota_user_lock_test_scope()
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn positive_writer_progresses_after_contention_release_without_external_wake() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_lock_cache();
|
|
||||||
let user = "quota-liveness-writer-positive";
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
|
|
||||||
let lock = quota_user_lock(user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold user quota lock before write");
|
|
||||||
|
|
||||||
let counters = Arc::new(SharedCounters::new());
|
|
||||||
let quota_exceeded = Arc::new(AtomicBool::new(false));
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
counters,
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.to_string(),
|
|
||||||
Some(1024),
|
|
||||||
quota_exceeded,
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let writer = tokio::spawn(async move { io.write_all(&[0x11]).await });
|
|
||||||
|
|
||||||
// Let the initial deferred wake fire while contention is still active.
|
|
||||||
tokio::time::sleep(Duration::from_millis(4)).await;
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
let completed = timeout(Duration::from_millis(250), writer)
|
|
||||||
.await
|
|
||||||
.expect("writer must be re-polled and complete after lock release")
|
|
||||||
.expect("writer task must not panic");
|
|
||||||
assert!(completed.is_ok(), "writer must complete after lock release");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn edge_reader_progresses_after_contention_release_without_external_wake() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_lock_cache();
|
|
||||||
let user = "quota-liveness-reader-edge";
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
|
|
||||||
let lock = quota_user_lock(user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold user quota lock before read");
|
|
||||||
|
|
||||||
let counters = Arc::new(SharedCounters::new());
|
|
||||||
let quota_exceeded = Arc::new(AtomicBool::new(false));
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::empty(),
|
|
||||||
counters,
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.to_string(),
|
|
||||||
Some(1024),
|
|
||||||
quota_exceeded,
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let reader = tokio::spawn(async move {
|
|
||||||
let mut one = [0u8; 1];
|
|
||||||
io.read(&mut one).await
|
|
||||||
});
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(4)).await;
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
let completed = timeout(Duration::from_millis(250), reader)
|
|
||||||
.await
|
|
||||||
.expect("reader must be re-polled and complete after lock release")
|
|
||||||
.expect("reader task must not panic");
|
|
||||||
assert!(completed.is_ok(), "reader must complete after lock release");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn adversarial_early_deferred_wake_consumption_does_not_deadlock_writer() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_lock_cache();
|
|
||||||
let user = "quota-liveness-adversarial";
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
|
|
||||||
let lock = quota_user_lock(user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold user quota lock before adversarial write");
|
|
||||||
|
|
||||||
let counters = Arc::new(SharedCounters::new());
|
|
||||||
let quota_exceeded = Arc::new(AtomicBool::new(false));
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
counters,
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.to_string(),
|
|
||||||
Some(1024),
|
|
||||||
quota_exceeded,
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let writer = tokio::spawn(async move { io.write_all(&[0x22]).await });
|
|
||||||
|
|
||||||
// Force multiple scheduler rounds while lock remains held so the first
|
|
||||||
// deferred wake has already been consumed under contention.
|
|
||||||
for _ in 0..32 {
|
|
||||||
tokio::task::yield_now().await;
|
|
||||||
}
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
let completed = timeout(Duration::from_millis(300), writer)
|
|
||||||
.await
|
|
||||||
.expect("writer must not stay parked forever after release")
|
|
||||||
.expect("writer task must not panic");
|
|
||||||
assert!(completed.is_ok());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn integration_parallel_waiters_resume_after_single_release_event() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_lock_cache();
|
|
||||||
let user = format!("quota-liveness-integration-{}", std::process::id());
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let barrier = Arc::new(Barrier::new(13));
|
|
||||||
|
|
||||||
let lock = quota_user_lock(&user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold user quota lock before launching waiters");
|
|
||||||
|
|
||||||
let mut waiters = Vec::new();
|
|
||||||
for _ in 0..12 {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let user = user.clone();
|
|
||||||
let barrier = Arc::clone(&barrier);
|
|
||||||
waiters.push(tokio::spawn(async move {
|
|
||||||
let counters = Arc::new(SharedCounters::new());
|
|
||||||
let quota_exceeded = Arc::new(AtomicBool::new(false));
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
counters,
|
|
||||||
stats,
|
|
||||||
user,
|
|
||||||
Some(4096),
|
|
||||||
quota_exceeded,
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
barrier.wait().await;
|
|
||||||
io.write_all(&[0x33]).await
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
barrier.wait().await;
|
|
||||||
tokio::time::sleep(Duration::from_millis(4)).await;
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
timeout(Duration::from_secs(1), async {
|
|
||||||
for waiter in waiters {
|
|
||||||
let outcome = waiter.await.expect("waiter must not panic");
|
|
||||||
assert!(
|
|
||||||
outcome.is_ok(),
|
|
||||||
"waiter must resume and complete after release"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("all waiters must complete in bounded time");
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn light_fuzz_release_timing_matrix_preserves_liveness() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_lock_cache();
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
|
|
||||||
let mut seed = 0xD1CE_F00D_0123_4567u64;
|
|
||||||
for round in 0..64u32 {
|
|
||||||
seed ^= seed << 7;
|
|
||||||
seed ^= seed >> 9;
|
|
||||||
seed ^= seed << 8;
|
|
||||||
|
|
||||||
let delay_ms = 1 + (seed & 0x7) as u64;
|
|
||||||
let user = format!("quota-liveness-fuzz-{}-{round}", std::process::id());
|
|
||||||
|
|
||||||
let lock = quota_user_lock(&user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold user quota lock in fuzz round");
|
|
||||||
|
|
||||||
let counters = Arc::new(SharedCounters::new());
|
|
||||||
let quota_exceeded = Arc::new(AtomicBool::new(false));
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
counters,
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user,
|
|
||||||
Some(2048),
|
|
||||||
quota_exceeded,
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let writer = tokio::spawn(async move { io.write_all(&[0x44]).await });
|
|
||||||
|
|
||||||
tokio::time::sleep(Duration::from_millis(delay_ms)).await;
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
let done = timeout(Duration::from_millis(300), writer)
|
|
||||||
.await
|
|
||||||
.expect("fuzz round writer must complete")
|
|
||||||
.expect("fuzz writer task must not panic");
|
|
||||||
assert!(
|
|
||||||
done.is_ok(),
|
|
||||||
"fuzz round writer must not stall after release"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
async fn stress_repeated_contention_cycles_remain_live() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_lock_cache();
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
|
|
||||||
for cycle in 0..40u32 {
|
|
||||||
let user = format!("quota-liveness-stress-{}-{cycle}", std::process::id());
|
|
||||||
let lock = quota_user_lock(&user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold lock before stress cycle");
|
|
||||||
|
|
||||||
let mut tasks = Vec::new();
|
|
||||||
for _ in 0..6 {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let user = user.clone();
|
|
||||||
tasks.push(tokio::spawn(async move {
|
|
||||||
let counters = Arc::new(SharedCounters::new());
|
|
||||||
let quota_exceeded = Arc::new(AtomicBool::new(false));
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
counters,
|
|
||||||
stats,
|
|
||||||
user,
|
|
||||||
Some(2048),
|
|
||||||
quota_exceeded,
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
io.write_all(&[0x55]).await
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
tokio::task::yield_now().await;
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
timeout(Duration::from_millis(700), async {
|
|
||||||
for task in tasks {
|
|
||||||
let outcome = task.await.expect("stress task must not panic");
|
|
||||||
assert!(outcome.is_ok(), "stress writer must complete");
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("stress cycle must finish in bounded time");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
@ -1,310 +0,0 @@
|
||||||
use super::*;
|
|
||||||
use crate::stats::Stats;
|
|
||||||
use dashmap::DashMap;
|
|
||||||
use std::pin::Pin;
|
|
||||||
use std::sync::Arc;
|
|
||||||
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
|
||||||
use std::task::{Context, Waker};
|
|
||||||
use tokio::io::{AsyncWriteExt, ReadBuf};
|
|
||||||
use tokio::time::{Duration, timeout};
|
|
||||||
|
|
||||||
#[derive(Default)]
|
|
||||||
struct WakeCounter {
|
|
||||||
wakes: AtomicUsize,
|
|
||||||
}
|
|
||||||
|
|
||||||
impl std::task::Wake for WakeCounter {
|
|
||||||
fn wake(self: Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
|
|
||||||
fn wake_by_ref(self: &Arc<Self>) {
|
|
||||||
self.wakes.fetch_add(1, Ordering::Relaxed);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn quota_test_guard() -> impl Drop {
|
|
||||||
super::quota_user_lock_test_scope()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn saturate_quota_user_locks() -> Vec<Arc<std::sync::Mutex<()>>> {
|
|
||||||
let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new);
|
|
||||||
map.clear();
|
|
||||||
|
|
||||||
let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX);
|
|
||||||
for idx in 0..QUOTA_USER_LOCKS_MAX {
|
|
||||||
retained.push(quota_user_lock(&format!("quota-waker-saturate-{idx}")));
|
|
||||||
}
|
|
||||||
retained
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn positive_contended_writer_emits_deferred_wake_for_liveness() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_quota_user_locks();
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let user = "quota-waker-positive-user";
|
|
||||||
|
|
||||||
let lock = quota_user_lock(user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold overflow lock before polling writer");
|
|
||||||
|
|
||||||
let counters = Arc::new(SharedCounters::new());
|
|
||||||
let quota_exceeded = Arc::new(AtomicBool::new(false));
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
counters,
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.to_string(),
|
|
||||||
Some(1024),
|
|
||||||
quota_exceeded,
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
|
|
||||||
let pending = Pin::new(&mut io).poll_write(&mut cx, &[0xA1]);
|
|
||||||
assert!(pending.is_pending());
|
|
||||||
|
|
||||||
timeout(Duration::from_millis(100), async {
|
|
||||||
loop {
|
|
||||||
if wake_counter.wakes.load(Ordering::Relaxed) >= 1 {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
tokio::task::yield_now().await;
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("contended writer must receive deferred wake");
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
let ready = Pin::new(&mut io).poll_write(&mut cx, &[0xA2]);
|
|
||||||
assert!(
|
|
||||||
ready.is_ready(),
|
|
||||||
"writer must progress after contention release"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn adversarial_blackhat_writer_contention_does_not_create_waker_storm() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_quota_user_locks();
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let user = "quota-waker-blackhat-writer";
|
|
||||||
|
|
||||||
let lock = quota_user_lock(user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold overflow lock before polling writer");
|
|
||||||
|
|
||||||
let counters = Arc::new(SharedCounters::new());
|
|
||||||
let quota_exceeded = Arc::new(AtomicBool::new(false));
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
counters,
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.to_string(),
|
|
||||||
Some(1024),
|
|
||||||
quota_exceeded,
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
|
|
||||||
for _ in 0..512 {
|
|
||||||
let poll = Pin::new(&mut io).poll_write(&mut cx, &[0xBE]);
|
|
||||||
assert!(
|
|
||||||
poll.is_pending(),
|
|
||||||
"writer must stay pending while lock is held"
|
|
||||||
);
|
|
||||||
tokio::task::yield_now().await;
|
|
||||||
}
|
|
||||||
|
|
||||||
let wakes = wake_counter.wakes.load(Ordering::Relaxed);
|
|
||||||
assert!(
|
|
||||||
wakes <= 128,
|
|
||||||
"pending writer retries must not trigger wake storm; observed wakes={wakes}"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
let ready = Pin::new(&mut io).poll_write(&mut cx, &[0xEF]);
|
|
||||||
assert!(ready.is_ready());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn edge_read_path_contention_keeps_wake_budget_bounded() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_quota_user_locks();
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let user = "quota-waker-read-edge";
|
|
||||||
|
|
||||||
let lock = quota_user_lock(user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold overflow lock before polling reader");
|
|
||||||
|
|
||||||
let counters = Arc::new(SharedCounters::new());
|
|
||||||
let quota_exceeded = Arc::new(AtomicBool::new(false));
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::empty(),
|
|
||||||
counters,
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.to_string(),
|
|
||||||
Some(1024),
|
|
||||||
quota_exceeded,
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
let mut storage = [0u8; 1];
|
|
||||||
|
|
||||||
for _ in 0..512 {
|
|
||||||
let mut buf = ReadBuf::new(&mut storage);
|
|
||||||
let poll = Pin::new(&mut io).poll_read(&mut cx, &mut buf);
|
|
||||||
assert!(poll.is_pending());
|
|
||||||
tokio::task::yield_now().await;
|
|
||||||
}
|
|
||||||
|
|
||||||
let wakes = wake_counter.wakes.load(Ordering::Relaxed);
|
|
||||||
assert!(
|
|
||||||
wakes <= 128,
|
|
||||||
"pending reader retries must not trigger wake storm; observed wakes={wakes}"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
let mut buf = ReadBuf::new(&mut storage);
|
|
||||||
let ready = Pin::new(&mut io).poll_read(&mut cx, &mut buf);
|
|
||||||
assert!(ready.is_ready());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn light_fuzz_mixed_poll_schedule_under_contention_stays_bounded() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_quota_user_locks();
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
let user = "quota-waker-fuzz-user";
|
|
||||||
|
|
||||||
let lock = quota_user_lock(user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold overflow lock before fuzz polling");
|
|
||||||
|
|
||||||
let counters_w = Arc::new(SharedCounters::new());
|
|
||||||
let mut writer_io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
counters_w,
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.to_string(),
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let counters_r = Arc::new(SharedCounters::new());
|
|
||||||
let mut reader_io = StatsIo::new(
|
|
||||||
tokio::io::empty(),
|
|
||||||
counters_r,
|
|
||||||
Arc::clone(&stats),
|
|
||||||
user.to_string(),
|
|
||||||
Some(1024),
|
|
||||||
Arc::new(AtomicBool::new(false)),
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
let wake_counter = Arc::new(WakeCounter::default());
|
|
||||||
let waker = Waker::from(Arc::clone(&wake_counter));
|
|
||||||
let mut cx = Context::from_waker(&waker);
|
|
||||||
let mut seed = 0xBADC_0FFE_EE11_2211u64;
|
|
||||||
let mut storage = [0u8; 1];
|
|
||||||
|
|
||||||
for _ in 0..1024 {
|
|
||||||
seed ^= seed << 7;
|
|
||||||
seed ^= seed >> 9;
|
|
||||||
seed ^= seed << 8;
|
|
||||||
|
|
||||||
if (seed & 1) == 0 {
|
|
||||||
let poll = Pin::new(&mut writer_io).poll_write(&mut cx, &[0x44]);
|
|
||||||
assert!(poll.is_pending());
|
|
||||||
} else {
|
|
||||||
let mut buf = ReadBuf::new(&mut storage);
|
|
||||||
let poll = Pin::new(&mut reader_io).poll_read(&mut cx, &mut buf);
|
|
||||||
assert!(poll.is_pending());
|
|
||||||
}
|
|
||||||
tokio::task::yield_now().await;
|
|
||||||
}
|
|
||||||
|
|
||||||
assert!(
|
|
||||||
wake_counter.wakes.load(Ordering::Relaxed) <= 192,
|
|
||||||
"mixed contention fuzz must keep deferred wake count tightly bounded"
|
|
||||||
);
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
let ready_w = Pin::new(&mut writer_io).poll_write(&mut cx, &[0x55]);
|
|
||||||
assert!(ready_w.is_ready());
|
|
||||||
|
|
||||||
let mut buf = ReadBuf::new(&mut storage);
|
|
||||||
let ready_r = Pin::new(&mut reader_io).poll_read(&mut cx, &mut buf);
|
|
||||||
assert!(ready_r.is_ready());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
|
||||||
#[ignore = "red-team detector: reveals possible starvation if deferred wake fires before contention release"]
|
|
||||||
async fn stress_many_contended_writers_complete_after_release() {
|
|
||||||
let _guard = quota_test_guard();
|
|
||||||
|
|
||||||
let _retained = saturate_quota_user_locks();
|
|
||||||
let user = "quota-waker-stress-user".to_string();
|
|
||||||
let stats = Arc::new(Stats::new());
|
|
||||||
|
|
||||||
let lock = quota_user_lock(&user);
|
|
||||||
let held_guard = lock
|
|
||||||
.try_lock()
|
|
||||||
.expect("test must hold overflow lock before launching contended tasks");
|
|
||||||
|
|
||||||
let mut tasks = Vec::new();
|
|
||||||
for _ in 0..32 {
|
|
||||||
let stats = Arc::clone(&stats);
|
|
||||||
let user = user.clone();
|
|
||||||
tasks.push(tokio::spawn(async move {
|
|
||||||
let counters = Arc::new(SharedCounters::new());
|
|
||||||
let quota_exceeded = Arc::new(AtomicBool::new(false));
|
|
||||||
let mut io = StatsIo::new(
|
|
||||||
tokio::io::sink(),
|
|
||||||
counters,
|
|
||||||
stats,
|
|
||||||
user,
|
|
||||||
Some(2048),
|
|
||||||
quota_exceeded,
|
|
||||||
tokio::time::Instant::now(),
|
|
||||||
);
|
|
||||||
|
|
||||||
io.write_all(&[0xAA]).await
|
|
||||||
}));
|
|
||||||
}
|
|
||||||
|
|
||||||
for _ in 0..8 {
|
|
||||||
tokio::task::yield_now().await;
|
|
||||||
}
|
|
||||||
|
|
||||||
drop(held_guard);
|
|
||||||
|
|
||||||
timeout(Duration::from_secs(2), async {
|
|
||||||
for task in tasks {
|
|
||||||
let result = task.await.expect("stress task must not panic");
|
|
||||||
assert!(result.is_ok(), "task must complete after lock release");
|
|
||||||
}
|
|
||||||
})
|
|
||||||
.await
|
|
||||||
.expect("all contended writer tasks must finish in bounded time after release");
|
|
||||||
}
|
|
||||||
File diff suppressed because it is too large
Load Diff
Loading…
Reference in New Issue