1 // Copyright 2022 The ChromiumOS Authors 2 // Use of this source code is governed by a BSD-style license that can be 3 // found in the LICENSE file. 4 5 //! Utility file to provide a ratelimit object 6 7 #![cfg_attr(not(any(target_os = "android", target_os = "linux")), allow(dead_code))] 8 9 use std::cmp; 10 use std::time::Duration; 11 use std::time::Instant; 12 13 const BUS_LOCK_SLICE_TIME: u32 = 1000000000; 14 15 pub struct Ratelimit { 16 slice_start_time: Instant, 17 slice_end_time: Instant, 18 slice_quota: u64, 19 slice_ns: Duration, 20 dispatched: u64, 21 } 22 23 impl Ratelimit { new() -> Self24 pub fn new() -> Self { 25 Ratelimit { 26 slice_start_time: Instant::now(), 27 slice_end_time: Instant::now(), 28 slice_quota: 0, 29 slice_ns: Duration::new(0, BUS_LOCK_SLICE_TIME), 30 dispatched: 0, 31 } 32 } ratelimit_set_speed(&mut self, speed: u64)33 pub fn ratelimit_set_speed(&mut self, speed: u64) { 34 if speed == 0 { 35 self.slice_quota = 0; 36 } else { 37 self.slice_quota = cmp::max( 38 speed * self.slice_ns.as_nanos() as u64 / BUS_LOCK_SLICE_TIME as u64, 39 1, 40 ); 41 } 42 } 43 /// Calculate and return delay for next request in ns 44 /// Record that we sent n data units (where n matches the scale chosen 45 /// during ratelimit_set_speed).If we may send more data units in the 46 /// current time slice, return 0 (i.e. no delay). 47 /// Otherwise return the amount of time (in ns) until the start of 48 /// the next time slice that will permit sending the next chunk of data. 49 /// 50 /// Recording sent data units even after exceeding the quota is permitted; 51 /// the time slice will be extended accordingly. ratelimit_calculate_delay(&mut self, n: u64) -> u6452 pub fn ratelimit_calculate_delay(&mut self, n: u64) -> u64 { 53 let now: Instant = Instant::now(); 54 if self.slice_quota == 0 { 55 return 0; 56 } 57 if self.slice_end_time < now { 58 // Previous, possibly extended, time slice finished; reset the accounting. 59 self.slice_start_time = now; 60 self.slice_end_time = now + self.slice_ns; 61 self.dispatched = 0; 62 } 63 64 self.dispatched += n; 65 if self.dispatched < self.slice_quota { 66 // We may send further data within the current time slice, 67 // no need to delay the next request. 68 return 0; 69 } 70 // Quota exceeded. Wait based on the excess amount and then start a new slice. 71 let delay_slices: f64 = self.dispatched as f64 / self.slice_quota as f64; 72 self.slice_end_time = self.slice_start_time + delay_slices as u32 * self.slice_ns; 73 (self.slice_end_time - now).as_nanos() as u64 74 } 75 } 76