mirror of https://github.com/telemt/telemt.git
136 lines
4.0 KiB
Rust
136 lines
4.0 KiB
Rust
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"
|
|
);
|
|
}
|