mf_persistence/subscriber/
mod.rs

1//! 基于事件总线的持久化订阅者：在 TrApply 事件中持久化事务并按策略写入快照
2use std::{collections::HashSet, fmt, sync::Arc};
3
4use async_trait::async_trait;
5use dashmap::DashMap;
6use mf_core::{
7    event::{Event, EventHandler},
8    ForgeResult,
9};
10use mf_state::transaction::Transaction;
11use serde_json::json;
12use uuid::Uuid;
13
14use crate::api::{CommitMode, EventStore, PersistOptions, PersistedEvent, Snapshot};
15use crate::ser::{
16    checksum32, compress_if_needed, frame_steps, SnapshotData, TypeWrapper,
17};
18
19#[derive(Default)]
20struct SnapshotCounters {
21    last_snapshot_ms: i64,
22    last_snapshot_upto_lsn: i64,
23    events_since: u32,
24    bytes_since: u64,
25}
26
27pub struct SnapshotSubscriber<E: EventStore + 'static> {
28    store: Arc<E>,
29    options: PersistOptions,
30    commit_mode: CommitMode,
31    default_doc_id: String,
32    // 进程内已持久化的事务ID集合，避免重复写入（进程重启后靠数据库幂等键）
33    persisted: DashMap<Uuid, ()>, // ✅ 改用 UUID 作为 key
34    // 每个文档的快照计数器
35    snap_counters: DashMap<String, SnapshotCounters>,
36    // 待写快照的 upto_lsn
37    pending_snapshot_lsn: DashMap<String, i64>,
38}
39
40impl<E: EventStore + 'static> fmt::Debug for SnapshotSubscriber<E> {
41    fn fmt(
42        &self,
43        f: &mut fmt::Formatter<'_>,
44    ) -> fmt::Result {
45        f.write_str("SnapshotSubscriber")
46    }
47}
48
49impl<E: EventStore + 'static> Clone for SnapshotSubscriber<E> {
50    fn clone(&self) -> Self {
51        Self {
52            store: self.store.clone(),
53            options: self.options.clone(),
54            commit_mode: self.commit_mode,
55            default_doc_id: self.default_doc_id.clone(),
56            persisted: DashMap::new(),
57            snap_counters: DashMap::new(),
58            pending_snapshot_lsn: DashMap::new(),
59        }
60    }
61}
62
63impl<E: EventStore + 'static> SnapshotSubscriber<E> {
64    pub fn new(
65        store: Arc<E>,
66        options: PersistOptions,
67        default_doc_id: impl Into<String>,
68    ) -> Self {
69        Self {
70            commit_mode: options.commit_mode,
71            store,
72            options,
73            default_doc_id: default_doc_id.into(),
74            persisted: DashMap::new(),
75            snap_counters: DashMap::new(),
76            pending_snapshot_lsn: DashMap::new(),
77        }
78    }
79
80    async fn persist_one(
81        &self,
82        tr_id: Uuid, // ✅ 改用 UUID
83        doc_id: &str,
84        transaction: &Transaction,
85    ) -> ForgeResult<Option<(i64, usize)>> {
86        // 将每个步骤包装为 {type_id, data}
87        let frames: Vec<TypeWrapper> = frame_steps(transaction);
88        let framed = serde_json::to_vec(&frames).map_err(|e| {
89            mf_core::error::error_utils::middleware_error(format!(
90                "步骤编码失败: {e}"
91            ))
92        })?;
93        let payload = compress_if_needed(&framed, self.options.compression)
94            .map_err(|e| {
95                mf_core::error::error_utils::middleware_error(e.to_string())
96            })?;
97        let checksum = checksum32(&payload);
98
99        let ev = PersistedEvent {
100            lsn: 0,
101            tr_id,
102            doc_id: doc_id.to_string(),
103            ts: chrono::Utc::now().timestamp_millis(),
104            actor: None,
105            idempotency_key: format!("tr:{tr_id}"),
106            payload,
107            meta: json!({}),
108            checksum,
109        };
110
111        match self.commit_mode {
112            CommitMode::MemoryOnly => Ok(None),
113            _ => {
114                let bytes = ev.payload.len();
115                self.store
116                    .append(ev)
117                    .await
118                    .map(|lsn| Some((lsn, bytes)))
119                    .map_err(|e| {
120                        mf_core::error::error_utils::middleware_error(
121                            e.to_string(),
122                        )
123                    })
124            },
125        }
126    }
127
128    fn should_snapshot(
129        &self,
130        doc_id: &str,
131    ) -> bool {
132        if let Some(counters) = self.snap_counters.get(doc_id) {
133            let now = chrono::Utc::now().timestamp_millis();
134            let time_ok = self.options.snapshot_every_ms > 0
135                && now - counters.last_snapshot_ms
136                    >= self.options.snapshot_every_ms as i64;
137            let count_ok = self.options.snapshot_every_n_events > 0
138                && counters.events_since
139                    >= self.options.snapshot_every_n_events;
140            let bytes_ok = self.options.snapshot_every_bytes > 0
141                && counters.bytes_since >= self.options.snapshot_every_bytes;
142            time_ok || count_ok || bytes_ok
143        } else {
144            false
145        }
146    }
147
148    async fn write_snapshot(
149        &self,
150        doc_id: &str,
151        upto_lsn: i64,
152        state: &mf_state::State,
153    ) -> ForgeResult<()> {
154        let mut ser = state.serialize().await.map_err(|e| {
155            mf_core::error::error_utils::middleware_error(format!(
156                "状态序列化失败: {e}"
157            ))
158        })?;
159        let snap = SnapshotData {
160            node_pool: std::mem::take(&mut ser.node_pool),
161            state_fields: std::mem::take(&mut ser.state_fields),
162        };
163        let blob = serde_json::to_vec(&snap).map_err(|e| {
164            mf_core::error::error_utils::middleware_error(format!(
165                "快照编码失败: {e}"
166            ))
167        })?;
168        let blob = compress_if_needed(&blob, self.options.compression)
169            .map_err(|e| {
170                mf_core::error::error_utils::middleware_error(format!(
171                    "快照压缩失败: {e}"
172                ))
173            })?;
174
175        let snap = Snapshot {
176            doc_id: doc_id.to_string(),
177            upto_lsn,
178            created_at: chrono::Utc::now().timestamp_millis(),
179            state_blob: blob,
180            version: 1,
181        };
182
183        self.store.write_snapshot(snap).await.map_err(|e| {
184            mf_core::error::error_utils::middleware_error(e.to_string())
185        })?;
186
187        // 重置计数器
188        let mut entry =
189            self.snap_counters.entry(doc_id.to_string()).or_default();
190        entry.last_snapshot_ms = chrono::Utc::now().timestamp_millis();
191        entry.last_snapshot_upto_lsn = upto_lsn;
192        entry.events_since = 0;
193        entry.bytes_since = 0;
194        Ok(())
195    }
196}
197
198#[async_trait]
199impl<E: EventStore + 'static> EventHandler<Event> for SnapshotSubscriber<E> {
200    async fn handle(
201        &self,
202        event: &Event,
203    ) -> ForgeResult<()> {
204        if let Event::Create(state) = event {
205            // 初始化时若不存在任何快照，则写入基础快照（upto_lsn = 0）
206            let doc_id = state.doc().root_id().to_string();
207            let has_snapshot = self
208                .store
209                .latest_snapshot(&doc_id)
210                .await
211                .map(|s| s.is_some())
212                .unwrap_or(false);
213            if !has_snapshot {
214                self.write_snapshot(&doc_id, 0, state).await?;
215            }
216        } else if let Event::TrApply { old_state: _, new_state, transactions } = event {
217            let state = new_state;
218            let trs = transactions;
219            let mut touched_docs: HashSet<String> = HashSet::new();
220            let max_lsn_by_doc: DashMap<String, i64> = DashMap::new();
221
222            for tr in trs.iter() {
223                let tr_id = tr.id;
224                if self.persisted.contains_key(&tr_id) {
225                    continue;
226                }
227                // 选择 doc_id：优先 meta 的 doc_id，否则使用当前状态根ID
228                let doc_id = tr
229                    .get_meta::<String>("doc_id")
230                    .unwrap_or_else(|| state.doc().root_id().to_string());
231                let doc_id = if doc_id.is_empty() {
232                    self.default_doc_id.clone()
233                } else {
234                    doc_id
235                };
236
237                match self.persist_one(tr_id, &doc_id, tr).await {
238                    Ok(Some((lsn, bytes))) => {
239                        self.persisted.insert(tr_id, ());
240                        touched_docs.insert(doc_id.clone());
241                        let mut entry = self
242                            .snap_counters
243                            .entry(doc_id.clone())
244                            .or_default();
245                        entry.events_since += 1;
246                        entry.bytes_since += bytes as u64;
247                        // 更新 up-to LSN
248                        let cur = self
249                            .pending_snapshot_lsn
250                            .get(&doc_id)
251                            .map(|v| *v.value())
252                            .unwrap_or(-1);
253                        if lsn > cur {
254                            self.pending_snapshot_lsn
255                                .insert(doc_id.clone(), lsn);
256                        }
257                        // 记录每个文档的最大 LSN（备用）
258                        let m = max_lsn_by_doc
259                            .get(&doc_id)
260                            .map(|v| *v.value())
261                            .unwrap_or(-1);
262                        if lsn > m {
263                            max_lsn_by_doc.insert(doc_id.clone(), lsn);
264                        }
265                    },
266                    Ok(None) => {},
267                    Err(e) => {
268                        return Err(
269                            mf_core::error::error_utils::middleware_error(
270                                format!("持久化失败: {e}"),
271                            ),
272                        );
273                    },
274                }
275            }
276
277            // 对触达的文档执行快照策略
278            for doc_id in touched_docs.into_iter() {
279                if let Some(upto) =
280                    self.pending_snapshot_lsn.get(&doc_id).map(|v| *v.value())
281                {
282                    if upto >= 0 {
283                        let has_snapshot = self
284                            .store
285                            .latest_snapshot(&doc_id)
286                            .await
287                            .map(|s| s.is_some())
288                            .unwrap_or(false);
289                        if !has_snapshot || self.should_snapshot(&doc_id) {
290                            self.write_snapshot(&doc_id, upto, state).await?;
291                            self.pending_snapshot_lsn.remove(&doc_id);
292                        }
293                    }
294                }
295            }
296        }
297        Ok(())
298    }
299}
mf_persistence/subscriber/mod.rs

mf_persistence/subscriber/
mod.rs
