5620643772

Committed 21 Jul 2023 09:09AM UTC coverage: 89.043% (-0.2%) from 89.194%

Build # 5620643772

Build Type

Pull #833

github

Committed by

web-flow

Commit Message

Merge 1ee0a296a into 2852314aa

Pull Request Pull Request #833: Shared-cache

Run Details

1128 of 1128 new or added lines in 9 files covered. (100.0%)

102925 of 115590 relevant lines covered (89.04%)

62633.59 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

80.74

/operators/src/pro/cache/shared_cache.rs

use super::{
    cache_chunks::{CachedFeatures, LandingZoneQueryFeatures},
    cache_tiles::{CachedTiles, LandingZoneQueryTiles, TypedCacheTileStream},
    error::CacheError,
    util::CacheSize,
};
use crate::engine::CanonicOperatorName;
use crate::util::Result;
use async_trait::async_trait;
use futures::Stream;
use geoengine_datatypes::{
    identifier,
    primitives::{CacheHint, RasterQueryRectangle, VectorQueryRectangle},
    raster::{Pixel, RasterTile2D},
    util::{test::TestDefault, ByteSize, Identifier},
};
use log::debug;
use lru::LruCache;
use std::{collections::HashMap, hash::Hash};
use tokio::sync::RwLock;

/// The tile cache caches all tiles of a query and is able to answer queries that are fully contained in the cache.
/// New tiles are inserted into the cache on-the-fly as they are produced by query processors.
/// The tiles are first inserted into a landing zone, until the query in completely finished and only then moved to the cache.
/// Both the landing zone and the cache have a maximum size.
/// If the landing zone is full, the caching of the current query will be aborted.
/// If the cache is full, the least recently used entries will be evicted if necessary to make room for the new entry.
#[derive(Debug)]
pub struct CacheBackend {
    // TODO: more fine granular locking?
    // for each operator graph, we have a cache, that can efficiently be accessed
    raster_caches: HashMap<CanonicOperatorName, RasterOperatorCacheEntry>,
    vector_caches: HashMap<CanonicOperatorName, VectorOperatorCacheEntry>,

    cache_size: CacheSize,
    landing_zone_size: CacheSize,

    // we only use the LruCache for determining the least recently used elements and evict as many entries as needed to fit the new one
    lru: LruCache<CacheEntryId, TypedCanonicOperatorName>,
}

