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//! - 存在进程内存里，服务重启后附件不在了（任务结果若持久化，重启后仍在，附件需要重抓）。
7//! - 下载接口和其他抓取接口一样在鉴权之后；id 是 128 位哈希前缀，猜不到。
8
9use std::collections::{BTreeMap, HashMap};
10use std::sync::{Arc, Mutex, OnceLock};
11use std::time::Instant;
12
13use axum::extract::Path;
14use axum::http::header;
15use axum::response::{IntoResponse, Response};
16
17use crate::error::{ApiError, ErrorCode};
18use crate::types::Document;
19
20struct Entry {
21    data: Arc<[u8]>,
22    mime: &'static str,
23    expires: Instant,
24}
25
26#[derive(Default)]
27struct Store {
28    map: HashMap<String, Entry>,
29    /// （过期时刻，id）：按过期先后淘汰。
30    order: BTreeMap<(Instant, u64), String>,
31    seq: u64,
32    bytes: usize,
33}
34
35fn store() -> &'static Mutex<Store> {
36    static STORE: OnceLock<Mutex<Store>> = OnceLock::new();
37    STORE.get_or_init(Default::default)
38}
39
40/// 存储上限：内存预算的 δ（与统计表同一份额）。
41fn cap() -> usize {
42    // [δ] 全局风险水平的默认值 0.05，作为愿意分给附件的内存份额，与统计表一致
43    let delta = 0.05;
44    // [策略] 读不到内存预算（非 Linux、刷新之前）时的回退：64 MiB，约 300 张 200 KiB 的图片
45    let fallback = 64 * 1024 * 1024;
46    pagesift_net::resources::budget()
47        .map_or(fallback, |b| (b as f64 * delta) as usize)
48        .max(1)
49}
50
51/// 记下文档里的附件内容（没有内容的跳过）。
52pub fn keep(doc: &Document) {
53    if doc.attachments.iter().all(|a| a.data.is_none()) {
54        return;
55    }
56    let expires = Instant::now() + crate::jobs::JOB_TTL;
57    let cap = cap();
58    let mut st = store().lock().unwrap_or_else(|e| e.into_inner());
59    let now = Instant::now();
60    // 先清掉过期的
61    while let Some((&key, _)) = st.order.first_key_value() {
62        if key.0 > now {
63            break;
64        }
65        drop_oldest(&mut st);
66    }
67    for a in &doc.attachments {
68        let Some(data) = &a.data else { continue };
69        if data.len() > cap {
70            continue;
71        }
72        if let Some(old) = st.map.remove(&a.id) {
73            st.bytes -= old.data.len();
74            let id = a.id.clone();
75            st.order.retain(|_, v| *v != id);
76        }
77        while st.bytes + data.len() > cap && !st.order.is_empty() {
78            drop_oldest(&mut st);
79        }
80        st.seq += 1;
81        let seq = st.seq;
82        st.order.insert((expires, seq), a.id.clone());
83        st.bytes += data.len();
84        st.map.insert(
85            a.id.clone(),
86            Entry {
87                data: data.clone(),
88                mime: a.mime(),
89                expires,
90            },
91        );
92    }
93}
94
95fn drop_oldest(st: &mut Store) {
96    if let Some((_, id)) = st.order.pop_first()
97        && let Some(e) = st.map.remove(&id)
98    {
99        st.bytes -= e.data.len();
100    }
101}
102
103/// 取一个附件：（内容，MIME 类型）。过期或不存在时为 `None`。
104pub fn get(id: &str) -> Option<(Arc<[u8]>, &'static str)> {
105    let st = store().lock().unwrap_or_else(|e| e.into_inner());
106    st.map
107        .get(id)
108        .filter(|e| e.expires > Instant::now())
109        .map(|e| (e.data.clone(), e.mime))
110}
111
112/// `GET /v2/attachments/{id}`。
113pub(crate) async fn download(Path(id): Path<String>) -> Result<Response, ApiError> {
114    let Some((data, mime)) = get(&id) else {
115        return Err(ApiError::new(
116            ErrorCode::NotFound,
117            format!("attachment {id} not found or expired"),
118        ));
119    };
120    Ok((
121        [
122            (header::CONTENT_TYPE, mime),
123            (header::CACHE_CONTROL, "private, max-age=86400, immutable"),
124        ],
125        data.to_vec(),
126    )
127        .into_response())
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133    use crate::types::Attachment;
134
135    #[test]
136    fn kept_attachments_can_be_fetched_by_id() {
137        let mut doc = Document::default();
138        doc.attachments
139            .push(Attachment::new(b"GIF89a-test".to_vec(), "gif", None));
140        keep(&doc);
141        let id = &doc.attachments[0].id;
142        let (data, mime) = get(id).unwrap();
143        assert_eq!(&*data, b"GIF89a-test");
144        assert_eq!(mime, "image/gif");
145        assert!(get("0000.png").is_none());
146    }
147}
pagesift_api/attachments.rs

pagesift_api/
attachments.rs
