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}