Line data    Source code

       1              : use std::sync::Arc;
       2              : 
       3              : use rand::Rng;
       4              : 
       5              : use crate::metrics::MISC_OPERATION_SECONDS;
       6              : 
       7              : /// Global rate limiter for background tasks.
       8              : #[derive(Clone)]
       9              : pub struct RateLimiter {
      10              :     partial_backup: Arc<tokio::sync::Semaphore>,
      11              :     eviction: Arc<tokio::sync::Semaphore>,
      12              : }
      13              : 
      14              : impl RateLimiter {
      15              :     /// Create a new rate limiter.
      16              :     /// - `partial_backup_max`: maximum number of concurrent partial backups.
      17              :     /// - `eviction_max`: maximum number of concurrent timeline evictions.
      18            0 :     pub fn new(partial_backup_max: usize, eviction_max: usize) -> Self {
      19            0 :         Self {
      20            0 :             partial_backup: Arc::new(tokio::sync::Semaphore::new(partial_backup_max)),
      21            0 :             eviction: Arc::new(tokio::sync::Semaphore::new(eviction_max)),
      22            0 :         }
      23            0 :     }
      24              : 
      25              :     /// Get a permit for partial backup. This will block if the maximum number of concurrent
      26              :     /// partial backups is reached.
      27            0 :     pub async fn acquire_partial_backup(&self) -> tokio::sync::OwnedSemaphorePermit {
      28            0 :         let _timer = MISC_OPERATION_SECONDS
      29            0 :             .with_label_values(&["partial_permit_acquire"])
      30            0 :             .start_timer();
      31            0 :         self.partial_backup
      32            0 :             .clone()
      33            0 :             .acquire_owned()
      34            0 :             .await
      35            0 :             .expect("semaphore is closed")
      36            0 :     }
      37              : 
      38              :     /// Try to get a permit for timeline eviction. This will return None if the maximum number of
      39              :     /// concurrent timeline evictions is reached.
      40            0 :     pub fn try_acquire_eviction(&self) -> Option<tokio::sync::OwnedSemaphorePermit> {
      41            0 :         self.eviction.clone().try_acquire_owned().ok()
      42            0 :     }
      43              : }
      44              : 
      45              : /// Generate a random duration that is a fraction of the given duration.
      46            0 : pub fn rand_duration(duration: &std::time::Duration) -> std::time::Duration {
      47            0 :     let randf64 = rand::thread_rng().gen_range(0.0..1.0);
      48            0 :     duration.mul_f64(randf64)
      49            0 : }