grafbase_sdk/host_io/
cache.rs

1//! A key-value cache shared with all instances of the extension.
2
3use std::{borrow::Cow, convert::Infallible, 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. If you need to additional validation on the bytes before storing it in the cache, consider using
65    /// [get_or_insert](Cache::get_or_insert) can keep additional state.
66    /// If you need error handling, use [try_get_or_insert_bytes](Cache::try_get_or_insert_bytes).
67    ///
68    /// See [try_get_or_insert_with_timeout](Cache::try_get_or_insert_with_timeout) for more details
69    pub fn get_or_insert_bytes(&self, key: &str, f: impl FnOnce() -> Vec<u8>) -> Vec<u8> {
70        self.try_get_or_insert_with_timeout(key, self.timeout, || {
71            let bytes = f();
72            Result::<_, Infallible>::Ok(((), bytes))
73        })
74        .unwrap()
75        .1
76    }
77
78    /// Retrieves a value from the cache by key or initialize it with the provided function using
79    /// the default timeout. If you need to additional validation on the bytes before storing it in the cache, consider using
80    /// [get_or_insert_with_timeout](Cache::get_or_insert_with_timeout) can keep additional state.
81    /// If you need error handling, use [try_get_or_insert_bytes_with_timeout](Cache::try_get_or_insert_bytes_with_timeout).
82    ///
83    /// See [try_get_or_insert_with_timeout](Cache::try_get_or_insert_with_timeout) for more details
84    pub fn get_or_insert_bytes_with_timeout(
85        &self,
86        key: &str,
87        timeout: Duration,
88        f: impl FnOnce() -> Vec<u8>,
89    ) -> Vec<u8> {
90        self.try_get_or_insert_with_timeout(key, timeout, || {
91            let bytes = f();
92            Result::<_, Infallible>::Ok(((), bytes))
93        })
94        .unwrap()
95        .1
96    }
97
98    /// Retrieves a value from the cache by key or initialize it with the provided function using
99    /// the default timeout.[get_or_insert_bytes](Cache::get_or_insert_bytes) is a simpler variant if you only handle bytes.
100    /// If you need error handling, use [try_get_or_insert](Cache::try_get_or_insert).
101    ///
102    /// See [try_get_or_insert_with_timeout](Cache::try_get_or_insert_with_timeout) for more details
103    pub fn get_or_insert<T>(&self, key: &str, f: impl FnOnce() -> (T, Vec<u8>)) -> (Option<T>, Vec<u8>) {
104        self.try_get_or_insert_with_timeout(key, self.timeout, || {
105            let value = f();
106            Result::<_, Infallible>::Ok(value)
107        })
108        .unwrap()
109    }
110
111    /// Retrieves a value from the cache by key or initialize it with the provided function using
112    /// the default timeout.[get_or_insert_bytes_with_timeout](Cache::get_or_insert_bytes_with_timeout) is a simpler variant if you only handle bytes.
113    /// If you need error handling, use [try_get_or_insert_with_timeout](Cache::try_get_or_insert_with_timeout).
114    ///
115    /// See [try_get_or_insert_with_timeout](Cache::try_get_or_insert_with_timeout) for more details
116    pub fn get_or_insert_with_timeout<T>(
117        &self,
118        key: &str,
119        timeout: Duration,
120        f: impl FnOnce() -> (T, Vec<u8>),
121    ) -> (Option<T>, Vec<u8>) {
122        self.try_get_or_insert_with_timeout(key, timeout, || {
123            let value = f();
124            Result::<_, Infallible>::Ok(value)
125        })
126        .unwrap()
127    }
128
129    /// Retrieves a value from the cache by key or initialize it with the provided function using
130    /// the default timeout. See [try_get_or_insert_with_timeout](Cache::try_get_or_insert_with_timeout) for more details
131    pub fn try_get_or_insert_bytes<E>(&self, key: &str, f: impl FnOnce() -> Result<Vec<u8>, E>) -> Result<Vec<u8>, E> {
132        self.try_get_or_insert_with_timeout(key, self.timeout, || f().map(|bytes| ((), bytes)))
133            .map(|(_, bytes)| bytes)
134    }
135
136    /// Retrieves a value from the cache by key or initialize it with the provided function using
137    /// the default timeout. See [try_get_or_insert_with_timeout](Cache::try_get_or_insert_with_timeout) for more details
138    pub fn try_get_or_insert_bytes_with_timeout<E>(
139        &self,
140        key: &str,
141        timeout: Duration,
142        f: impl FnOnce() -> Result<Vec<u8>, E>,
143    ) -> Result<Vec<u8>, E> {
144        self.try_get_or_insert_with_timeout(key, timeout, || f().map(|bytes| ((), bytes)))
145            .map(|(_, bytes)| bytes)
146    }
147
148    /// Retrieves a value from the cache by key or initialize it with the provided function using
149    /// the default timeout. See [try_get_or_insert_with_timeout](Cache::try_get_or_insert_with_timeout) for more details
150    pub fn try_get_or_insert<T, E>(
151        &self,
152        key: &str,
153        f: impl FnOnce() -> Result<(T, Vec<u8>), E>,
154    ) -> Result<(Option<T>, Vec<u8>), E> {
155        self.try_get_or_insert_with_timeout(key, self.timeout, f)
156    }
157
158    /// Retrieves a value from the cache by key or initialize it with the provided function.
159    /// If there is no existing value in the cache, the callback function will be immediately
160    /// called to fill the cache. All further calls during the callback execution will wait for
161    /// the value to be computed. As the callback might crash, a timeout limits how long this
162    /// function will wait. Unfortunately it does result in a thundering herd problem where all
163    /// Wasm instances will try to compute the value at the same time.
164    pub fn try_get_or_insert_with_timeout<T, E>(
165        &self,
166        key: &str,
167        timeout: Duration,
168        f: impl FnOnce() -> Result<(T, Vec<u8>), E>,
169    ) -> Result<(Option<T>, Vec<u8>), E> {
170        if let Some(value) = self.inner.get_or_reserve(key, timeout.as_millis() as u64) {
171            return Ok((None, value));
172        }
173        let (value, bytes) = f()?;
174        self.inner.insert(key, &bytes);
175        Ok((Some(value), bytes))
176    }
177}