teaql_runtime/context/
pagination.rs

1use std::collections::HashMap;
2use std::sync::{Arc, Mutex, OnceLock};
3use std::time::SystemTime;
4
5use teaql_core::Value;
6
7use super::UserContext;
8
9#[derive(Debug, Clone, PartialEq)]
10pub struct ContinuousPageCursor {
11    pub cursor_id: String,
12    pub query_key: String,
13    pub entity: String,
14    pub direction: teaql_core::SortDirection,
15    pub boundary: Value,
16    pub page_size: u64,
17    pub next_offset: u64,
18    pub expires_at: SystemTime,
19}
20
21#[async_trait::async_trait]
22pub trait ContinuousPageCursorStore: Send + Sync + 'static {
23    async fn get(
24        &self,
25        query_key: &str,
26        target_offset: u64,
27    ) -> Result<Option<ContinuousPageCursor>, String>;
28    async fn put(&self, cursor: ContinuousPageCursor) -> Result<(), String>;
29    async fn invalidate(&self, query_key: &str) -> Result<(), String>;
30}
31
32pub struct InMemoryContinuousPageCursorStore {
33    cursors: Mutex<HashMap<String, ContinuousPageCursor>>,
34    max_entries: usize,
35}
36
37impl Default for InMemoryContinuousPageCursorStore {
38    fn default() -> Self {
39        Self {
40            cursors: Mutex::new(HashMap::new()),
41            max_entries: 4096,
42        }
43    }
44}
45
46#[async_trait::async_trait]
47impl ContinuousPageCursorStore for InMemoryContinuousPageCursorStore {
48    async fn get(
49        &self,
50        query_key: &str,
51        target_offset: u64,
52    ) -> Result<Option<ContinuousPageCursor>, String> {
53        let key = format!("{query_key}:{target_offset}");
54        let mut cursors = self.cursors.lock().map_err(|error| error.to_string())?;
55        if cursors
56            .get(&key)
57            .is_some_and(|cursor| cursor.expires_at <= SystemTime::now())
58        {
59            cursors.remove(&key);
60        }
61        Ok(cursors.get(&key).cloned())
62    }
63
64    async fn put(&self, cursor: ContinuousPageCursor) -> Result<(), String> {
65        let key = format!("{}:{}", cursor.query_key, cursor.next_offset);
66        let mut cursors = self.cursors.lock().map_err(|error| error.to_string())?;
67        if cursors.len() >= self.max_entries {
68            if let Some(oldest) = cursors
69                .iter()
70                .min_by_key(|(_, value)| value.expires_at)
71                .map(|(key, _)| key.clone())
72            {
73                cursors.remove(&oldest);
74            }
75        }
76        cursors.insert(key, cursor);
77        Ok(())
78    }
79
80    async fn invalidate(&self, query_key: &str) -> Result<(), String> {
81        let prefix = format!("{query_key}:");
82        self.cursors
83            .lock()
84            .map_err(|error| error.to_string())?
85            .retain(|key, _| !key.starts_with(&prefix));
86        Ok(())
87    }
88}
89
90#[derive(Debug, Clone)]
91pub struct RetainedIdSet {
92    pub query_key: String,
93    pub ids: Arc<Vec<u64>>,
94    pub expires_at: SystemTime,
95}
96
97#[async_trait::async_trait]
98pub trait IdSetStore: Send + Sync + 'static {
99    async fn get(&self, query_key: &str) -> Result<Option<RetainedIdSet>, String>;
100    async fn put(&self, id_set: RetainedIdSet) -> Result<(), String>;
101    async fn invalidate(&self, query_key: &str) -> Result<(), String>;
102}
103
104pub struct InMemoryIdSetStore {
105    sets: Mutex<HashMap<String, RetainedIdSet>>,
106    max_entries: usize,
107    max_bytes: usize,
108}
109
110impl Default for InMemoryIdSetStore {
111    fn default() -> Self {
112        Self {
113            sets: Mutex::new(HashMap::new()),
114            max_entries: 64,
115            max_bytes: 256 * 1024 * 1024,
116        }
117    }
118}
119
120impl InMemoryIdSetStore {
121    fn retained_bytes(sets: &HashMap<String, RetainedIdSet>) -> usize {
122        sets.values()
123            .map(|value| value.ids.len().saturating_mul(std::mem::size_of::<u64>()))
124            .sum()
125    }
126}
127
128#[async_trait::async_trait]
129impl IdSetStore for InMemoryIdSetStore {
130    async fn get(&self, query_key: &str) -> Result<Option<RetainedIdSet>, String> {
131        let mut sets = self.sets.lock().map_err(|error| error.to_string())?;
132        if sets
133            .get(query_key)
134            .is_some_and(|value| value.expires_at <= SystemTime::now())
135        {
136            sets.remove(query_key);
137        }
138        Ok(sets.get(query_key).cloned())
139    }
140
141    async fn put(&self, id_set: RetainedIdSet) -> Result<(), String> {
142        let incoming_bytes = id_set.ids.len().saturating_mul(std::mem::size_of::<u64>());
143        if incoming_bytes > self.max_bytes {
144            return Err("ID set exceeds the process-local store memory ceiling".to_owned());
145        }
146        let mut sets = self.sets.lock().map_err(|error| error.to_string())?;
147        sets.retain(|_, value| value.expires_at > SystemTime::now());
148        while sets.len() >= self.max_entries
149            || Self::retained_bytes(&sets).saturating_add(incoming_bytes) > self.max_bytes
150        {
151            let Some(oldest) = sets
152                .iter()
153                .min_by_key(|(_, value)| value.expires_at)
154                .map(|(key, _)| key.clone())
155            else {
156                break;
157            };
158            sets.remove(&oldest);
159        }
160        sets.insert(id_set.query_key.clone(), id_set);
161        Ok(())
162    }
163
164    async fn invalidate(&self, query_key: &str) -> Result<(), String> {
165        self.sets
166            .lock()
167            .map_err(|error| error.to_string())?
168            .remove(query_key);
169        Ok(())
170    }
171}
172
173pub(super) fn id_set_build_lock(query_key: &str) -> Arc<futures_util::lock::Mutex<()>> {
174    static LOCKS: OnceLock<Mutex<HashMap<String, std::sync::Weak<futures_util::lock::Mutex<()>>>>> =
175        OnceLock::new();
176    let mut locks = LOCKS
177        .get_or_init(|| Mutex::new(HashMap::new()))
178        .lock()
179        .expect("ID set build lock registry poisoned");
180    locks.retain(|_, lock| lock.strong_count() > 0);
181    if let Some(lock) = locks.get(query_key).and_then(std::sync::Weak::upgrade) {
182        return lock;
183    }
184    let lock = Arc::new(futures_util::lock::Mutex::new(()));
185    locks.insert(query_key.to_owned(), Arc::downgrade(&lock));
186    lock
187}
188
189impl UserContext {
190    pub fn set_continuous_page_cursor_store(&mut self, store: Arc<dyn ContinuousPageCursorStore>) {
191        self.continuous_page_cursor_store = store;
192    }
193
194    pub fn continuous_page_plan(&self) -> Option<String> {
195        self.continuous_page_observation
196            .lock()
197            .ok()
198            .map(|value| value.0.clone())
199    }
200
201    pub fn continuous_page_cursor_id(&self) -> Option<String> {
202        self.continuous_page_observation
203            .lock()
204            .ok()
205            .and_then(|value| value.1.clone())
206    }
207
208    pub(crate) fn observe_continuous_page(
209        &self,
210        plan: impl Into<String>,
211        cursor_id: Option<String>,
212    ) {
213        if let Ok(mut observation) = self.continuous_page_observation.lock() {
214            *observation = (plan.into(), cursor_id);
215        }
216    }
217
218    pub(crate) fn continuous_page_cursor_store(&self) -> &dyn ContinuousPageCursorStore {
219        self.continuous_page_cursor_store.as_ref()
220    }
221
222    pub fn set_id_set_store(&mut self, store: Arc<dyn IdSetStore>) {
223        self.id_set_store = store;
224    }
225
226    pub fn id_set_plan(&self) -> Option<String> {
227        self.id_set_observation
228            .lock()
229            .ok()
230            .map(|observation| observation.0.clone())
231    }
232
233    pub fn id_set_count(&self) -> Option<u64> {
234        self.id_set_observation
235            .lock()
236            .ok()
237            .and_then(|observation| observation.1)
238    }
239
240    pub(crate) fn observe_id_set(&self, plan: impl Into<String>, count: Option<u64>) {
241        if let Ok(mut observation) = self.id_set_observation.lock() {
242            *observation = (plan.into(), count);
243        }
244    }
245
246    pub(crate) fn id_set_store(&self) -> &dyn IdSetStore {
247        self.id_set_store.as_ref()
248    }
249
250    pub(crate) fn id_set_build_lock(&self, query_key: &str) -> Arc<futures_util::lock::Mutex<()>> {
251        id_set_build_lock(query_key)
252    }
253}
teaql_runtime/context/pagination.rs

teaql_runtime/context/
pagination.rs
