This commit is contained in:
Alexey
2026-09-25 19:46:31 +03:00
parent 2d63fcf376
commit 326c0ecdb9
128 changed files with 914 additions and 1248 deletions
+12 -25
View File
@@ -67,12 +67,8 @@ impl DirectionBucket {
if self.should_force_reserve_failure() {
return Err(current);
}
self.state.compare_exchange(
current,
next,
Ordering::Relaxed,
Ordering::Relaxed,
)
self.state
.compare_exchange(current, next, Ordering::Relaxed, Ordering::Relaxed)
}
#[inline(always)]
@@ -81,12 +77,8 @@ impl DirectionBucket {
if self.should_force_refund_failure() {
return Err(current);
}
self.state.compare_exchange(
current,
next,
Ordering::Relaxed,
Ordering::Relaxed,
)
self.state
.compare_exchange(current, next, Ordering::Relaxed, Ordering::Relaxed)
}
fn unpack(state: u64) -> (u64, u64) {
@@ -355,9 +347,9 @@ impl CidrDirectionBucket {
});
};
let user_granted = user_debit.granted();
let Some(aggregate_debit) = self
.used
.try_reserve_at(epoch, cap_epoch, user_granted, budget)?
let Some(aggregate_debit) =
self.used
.try_reserve_at(epoch, cap_epoch, user_granted, budget)?
else {
return Ok(CidrReservation {
granted: 0,
@@ -410,12 +402,8 @@ impl CidrUserDirectionState {
if observed_epoch > epoch {
return Err(BucketReserveError::StaleEpoch);
}
let Some(mut active_debit) = active_users.try_reserve_at(
epoch,
PACKED_USAGE_MASK,
1,
budget,
)?
let Some(mut active_debit) =
active_users.try_reserve_at(epoch, PACKED_USAGE_MASK, 1, budget)?
else {
return Ok(false);
};
@@ -532,10 +520,9 @@ impl CidrBucket {
}
let cap_epoch = bytes_per_epoch(cap_bps);
match direction {
RateDirection::Up => {
self.up
.try_reserve(&share.up, epoch, cap_epoch, requested, budget)
}
RateDirection::Up => self
.up
.try_reserve(&share.up, epoch, cap_epoch, requested, budget),
RateDirection::Down => {
self.down
.try_reserve(&share.down, epoch, cap_epoch, requested, budget)
+25 -28
View File
@@ -61,32 +61,28 @@ impl TrafficLease {
let mut granted = requested;
let mut user_debit = None;
if let Some(user_bucket) = binding.user_bucket.as_ref() {
let user_reservation = match user_bucket.try_reserve(
direction,
epoch,
granted,
&mut budget,
) {
Ok(reservation) => reservation,
Err(error) => {
if error.exhausted_reserve_budget() {
self.limiter
.user_scope
.reserve_cas_retry_exhausted(direction);
let user_reservation =
match user_bucket.try_reserve(direction, epoch, granted, &mut budget) {
Ok(reservation) => reservation,
Err(error) => {
if error.exhausted_reserve_budget() {
self.limiter
.user_scope
.reserve_cas_retry_exhausted(direction);
}
return TrafficReservation {
result: TrafficConsumeResult {
granted: 0,
blocked_user: false,
blocked_cidr: false,
},
_binding: binding,
user: None,
cidr: None,
cidr_user: None,
};
}
return TrafficReservation {
result: TrafficConsumeResult {
granted: 0,
blocked_user: false,
blocked_cidr: false,
},
_binding: binding,
user: None,
cidr: None,
cidr_user: None,
};
}
};
};
user_debit = user_reservation.debit;
if user_reservation.granted == 0 {
self.limiter.observe_throttle(direction, true, false);
@@ -107,9 +103,10 @@ impl TrafficLease {
let mut cidr_debit = None;
let mut cidr_user_debit = None;
if let (Some(cidr_bucket), Some(cidr_user_share)) =
(binding.cidr_bucket.as_ref(), binding.cidr_user_share.as_ref())
{
if let (Some(cidr_bucket), Some(cidr_user_share)) = (
binding.cidr_bucket.as_ref(),
binding.cidr_user_share.as_ref(),
) {
let cidr_reservation = match cidr_bucket.try_reserve_for_user(
direction,
cidr_user_share,
+32 -8
View File
@@ -330,14 +330,38 @@ impl TrafficLimiter {
cidr_refund_down,
] = values;
for (counter, value) in [
(&self.user_scope.contention_up.reserve_exhausted_total, user_reserve_up),
(&self.user_scope.contention_down.reserve_exhausted_total, user_reserve_down),
(&self.user_scope.contention_up.refund_exhausted_total, user_refund_up),
(&self.user_scope.contention_down.refund_exhausted_total, user_refund_down),
(&self.cidr_scope.contention_up.reserve_exhausted_total, cidr_reserve_up),
(&self.cidr_scope.contention_down.reserve_exhausted_total, cidr_reserve_down),
(&self.cidr_scope.contention_up.refund_exhausted_total, cidr_refund_up),
(&self.cidr_scope.contention_down.refund_exhausted_total, cidr_refund_down),
(
&self.user_scope.contention_up.reserve_exhausted_total,
user_reserve_up,
),
(
&self.user_scope.contention_down.reserve_exhausted_total,
user_reserve_down,
),
(
&self.user_scope.contention_up.refund_exhausted_total,
user_refund_up,
),
(
&self.user_scope.contention_down.refund_exhausted_total,
user_refund_down,
),
(
&self.cidr_scope.contention_up.reserve_exhausted_total,
cidr_reserve_up,
),
(
&self.cidr_scope.contention_down.reserve_exhausted_total,
cidr_reserve_down,
),
(
&self.cidr_scope.contention_up.refund_exhausted_total,
cidr_refund_up,
),
(
&self.cidr_scope.contention_down.refund_exhausted_total,
cidr_refund_down,
),
] {
counter.store(value, Ordering::Relaxed);
}
+2 -11
View File
@@ -65,12 +65,7 @@ fn reserve_at(
cap: u64,
requested: u64,
) -> Result<Option<DirectionDebit>, BucketReserveError> {
bucket.try_reserve_at(
epoch,
cap,
requested,
&mut ReserveCasBudget::new(),
)
bucket.try_reserve_at(epoch, cap, requested, &mut ReserveCasBudget::new())
}
#[test]
@@ -346,11 +341,7 @@ fn concurrent_first_use_counts_one_active_cidr_user() {
let barrier = Arc::clone(&barrier);
threads.push(std::thread::spawn(move || {
barrier.wait();
user.ensure_active(
13,
&bucket.active_users,
&mut ReserveCasBudget::new(),
)
user.ensure_active(13, &bucket.active_users, &mut ReserveCasBudget::new())
}));
}
let results: Vec<_> = threads
@@ -9,10 +9,7 @@ fn reserve_stops_after_the_attempt_limit() {
let reservation = bucket.try_reserve_at(1, 100, 1, &mut budget);
assert!(matches!(
reservation,
Err(BucketReserveError::Contended)
));
assert!(matches!(reservation, Err(BucketReserveError::Contended)));
assert_eq!(
bucket.reserve_cas_attempts(),
RESERVE_CAS_ATTEMPT_LIMIT as u64
@@ -32,7 +29,10 @@ fn reserve_succeeds_on_the_last_allowed_attempt() {
.unwrap();
assert_eq!(debit.commit_all(), 80);
assert_eq!(bucket.reserve_cas_attempts(), RESERVE_CAS_ATTEMPT_LIMIT as u64);
assert_eq!(
bucket.reserve_cas_attempts(),
RESERVE_CAS_ATTEMPT_LIMIT as u64
);
assert!(budget.is_exhausted());
assert_eq!(bucket.used_at(1), Some(80));
}
@@ -121,15 +121,7 @@ fn lease_contention_is_not_reported_as_throttling() {
let lease = limiter
.acquire_lease("alice", "203.0.113.7".parse().unwrap())
.unwrap();
let bucket = Arc::clone(
&lease
.binding
.load_full()
.user_bucket
.as_ref()
.unwrap()
.down,
);
let bucket = Arc::clone(&lease.binding.load_full().user_bucket.as_ref().unwrap().down);
bucket.force_reserve_failures(RESERVE_CAS_ATTEMPT_LIMIT);
let result = lease.try_consume(RateDirection::Down, 1);
@@ -187,12 +179,7 @@ fn cidr_contention_rolls_back_provisional_user_debits() {
assert!(!result.blocked_user);
assert!(!result.blocked_cidr);
assert_eq!(
binding
.user_bucket
.as_ref()
.unwrap()
.down
.used_at(epoch),
binding.user_bucket.as_ref().unwrap().down.used_at(epoch),
Some(0)
);
assert_eq!(cidr_bucket.down.used.used_at(epoch), None);
@@ -275,8 +262,7 @@ fn contention_snapshot_preserves_scope_direction_and_operation() {
fn cidr_activation_consumes_one_shared_attempt_budget() {
let bucket = CidrDirectionBucket::default();
let user = CidrUserDirectionState::default();
user.used
.force_reserve_failures(RESERVE_CAS_ATTEMPT_LIMIT);
user.used.force_reserve_failures(RESERVE_CAS_ATTEMPT_LIMIT);
let mut budget = ReserveCasBudget::new();
let activation = user.ensure_active(13, &bucket.active_users, &mut budget);
@@ -360,19 +346,11 @@ fn cidr_first_grants_preserve_the_current_soft_fair_share() {
let first = CidrUserDirectionState::default();
let second = CidrUserDirectionState::default();
assert_eq!(
first.ensure_active(
17,
&bucket.active_users,
&mut ReserveCasBudget::new(),
),
first.ensure_active(17, &bucket.active_users, &mut ReserveCasBudget::new(),),
Ok(true)
);
assert_eq!(
second.ensure_active(
17,
&bucket.active_users,
&mut ReserveCasBudget::new(),
),
second.ensure_active(17, &bucket.active_users, &mut ReserveCasBudget::new(),),
Ok(true)
);
@@ -391,21 +369,13 @@ fn cidr_first_grants_preserve_the_current_soft_fair_share() {
.as_mut()
.unwrap()
.commit_all();
first_reservation
.user_debit
.as_mut()
.unwrap()
.commit_all();
first_reservation.user_debit.as_mut().unwrap().commit_all();
second_reservation
.aggregate_debit
.as_mut()
.unwrap()
.commit_all();
second_reservation
.user_debit
.as_mut()
.unwrap()
.commit_all();
second_reservation.user_debit.as_mut().unwrap().commit_all();
assert_eq!(bucket.used.used_at(17), Some(100));
assert_eq!(first.used.used_at(17), Some(50));
assert_eq!(second.used.used_at(17), Some(50));