Line data    Source code

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