grafbase_sdk/host_io/
cache.rs

1//! A key-value cache shared with all instances of the extension.
2
3use std::{borrow::Cow, time::Duration};
4
5use crate::wit;
6
7/// The cache is a key-value store shared across Wasm instances. As Wasm is single threaded, the
8/// gateway uses a pool of Wasm instances to execute extensions. Cache with the same name will be
9/// the same across those instances and share the same data.
10pub struct Cache {
11    inner: wit::Cache,
12    timeout: Duration,
13}
14
15/// The builder for the cache. It allows to set the name, size, ttl and timeout of the cache.
16pub struct CacheBuilder {
17    name: Cow<'static, str>,
18    size: usize,
19    ttl: Option<Duration>,
20    timeout: Duration,
21}
22
23impl CacheBuilder {
24    /// Time to live for cached entries.
25    pub fn time_to_live(mut self, ttl: Option<Duration>) -> Self {
26        self.ttl = ttl;
27        self
28    }
29
30    /// The default timeout to use when retrieving data from the cache.
31    pub fn timeout(mut self, timeout: Duration) -> Self {
32        self.timeout = timeout;
33        self
34    }
35
36    /// Builds the cache
37    pub fn build(self) -> Cache {
38        let Self {
39            name,
40            size,
41            ttl,
42            timeout,
43        } = self;
44        Cache {
45            inner: wit::Cache::init(&name, size as u32, ttl.map(|d| d.as_millis() as u64)),
46            timeout,
47        }
48    }
49}
50
51impl Cache {
52    /// Creates a new cache builder with the given name and size.
53    /// Caches are unique for a given name and extension.
54    pub fn builder(name: impl Into<Cow<'static, str>>, size: usize) -> CacheBuilder {
55        CacheBuilder {
56            name: name.into(),
57            size,
58            ttl: None,
59            timeout: Duration::from_secs(5),
60        }
61    }
62
63    /// Retrieves a value from the cache by key or initialize it with the provided function using
64    /// the default timeout. See [get_or_init_with_timeout](Cache::get_or_init_with_timeout) for more details
65    pub fn get_or_insert<E>(&self, key: &str, f: impl FnOnce() -> Result<Vec<u8>, E>) -> Result<Vec<u8>, E> {
66        self.get_or_insert_with_timeout(key, self.timeout, f)
67    }
68
69    /// Retrieves a value from the cache by key or initialize it with the provided function.
70    /// If there is no existing value in the cache, the callback function will be immediately
71    /// called to fill the cache. All further calls during the callback execution will wait for
72    /// the value to be computed. As the callback might crash, a timeout limits how long this
73    /// function will wait. Unfortunately it does result in a thundering herd problem where all
74    /// Wasm instances will try to compute the value at the same time.
75    pub fn get_or_insert_with_timeout<E>(
76        &self,
77        key: &str,
78        timeout: Duration,
79        f: impl FnOnce() -> Result<Vec<u8>, E>,
80    ) -> Result<Vec<u8>, E> {
81        if let Some(value) = self.inner.get_or_reserve(key, timeout.as_millis() as u64) {
82            return Ok(value);
83        }
84        let value = f()?;
85        self.inner.insert(key, &value);
86        Ok(value)
87    }
88}