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            0 :     fn default() -> Self {
      19            0 :         Self {
      20            0 :             timelines: std::sync::Mutex::new(HashMap::new()),
      21            0 :         }
      22            0 :     }
      23              : }
      24              : 
      25              : impl TimelinesSet {
      26            0 :     pub fn insert(&self, tli: Arc<Timeline>) {
      27            0 :         self.timelines.lock().unwrap().insert(tli.ttid, tli);
      28            0 :     }
      29              : 
      30            0 :     pub fn delete(&self, ttid: &TenantTimelineId) {
      31            0 :         self.timelines.lock().unwrap().remove(ttid);
      32            0 :     }
      33              : 
      34              :     /// If present is true, adds timeline to the set, otherwise removes it.
      35            0 :     pub fn set_present(&self, tli: Arc<Timeline>, present: bool) {
      36            0 :         if present {
      37            0 :             self.insert(tli);
      38            0 :         } else {
      39            0 :             self.delete(&tli.ttid);
      40            0 :         }
      41            0 :     }
      42              : 
      43            0 :     pub fn is_present(&self, ttid: &TenantTimelineId) -> bool {
      44            0 :         self.timelines.lock().unwrap().contains_key(ttid)
      45            0 :     }
      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            0 :     pub fn guard(self: &Arc<Self>, tli: Arc<Timeline>) -> TimelineSetGuard {
      54            0 :         let is_present = self.is_present(&tli.ttid);
      55            0 :         TimelineSetGuard {
      56            0 :             timelines_set: self.clone(),
      57            0 :             tli,
      58            0 :             is_present,
      59            0 :         }
      60            0 :     }
      61              : }
      62              : 
      63              : /// Guard is used to add or remove timeline from the set.
      64              : /// If the timeline present in set, it will be removed from it on drop.
      65              : /// Note: do not use more than one guard for the same timeline, it caches the presence state.
      66              : /// It is designed to be used in the manager task only.
      67              : pub struct TimelineSetGuard {
      68              :     timelines_set: Arc<TimelinesSet>,
      69              :     tli: Arc<Timeline>,
      70              :     is_present: bool,
      71              : }
      72              : 
      73              : impl TimelineSetGuard {
      74              :     /// Returns true if the state was changed.
      75            0 :     pub fn set(&mut self, present: bool) -> bool {
      76            0 :         if present == self.is_present {
      77            0 :             return false;
      78            0 :         }
      79            0 :         self.is_present = present;
      80            0 :         self.timelines_set.set_present(self.tli.clone(), present);
      81            0 :         true
      82            0 :     }
      83              : }
      84              : 
      85              : impl Drop for TimelineSetGuard {
      86            0 :     fn drop(&mut self) {
      87            0 :         // remove timeline from the map on drop
      88            0 :         self.timelines_set.delete(&self.tli.ttid);
      89            0 :     }
      90              : }