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            && 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        cursors.insert(key, cursor);
76        Ok(())
77    }
78
79    async fn invalidate(&self, query_key: &str) -> Result<(), String> {
80        let prefix = format!("{query_key}:");
81        self.cursors
82            .lock()
83            .map_err(|error| error.to_string())?
84            .retain(|key, _| !key.starts_with(&prefix));
85        Ok(())
86    }
87}
88
89#[derive(Debug, Clone)]
90pub struct RetainedIdSet {
91    pub query_key: String,
92    pub ids: Arc<Vec<u64>>,
93    pub expires_at: SystemTime,
94}
95
96#[async_trait::async_trait]
97pub trait IdSetStore: Send + Sync + 'static {
98    async fn get(&self, query_key: &str) -> Result<Option<RetainedIdSet>, String>;
99    async fn put(&self, id_set: RetainedIdSet) -> Result<(), String>;
100    async fn invalidate(&self, query_key: &str) -> Result<(), String>;
101}
102
103pub struct InMemoryIdSetStore {
104    sets: Mutex<HashMap<String, RetainedIdSet>>,
105    max_entries: usize,
106    max_bytes: usize,
107}
108
109impl Default for InMemoryIdSetStore {
110    fn default() -> Self {
111        Self {
112            sets: Mutex::new(HashMap::new()),
113            max_entries: 64,
114            max_bytes: 256 * 1024 * 1024,
115        }
116    }
117}
118
119impl InMemoryIdSetStore {
120    fn retained_bytes(sets: &HashMap<String, RetainedIdSet>) -> usize {
121        sets.values()
122            .map(|value| value.ids.len().saturating_mul(std::mem::size_of::<u64>()))
123            .sum()
124    }
125}
126
127#[async_trait::async_trait]
128impl IdSetStore for InMemoryIdSetStore {
129    async fn get(&self, query_key: &str) -> Result<Option<RetainedIdSet>, String> {
130        let mut sets = self.sets.lock().map_err(|error| error.to_string())?;
131        if sets
132            .get(query_key)
133            .is_some_and(|value| value.expires_at <= SystemTime::now())
134        {
135            sets.remove(query_key);
136        }
137        Ok(sets.get(query_key).cloned())
138    }
139
140    async fn put(&self, id_set: RetainedIdSet) -> Result<(), String> {
141        let incoming_bytes = id_set.ids.len().saturating_mul(std::mem::size_of::<u64>());
142        if incoming_bytes > self.max_bytes {
143            return Err("ID set exceeds the process-local store memory ceiling".to_owned());
144        }
145        let mut sets = self.sets.lock().map_err(|error| error.to_string())?;
146        sets.retain(|_, value| value.expires_at > SystemTime::now());
147        while sets.len() >= self.max_entries
148            || Self::retained_bytes(&sets).saturating_add(incoming_bytes) > self.max_bytes
149        {
150            let Some(oldest) = sets
151                .iter()
152                .min_by_key(|(_, value)| value.expires_at)
153                .map(|(key, _)| key.clone())
154            else {
155                break;
156            };
157            sets.remove(&oldest);
158        }
159        sets.insert(id_set.query_key.clone(), id_set);
160        Ok(())
161    }
162
163    async fn invalidate(&self, query_key: &str) -> Result<(), String> {
164        self.sets
165            .lock()
166            .map_err(|error| error.to_string())?
167            .remove(query_key);
168        Ok(())
169    }
170}
171
172pub(super) fn id_set_build_lock(query_key: &str) -> Arc<futures_util::lock::Mutex<()>> {
173    static LOCKS: OnceLock<Mutex<HashMap<String, std::sync::Weak<futures_util::lock::Mutex<()>>>>> =
174        OnceLock::new();
175    let mut locks = LOCKS
176        .get_or_init(|| Mutex::new(HashMap::new()))
177        .lock()
178        .expect("ID set build lock registry poisoned");
179    locks.retain(|_, lock| lock.strong_count() > 0);
180    if let Some(lock) = locks.get(query_key).and_then(std::sync::Weak::upgrade) {
181        return lock;
182    }
183    let lock = Arc::new(futures_util::lock::Mutex::new(()));
184    locks.insert(query_key.to_owned(), Arc::downgrade(&lock));
185    lock
186}
187
188impl UserContext {
189    pub fn set_continuous_page_cursor_store(&mut self, store: Arc<dyn ContinuousPageCursorStore>) {
190        self.continuous_page_cursor_store = store;
191    }
192
193    pub fn continuous_page_plan(&self) -> Option<String> {
194        self.continuous_page_observation
195            .lock()
196            .ok()
197            .map(|value| value.0.clone())
198    }
199
200    pub fn continuous_page_cursor_id(&self) -> Option<String> {
201        self.continuous_page_observation
202            .lock()
203            .ok()
204            .and_then(|value| value.1.clone())
205    }
206
207    pub(crate) fn observe_continuous_page(
208        &self,
209        plan: impl Into<String>,
210        cursor_id: Option<String>,
211    ) {
212        if let Ok(mut observation) = self.continuous_page_observation.lock() {
213            *observation = (plan.into(), cursor_id);
214        }
215    }
216
217    pub(crate) fn continuous_page_cursor_store(&self) -> &dyn ContinuousPageCursorStore {
218        self.continuous_page_cursor_store.as_ref()
219    }
220
221    pub fn set_id_set_store(&mut self, store: Arc<dyn IdSetStore>) {
222        self.id_set_store = store;
223    }
224
225    pub fn id_set_plan(&self) -> Option<String> {
226        self.id_set_observation
227            .lock()
228            .ok()
229            .map(|observation| observation.0.clone())
230    }
231
232    pub fn id_set_count(&self) -> Option<u64> {
233        self.id_set_observation
234            .lock()
235            .ok()
236            .and_then(|observation| observation.1)
237    }
238
239    pub(crate) fn observe_id_set(&self, plan: impl Into<String>, count: Option<u64>) {
240        if let Ok(mut observation) = self.id_set_observation.lock() {
241            *observation = (plan.into(), count);
242        }
243    }
244
245    pub(crate) fn id_set_store(&self) -> &dyn IdSetStore {
246        self.id_set_store.as_ref()
247    }
248
249    pub(crate) fn id_set_build_lock(&self, query_key: &str) -> Arc<futures_util::lock::Mutex<()>> {
250        id_set_build_lock(query_key)
251    }
252}
teaql_runtime/context/pagination.rs

teaql_runtime/context/
pagination.rs
