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<T, E>(
66        &self,
67        key: &str,
68        f: impl FnOnce() -> Result<(T, Vec<u8>), E>,
69    ) -> Result<(Option<T>, Vec<u8>), E> {
70        self.get_or_insert_with_timeout(key, self.timeout, f)
71    }
72
73    /// Retrieves a value from the cache by key or initialize it with the provided function.
74    /// If there is no existing value in the cache, the callback function will be immediately
75    /// called to fill the cache. All further calls during the callback execution will wait for
76    /// the value to be computed. As the callback might crash, a timeout limits how long this
77    /// function will wait. Unfortunately it does result in a thundering herd problem where all
78    /// Wasm instances will try to compute the value at the same time.
79    pub fn get_or_insert_with_timeout<T, E>(
80        &self,
81        key: &str,
82        timeout: Duration,
83        f: impl FnOnce() -> Result<(T, Vec<u8>), E>,
84    ) -> Result<(Option<T>, Vec<u8>), E> {
85        if let Some(value) = self.inner.get_or_reserve(key, timeout.as_millis() as u64) {
86            return Ok((None, value));
87        }
88        let (value, bytes) = f()?;
89        self.inner.insert(key, &bytes);
90        Ok((Some(value), bytes))
91    }
92}
grafbase_sdk/host_io/cache.rs

grafbase_sdk/host_io/
cache.rs
