pagesift_api/
attachments.rs

1//! 附件的存储与下载：响应里只列 `{id, format, bytes, hash}`，内容用 `GET /v2/attachments/{id}` 取。
2//!
3//! - 保留期与结果相同（[`crate::jobs::JOB_TTL`]，任务结果的保留期）；每次再出现同一内容就续期。
4//! - 内容寻址（id 是 SHA-256 前缀加扩展名），同一内容只存一份。
5//! - 占用有上限：内存预算（`pagesift_net::resources::budget`）的 `δ`，超出时先淘汰最早过期的。
6//! - 先存在进程内存里。服务开了任务持久化（有存储）时同时写进存储（HTTP API；MCP 只用内存），内存里淘汰了、
7//!   或者服务重启了，下载接口从存储里取，到期前链接一直有效；0.90 只存内存，重启或被新图片挤掉后
8//!   链接就 404，而任务结果还在。到期的在启动时和下次访问时从存储删掉。
9//! - 有温备时（网关的任务镜像），附件的写入和删除同任务记录一样复制到温备，快照里也带着；
10//!   单个附件超过一条复制记录的上限（帧上限的一半，[`crate::JobLog::max_op`]）时只存本地、记警告，
11//!   温备接管后这个附件的链接 404。
12//! - 没有存储时只在内存里：重启后没了，内存满了先淘汰最早过期的（易失缓存，不是可靠交付）。
13//! - 下载接口和其他抓取接口一样在鉴权之后；id 是 128 位哈希前缀，猜不到。
14
15use std::collections::{BTreeMap, HashMap};
16use std::sync::{Arc, Mutex, OnceLock};
17use std::time::Instant;
18
19use axum::extract::{Path, State};
20use axum::http::header;
21use axum::response::{IntoResponse, Response};
22
23use crate::error::{ApiError, ErrorCode};
24use crate::types::Document;
25use crate::types::attachment::mime_of;
26
27/// 存储里附件内容的键（`attachment/<id>`）和到期时刻的键（`attachment-exp/<id>`，Unix 毫秒）。
28/// 先写内容再写到期时刻：有到期时刻就一定有内容。
29fn data_key(id: &str) -> String {
30    format!("attachment/{id}")
31}
32
33fn exp_key(id: &str) -> String {
34    format!("attachment-exp/{id}")
35}
36
37fn now_ms() -> u64 {
38    std::time::SystemTime::now()
39        .duration_since(std::time::UNIX_EPOCH)
40        .map_or(0, |d| d.as_millis() as u64)
41}
42
43/// 附件写进的地方：任务存储，和（网关上）把写入复制到温备的镜像。
44#[derive(Clone)]
45pub(crate) struct Disk {
46    pub store: pagesift_store::Store,
47    pub mirror: Option<Arc<crate::jobs::mirror::JobMirror>>,
48}
49
50/// 删掉一个附件（经镜像，温备上也删）。
51fn remove(d: &Disk, id: &str) {
52    let r = crate::jobs::mirror::write(&d.store, d.mirror.as_deref(), |w| {
53        w.kv_delete(&data_key(id))?;
54        w.kv_delete(&exp_key(id))
55    });
56    if let Err(e) = r {
57        tracing::warn!(attachment = %id, error = %e, "deleting an expired attachment failed");
58    }
59}
60
61/// （键，值）。
62type Kv = (String, Vec<u8>);
63
64/// 存储里的全部附件：（（内容键，内容），（到期键，到期时刻））。快照用。
65pub(crate) fn stored(t: &pagesift_store::Tasks<'_>) -> Vec<(Kv, Kv)> {
66    let mut out = Vec::new();
67    for k in t.kv_keys("attachment-exp/").unwrap_or_default() {
68        let Ok(k) = String::from_utf8(k) else {
69            continue;
70        };
71        let id = &k["attachment-exp/".len()..];
72        if let (Ok(Some(exp)), Ok(Some(data))) = (t.kv_get(&k), t.kv_get(data_key(id))) {
73            out.push(((data_key(id), data), (k.clone(), exp)));
74        }
75    }
76    out
77}
78
79/// 清掉存储里已经到期的附件（建服务时调用，在后台做）。要在 tokio 运行时里调用。
80pub(crate) fn sweep_expired(d: Disk) {
81    tokio::spawn(async move {
82        let s = d.store.clone();
83        s.blocking(move |s| {
84            let now = now_ms();
85            for k in s.tasks().kv_keys("attachment-exp/").unwrap_or_default() {
86                let Ok(k) = String::from_utf8(k) else {
87                    continue;
88                };
89                let id = &k["attachment-exp/".len()..];
90                if read_exp(s, id).is_none_or(|e| e <= now) {
91                    remove(&d, id);
92                }
93            }
94        })
95        .await;
96    });
97}
98
99fn read_exp(s: &pagesift_store::Store, id: &str) -> Option<u64> {
100    let v = s.tasks().kv_get(exp_key(id)).ok()??;
101    Some(u64::from_le_bytes(v.as_slice().try_into().ok()?))
102}
103
104/// 写进存储（已有就只续期），有镜像时同时复制到温备。在后台做，不等。
105fn save(disk: Option<&Disk>, id: String, data: Arc<[u8]>) {
106    let (Some(d), Ok(rt)) = (disk.cloned(), tokio::runtime::Handle::try_current()) else {
107        return;
108    };
109    let exp = now_ms() + crate::jobs::JOB_TTL.as_millis() as u64;
110    rt.spawn(async move {
111        let s = d.store.clone();
112        s.blocking(move |s| {
113            // 超过一条复制记录上限的只存本地
114            let mirror = d
115                .mirror
116                .as_deref()
117                .filter(|m| data.len() + data_key(&id).len() < m.max_op());
118            if d.mirror.is_some() && mirror.is_none() {
119                tracing::warn!(attachment = %id, bytes = data.len(), "attachment too large to replicate to the standby; stored locally only");
120            }
121            let fresh = read_exp(s, &id).is_none();
122            let r = crate::jobs::mirror::write(s, mirror, |w| {
123                if fresh {
124                    w.kv_put(&data_key(&id), &data)?;
125                }
126                w.kv_put(&exp_key(&id), &exp.to_le_bytes())
127            });
128            if let Err(e) = r {
129                tracing::warn!(attachment = %id, error = %e, "saving an attachment failed; it stays in memory only");
130            }
131        })
132        .await;
133    });
134}
135
136/// 从存储取（内存里没有时）。到期的顺手删掉。
137async fn load(disk: Option<Disk>, id: String) -> Option<(Arc<[u8]>, &'static str)> {
138    let d = disk?;
139    let store = d.store.clone();
140    store
141        .blocking(move |s| {
142            let exp = read_exp(s, &id)?;
143            if exp <= now_ms() {
144                remove(&d, &id);
145                return None;
146            }
147            let data = s.tasks().kv_get(data_key(&id)).ok()??;
148            let ext = id.rsplit_once('.').map_or("", |(_, e)| e);
149            Some((Arc::from(data), mime_of(ext)))
150        })
151        .await
152}
153
154struct Entry {
155    data: Arc<[u8]>,
156    mime: &'static str,
157    expires: Instant,
158}
159
160#[derive(Default)]
161struct Store {
162    map: HashMap<String, Entry>,
163    /// （过期时刻，id）：按过期先后淘汰。
164    order: BTreeMap<(Instant, u64), String>,
165    seq: u64,
166    bytes: usize,
167}
168
169fn store() -> &'static Mutex<Store> {
170    static STORE: OnceLock<Mutex<Store>> = OnceLock::new();
171    STORE.get_or_init(Default::default)
172}
173
174/// 存储上限：内存预算的 δ（与统计表同一份额）。
175fn cap() -> usize {
176    // [δ] 全局风险水平的默认值 0.05，作为愿意分给附件的内存份额，与统计表一致
177    let delta = 0.05;
178    // [策略] 读不到内存预算（非 Linux、刷新之前）时的回退：64 MiB，约 300 张 200 KiB 的图片
179    let fallback = 64 * 1024 * 1024;
180    pagesift_net::resources::budget()
181        .map_or(fallback, |b| (b as f64 * delta) as usize)
182        .max(1)
183}
184
185/// 记下文档里的附件内容（没有内容的跳过）；有 `disk`（任务持久化的存储）时也写进去。
186pub fn keep(doc: &Document, disk: Option<&pagesift_store::Store>) {
187    let disk = disk.map(|store| Disk {
188        store: store.clone(),
189        mirror: None,
190    });
191    keep_in(doc, disk.as_ref());
192}
193
194/// 同 [`keep`]，写进任务存储并经镜像复制到温备（HTTP API 用）。
195pub(crate) fn keep_in(doc: &Document, disk: Option<&Disk>) {
196    if doc.attachments.iter().all(|a| a.data.is_none()) {
197        return;
198    }
199    let expires = Instant::now() + crate::jobs::JOB_TTL;
200    let cap = cap();
201    let mut st = store().lock().unwrap_or_else(|e| e.into_inner());
202    let now = Instant::now();
203    // 先清掉过期的
204    while let Some((&key, _)) = st.order.first_key_value() {
205        if key.0 > now {
206            break;
207        }
208        drop_oldest(&mut st);
209    }
210    for a in &doc.attachments {
211        let Some(data) = &a.data else { continue };
212        if data.len() > cap {
213            // 超过内存上限的只进存储（有的话）
214            save(disk, a.id.clone(), data.clone());
215            continue;
216        }
217        if let Some(old) = st.map.remove(&a.id) {
218            st.bytes -= old.data.len();
219            let id = a.id.clone();
220            st.order.retain(|_, v| *v != id);
221        }
222        while st.bytes + data.len() > cap && !st.order.is_empty() {
223            drop_oldest(&mut st);
224        }
225        st.seq += 1;
226        let seq = st.seq;
227        st.order.insert((expires, seq), a.id.clone());
228        st.bytes += data.len();
229        st.map.insert(
230            a.id.clone(),
231            Entry {
232                data: data.clone(),
233                mime: a.mime(),
234                expires,
235            },
236        );
237        save(disk, a.id.clone(), data.clone());
238    }
239}
240
241fn drop_oldest(st: &mut Store) {
242    if let Some((_, id)) = st.order.pop_first()
243        && let Some(e) = st.map.remove(&id)
244    {
245        st.bytes -= e.data.len();
246    }
247}
248
249/// 取一个附件：（内容，MIME 类型）。过期或不存在时为 `None`。
250pub fn get(id: &str) -> Option<(Arc<[u8]>, &'static str)> {
251    let st = store().lock().unwrap_or_else(|e| e.into_inner());
252    st.map
253        .get(id)
254        .filter(|e| e.expires > Instant::now())
255        .map(|e| (e.data.clone(), e.mime))
256}
257
258/// `GET /v2/attachments/{id}`。
259pub(crate) async fn download(
260    State(s): State<crate::server::AppState>,
261    Path(id): Path<String>,
262) -> Result<Response, ApiError> {
263    let found = match get(&id) {
264        Some(x) => Some(x),
265        None => load(s.jobs.disk(), id.clone()).await,
266    };
267    let Some((data, mime)) = found else {
268        return Err(ApiError::new(
269            ErrorCode::NotFound,
270            format!("attachment {id} not found or expired"),
271        ));
272    };
273    Ok((
274        [
275            (header::CONTENT_TYPE, mime),
276            (header::CACHE_CONTROL, "private, max-age=86400, immutable"),
277        ],
278        data.to_vec(),
279    )
280        .into_response())
281}
282
283#[cfg(test)]
284mod tests {
285    use super::*;
286    use crate::types::Attachment;
287
288    #[test]
289    fn kept_attachments_can_be_fetched_by_id() {
290        let mut doc = Document::default();
291        doc.attachments
292            .push(Attachment::new(b"GIF89a-test".to_vec(), "gif", None));
293        keep(&doc, None);
294        let id = &doc.attachments[0].id;
295        let (data, mime) = get(id).unwrap();
296        assert_eq!(&*data, b"GIF89a-test");
297        assert_eq!(mime, "image/gif");
298        assert!(get("0000.png").is_none());
299    }
300
301    /// 有存储时，内存里没有了（重启、被挤掉）也能从存储读回；0.90 这时 404。
302    #[tokio::test(flavor = "multi_thread")]
303    async fn persisted_attachments_survive_losing_the_memory_copy() {
304        let dir = tempfile::tempdir().unwrap();
305        let store = pagesift_store::Store::open_dir(dir.path()).unwrap();
306        let mut doc = Document::default();
307        doc.attachments
308            .push(Attachment::new(b"\x89PNG-persisted".to_vec(), "png", None));
309        keep(&doc, Some(&store));
310        let id = doc.attachments[0].id.clone();
311        // 等后台写完
312        for _ in 0..500 {
313            if read_exp(&store, &id).is_some() {
314                break;
315            }
316            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
317        }
318        // 模拟重启：清空内存
319        {
320            let mut st = super::store().lock().unwrap();
321            let e = st.map.remove(&id).unwrap();
322            st.bytes -= e.data.len();
323            st.order.retain(|_, v| *v != id);
324        }
325        assert!(get(&id).is_none());
326        let (data, mime) = load(
327            Some(Disk {
328                store: store.clone(),
329                mirror: None,
330            }),
331            id.clone(),
332        )
333        .await
334        .unwrap();
335        assert_eq!((&*data, mime), (&b"\x89PNG-persisted"[..], "image/png"));
336        // 到期的不再给出，并从存储删掉
337        store
338            .tasks()
339            .kv_put(exp_key(&id), &1u64.to_le_bytes())
340            .unwrap();
341        assert!(
342            load(
343                Some(Disk {
344                    store: store.clone(),
345                    mirror: None
346                }),
347                id.clone()
348            )
349            .await
350            .is_none()
351        );
352        assert!(store.tasks().kv_get(data_key(&id)).unwrap().is_none());
353    }
354
355    struct Ops(Mutex<Vec<Vec<u8>>>);
356
357    impl crate::JobLog for Ops {
358        fn append(&self, op: Vec<u8>) -> Result<(), String> {
359            self.0.lock().unwrap().push(op);
360            Ok(())
361        }
362    }
363
364    /// 网关上的附件随任务记录复制：温备回放日志后能下载，从快照重建的副本也能（0.90 不复制附件）。
365    #[tokio::test(flavor = "multi_thread")]
366    async fn attachments_replicate_with_the_job_records() {
367        let dirs: Vec<_> = (0..3).map(|_| tempfile::tempdir().unwrap()).collect();
368        let open = |i: usize| pagesift_store::Store::open_dir(dirs[i].path()).unwrap();
369        let (src, dst, from_snapshot) = (open(0), open(1), open(2));
370        let log = Arc::new(Ops(Mutex::default()));
371        let mirror = crate::JobMirror::new(log.clone());
372        let mut doc = Document::default();
373        doc.attachments
374            .push(Attachment::new(b"GIF89a-replicated".to_vec(), "gif", None));
375        let disk = Disk {
376            store: src.clone(),
377            mirror: Some(mirror.clone()),
378        };
379        keep_in(&doc, Some(&disk));
380        let id = doc.attachments[0].id.clone();
381        for _ in 0..500 {
382            if log.0.lock().unwrap().len() >= 2 {
383                break;
384            }
385            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
386        }
387        for op in log.0.lock().unwrap().iter() {
388            crate::apply_op(&dst, op).unwrap();
389        }
390        for op in mirror.dump(&src) {
391            crate::apply_op(&from_snapshot, &op).unwrap();
392        }
393        for s in [dst, from_snapshot] {
394            let d = Disk {
395                store: s,
396                mirror: None,
397            };
398            let (data, mime) = load(Some(d), id.clone()).await.expect("replicated");
399            assert_eq!((&*data, mime), (&b"GIF89a-replicated"[..], "image/gif"));
400        }
401    }
402}
pagesift_api/attachments.rs

pagesift_api/
attachments.rs
