Line data    Source code

       1              : use std::{
       2              :     ops::{Deref, DerefMut},
       3              :     time::{Duration, Instant},
       4              : };
       5              : 
       6              : use moka::Expiry;
       7              : 
       8              : use crate::control_plane::messages::ControlPlaneErrorMessage;
       9              : 
      10              : /// Default TTL used when caching errors from control plane.
      11              : pub const DEFAULT_ERROR_TTL: Duration = Duration::from_secs(30);
      12              : 
      13              : /// A generic trait which exposes types of cache's key and value,
      14              : /// as well as the notion of cache entry invalidation.
      15              : /// This is useful for [`Cached`].
      16              : pub(crate) trait Cache {
      17              :     /// Entry's key.
      18              :     type Key;
      19              : 
      20              :     /// Entry's value.
      21              :     type Value;
      22              : 
      23              :     /// Used for entry invalidation.
      24              :     type LookupInfo<Key>;
      25              : 
      26              :     /// Invalidate an entry using a lookup info.
      27              :     /// We don't have an empty default impl because it's error-prone.
      28              :     fn invalidate(&self, _: &Self::LookupInfo<Self::Key>);
      29              : }
      30              : 
      31              : impl<C: Cache> Cache for &C {
      32              :     type Key = C::Key;
      33              :     type Value = C::Value;
      34              :     type LookupInfo<Key> = C::LookupInfo<Key>;
      35              : 
      36            6 :     fn invalidate(&self, info: &Self::LookupInfo<Self::Key>) {
      37            6 :         C::invalidate(self, info);
      38            6 :     }
      39              : }
      40              : 
      41              : /// Wrapper for convenient entry invalidation.
      42              : pub(crate) struct Cached<C: Cache, V = <C as Cache>::Value> {
      43              :     /// Cache + lookup info.
      44              :     pub(crate) token: Option<(C, C::LookupInfo<C::Key>)>,
      45              : 
      46              :     /// The value itself.
      47              :     pub(crate) value: V,
      48              : }
      49              : 
      50              : impl<C: Cache, V> Cached<C, V> {
      51              :     /// Place any entry into this wrapper; invalidation will be a no-op.
      52            1 :     pub(crate) fn new_uncached(value: V) -> Self {
      53            1 :         Self { token: None, value }
      54            1 :     }
      55              : 
      56            0 :     pub(crate) fn take_value(self) -> (Cached<C, ()>, V) {
      57            0 :         (
      58            0 :             Cached {
      59            0 :                 token: self.token,
      60            0 :                 value: (),
      61            0 :             },
      62            0 :             self.value,
      63            0 :         )
      64            0 :     }
      65              : 
      66           17 :     pub(crate) fn map<U>(self, f: impl FnOnce(V) -> U) -> Cached<C, U> {
      67           17 :         Cached {
      68           17 :             token: self.token,
      69           17 :             value: f(self.value),
      70           17 :         }
      71           17 :     }
      72              : 
      73              :     /// Drop this entry from a cache if it's still there.
      74            6 :     pub(crate) fn invalidate(self) -> V {
      75            6 :         if let Some((cache, info)) = &self.token {
      76            6 :             cache.invalidate(info);
      77            6 :         }
      78            6 :         self.value
      79            6 :     }
      80              : 
      81              :     /// Tell if this entry is actually cached.
      82           16 :     pub(crate) fn cached(&self) -> bool {
      83           16 :         self.token.is_some()
      84           16 :     }
      85              : }
      86              : 
      87              : impl<C: Cache, V> Deref for Cached<C, V> {
      88              :     type Target = V;
      89              : 
      90            0 :     fn deref(&self) -> &Self::Target {
      91            0 :         &self.value
      92            0 :     }
      93              : }
      94              : 
      95              : impl<C: Cache, V> DerefMut for Cached<C, V> {
      96            0 :     fn deref_mut(&mut self) -> &mut Self::Target {
      97            0 :         &mut self.value
      98            0 :     }
      99              : }
     100              : 
     101              : pub type ControlPlaneResult<T> = Result<T, Box<ControlPlaneErrorMessage>>;
     102              : 
     103              : #[derive(Clone, Copy)]
     104              : pub struct CplaneExpiry {
     105              :     pub error: Duration,
     106              : }
     107              : 
     108              : impl Default for CplaneExpiry {
     109            2 :     fn default() -> Self {
     110            2 :         Self {
     111            2 :             error: DEFAULT_ERROR_TTL,
     112            2 :         }
     113            2 :     }
     114              : }
     115              : 
     116              : impl CplaneExpiry {
     117           12 :     pub fn expire_early<V>(
     118           12 :         &self,
     119           12 :         value: &ControlPlaneResult<V>,
     120           12 :         updated: Instant,
     121           12 :     ) -> Option<Duration> {
     122           12 :         match value {
     123            8 :             Ok(_) => None,
     124            4 :             Err(err) => Some(self.expire_err_early(err, updated)),
     125              :         }
     126           12 :     }
     127              : 
     128            4 :     pub fn expire_err_early(&self, err: &ControlPlaneErrorMessage, updated: Instant) -> Duration {
     129            4 :         err.status
     130            4 :             .as_ref()
     131            4 :             .and_then(|s| s.details.retry_info.as_ref())
     132            4 :             .map_or(self.error, |r| r.retry_at.into_std() - updated)
     133            4 :     }
     134              : }
     135              : 
     136              : impl<K, V> Expiry<K, ControlPlaneResult<V>> for CplaneExpiry {
     137            7 :     fn expire_after_create(
     138            7 :         &self,
     139            7 :         _key: &K,
     140            7 :         value: &ControlPlaneResult<V>,
     141            7 :         created_at: Instant,
     142            7 :     ) -> Option<Duration> {
     143            7 :         self.expire_early(value, created_at)
     144            7 :     }
     145              : 
     146            5 :     fn expire_after_update(
     147            5 :         &self,
     148            5 :         _key: &K,
     149            5 :         value: &ControlPlaneResult<V>,
     150            5 :         updated_at: Instant,
     151            5 :         _duration_until_expiry: Option<Duration>,
     152            5 :     ) -> Option<Duration> {
     153            5 :         self.expire_early(value, updated_at)
     154            5 :     }
     155              : }