pagesift_store/
tasks.rs

1//! 爬取任务状态的持久化原语（F-6.4）：KV、追加列表（队列）、集合。
2//!
3//! 每次调用是一条独立记录，单条写入是原子的，多条之间不是。
4//! 例如爬虫应先 `set_insert` 已完成，再 `list_pop_front`，崩溃时最多重复抓一个页面。
5
6use std::time::SystemTime;
7
8use crate::engine::{Tx, group_key, ns};
9use crate::{Result, Store, to_ms};
10
11pub struct Tasks<'a> {
12    store: &'a Store,
13}
14
15impl<'a> Tasks<'a> {
16    pub(crate) fn new(store: &'a Store) -> Self {
17        Self { store }
18    }
19
20    // ---- KV ----
21
22    pub fn kv_get(&self, key: impl AsRef<[u8]>) -> Result<Option<Vec<u8>>> {
23        Ok(self.store.get(ns::KV, key.as_ref())?.map(|(_, v)| v))
24    }
25
26    pub fn kv_put(&self, key: impl AsRef<[u8]>, value: &[u8]) -> Result<()> {
27        self.store
28            .put(ns::KV, key.as_ref(), value, now_ms())
29            .map(drop)
30    }
31
32    /// 返回键之前是否存在。
33    pub fn kv_delete(&self, key: impl AsRef<[u8]>) -> Result<bool> {
34        self.store.delete(ns::KV, key.as_ref())
35    }
36
37    /// 以 `prefix` 开头的所有键，按字节序排序。要逐条读盘，适合启动时列出未完成任务这类低频操作。
38    pub fn kv_keys(&self, prefix: impl AsRef<[u8]>) -> Result<Vec<Vec<u8>>> {
39        let p = prefix.as_ref();
40        self.store.keys(ns::KV, |k| k.starts_with(p))
41    }
42
43    // ---- 列表 ----
44
45    /// 追加到列表尾部，返回序号。
46    pub fn list_push(&self, list: impl AsRef<[u8]>, item: &[u8]) -> Result<u64> {
47        let r = list_push(
48            &mut self.store.tx(),
49            ns::LIST_ITEM,
50            list.as_ref(),
51            item,
52            now_ms(),
53        );
54        self.store.finish();
55        r
56    }
57
58    /// 弹出列表头部。
59    pub fn list_pop_front(&self, list: impl AsRef<[u8]>) -> Result<Option<Vec<u8>>> {
60        let list = list.as_ref();
61        let mut tx = self.store.tx();
62        let Some((seq, item)) = list_front(&tx, ns::LIST_ITEM, list)? else {
63            return Ok(None);
64        };
65        tx.delete(ns::LIST_ITEM, &group_key(list, &seq.to_be_bytes())?)?;
66        drop(tx);
67        self.store.finish();
68        Ok(Some(item))
69    }
70
71    pub fn list_front(&self, list: impl AsRef<[u8]>) -> Result<Option<Vec<u8>>> {
72        let tx = self.store.tx();
73        Ok(list_front(&tx, ns::LIST_ITEM, list.as_ref())?.map(|(_, v)| v))
74    }
75
76    pub fn list_len(&self, list: impl AsRef<[u8]>) -> u64 {
77        self.store
78            .group(ns::LIST_ITEM, list.as_ref())
79            .map_or(0, |g| g.live)
80    }
81
82    /// 从头部开始最多 `limit` 个元素。
83    pub fn list_items(&self, list: impl AsRef<[u8]>, limit: usize) -> Result<Vec<Vec<u8>>> {
84        Ok(list_scan(self.store, ns::LIST_ITEM, list.as_ref(), limit)?
85            .into_iter()
86            .map(|(_, _, v)| v)
87            .collect())
88    }
89
90    /// 从序号 `from` 开始最多 `limit` 个元素，跳过已弹出的。只追加、不弹出的列表里序号就是下标，
91    /// 可以用来分页。
92    pub fn list_range(
93        &self,
94        list: impl AsRef<[u8]>,
95        from: u64,
96        limit: usize,
97    ) -> Result<Vec<Vec<u8>>> {
98        let list = list.as_ref();
99        let mut out = Vec::new();
100        let Some(g) = self.store.group(ns::LIST_ITEM, list) else {
101            return Ok(out);
102        };
103        for seq in from.max(g.head)..g.next {
104            if out.len() >= limit {
105                break;
106            }
107            if let Some((_, v)) = self
108                .store
109                .get(ns::LIST_ITEM, &group_key(list, &seq.to_be_bytes())?)?
110            {
111                out.push(v);
112            }
113        }
114        Ok(out)
115    }
116
117    /// 清空列表，返回删除的元素个数。
118    pub fn list_clear(&self, list: impl AsRef<[u8]>) -> Result<u64> {
119        let list = list.as_ref();
120        let mut tx = self.store.tx();
121        let mut n = 0;
122        if let Some(g) = tx.group(ns::LIST_ITEM, list) {
123            for seq in g.head..g.next {
124                if tx.delete(ns::LIST_ITEM, &group_key(list, &seq.to_be_bytes())?)? {
125                    n += 1;
126                }
127            }
128        }
129        drop(tx);
130        self.store.finish();
131        Ok(n)
132    }
133
134    // ---- 集合 ----
135
136    /// 返回是否是新加入的成员。
137    pub fn set_insert(&self, set: impl AsRef<[u8]>, member: impl AsRef<[u8]>) -> Result<bool> {
138        let key = group_key(set.as_ref(), member.as_ref())?;
139        let mut tx = self.store.tx();
140        if tx.contains(ns::SET_MEMBER, &key)? {
141            return Ok(false);
142        }
143        tx.put(ns::SET_MEMBER, &key, &[], now_ms())?;
144        drop(tx);
145        self.store.finish();
146        Ok(true)
147    }
148
149    pub fn set_contains(&self, set: impl AsRef<[u8]>, member: impl AsRef<[u8]>) -> Result<bool> {
150        let key = group_key(set.as_ref(), member.as_ref())?;
151        self.store.contains(ns::SET_MEMBER, &key)
152    }
153
154    pub fn set_remove(&self, set: impl AsRef<[u8]>, member: impl AsRef<[u8]>) -> Result<bool> {
155        let key = group_key(set.as_ref(), member.as_ref())?;
156        self.store.delete(ns::SET_MEMBER, &key)
157    }
158
159    pub fn set_len(&self, set: impl AsRef<[u8]>) -> u64 {
160        self.store
161            .group(ns::SET_MEMBER, set.as_ref())
162            .map_or(0, |g| g.live)
163    }
164
165    /// 所有成员（逐条读盘，低频使用）。
166    pub fn set_members(&self, set: impl AsRef<[u8]>) -> Result<Vec<Vec<u8>>> {
167        let prefix = group_key(set.as_ref(), &[])?;
168        Ok(self
169            .store
170            .keys(ns::SET_MEMBER, |k| k.starts_with(&prefix))?
171            .into_iter()
172            .map(|k| k[prefix.len()..].to_vec())
173            .collect())
174    }
175
176    /// 清空集合，返回删除的成员个数。
177    pub fn set_clear(&self, set: impl AsRef<[u8]>) -> Result<u64> {
178        let set = set.as_ref();
179        let mut n = 0;
180        for m in self.set_members(set)? {
181            if self.set_remove(set, &m)? {
182                n += 1;
183            }
184        }
185        Ok(n)
186    }
187}
188
189fn now_ms() -> u64 {
190    to_ms(SystemTime::now())
191}
192
193pub(crate) fn list_push(tx: &mut Tx<'_>, n: u8, list: &[u8], item: &[u8], ts: u64) -> Result<u64> {
194    let seq = tx.group(n, list).map_or(0, |g| g.next);
195    tx.put(n, &group_key(list, &seq.to_be_bytes())?, item, ts)?;
196    Ok(seq)
197}
198
199/// 头部元素 (序号, 值)。
200fn list_front(tx: &Tx<'_>, n: u8, list: &[u8]) -> Result<Option<(u64, Vec<u8>)>> {
201    let Some(g) = tx.group(n, list) else {
202        return Ok(None);
203    };
204    for seq in g.head..g.next {
205        if let Some((_, v)) = tx.get(n, &group_key(list, &seq.to_be_bytes())?)? {
206            if seq != g.head {
207                tx.set_head(n, list, seq);
208            }
209            return Ok(Some((seq, v)));
210        }
211    }
212    Ok(None)
213}
214
215/// 从头部开始顺序读，返回 (序号, 写入时间, 值)。
216pub(crate) fn list_scan(
217    store: &Store,
218    n: u8,
219    list: &[u8],
220    limit: usize,
221) -> Result<Vec<(u64, u64, Vec<u8>)>> {
222    let mut out = Vec::new();
223    let Some(g) = store.group(n, list) else {
224        return Ok(out);
225    };
226    for seq in g.head..g.next {
227        if out.len() >= limit || out.len() as u64 >= g.live {
228            break;
229        }
230        if let Some((ts, v)) = store.get(n, &group_key(list, &seq.to_be_bytes())?)? {
231            out.push((seq, ts, v));
232        }
233    }
234    Ok(out)
235}
pagesift_store/tasks.rs

pagesift_store/
tasks.rs
