Line data    Source code

       1              : use std::collections::HashMap;
       2              : use std::sync::Arc;
       3              : 
       4              : use utils::id::TenantTimelineId;
       5              : 
       6              : use crate::timeline::Timeline;
       7              : 
       8              : /// Set of timelines, supports operations:
       9              : /// - add timeline
      10              : /// - remove timeline
      11              : /// - clone the set
      12              : ///
      13              : /// Usually used for keeping subset of timelines. For example active timelines that require broker push.
      14              : pub struct TimelinesSet {
      15              :     timelines: std::sync::Mutex<HashMap<TenantTimelineId, Arc<Timeline>>>,
      16              : }
      17              : 
      18              : impl Default for TimelinesSet {
      19            5 :     fn default() -> Self {
      20            5 :         Self {
      21            5 :             timelines: std::sync::Mutex::new(HashMap::new()),
      22            5 :         }
      23            5 :     }
      24              : }
      25              : 
      26              : impl TimelinesSet {
      27            5 :     pub fn insert(&self, tli: Arc<Timeline>) {
      28            5 :         self.timelines.lock().unwrap().insert(tli.ttid, tli);
      29            5 :     }
      30              : 
      31            5 :     pub fn delete(&self, ttid: &TenantTimelineId) {
      32            5 :         self.timelines.lock().unwrap().remove(ttid);
      33            5 :     }
      34              : 
      35              :     /// If present is true, adds timeline to the set, otherwise removes it.
      36            5 :     pub fn set_present(&self, tli: Arc<Timeline>, present: bool) {
      37            5 :         if present {
      38            5 :             self.insert(tli);
      39            5 :         } else {
      40            0 :             self.delete(&tli.ttid);
      41            0 :         }
      42            5 :     }
      43              : 
      44            5 :     pub fn is_present(&self, ttid: &TenantTimelineId) -> bool {
      45            5 :         self.timelines.lock().unwrap().contains_key(ttid)
      46            5 :     }
      47              : 
      48              :     /// Returns all timelines in the set.
      49            0 :     pub fn get_all(&self) -> Vec<Arc<Timeline>> {
      50            0 :         self.timelines.lock().unwrap().values().cloned().collect()
      51            0 :     }
      52              : 
      53              :     /// Returns a timeline guard for easy presence control.
      54            5 :     pub fn guard(self: &Arc<Self>, tli: Arc<Timeline>) -> TimelineSetGuard {
      55            5 :         let is_present = self.is_present(&tli.ttid);
      56            5 :         TimelineSetGuard {
      57            5 :             timelines_set: self.clone(),
      58            5 :             tli,
      59            5 :             is_present,
      60            5 :         }
      61            5 :     }
      62              : }
      63              : 
      64              : /// Guard is used to add or remove timelines from the set.
      65              : ///
      66              : /// If the timeline present in set, it will be removed from it on drop.
      67              : /// Note: do not use more than one guard for the same timeline, it caches the presence state.
      68              : /// It is designed to be used in the manager task only.
      69              : pub struct TimelineSetGuard {
      70              :     timelines_set: Arc<TimelinesSet>,
      71              :     tli: Arc<Timeline>,
      72              :     is_present: bool,
      73              : }
      74              : 
      75              : impl TimelineSetGuard {
      76              :     /// Returns true if the state was changed.
      77           35 :     pub fn set(&mut self, present: bool) -> bool {
      78           35 :         if present == self.is_present {
      79           30 :             return false;
      80            5 :         }
      81            5 :         self.is_present = present;
      82            5 :         self.timelines_set.set_present(self.tli.clone(), present);
      83            5 :         true
      84           35 :     }
      85              : 
      86            0 :     pub fn get(&self) -> bool {
      87            0 :         self.is_present
      88            0 :     }
      89              : }
      90              : 
      91              : impl Drop for TimelineSetGuard {
      92            5 :     fn drop(&mut self) {
      93            5 :         // remove timeline from the map on drop
      94            5 :         self.timelines_set.delete(&self.tli.ttid);
      95            5 :     }
      96              : }