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;
12
13use crate::api::{CommitMode, EventStore, PersistOptions, PersistedEvent, Snapshot};
14use crate::ser::{
15    checksum32, compress_if_needed, frame_steps, SnapshotData, TypeWrapper,
16};
17
18#[derive(Default)]
19struct SnapshotCounters {
20    last_snapshot_ms: i64,
21    last_snapshot_upto_lsn: i64,
22    events_since: u32,
23    bytes_since: u64,
24}
25
26pub struct SnapshotSubscriber<E: EventStore + 'static> {
27    store: Arc<E>,
28    options: PersistOptions,
29    commit_mode: CommitMode,
30    default_doc_id: String,
31    // 进程内已持久化的事务ID集合，避免重复写入（进程重启后靠数据库幂等键）
32    persisted: DashMap<u64, ()>,
33    // 每个文档的快照计数器
34    snap_counters: DashMap<String, SnapshotCounters>,
35    // 待写快照的 upto_lsn
36    pending_snapshot_lsn: DashMap<String, i64>,
37}
38
39impl<E: EventStore + 'static> fmt::Debug for SnapshotSubscriber<E> {
40    fn fmt(
41        &self,
42        f: &mut fmt::Formatter<'_>,
43    ) -> fmt::Result {
44        f.write_str("SnapshotSubscriber")
45    }
46}
47
48impl<E: EventStore + 'static> Clone for SnapshotSubscriber<E> {
49    fn clone(&self) -> Self {
50        Self {
51            store: self.store.clone(),
52            options: self.options.clone(),
53            commit_mode: self.commit_mode,
54            default_doc_id: self.default_doc_id.clone(),
55            persisted: DashMap::new(),
56            snap_counters: DashMap::new(),
57            pending_snapshot_lsn: DashMap::new(),
58        }
59    }
60}
61
62impl<E: EventStore + 'static> SnapshotSubscriber<E> {
63    pub fn new(
64        store: Arc<E>,
65        options: PersistOptions,
66        default_doc_id: impl Into<String>,
67    ) -> Self {
68        Self {
69            commit_mode: options.commit_mode,
70            store,
71            options,
72            default_doc_id: default_doc_id.into(),
73            persisted: DashMap::new(),
74            snap_counters: DashMap::new(),
75            pending_snapshot_lsn: DashMap::new(),
76        }
77    }
78
79    async fn persist_one(
80        &self,
81        tr_id: u64,
82        doc_id: &str,
83        transaction: &Transaction,
84    ) -> ForgeResult<Option<(i64, usize)>> {
85        // 将每个步骤包装为 {type_id, data}
86        let frames: Vec<TypeWrapper> = frame_steps(transaction);
87        let framed = serde_json::to_vec(&frames).map_err(|e| {
88            mf_core::error::error_utils::middleware_error(format!(
89                "步骤编码失败: {}",
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                "状态序列化失败: {}",
157                e
158            ))
159        })?;
160        let snap = SnapshotData {
161            node_pool: std::mem::take(&mut ser.node_pool),
162            state_fields: std::mem::take(&mut ser.state_fields),
163        };
164        let blob = serde_json::to_vec(&snap).map_err(|e| {
165            mf_core::error::error_utils::middleware_error(format!(
166                "快照编码失败: {}",
167                e
168            ))
169        })?;
170        let blob = compress_if_needed(&blob, self.options.compression)
171            .map_err(|e| {
172                mf_core::error::error_utils::middleware_error(format!(
173                    "快照压缩失败: {}",
174                    e
175                ))
176            })?;
177
178        let snap = Snapshot {
179            doc_id: doc_id.to_string(),
180            upto_lsn,
181            created_at: chrono::Utc::now().timestamp_millis(),
182            state_blob: blob,
183            version: 1,
184        };
185
186        self.store.write_snapshot(snap).await.map_err(|e| {
187            mf_core::error::error_utils::middleware_error(e.to_string())
188        })?;
189
190        // 重置计数器
191        let mut entry = self
192            .snap_counters
193            .entry(doc_id.to_string())
194            .or_insert_with(SnapshotCounters::default);
195        entry.last_snapshot_ms = chrono::Utc::now().timestamp_millis();
196        entry.last_snapshot_upto_lsn = upto_lsn;
197        entry.events_since = 0;
198        entry.bytes_since = 0;
199        Ok(())
200    }
201}
202
203#[async_trait]
204impl<E: EventStore + 'static> EventHandler<Event> for SnapshotSubscriber<E> {
205    async fn handle(
206        &self,
207        event: &Event,
208    ) -> ForgeResult<()> {
209        if let Event::Create(state) = event {
210            // 初始化时若不存在任何快照，则写入基础快照（upto_lsn = 0）
211            let doc_id = state.doc().root_id().to_string();
212            let has_snapshot = self
213                .store
214                .latest_snapshot(&doc_id)
215                .await
216                .map(|s| s.is_some())
217                .unwrap_or(false);
218            if !has_snapshot {
219                self.write_snapshot(&doc_id, 0, state).await?;
220            }
221        } else if let Event::TrApply(_old, trs, state) = event {
222            let mut touched_docs: HashSet<String> = HashSet::new();
223            let max_lsn_by_doc: DashMap<String, i64> = DashMap::new();
224
225            for tr in trs.iter() {
226                let tr_id = tr.id;
227                if self.persisted.contains_key(&tr_id) {
228                    continue;
229                }
230                // 选择 doc_id：优先 meta 的 doc_id，否则使用当前状态根ID
231                let doc_id = tr
232                    .get_meta::<String>("doc_id")
233                    .unwrap_or_else(|| state.doc().root_id().to_string());
234                let doc_id = if doc_id.is_empty() {
235                    self.default_doc_id.clone()
236                } else {
237                    doc_id
238                };
239
240                match self.persist_one(tr_id, &doc_id, tr).await {
241                    Ok(Some((lsn, bytes))) => {
242                        self.persisted.insert(tr_id, ());
243                        touched_docs.insert(doc_id.clone());
244                        let mut entry = self
245                            .snap_counters
246                            .entry(doc_id.clone())
247                            .or_insert_with(SnapshotCounters::default);
248                        entry.events_since += 1;
249                        entry.bytes_since += bytes as u64;
250                        // 更新 up-to LSN
251                        let cur = self
252                            .pending_snapshot_lsn
253                            .get(&doc_id)
254                            .map(|v| *v.value())
255                            .unwrap_or(-1);
256                        if lsn > cur {
257                            self.pending_snapshot_lsn
258                                .insert(doc_id.clone(), lsn);
259                        }
260                        // 记录每个文档的最大 LSN（备用）
261                        let m = max_lsn_by_doc
262                            .get(&doc_id)
263                            .map(|v| *v.value())
264                            .unwrap_or(-1);
265                        if lsn > m {
266                            max_lsn_by_doc.insert(doc_id.clone(), lsn);
267                        }
268                    },
269                    Ok(None) => {},
270                    Err(e) => {
271                        return Err(
272                            mf_core::error::error_utils::middleware_error(
273                                format!("持久化失败: {}", e),
274                            ),
275                        );
276                    },
277                }
278            }
279
280            // 对触达的文档执行快照策略
281            for doc_id in touched_docs.into_iter() {
282                if let Some(upto) =
283                    self.pending_snapshot_lsn.get(&doc_id).map(|v| *v.value())
284                {
285                    if upto >= 0 {
286                        let has_snapshot = self
287                            .store
288                            .latest_snapshot(&doc_id)
289                            .await
290                            .map(|s| s.is_some())
291                            .unwrap_or(false);
292                        if !has_snapshot || self.should_snapshot(&doc_id) {
293                            self.write_snapshot(&doc_id, upto, state).await?;
294                            self.pending_snapshot_lsn.remove(&doc_id);
295                        }
296                    }
297                }
298            }
299        }
300        Ok(())
301    }
302}
mf_persistence/subscriber/mod.rs

mf_persistence/subscriber/
mod.rs