impl CacheBackend {
    /// This method removes entries from the cache until it can fit the given amount of bytes.
    pub fn evict_until_can_fit_bytes(&mut self, bytes: usize) {
        while !self.cache_size.can_fit_bytes(bytes) {
            if let Some((pop_id, pop_key)) = self.lru.pop_lru() {
                match pop_key {
                    TypedCanonicOperatorName::Raster(raster_pop_key) => {
                        let op_cache = self
                            .raster_caches
                            .get_mut(&raster_pop_key)
                            .expect("LRU entry must exist in the cache!");
                        let query_element = op_cache
                            .remove_cache_entry(&pop_id)
                            .expect("LRU entry must exist in the cache!");
                        self.cache_size.remove_element_bytes(&query_element);
                    }
                    TypedCanonicOperatorName::Vector(vector_pop_key) => {
                        let op_cache = self
                            .vector_caches
                            .get_mut(&vector_pop_key)
                            .expect("LRU entry must exist in the cache!");
                        let query_element = op_cache
                            .remove_cache_entry(&pop_id)
                            .expect("LRU entry must exist in the cache!");
                        self.cache_size.remove_element_bytes(&query_element);
                    }
                };
                self.cache_size.remove_element_bytes(&pop_id);

                debug!(
                    "Evicted query {}. Cache size: {}. Cache size used: {}, Cache used percentage: {}.",
                    pop_id,
                    self.cache_size.total_byte_size(),
                    self.cache_size.byte_size_used(),
                    self.cache_size.size_used_fraction()
                );
            }
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum TypedCanonicOperatorName {
    Raster(CanonicOperatorName),
    Vector(CanonicOperatorName),
}

impl TypedCanonicOperatorName {
    pub fn as_raster(&self) -> Option<&CanonicOperatorName> {
        match self {
            Self::Raster(name) => Some(name),
            Self::Vector(_) => None,
        }
    }

    pub fn as_vector(&self) -> Option<&CanonicOperatorName> {
        match self {
            Self::Raster(_) => None,
            Self::Vector(name) => Some(name),
        }
    }
}

pub trait CacheEvictUntilFit {
    fn evict_entries_until_can_fit_bytes(&mut self, bytes: usize);
}

impl CacheEvictUntilFit for CacheBackend {
    fn evict_entries_until_can_fit_bytes(&mut self, bytes: usize) {
        self.evict_until_can_fit_bytes(bytes);
    }
}

pub trait CacheView<C, L>: CacheEvictUntilFit {
    fn operator_caches_mut(
        &mut self,
    ) -> &mut HashMap<CanonicOperatorName, OperatorCacheEntry<C, L>>;

    fn create_operator_cache_if_needed(&mut self, key: CanonicOperatorName) {
        self.operator_caches_mut()
            .entry(key)
            .or_insert_with(|| OperatorCacheEntry::new());
    }

    fn remove_operator_cache(
        &mut self,
        key: &CanonicOperatorName,
    ) -> Option<OperatorCacheEntry<C, L>> {
        // TODO: maybe remove the size of the OperatorCacheEntry to the cache size?
        self.operator_caches_mut().remove(key)
    }
}

#[allow(clippy::type_complexity)]
pub struct OperatorCacheEntryView<'a, C: CacheElement> {
    operator_cache: &'a mut OperatorCacheEntry<
        CacheQueryEntry<C::Query, C::CacheContainer>,
        CacheQueryEntry<C::Query, C::LandingZoneContainer>,
    >,
    cache_size: &'a mut CacheSize,
    landing_zone_size: &'a mut CacheSize,
    lru: &'a mut LruCache<CacheEntryId, TypedCanonicOperatorName>,
}

impl<'a, C> OperatorCacheEntryView<'a, C>
where
    C: CacheElement + ByteSize,
    C::Query: Clone + CacheQueryMatch,
    CacheQueryEntry<C::Query, C::LandingZoneContainer>: ByteSize,
    CacheQueryEntry<C::Query, C::CacheContainer>: ByteSize,
{
    fn is_empty(&self) -> bool {
        self.operator_cache.is_empty()
    }

    /// This method removes a query from the landing zone.
    ///
    /// If the query is not in the landing zone, this method returns None.
    ///
    fn remove_query_from_landing_zone(
        &mut self,
        query_id: &QueryId,
    ) -> Option<CacheQueryEntry<C::Query, C::LandingZoneContainer>> {
        if let Some(entry) = self.operator_cache.remove_landing_zone_entry(query_id) {
            self.landing_zone_size.remove_element_bytes(query_id);
            self.landing_zone_size.remove_element_bytes(&entry);

            // debug output
            log::debug!(
                "Removed query {}. Landing zone size: {}. Landing zone size used: {}, Landing zone used percentage: {}.",
                query_id, self.landing_zone_size.total_byte_size(), self.landing_zone_size.byte_size_used(), self.landing_zone_size.size_used_fraction()
            );

            Some(entry)
        } else {
            None
        }
    }

    /// This method removes a query from the cache and the LRU.
    /// It will remove a queries cache entry from the cache and the LRU.
    ///
    /// If the query is not in the cache, this method returns None.
    ///
    fn remove_query_from_cache_and_lru(
        &mut self,
        cache_entry_id: &CacheEntryId,
    ) -> Option<CacheQueryEntry<C::Query, C::CacheContainer>> {
        if let Some(entry) = self.operator_cache.remove_cache_entry(cache_entry_id) {
            let old_lru_entry = self.lru.pop_entry(cache_entry_id);
            debug_assert!(old_lru_entry.is_some(), "CacheEntryId not found in LRU");
            self.cache_size.remove_element_bytes(cache_entry_id);
            self.cache_size.remove_element_bytes(&entry);

            log::debug!(
                "Removed cache entry {}. Cache size: {}. Cache size used: {}, Cache used percentage: {}.",
                cache_entry_id, self.cache_size.total_byte_size(), self.cache_size.byte_size_used(), self.cache_size.size_used_fraction()
            );

            Some(entry)
        } else {
            None
        }
    }

    /// This method removes a list of queries from the cache and the LRU.
    fn discard_querys_from_cache_and_lru(&mut self, cache_entry_ids: &[CacheEntryId]) {
        for cache_entry_id in cache_entry_ids {
            let old_entry = self.remove_query_from_cache_and_lru(cache_entry_id);
            debug_assert!(
                old_entry.is_some(),
                "CacheEntryId not found in OperatorCacheEntry"
            );
        }
    }

    /// This method adds a query element to the landing zone.
    /// It will add the element to the landing zone entry of the query.
    ///
    /// # Errors
    ///
    /// This method returns an error if the query is not in the landing zone.
    /// This method returns an error if the element is already expired.
    /// This method returns an error if the landing zone is full or the new element would cause the landing zone to overflow.
    ///
    fn add_query_element_to_landing_zone(
        &mut self,
        query_id: &QueryId,
        landing_zone_element: C,
    ) -> Result<(), CacheError> {
        let landing_zone_entry = self
            .operator_cache
            .landing_zone_entry_mut(query_id)
            .ok_or(CacheError::QueryNotFoundInLandingZone)?;

        if landing_zone_element.cache_hint().is_expired() {
            return Err(CacheError::TileExpiredBeforeInsertion);
        };

        let element_bytes_size = landing_zone_element.byte_size();

        if !self.landing_zone_size.can_fit_bytes(element_bytes_size) {
            return Err(CacheError::NotEnoughSpaceInLandingZone);
        }

        landing_zone_entry.insert_element(landing_zone_element)?;

        // we add the bytes size of the element to the landing zone size after we have inserted it.
        self.landing_zone_size
            .try_add_bytes(element_bytes_size)
            .expect(
            "The Landing Zone must have enough space for the element since we checked it before",
        );

        log::trace!(
            "Inserted tile for query {} into landing zone. Landing zone size: {}. Landing zone size used: {}. Landing zone used percentage: {}",
            query_id, self.landing_zone_size.total_byte_size(), self.landing_zone_size.byte_size_used(), self.landing_zone_size.size_used_fraction()
        );

        Ok(())
    }

    /// This method inserts a query into the landing zone.
    /// It will cause the operator cache to create a new landing zone entry.
    /// Therefore, the size of the query and the size of the landing zone entry will be added to the landing zone size.
    ///
    /// # Errors
    ///
    /// This method returns an error if the query is already in the landing zone.
    /// This method returns an error if the landing zone is full or the new query would cause the landing zone to overflow.
    ///
    fn insert_query_into_landing_zone(&mut self, query: &C::Query) -> Result<QueryId, CacheError> {
        let landing_zone_entry = CacheQueryEntry::create_empty::<C>(query.clone());
        let query_id = QueryId::new();

        self.landing_zone_size.try_add_element_bytes(&query_id)?;
        self.landing_zone_size
            .try_add_element_bytes(&landing_zone_entry)?;

        self.operator_cache
            .insert_landing_zone_entry(query_id, landing_zone_entry)?;

        // debug output
        log::trace!(
            "Added query {} to landing zone. Landing zone size: {}. Landing zone size used: {}, Landing zone used percentage: {}.",
            query_id, self.landing_zone_size.total_byte_size(), self.landing_zone_size.byte_size_used(), self.landing_zone_size.size_used_fraction()
        );

        Ok(query_id)
    }

    /// This method inserts a cache entry into the cache and the LRU.
    /// It allows the element cache to overflow the cache size.
    /// This is done because the total cache size is the cache size + the landing zone size.
    /// This method is used when moving an element from the landing zone to the cache.
    ///
    /// # Errors
    ///
    /// This method returns an error if the cache entry is already in the cache.
    ///
    fn insert_cache_entry_allow_overflow(
        &mut self,
        cache_entry: CacheQueryEntry<C::Query, C::CacheContainer>,
        key: &CanonicOperatorName,
    ) -> Result<CacheEntryId, CacheError> {
        let cache_entry_id = CacheEntryId::new();
        let bytes = cache_entry.byte_size() + cache_entry_id.byte_size();
        // When inserting data from the landing zone into the cache, we allow the cache to overflow.
        // This is done because the total cache size is the cache size + the landing zone size.
        self.cache_size.add_bytes_allow_overflow(bytes);
        self.operator_cache
            .insert_cache_entry(cache_entry_id, cache_entry)?;
        // we have to wrap the key in a TypedCanonicOperatorName to be able to insert it into the LRU
        self.lru.push(
            cache_entry_id,
            C::typed_canonical_operator_name(key.clone()),
        );

        // debug output
        log::trace!(
            "Added cache entry {}. Cache size: {}. Cache size used: {}, Cache used percentage: {}.",
            cache_entry_id,
            self.cache_size.total_byte_size(),
            self.cache_size.byte_size_used(),
            self.cache_size.size_used_fraction()
        );

        Ok(cache_entry_id)
    }

    /*
    fn insert_cache_entry(
        &mut self,
        cache_entry: CacheQueryEntry<C::Query, C::CacheContainer>,
        key: CanonicOperatorName,
    ) -> Result<CacheEntryId, CacheError> {
        let cache_entry_id = CacheEntryId::new();
        self.cache_size.try_add_element_bytes(&cache_entry)?;
        self.cache_size.try_add_element_bytes(&cache_entry_id)?;
        self.operator_cache
            .insert_cache_entry(cache_entry_id, cache_entry)?;
        self.lru
            .push(cache_entry_id, C::typed_canonical_operator_name(key));

        Ok(cache_entry_id)
    }
    */

    /// This method finds a cache entry in the cache that matches the query.
    /// It will also collect all expired cache entries.
    /// The cache entry is returned together with the expired ids.
    fn find_matching_cache_entry_and_collect_expired_entries(
        &mut self,
        query: &C::Query,
    ) -> CacheQueryResult<C::Query, C::CacheContainer> {
        let mut expired_cache_entry_ids = vec![];

        let x = self.operator_cache.iter().find(|&(id, entry)| {
            if entry.elements.is_expired() {
                expired_cache_entry_ids.push(*id);
                return false;
            }
            entry.query.is_match(query)
        });

        CacheQueryResult {
            cache_hit: x.map(|(id, entry)| (*id, entry)),
            expired_cache_entry_ids,
        }
    }
}

struct CacheQueryResult<'a, Query, CE> {
    cache_hit: Option<(CacheEntryId, &'a CacheQueryEntry<Query, CE>)>,
    expired_cache_entry_ids: Vec<CacheEntryId>,
}

pub trait Cache<C: CacheElement>:
    CacheView<
    CacheQueryEntry<C::Query, C::CacheContainer>,
    CacheQueryEntry<C::Query, C::LandingZoneContainer>,
>
where
    C::Query: Clone + CacheQueryMatch,
    CacheQueryEntry<C::Query, C::LandingZoneContainer>: ByteSize,
    CacheQueryEntry<C::Query, C::CacheContainer>: ByteSize,
    CacheQueryEntry<C::Query, C::CacheContainer>:
        From<CacheQueryEntry<C::Query, C::LandingZoneContainer>>,
{
    /// This method returns a mutable reference to the cache entry of an operator.
    /// If there is no cache entry for the operator, this method returns None.
    fn operator_cache_view_mut(
        &mut self,
        key: &CanonicOperatorName,
    ) -> Option<OperatorCacheEntryView<C>>;

    /// This method queries the cache for a given query.
    /// If the query matches an entry in the cache, the cache entry is returned and it is promoted in the LRU.
    /// If the query does not match an entry in the cache, None is returned. Also if a cache entry is found but it is expired, None is returned.
    ///
    /// # Errors
    /// This method returns an error if the cache entry is not found.
    ///
    fn query_and_promote(
        &mut self,
        key: &CanonicOperatorName,
        query: &C::Query,
    ) -> Result<Option<C::ResultStream>, CacheError> {
        let mut cache = self
            .operator_cache_view_mut(key)
            .ok_or(CacheError::OperatorCacheEntryNotFound)?;

        let CacheQueryResult {
            cache_hit,
            expired_cache_entry_ids,
        } = cache.find_matching_cache_entry_and_collect_expired_entries(query);

        let res = if let Some((cache_entry_id, cache_entry)) = cache_hit {
            let stream = cache_entry.elements.result_stream(query);

            // promote the cache entry in the LRU
            cache.lru.promote(&cache_entry_id);
            Some(stream)
        } else {
            None
        };

        // discard expired cache entries
        cache.discard_querys_from_cache_and_lru(&expired_cache_entry_ids);

        Ok(res.flatten())
    }

    /// This method inserts a query into the cache.
    ///
    /// # Errors
    /// This method returns an error if the query is already in the cache.
    ///
    fn insert_query_into_landing_zone(
        &mut self,
        key: &CanonicOperatorName,
        query: &C::Query,
    ) -> Result<QueryId, CacheError> {
        self.create_operator_cache_if_needed(key.clone());
        self.operator_cache_view_mut(key)
            .expect("OperatorCache was Just created ")
            .insert_query_into_landing_zone(query)
    }

    fn insert_query_element_into_landing_zone(
        &mut self,
        key: &CanonicOperatorName,
        query_id: &QueryId,
        landing_zone_element: C,
    ) -> Result<(), CacheError> {
        let mut cache = self
            .operator_cache_view_mut(key)
            .ok_or(CacheError::QueryNotFoundInLandingZone)?;
        let res = cache.add_query_element_to_landing_zone(query_id, landing_zone_element);

        // if we cant add the element to the landing zone, we remove the query from the landing zone
        if res.is_err() {
            let _old_entry = cache.remove_query_from_landing_zone(query_id);

            // if the operator cache is empty, we remove it from the cache
            if cache.is_empty() {
                self.remove_operator_cache(key);
            }
        }

        res
    }

    /// This method discards a query from the landing zone.
    /// If the query is not in the landing zone, this method does nothing.
    fn discard_query_from_landing_zone(&mut self, key: &CanonicOperatorName, query_id: &QueryId) {
        if let Some(mut cache) = self.operator_cache_view_mut(key) {
            cache.remove_query_from_landing_zone(query_id);
            if cache.is_empty() {
                self.remove_operator_cache(key);
            }
        }
    }

    /// This method discards a query from the cache and the LRU.
    /// If the query is not in the cache, this method does nothing.
    fn discard_querys_from_cache_and_lru(
        &mut self,
        key: &CanonicOperatorName,
        cache_entry_ids: &[CacheEntryId],
    ) {
        if let Some(mut cache) = self.operator_cache_view_mut(key) {
            cache.discard_querys_from_cache_and_lru(cache_entry_ids);
            if cache.is_empty() {
                self.remove_operator_cache(key);
            }
        }
    }

    /// This method moves a query from the landing zone to the cache.
    /// It will remove the query from the landing zone and insert it into the cache.
    /// If the cache is full, the least recently used entries will be evicted if necessary to make room for the new entry.
    /// This method returns the cache entry id of the inserted cache entry.
    ///
    /// # Errors
    /// This method returns an error if the query is not in the landing zone.
    /// This method returns an error if the cache entry is already in the cache.
    /// This method returns an error if the cache is full and the least recently used entries cannot be evicted to make room for the new entry.
    ///
    fn move_query_from_landing_to_cache(
        &mut self,
        key: &CanonicOperatorName,
        query_id: &QueryId,
    ) -> Result<CacheEntryId, CacheError> {
        let mut operator_cache = self
            .operator_cache_view_mut(key)
            .ok_or(CacheError::OperatorCacheEntryNotFound)?;
        let landing_zone_entry = operator_cache
            .remove_query_from_landing_zone(query_id)
            .ok_or(CacheError::QueryNotFoundInLandingZone)?;
        let cache_entry: CacheQueryEntry<
            <C as CacheElement>::Query,
            <C as CacheElement>::CacheContainer,
        > = landing_zone_entry.into();
        // when moving an element from the landing zone to the cache, we allow the cache size to overflow.
        // This is done because the total cache size is the cache size + the landing zone size.
        let cache_entry_id = operator_cache.insert_cache_entry_allow_overflow(cache_entry, key)?;
        // We could also first try to evict until the cache can hold the entry.
        // However, then we would need to lookup the cache entry twice.
        // To avoid that, we just evict after we moved the entry from the landing zone to the cache.
        // This is also not a problem since the total cache size is the cache size + the landing zone size.
        self.evict_entries_until_can_fit_bytes(0);

        Ok(cache_entry_id)
    }
}

impl<T> Cache<RasterTile2D<T>> for CacheBackend
where
    T: Pixel + CachableSubType<CacheElementType = RasterTile2D<T>>,
{
    fn operator_cache_view_mut(
        &mut self,
        key: &CanonicOperatorName,
    ) -> Option<OperatorCacheEntryView<RasterTile2D<T>>> {
        self.raster_caches
            .get_mut(key)
            .map(|op| OperatorCacheEntryView {
                operator_cache: op,
                cache_size: &mut self.cache_size,
                landing_zone_size: &mut self.landing_zone_size,
                lru: &mut self.lru,
            })
    }
}

impl CacheView<RasterCacheQueryEntry, RasterLandingQueryEntry> for CacheBackend {
    fn operator_caches_mut(
        &mut self,
    ) -> &mut HashMap<CanonicOperatorName, RasterOperatorCacheEntry> {
        &mut self.raster_caches
    }
}

impl CacheView<VectorCacheQueryEntry, VectorLandingQueryEntry> for CacheBackend {
    fn operator_caches_mut(
        &mut self,
    ) -> &mut HashMap<CanonicOperatorName, VectorOperatorCacheEntry> {
        &mut self.vector_caches
    }
}

pub trait CacheElement: ByteSize + Send + ByteSize
where
    Self: Sized,
{
    type Query: CacheQueryMatch + Clone + Send + Sync;
    type LandingZoneContainer: LandingZoneElementsContainer<Self>;
    type CacheContainer: CacheElementsContainer<Self::Query, Self, ResultStream = Self::ResultStream>
        + From<Self::LandingZoneContainer>;
    type ResultStream;

    fn move_into_landing_zone(
        self,
        landing_zone: &mut Self::LandingZoneContainer,
    ) -> Result<(), CacheError> {
        landing_zone.insert_element(self)
    }

    fn cache_hint(&self) -> CacheHint;

    fn typed_canonical_operator_name(key: CanonicOperatorName) -> TypedCanonicOperatorName;
}

pub trait CachableSubType {
    type CacheElementType: CacheElement;

    fn insert_element_into_landing_zone(
        landing_zone: &mut <Self::CacheElementType as CacheElement>::LandingZoneContainer,
        element: Self::CacheElementType,
    ) -> Result<(), super::error::CacheError>;

    fn create_empty_landing_zone() -> <Self::CacheElementType as CacheElement>::LandingZoneContainer;

    fn result_stream(
        cache_elements_container: &<Self::CacheElementType as CacheElement>::CacheContainer,
        query: &<Self::CacheElementType as CacheElement>::Query,
    ) -> Option<<Self::CacheElementType as CacheElement>::ResultStream>;
}

#[derive(Debug)]
pub struct SharedCache {
    backend: RwLock<CacheBackend>,
}

impl SharedCache {
    pub fn new(cache_size_in_mb: usize, landing_zone_ratio: f64) -> Result<Self> {
        if landing_zone_ratio <= 0.0 {
            return Err(crate::error::Error::QueryingProcessorFailed {
                source: Box::new(CacheError::LandingZoneRatioMustBeLargerThanZero),
            });
        }

        if landing_zone_ratio >= 1.0 {
            return Err(crate::error::Error::QueryingProcessorFailed {
                source: Box::new(CacheError::LandingZoneRatioMustBeSmallerThanOne),
            });
        }

        // TODO: landing zone < 50% of cache size

        let cache_size_bytes =
            (cache_size_in_mb as f64 * (1.0 - landing_zone_ratio) * 1024.0 * 1024.0) as usize;

        let landing_zone_size_bytes =
            (cache_size_in_mb as f64 * landing_zone_ratio * 1024.0 * 1024.0) as usize;

        Ok(Self {
            backend: RwLock::new(CacheBackend {
                vector_caches: Default::default(),
                raster_caches: Default::default(),
                lru: LruCache::unbounded(), // we need no cap because we evict manually
                cache_size: CacheSize::new(cache_size_bytes),
                landing_zone_size: CacheSize::new(landing_zone_size_bytes),
            }),
        })
    }
}

impl TestDefault for SharedCache {
    fn test_default() -> Self {
        Self {
            backend: RwLock::new(CacheBackend {
                vector_caches: Default::default(),
                raster_caches: Default::default(),
                lru: LruCache::unbounded(), // we need no cap because we evict manually
                cache_size: CacheSize::new(usize::MAX),
                landing_zone_size: CacheSize::new(usize::MAX),
            }),
        }
    }
}

/// Holds all the cached results for an operator graph (workflow)
#[derive(Debug, Default)]
pub struct OperatorCacheEntry<C, L> {
    // for a given operator and query we need to look through all entries to find one that matches
    // TODO: use a multi-dimensional index to speed up the lookup
    entries: HashMap<CacheEntryId, C>,

    // running queries insert their tiles as they are produced. The entry will be created once the query is done.
    // The query is identified by a Uuid instead of the query rectangle to avoid confusions with other queries
    landing_zone: HashMap<QueryId, L>,
}

impl<C, L> OperatorCacheEntry<C, L> {
    pub fn new() -> Self {
        Self {
            entries: Default::default(),
            landing_zone: Default::default(),
        }
    }

    fn insert_landing_zone_entry(
        &mut self,
        query_id: QueryId,
        landing_zone_entry: L,
    ) -> Result<(), CacheError> {
        let old_entry = self.landing_zone.insert(query_id, landing_zone_entry);

        if old_entry.is_some() {
            Err(CacheError::QueryIdAlreadyInLandingZone)
        } else {
            Ok(())
        }
    }

    fn remove_landing_zone_entry(&mut self, query_id: &QueryId) -> Option<L> {
        self.landing_zone.remove(query_id)
    }

    fn landing_zone_entry_mut(&mut self, query_id: &QueryId) -> Option<&mut L> {
        self.landing_zone.get_mut(query_id)
    }

    fn insert_cache_entry(
        &mut self,
        cache_entry_id: CacheEntryId,
        cache_entry: C,
    ) -> Result<(), CacheError> {
        let old_entry = self.entries.insert(cache_entry_id, cache_entry);

        if old_entry.is_some() {
            Err(CacheError::CacheEntryIdAlreadyInCache)
        } else {
            Ok(())
        }
    }

    fn remove_cache_entry(&mut self, cache_entry_id: &CacheEntryId) -> Option<C> {
        self.entries.remove(cache_entry_id)
    }

    fn is_empty(&self) -> bool {
        self.entries.is_empty() && self.landing_zone.is_empty()
    }

    fn iter(&self) -> impl Iterator<Item = (&CacheEntryId, &C)> {
        self.entries.iter()
    }
}

identifier!(QueryId);

impl ByteSize for QueryId {}

identifier!(CacheEntryId);

impl ByteSize for CacheEntryId {}

/// Holds all the tiles for a given query and is able to answer queries that are fully contained
#[derive(Debug, Hash)]
pub struct CacheQueryEntry<Query, Elements> {
    query: Query,
    elements: Elements,
}
type RasterOperatorCacheEntry = OperatorCacheEntry<RasterCacheQueryEntry, RasterLandingQueryEntry>;
type RasterCacheQueryEntry = CacheQueryEntry<RasterQueryRectangle, CachedTiles>;
type RasterLandingQueryEntry = CacheQueryEntry<RasterQueryRectangle, LandingZoneQueryTiles>;

type VectorOperatorCacheEntry = OperatorCacheEntry<VectorCacheQueryEntry, VectorLandingQueryEntry>;
type VectorCacheQueryEntry = CacheQueryEntry<VectorQueryRectangle, CachedFeatures>;
type VectorLandingQueryEntry = CacheQueryEntry<VectorQueryRectangle, LandingZoneQueryFeatures>;

impl<Query, Elements> CacheQueryEntry<Query, Elements> {
    pub fn create_empty<E>(query: Query) -> Self
    where
        Elements: LandingZoneElementsContainer<E>,
    {
        Self {
            query,
            elements: Elements::create_empty(),
        }
    }

    pub fn query(&self) -> &Query {
        &self.query
    }

    pub fn elements_mut(&mut self) -> &mut Elements {
        &mut self.elements
    }

    pub fn insert_element<E>(&mut self, element: E) -> Result<(), CacheError>
    where
        Elements: LandingZoneElementsContainer<E>,
    {
        self.elements.insert_element(element)
    }
}

impl<Query, Elements> ByteSize for CacheQueryEntry<Query, Elements>
where
    Elements: ByteSize,
{
    fn heap_byte_size(&self) -> usize {
        self.elements.heap_byte_size()
    }
}

pub trait CacheQueryMatch<RHS = Self> {
    fn is_match(&self, query: &RHS) -> bool;
}

impl CacheQueryMatch for RasterQueryRectangle {
    fn is_match(&self, query: &RasterQueryRectangle) -> bool {
        self.spatial_bounds.contains(&query.spatial_bounds)
            && self.time_interval.contains(&query.time_interval)
            && self.spatial_resolution == query.spatial_resolution
    }
}

impl CacheQueryMatch for VectorQueryRectangle {
    // TODO: check if that is what we need
    fn is_match(&self, query: &VectorQueryRectangle) -> bool {
        self.spatial_bounds.contains_bbox(&query.spatial_bounds)
            && self.time_interval.contains(&query.time_interval)
            && self.spatial_resolution == query.spatial_resolution
    }
}

pub trait LandingZoneElementsContainer<E> {
    fn insert_element(&mut self, element: E) -> Result<(), CacheError>;
    fn create_empty() -> Self;
}

pub trait CacheElementsContainerInfos<Query> {
    fn is_expired(&self) -> bool;
}

pub trait CacheElementsContainer<Query, E>: CacheElementsContainerInfos<Query> {
    type ResultStream: Stream<Item = Result<E>>;

    fn result_stream(&self, query: &Query) -> Option<Self::ResultStream>;
}

impl CacheQueryEntry<RasterQueryRectangle, CachedTiles> {
    /// Return true if the query can be answered in full by this cache entry
    /// For this, the bbox and time has to be fully contained, and the spatial resolution has to match
    pub fn matches(&self, query: &RasterQueryRectangle) -> bool {
        self.query.spatial_bounds.contains(&query.spatial_bounds)
            && self.query.time_interval.contains(&query.time_interval)
            && self.query.spatial_resolution == query.spatial_resolution
    }

    /// Produces a tile stream from the cache
    pub fn tile_stream(&self, query: &RasterQueryRectangle) -> TypedCacheTileStream {
        self.elements.tile_stream(query)
    }
}

impl From<RasterLandingQueryEntry> for RasterCacheQueryEntry {
    fn from(value: RasterLandingQueryEntry) -> Self {
        Self {
            query: value.query,
            elements: value.elements.into(),
        }
    }
}

impl From<VectorLandingQueryEntry> for VectorCacheQueryEntry {
    fn from(value: VectorLandingQueryEntry) -> Self {
        Self {
            query: value.query,
            elements: value.elements.into(),
        }
    }
}

#[async_trait]
pub trait AsyncCache<C: CacheElement> {
    async fn query_cache<S: CachableSubType<CacheElementType = C>>(
        &self,
        key: &CanonicOperatorName,
        query: &C::Query,
    ) -> Result<Option<C::ResultStream>, CacheError>;

    async fn insert_query<S: CachableSubType<CacheElementType = C>>(
        &self,
        key: &CanonicOperatorName,
        query: &C::Query,
    ) -> Result<QueryId, CacheError>;

    async fn insert_query_element<S: CachableSubType<CacheElementType = C>>(
        &self,
        key: &CanonicOperatorName,
        query_id: &QueryId,
        landing_zone_element: S::CacheElementType,
    ) -> Result<(), CacheError>;

    async fn abort_query<S: CachableSubType<CacheElementType = C>>(
        &self,
        key: &CanonicOperatorName,
        query_id: &QueryId,
    );

    async fn finish_query<S: CachableSubType<CacheElementType = C>>(
        &self,
        key: &CanonicOperatorName,
        query_id: &QueryId,
    ) -> Result<CacheEntryId, CacheError>;
}

#[async_trait]
impl<C> AsyncCache<C> for SharedCache
where
    C: CacheElement + ByteSize + Send + 'static,
    C::Query: Send + Sync,
    CacheBackend: Cache<C>,
    CacheQueryEntry<C::Query, C::LandingZoneContainer>: ByteSize,
    CacheQueryEntry<C::Query, C::CacheContainer>: ByteSize,
    CacheQueryEntry<C::Query, C::CacheContainer>:
        From<CacheQueryEntry<C::Query, C::LandingZoneContainer>>,
{
    /// Query the cache and on hit create a stream of cache elements
    async fn query_cache<S: CachableSubType<CacheElementType = C>>(
        &self,
        key: &CanonicOperatorName,
        query: &C::Query,
    ) -> Result<Option<C::ResultStream>, CacheError> {
        let mut backend = self.backend.write().await;
        backend.query_and_promote(key, query)
    }

    /// When inserting a new query, we first register the query and then insert the elements as they are produced
    /// This is to avoid confusing different queries on the same operator and query rectangle
    async fn insert_query<S: CachableSubType<CacheElementType = C>>(
        &self,
        key: &CanonicOperatorName,
        query: &C::Query,
    ) -> Result<QueryId, CacheError> {
        let mut backend = self.backend.write().await;
        backend.insert_query_into_landing_zone(key, query)
    }

    /// Insert a cachable element for a given query. The query has to be inserted first.
    /// The element is inserted into the landing zone and only moved to the cache when the query is finished.
    /// If the landing zone is full or the element size would cause the landing zone size to overflow, the caching of the query is aborted.
    async fn insert_query_element<S: CachableSubType<CacheElementType = C>>(
        &self,
        key: &CanonicOperatorName,
        query_id: &QueryId,
        landing_zone_element: C,
    ) -> Result<(), CacheError> {
        let mut backend = self.backend.write().await;
        backend.insert_query_element_into_landing_zone(key, query_id, landing_zone_element)
    }

    /// Abort the query and remove already inserted elements from the caches landing zone
    async fn abort_query<S: CachableSubType<CacheElementType = C>>(
        &self,
        key: &CanonicOperatorName,
        query_id: &QueryId,
    ) {
        let mut backend = self.backend.write().await;
        backend.discard_query_from_landing_zone(key, query_id);
    }

    /// Finish the query and make the inserted elements available in the cache
    async fn finish_query<S: CachableSubType<CacheElementType = C>>(
        &self,
        key: &CanonicOperatorName,
        query_id: &QueryId,
    ) -> Result<CacheEntryId, CacheError> {
        let mut backend = self.backend.write().await;
        backend.move_query_from_landing_to_cache(key, query_id)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use geoengine_datatypes::{
        primitives::{CacheHint, DateTime, SpatialPartition2D, SpatialResolution, TimeInterval},
        raster::{Grid, RasterProperties},
    };
    use serde_json::json;

    use super::*;

    async fn process_query(tile_cache: &mut SharedCache, op_name: CanonicOperatorName) {
        let query_id = tile_cache
            .insert_query::<u8>(&op_name, &query_rect())
            .await
            .unwrap();

        tile_cache
            .insert_query_element::<u8>(&op_name, &query_id, create_tile())
            .await
            .unwrap();

        tile_cache
            .finish_query::<u8>(&op_name, &query_id)
            .await
            .unwrap();
    }

    fn create_tile() -> RasterTile2D<u8> {
        RasterTile2D::<u8> {
            time: TimeInterval::new_instant(DateTime::new_utc(2014, 3, 1, 0, 0, 0)).unwrap(),
            tile_position: [-1, 0].into(),
            global_geo_transform: TestDefault::test_default(),
            grid_array: Grid::new([3, 2].into(), vec![1, 2, 3, 4, 5, 6])
                .unwrap()
                .into(),
            properties: RasterProperties::default(),
            cache_hint: CacheHint::max_duration(),
        }
    }

    fn query_rect() -> RasterQueryRectangle {
        RasterQueryRectangle {
            spatial_bounds: SpatialPartition2D::new_unchecked(
                (-180., 90.).into(),
                (180., -90.).into(),
            ),
            time_interval: TimeInterval::new_instant(DateTime::new_utc(2014, 3, 1, 0, 0, 0))
                .unwrap(),
            spatial_resolution: SpatialResolution::one(),
        }
    }

    fn op(idx: usize) -> CanonicOperatorName {
        CanonicOperatorName::new_unchecked(&json!({
            "type": "GdalSource",
            "params": {
                "data": idx
            }
        }))
    }

    #[tokio::test]
    async fn it_evicts_lru() {
        // Create cache entry and landing zone entry to geht the size of both
        let landing_zone_entry = RasterLandingQueryEntry {
            query: query_rect(),
            elements: LandingZoneQueryTiles::U8(vec![create_tile()]),
        };
        let query_id = QueryId::new();
        let size_of_landing_zone_entry = landing_zone_entry.byte_size() + query_id.byte_size();
        let cache_entry: RasterCacheQueryEntry = landing_zone_entry.into();
        let cache_entry_id = CacheEntryId::new();
        let size_of_cache_entry = cache_entry.byte_size() + cache_entry_id.byte_size();

        // Select the max of both sizes
        // This is done because the landing zone should not be smaller then the cache
        let m_size = size_of_cache_entry.max(size_of_landing_zone_entry);

        // set limits s.t. three tiles fit
        let mut tile_cache = SharedCache {
            backend: RwLock::new(CacheBackend {
                raster_caches: Default::default(),
                vector_caches: Default::default(),
                lru: LruCache::unbounded(),
                cache_size: CacheSize::new(m_size * 3),
                landing_zone_size: CacheSize::new(m_size * 3),
            }),
        };

        // process three different queries
        process_query(&mut tile_cache, op(1)).await;
        process_query(&mut tile_cache, op(2)).await;
        process_query(&mut tile_cache, op(3)).await;

        // query the first one s.t. it is the most recently used
        tile_cache
            .query_cache::<u8>(&op(1), &query_rect())
            .await
            .unwrap();
        // process a fourth query
        process_query(&mut tile_cache, op(4)).await;

        // assure the seconds query is evicted because it is the least recently used
        assert!(tile_cache
            .query_cache::<u8>(&op(2), &query_rect())
            .await
            .unwrap()
            .is_none());

        // assure that the other queries are still in the cache
        for i in [1, 3, 4] {
            assert!(tile_cache
                .query_cache::<u8>(&op(i), &query_rect())
                .await
                .unwrap()
                .is_some());
        }

        assert_eq!(
            tile_cache.backend.read().await.cache_size.byte_size_used(),
            3 * size_of_cache_entry
        );
    }

    #[test]
    fn cache_byte_size() {
        assert_eq!(create_tile().byte_size(), 284);
        assert_eq!(
            CachedTiles::U8(Arc::new(vec![create_tile()])).byte_size(),
            /* enum + arc */ 16 + /* vec */ 24  + /* tile */ 284
        );
        assert_eq!(
            CachedTiles::U8(Arc::new(vec![create_tile(), create_tile()])).byte_size(),
            /* enum + arc */ 16 + /* vec */ 24  + /* tile */ 2 * 284
        );
    }

    #[tokio::test]
    async fn it_checks_ttl() {
        let mut tile_cache = SharedCache {
            backend: RwLock::new(CacheBackend {
                raster_caches: Default::default(),
                vector_caches: Default::default(),
                lru: LruCache::unbounded(),
                cache_size: CacheSize::new(usize::MAX),
                landing_zone_size: CacheSize::new(usize::MAX),
            }),
        };

        process_query(&mut tile_cache, op(1)).await;

        // access works because no ttl is set
        tile_cache
            .query_cache::<u8>(&op(1), &query_rect())
            .await
            .unwrap()
            .unwrap();

        // manually expire entry
        {
            let mut backend = tile_cache.backend.write().await;
            let cache = backend.raster_caches.iter_mut().next().unwrap();

            let tiles = &mut cache.1.entries.iter_mut().next().unwrap().1.elements;
            match tiles {
                CachedTiles::U8(tiles) => {
                    let mut expired_tiles = (**tiles).clone();
                    expired_tiles[0].cache_hint = CacheHint::with_created_and_expires(
                        DateTime::new_utc(0, 1, 1, 0, 0, 0),
                        DateTime::new_utc(0, 1, 1, 0, 0, 1).into(),
                    );
                    *tiles = Arc::new(expired_tiles);
                }
                _ => panic!("wrong tile type"),
            }
        }

        // access fails because ttl is expired
        assert!(tile_cache
            .query_cache::<u8>(&op(1), &query_rect())
            .await
            .unwrap()
            .is_none());
    }

    #[tokio::test]
    async fn tile_cache_init_size() {
        let tile_cache = SharedCache::new(100, 0.1).unwrap();

        let backend = tile_cache.backend.read().await;

        let cache_size = 90 * 1024 * 1024;
        let landing_zone_size = 10 * 1024 * 1024;

        assert_eq!(backend.cache_size.total_byte_size(), cache_size);
        assert_eq!(
            backend.landing_zone_size.total_byte_size(),
            landing_zone_size
        );
    }
}

geo-engine / geoengine / 5620643772

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous