holographic_memory/core/engine/
compaction.rs

1// Copyright 2024-2026 WritersLogic Contributors
2// SPDX-License-Identifier: Apache-2.0
3
4use std::collections::HashSet;
5use std::sync::Arc;
6
7use anyhow::{ensure, Result};
8use serde::Serialize;
9
10use super::{mutation::Mutation, HmsCore};
11use crate::core::documents::StoredDocument;
12use crate::core::storage::StagedArena;
13
14/// Removes an abandoned staging generation; disarmed once publication begins.
15struct StagingDir {
16    path: std::path::PathBuf,
17    armed: std::cell::Cell<bool>,
18}
19
20impl StagingDir {
21    fn new(path: std::path::PathBuf) -> Self {
22        Self {
23            path,
24            armed: std::cell::Cell::new(true),
25        }
26    }
27    fn path(&self) -> &std::path::Path {
28        &self.path
29    }
30    fn disarm(&self) {
31        self.armed.set(false);
32    }
33}
34
35impl Drop for StagingDir {
36    fn drop(&mut self) {
37        if self.armed.get() {
38            let _ = std::fs::remove_dir_all(&self.path);
39        }
40    }
41}
42
43struct Snapshot {
44    checkpoint: (String, usize),
45    mutations: Vec<Mutation>,
46    documents: Vec<Arc<StoredDocument>>,
47    /// The store's text segment at snapshot time: chunk text moved there is written back
48    /// inline into the compaction frames.
49    text: Option<Arc<crate::core::documents::TextSegment>>,
50}
51
52#[derive(Serialize)]
53enum DocumentMutation<'a> {
54    Document(crate::core::documents::DocumentWithText<'a>),
55}
56
57impl HmsCore {
58    pub fn compact(&self) -> Result<()> {
59        self.compact_with_hook(|| {})
60    }
61
62    /// Compacts and then physically erases what the compaction replaced: the old generation
63    /// and any unreferenced leftovers are overwritten with zeros and synced before they are
64    /// unlinked. The text segment is scrubbed per deleted document at delete time.
65    ///
66    /// Afterwards nothing derived depends on the purged documents or on the order documents
67    /// arrived in (see `canonicalize_derived_state`), so every query, status field
68    /// other than the opaque counters, and file other than the generation's name match a store
69    /// that never held them and was purged at the same point.
70    pub fn compact_scrubbing(&self) -> Result<()> {
71        self.sync_text_segment()?;
72        self.compact_impl(|| {}, || {}, true)
73    }
74
75    /// Rebuilds the state a delete leaves shaped by history rather than by the live set: the
76    /// text segment (zeroed records of deleted text, offsets that count it) and the slots, in
77    /// id order as a reopen of the compacted journal lays them out; a published vector
78    /// snapshot, whose graph still holds the deleted rows' vertices and the edges they shaped
79    /// (rebuilt with the same graph choice); and the dense document index a delete dropped.
80    fn canonicalize_derived_state(&self) -> Result<()> {
81        {
82            let _transaction = self.mutation_gate.maintenance();
83            self.documents
84                .write()
85                .rebuild_canonical(&self.storage_path.join("documents.text"))?;
86        }
87        if let Some(snapshot) = self.vector_index.get() {
88            let graph_min_rows = if snapshot.has_graph() { 0 } else { usize::MAX };
89            self.rebuild_vector_index_waiting(graph_min_rows)?;
90        }
91        if self.documents.read().dense_index_recommended() {
92            self.train_dense_index()?;
93        }
94        Ok(())
95    }
96
97    fn compact_with_hook(&self, after_snapshot: impl FnOnce()) -> Result<()> {
98        self.compact_with_hooks(after_snapshot, || {})
99    }
100
101    fn compact_with_hooks(
102        &self,
103        after_snapshot: impl FnOnce(),
104        before_publication: impl FnOnce(),
105    ) -> Result<()> {
106        self.compact_impl(after_snapshot, before_publication, false)
107    }
108
109    fn compact_impl(
110        &self,
111        after_snapshot: impl FnOnce(),
112        before_publication: impl FnOnce(),
113        scrub: bool,
114    ) -> Result<()> {
115        let _compaction = self.compaction.lock();
116        let staging = StagingDir::new(self.arena.new_generation_dir()?);
117        let snapshot = self.compaction_snapshot()?;
118        let (generation, mut cursor) = snapshot.checkpoint;
119        after_snapshot();
120        let mut arena = StagedArena::new(staging.path(), Some(&self.arena))?;
121        for mutation in snapshot.mutations {
122            self.write_compaction_frame(&mut arena, &[mutation])?;
123        }
124        let mut payload = Vec::new();
125        for document in &snapshot.documents {
126            // A document deleted since the snapshot has had its text zeroed; the delete is in
127            // the tail, which the replay below applies, so it is simply not carried over.
128            {
129                let store = self.documents.read();
130                if store.is_deleted(document) {
131                    continue;
132                }
133                let binary = match self.journal {
134                    super::journal::JournalFormat::Binary => {
135                        super::journal::binary_document_payload(document, snapshot.text.as_deref())?
136                    }
137                    super::journal::JournalFormat::Json => None,
138                };
139                if let Some(binary) = binary {
140                    payload = binary;
141                } else {
142                    // The JSON form carries any document; a binary store reads both.
143                    payload.clear();
144                    payload.extend_from_slice(super::mutation::TRANSACTION_MAGIC);
145                    serde_json::to_writer(
146                        &mut payload,
147                        &[DocumentMutation::Document(
148                            crate::core::documents::DocumentWithText {
149                                document,
150                                segment: snapshot.text.as_deref(),
151                            },
152                        )],
153                    )?;
154                }
155            }
156            arena.write_slice(&self.maybe_encrypt(&payload)?)?;
157        }
158        let mut prepared = arena.finish()?.prepare_compacted()?;
159        for _ in 0..2 {
160            let end = {
161                let _transaction = self.mutation_gate.maintenance();
162                self.compaction_end(&generation)?
163            };
164            self.arena
165                .replay_compaction_tail(&mut prepared, &mut cursor, end)?;
166        }
167        before_publication();
168        let _transaction = self.mutation_gate.maintenance();
169        let end = self.compaction_end(&generation)?;
170        self.arena
171            .replay_compaction_tail(&mut prepared, &mut cursor, end)?;
172        staging.disarm();
173        if scrub {
174            self.arena.replace_with_compacted_scrubbing(prepared)?;
175        } else {
176            self.arena.replace_with_compacted(prepared)?;
177        }
178        self.revision
179            .fetch_add(1, std::sync::atomic::Ordering::Release);
180        // The arena pointer has moved, so writers may continue: the lease ends here. Index
181        // persistence binds the checkpoint read under the lease, holds the shards only while
182        // encoding, and writes the files with no lock held; the persist call that follows
183        // then finds identical metadata and writes nothing.
184        let checkpoint = self.arena.checkpoint();
185        _transaction.release();
186        self.persist_indices_detached(checkpoint)?;
187        self.persist_indices(&self.shards.read())?;
188        if let Some(audit) = &self.audit {
189            audit.record(
190                crate::core::audit::AuditOp::Compact,
191                "",
192                self.sign_fn().as_deref(),
193            )?;
194        }
195        if scrub {
196            // Still under the compaction lock: a concurrent compaction must not capture
197            // documents whose text offsets the canonical rebuild is replacing.
198            self.canonicalize_derived_state()?;
199        }
200        Ok(())
201    }
202
203    fn compaction_snapshot(&self) -> Result<Snapshot> {
204        let _transaction = self.mutation_gate.maintenance();
205        let shards = self.shards.read();
206        let mut vectors = Vec::new();
207        let composite_ids: HashSet<_> = self
208            .composite_memory
209            .as_ref()
210            .map(|memory| {
211                memory
212                    .inner()
213                    .all_vectors()
214                    .into_iter()
215                    .map(|(_, id, _)| id)
216                    .collect()
217            })
218            .unwrap_or_default();
219        shards.try_for_each_shard(|shard| {
220            for (id, vector) in shard.vectors.read().iter() {
221                if !id.starts_with(crate::core::documents::CHUNK_PREFIX)
222                    && !composite_ids.contains(id)
223                {
224                    vectors.push((id.clone(), vector.clone()));
225                }
226            }
227            Ok(())
228        })?;
229        // The shard map's iteration order depends on its insertion and removal history; id
230        // order makes the compacted journal a function of the live vectors alone.
231        vectors.sort_unstable_by(|a, b| a.0.cmp(&b.0));
232        let mut mutations: Vec<Mutation> = vectors
233            .into_iter()
234            .map(|(id, vector)| Mutation::Vector { id, vector })
235            .collect();
236        let (documents, text) = {
237            let store = self.documents.read();
238            (store.shared_values().collect(), store.text_segment())
239        };
240        mutations.extend(self.graph.snapshot().into_iter().map(Mutation::Relation));
241        mutations.extend(
242            self.graph
243                .type_snapshot()
244                .into_iter()
245                .map(Mutation::RelationType),
246        );
247        if let Some(atoms) = &self.atom_memory {
248            mutations.extend(
249                atoms
250                    .inner()
251                    .all_vectors()
252                    .into_iter()
253                    .map(|(_, id, vector)| Mutation::Atom { id, vector }),
254            );
255        }
256        if let Some(composites) = &self.composite_memory {
257            mutations.extend(
258                composites
259                    .inner()
260                    .all_vectors()
261                    .into_iter()
262                    .map(|(_, id, vector)| Mutation::Composite { id, vector }),
263            );
264        }
265        if let Some(triples) = &self.triple_store {
266            mutations.extend(triples.snapshot().into_iter().map(Mutation::StoredTriple));
267        }
268        if let Some(rules) = &self.rule_store {
269            mutations.extend(rules.all_rules().into_iter().map(Mutation::Rule));
270        }
271        mutations.extend(
272            self.derived
273                .read()
274                .iter()
275                .map(|(id, children)| Mutation::Derived {
276                    id: id.clone(),
277                    children: children.clone(),
278                }),
279        );
280        Ok(Snapshot {
281            checkpoint: self.arena.checkpoint(),
282            mutations,
283            documents,
284            text,
285        })
286    }
287
288    fn write_compaction_frame(
289        &self,
290        arena: &mut StagedArena,
291        mutations: &[Mutation],
292    ) -> Result<()> {
293        let payload = self.transaction_payload(mutations)?;
294        arena.write_slice(&self.maybe_encrypt(&payload)?)?;
295        Ok(())
296    }
297
298    fn compaction_end(&self, generation: &str) -> Result<usize> {
299        let checkpoint = self.arena.checkpoint();
300        ensure!(
301            checkpoint.0 == generation,
302            "arena generation changed during compaction"
303        );
304        Ok(checkpoint.1)
305    }
306
307    /// Persists the indices for `checkpoint` without holding the writer lease: the shards are
308    /// read only while the indices are encoded, and the files are written afterwards.
309    fn persist_indices_detached(&self, checkpoint: (String, usize)) -> Result<()> {
310        let encoded = self.encode_indices(&self.shards.read(), checkpoint)?;
311        self.write_indices(encoded)
312    }
313}
314
315#[cfg(test)]
316mod tests {
317    use super::*;
318    use std::sync::mpsc;
319    use std::time::Duration;
320
321    use crate::core::documents::{DocumentInput, SearchOptions};
322    use crate::core::{EmbeddingSpace, HmsConfig};
323    use serde_json::json;
324
325    fn config() -> HmsConfig {
326        HmsConfig {
327            embedding_space: Some(EmbeddingSpace {
328                model: "fixture".into(),
329                revision: "1".into(),
330                dimensions: 2,
331                normalization: "l2".into(),
332                metric: "cosine".into(),
333            }),
334            ..Default::default()
335        }
336    }
337
338    fn configurations() -> impl Iterator<Item = HmsConfig> {
339        [false, true]
340            .into_iter()
341            .filter(|encrypted| !encrypted || cfg!(feature = "security"))
342            .map(|encrypted| {
343                let mut config = config();
344                config.security.encryption_enabled = encrypted;
345                if encrypted {
346                    config.security.encryption_passphrase =
347                        Some("compaction-replay-fixture".into());
348                }
349                config
350            })
351    }
352
353    fn document(id: &str, version: &str) -> DocumentInput {
354        serde_json::from_value(json!({
355            "id": id, "version": version, "text": "needle durable source",
356            "sourceUri": format!("source:{version}"), "metadata": {"version": version},
357            "embeddings": [[1.0, 0.0]],
358        }))
359        .unwrap()
360    }
361
362    fn seed(core: &HmsCore) -> Result<()> {
363        core.memorize_document(document("replace", "old"))?;
364        core.memorize_document(document("delete", "old"))?;
365        core.memorize("vector".into(), core.encode_text("original"))?;
366        core.memorize("remove".into(), core.encode_text("removed"))?;
367        core.add_relation(&serde_json::from_value(json!({
368            "sourceId": "a", "relationType": "knows", "targetId": "b",
369            "properties": null, "validFrom": 0, "validTo": 0,
370        }))?)?;
371        Ok(())
372    }
373
374    fn mutate(core: &HmsCore) -> Result<()> {
375        core.memorize_documents(vec![document("replace", "new"), document("insert", "new")])?;
376        core.delete_document("delete")?;
377        core.memorize("vector".into(), core.encode_text("replacement"))?;
378        core.delete("remove")?;
379        core.remove_relation("a", "knows", "b")?;
380        Ok(())
381    }
382
383    fn assert_current(core: &HmsCore) -> Result<serde_json::Value> {
384        let options = SearchOptions {
385            embedding: Some(vec![1.0, 0.0]),
386            source_uri: Some("source:new".into()),
387            filter: Some(json!({"version": "new"})),
388            ..Default::default()
389        };
390        let results = core.search_documents("needle", &options)?;
391        assert_eq!(results.len(), 2);
392        assert!(results.iter().all(|r| r.version == "new"
393            && r.source_uri.as_deref() == Some("source:new")
394            && r.semantic_score == Some(1.0)));
395        assert_eq!(core.relation_count(), 0);
396        assert_eq!(core.vector_count(), 3);
397        assert_eq!(
398            core.query(&core.encode_text("replacement"), 1)[0].id,
399            "vector"
400        );
401        Ok(serde_json::to_value(results)?)
402    }
403
404    #[test]
405    fn compaction_replays_acknowledged_writes_and_preserves_batch_frames() -> Result<()> {
406        for config in configurations() {
407            let directory = tempfile::tempdir()?;
408            let path = directory.path().display().to_string();
409            let core = Arc::new(HmsCore::new(
410                1024,
411                Some(path.clone()),
412                Some(config.clone()),
413            )?);
414            seed(&core)?;
415            let writer_core = Arc::clone(&core);
416            let (start_tx, start_rx) = mpsc::channel();
417            let (done_tx, done_rx) = mpsc::channel();
418            let writer = std::thread::spawn(move || -> Result<()> {
419                start_rx.recv()?;
420                mutate(&writer_core)?;
421                assert_current(&writer_core)?;
422                done_tx.send(())?;
423                Ok(())
424            });
425            let compacted = core.compact_with_hook(|| {
426                start_tx.send(()).unwrap();
427                done_rx
428                    .recv_timeout(Duration::from_secs(5))
429                    .expect("compaction snapshot blocked mutations or queries");
430            });
431            writer.join().unwrap()?;
432            compacted?;
433            let expected = assert_current(&core)?;
434            let mut offset = 0;
435            let mut batch_frames = 0;
436            while offset < core.arena.stats().used_bytes {
437                let (data, _) = core.arena_read_frame(offset)?;
438                let operations = super::super::journal::decode(&data)?;
439                if operations.len() == 2 {
440                    batch_frames += 1;
441                    assert!(operations
442                        .iter()
443                        .all(|op| matches!(op, Mutation::Document(_))));
444                }
445                offset = core.arena.next_offset(offset)?;
446            }
447            assert_eq!(batch_frames, 1);
448            drop(core);
449            let restored = HmsCore::new(1024, Some(path), Some(config))?;
450            assert_eq!(assert_current(&restored)?, expected);
451        }
452        Ok(())
453    }
454
455    #[test]
456    fn compaction_snapshot_shares_document_payloads_after_replacement() -> Result<()> {
457        let directory = tempfile::tempdir()?;
458        let core = HmsCore::new(
459            1024,
460            Some(directory.path().display().to_string()),
461            Some(config()),
462        )?;
463        core.memorize_document(document("replace", "old"))?;
464        let snapshot = core.compaction_snapshot()?;
465        let owner = core.documents.read().shared_values().next().unwrap();
466        assert!(Arc::ptr_eq(&owner, &snapshot.documents[0]));
467        let dense_pointer = owner.chunks[0].embedding.as_ref().unwrap().as_ptr();
468        core.memorize_document(document("replace", "new"))?;
469        assert_eq!(snapshot.documents[0].version, "old");
470        assert_eq!(
471            snapshot.documents[0].chunks[0]
472                .embedding
473                .as_ref()
474                .unwrap()
475                .as_ptr(),
476            dense_pointer
477        );
478        Ok(())
479    }
480
481    #[test]
482    fn compaction_replay_crash_child() -> Result<()> {
483        let Ok(root) = std::env::var("HMS_TEST_CORE_REPLAY_ROOT") else {
484            return Ok(());
485        };
486        let encrypted = std::env::var("HMS_TEST_CORE_REPLAY_ENCRYPTED").as_deref() == Ok("true");
487        let config = configurations()
488            .find(|config| config.security.encryption_enabled == encrypted)
489            .unwrap();
490        let core = HmsCore::new(1024, Some(root), Some(config))?;
491        seed(&core)?;
492        core.compact_with_hook(|| mutate(&core).unwrap())
493    }
494
495    #[test]
496    fn compaction_replayed_writes_survive_each_publication_interruption() -> Result<()> {
497        for config in configurations() {
498            for phase in ["staged", "published", "mapped"] {
499                let directory = tempfile::tempdir()?;
500                let output = std::process::Command::new(std::env::current_exe()?)
501                    .args([
502                        "--exact",
503                        "core::engine::compaction::tests::compaction_replay_crash_child",
504                    ])
505                    .env("HMS_TEST_CORE_REPLAY_ROOT", directory.path())
506                    .env(
507                        "HMS_TEST_CORE_REPLAY_ENCRYPTED",
508                        config.security.encryption_enabled.to_string(),
509                    )
510                    .env("HMS_TEST_COMPACTION_PHASE", phase)
511                    .output()?;
512                assert_eq!(output.status.code(), Some(73), "{phase}: {output:?}");
513                let core = HmsCore::new(
514                    1024,
515                    Some(directory.path().display().to_string()),
516                    Some(config.clone()),
517                )?;
518                assert_current(&core)?;
519                core.memorize_document(document("after-recovery", "recovered"))?;
520                core.compact()?;
521            }
522        }
523        Ok(())
524    }
525
526    #[test]
527    fn concurrent_compactors_serialize_without_losing_intervening_mutations() -> Result<()> {
528        let directory = tempfile::tempdir()?;
529        let path = directory.path().display().to_string();
530        let core = Arc::new(HmsCore::new(1024, Some(path.clone()), Some(config()))?);
531        seed(&core)?;
532        let other_core = Arc::clone(&core);
533        let (start_tx, start_rx) = mpsc::channel();
534        let other = std::thread::spawn(move || -> Result<()> {
535            start_rx.recv()?;
536            other_core.compact()
537        });
538        core.compact_with_hook(|| {
539            start_tx.send(()).unwrap();
540            assert!(core.compaction.try_lock().is_none());
541            mutate(&core).unwrap();
542        })?;
543        other.join().unwrap()?;
544        let expected = assert_current(&core)?;
545        drop(core);
546        let restored = HmsCore::new(1024, Some(path), Some(config()))?;
547        assert_eq!(assert_current(&restored)?, expected);
548        Ok(())
549    }
550
551    #[test]
552    fn replay_rejects_partial_frames_and_accepts_exact_capture_boundaries() -> Result<()> {
553        for delta in [-1_isize, 0, 1] {
554            let directory = tempfile::tempdir()?;
555            let core = HmsCore::new(1024, Some(directory.path().display().to_string()), None)?;
556            let mut cursor = core.arena.stats().used_bytes;
557            core.memorize("after".into(), core.encode_text("acknowledged"))?;
558            let stage = tempfile::tempdir_in(directory.path())?;
559            let staged = StagedArena::new(stage.path(), None)?.finish()?;
560            let mut prepared = staged.prepare_compacted()?;
561            let end = core
562                .arena
563                .stats()
564                .used_bytes
565                .checked_add_signed(delta)
566                .unwrap();
567            let result = core
568                .arena
569                .replay_compaction_tail(&mut prepared, &mut cursor, end);
570            assert_eq!(result.is_ok(), delta == 0);
571            assert!(!core.storage_path.join("vectors_data.bin/CURRENT").exists());
572            assert_eq!(
573                core.query(&core.encode_text("acknowledged"), 1)[0].id,
574                "after"
575            );
576        }
577        Ok(())
578    }
579
580    #[test]
581    fn compaction_writes_after_catchup_survive_publication_and_restart() -> Result<()> {
582        for config in configurations() {
583            let directory = tempfile::tempdir()?;
584            let path = directory.path().display().to_string();
585            let core = HmsCore::new(1024, Some(path.clone()), Some(config.clone()))?;
586            let initial = crate::EntangledHVec::new_deterministic(1024, 1);
587            let during = crate::EntangledHVec::new_deterministic(1024, 2);
588            let late = crate::EntangledHVec::new_deterministic(1024, 3);
589            assert_ne!(initial.indices(), late.indices());
590            assert_ne!(during.indices(), late.indices());
591            core.memorize("vector".into(), initial)?;
592            core.compact_with_hooks(
593                || core.memorize("vector".into(), during).unwrap(),
594                || {
595                    core.memorize("vector".into(), late.clone()).unwrap();
596                    core.memorize_document(document("late", "late")).unwrap();
597                },
598            )?;
599            let expected =
600                serde_json::to_value(core.search_documents("needle", &Default::default())?)?;
601            assert_eq!(expected[0]["version"], "late");
602            drop(core);
603            let restored = HmsCore::new(1024, Some(path), Some(config))?;
604            assert_eq!(
605                restored
606                    .shards
607                    .read()
608                    .get_vector("vector")
609                    .unwrap()
610                    .indices(),
611                late.indices()
612            );
613            assert_eq!(
614                serde_json::to_value(restored.search_documents("needle", &Default::default())?)?,
615                expected
616            );
617        }
618        Ok(())
619    }
620}
holographic_memory/core/engine/compaction.rs

holographic_memory/core/engine/
compaction.rs
