Skip to main content

uni_store/storage/
compaction.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright 2024-2026 Dragonscale Team
3
4use crate::storage::delta::{ENTRY_SIZE_ESTIMATE, L1Entry, Op};
5use crate::storage::manager::StorageManager;
6use anyhow::{Result, anyhow};
7use arrow_array::Array;
8use arrow_array::builder::{ListBuilder, UInt64Builder};
9use arrow_array::{ListArray, RecordBatch, UInt64Array};
10use metrics;
11use std::collections::{HashMap, HashSet};
12use std::sync::Arc;
13use tracing::{error, info, instrument};
14use uni_common::core::id::{Eid, Vid};
15use uni_common::core::schema::DataType;
16use uni_common::{Properties, Value};
17use uni_crdt::Crdt;
18
19pub struct Compactor {
20    storage: Arc<StorageManager>,
21}
22
23impl Compactor {
24    pub fn new(storage: Arc<StorageManager>) -> Self {
25        Self { storage }
26    }
27
28    #[instrument(skip(self), level = "info")]
29    pub async fn compact_all(&self) -> Result<Vec<CompactionInfo>> {
30        let start = std::time::Instant::now();
31        let schema = self.storage.schema_manager().schema();
32        let mut compaction_results = Vec::new();
33
34        // Compact Vertices
35        for label in schema.labels.keys() {
36            info!("Compacting vertices for label {}", label);
37            if let Err(e) = self.compact_vertices(label).await {
38                error!("Failed to compact vertices for {}: {}", label, e);
39            }
40        }
41
42        // Compact Edges
43        for (edge_type, meta) in &schema.edge_types {
44            // Outgoing: src_labels
45            for label in &meta.src_labels {
46                info!("Compacting adjacency {} -> {} (fwd)", label, edge_type);
47                match self.compact_adjacency(edge_type, label, "fwd").await {
48                    Ok(info) => compaction_results.push(info),
49                    Err(e) => {
50                        error!(
51                            "Failed to compact adjacency {} -> {}: {}",
52                            label, edge_type, e
53                        );
54                    }
55                }
56            }
57
58            // Incoming: dst_labels
59            for label in &meta.dst_labels {
60                info!("Compacting adjacency {} <- {} (bwd)", label, edge_type);
61                match self.compact_adjacency(edge_type, label, "bwd").await {
62                    Ok(info) => compaction_results.push(info),
63                    Err(e) => {
64                        error!(
65                            "Failed to compact adjacency {} <- {}: {}",
66                            label, edge_type, e
67                        );
68                    }
69                }
70            }
71        }
72
73        metrics::counter!("uni_compaction_runs_total").increment(1);
74        metrics::histogram!("uni_compaction_duration_seconds")
75            .record(start.elapsed().as_secs_f64());
76
77        Ok(compaction_results)
78    }
79
80    #[instrument(skip(self), fields(rows_processed, duration_ms), level = "info")]
81    pub async fn compact_vertices(&self, label: &str) -> Result<()> {
82        let start = std::time::Instant::now();
83        let schema_manager = self.storage.schema_manager();
84        let schema = schema_manager.schema();
85
86        let label_props = schema
87            .properties
88            .get(label)
89            .ok_or_else(|| anyhow!("Label not found"))?;
90
91        // Identify CRDT properties
92        let crdt_props: HashSet<String> = label_props
93            .iter()
94            .filter(|(_, meta)| matches!(meta.r#type, DataType::Crdt(_)))
95            .map(|(name, _)| name.clone())
96            .collect();
97
98        let dataset = self.storage.vertex_dataset(label)?;
99        let backend = self.storage.backend();
100        let table_name = dataset.table_name();
101
102        // Check if table exists
103        if !backend.table_exists(&table_name).await.unwrap_or(false) {
104            info!("No vertex data to compact for label '{}'", label);
105            return Ok(());
106        }
107
108        // In-memory compaction for now (MVP).
109        // For large datasets, this needs to be streaming/chunked with external sort.
110        // Current approach: Read ALL, merge in map, write NEW.
111        // TODO(perf): This accumulates ALL vertices in memory, causing OOM for large
112        // labels (millions of vertices). Refactor to use streaming merge-sort with
113        // constant memory usage (e.g., external sort or Lance fragment-by-fragment merge).
114
115        let row_count = backend.count_rows(&table_name, None).await?;
116        crate::storage::delta::check_oom_guard(
117            row_count,
118            self.storage.config.max_compaction_rows,
119            label,
120            "vertices",
121        )?;
122
123        info!(
124            label = %label,
125            row_count,
126            estimated_bytes = row_count * 200,
127            "Starting vertex compaction"
128        );
129
130        // Serialize the whole read → merge → OVERWRITE against concurrent flushes.
131        // The flush path takes the same per-table lock in
132        // `merge_insert_batch_with_lance_conflict_retry`; without it, an unguarded
133        // `AddDataMode::Overwrite` below would silently discard rows a flush
134        // appended in the window between this scan and the overwrite commit.
135        let _write_guard = backend.lock_table_for_write(&table_name).await;
136
137        use crate::backend::types::ScanRequest;
138        let batches: Vec<RecordBatch> = backend.scan(ScanRequest::all(&table_name)).await?;
139
140        // Vid -> (Properties, Deleted)
141        let mut vertex_state: HashMap<Vid, (Properties, bool)> = HashMap::new();
142        let mut vertex_versions: HashMap<Vid, u64> = HashMap::new();
143        let mut vertex_labels: HashMap<Vid, Vec<String>> = HashMap::new();
144
145        let mut rows_processed = 0;
146
147        for batch in batches {
148            rows_processed += batch.num_rows();
149            let vid_col = batch
150                .column_by_name("_vid")
151                .unwrap()
152                .as_any()
153                .downcast_ref::<UInt64Array>()
154                .unwrap();
155            let ver_col = batch
156                .column_by_name("_version")
157                .unwrap()
158                .as_any()
159                .downcast_ref::<UInt64Array>()
160                .unwrap();
161            let del_col = batch
162                .column_by_name("_deleted")
163                .unwrap()
164                .as_any()
165                .downcast_ref::<arrow_array::BooleanArray>()
166                .unwrap();
167
168            // Read _labels column (List<Utf8>) if present
169            let labels_col = batch
170                .column_by_name("_labels")
171                .and_then(|c| c.as_any().downcast_ref::<arrow_array::ListArray>());
172
173            for i in 0..batch.num_rows() {
174                let vid = Vid::from(vid_col.value(i));
175                let version = ver_col.value(i);
176                let deleted = del_col.value(i);
177
178                // Extract labels from the _labels column (keep latest version's labels)
179                if let Some(list_arr) = labels_col
180                    && version >= *vertex_versions.entry(vid).or_insert(0)
181                {
182                    let labels = crate::storage::arrow_convert::labels_from_list_array(list_arr, i);
183                    if !labels.is_empty() {
184                        vertex_labels.insert(vid, labels);
185                    }
186                }
187
188                let current_entry = vertex_state
189                    .entry(vid)
190                    .or_insert((Properties::new(), false));
191                let current_version = vertex_versions.entry(vid).or_insert(0);
192
193                // If this row is newer than what we've seen (or same), we apply logic.
194                // Wait, if we process unordered, we need to be careful.
195                // For CRDTs, we MERGE regardless of version (commutative).
196                // For LWW, we take MAX version.
197
198                // If it's a deletion, and it's newer, it wins.
199                if deleted {
200                    if version >= *current_version {
201                        current_entry.1 = true;
202                        current_entry.0.clear(); // Clear properties on delete
203                        *current_version = version;
204                    }
205                    continue;
206                }
207
208                // It's an update/insert
209                // Extract props and track NULLs (property removals)
210                let mut row_props = Properties::new();
211                let mut null_props = Vec::new(); // Track explicitly NULL properties
212                for (name, meta) in label_props {
213                    if let Some(col) = batch.column_by_name(name) {
214                        if col.is_null(i) {
215                            // Property was explicitly removed (set to NULL)
216                            null_props.push(name.clone());
217                        } else {
218                            let val = crate::storage::value_codec::decode_column_value(
219                                col.as_ref(),
220                                &meta.r#type,
221                                i,
222                                crate::storage::value_codec::CrdtDecodeMode::Strict,
223                            )?;
224                            row_props.insert(name.clone(), val);
225                        }
226                    }
227                }
228
229                Self::merge_row_into_state(
230                    row_props,
231                    null_props,
232                    version,
233                    current_entry,
234                    current_version,
235                    &crdt_props,
236                    self.storage.plugin_registry().map(|a| a.as_ref()),
237                )?;
238            }
239        }
240
241        // Convert state to RecordBatch and write OVERWRITE
242        let mut valid_vertices = Vec::new();
243        let mut valid_versions = Vec::new();
244        let mut valid_deleted = Vec::new(); // Should be all false if we filter out tombstones?
245        // Or we keep tombstones if they are recent?
246        // Compaction usually removes tombstones.
247
248        for (vid, (props, deleted)) in vertex_state {
249            if !deleted {
250                let labels = vertex_labels.remove(&vid).unwrap_or_default();
251                valid_vertices.push((vid, labels, props));
252                valid_versions.push(vertex_versions[&vid]);
253                valid_deleted.push(false);
254            }
255        }
256
257        if !valid_vertices.is_empty() {
258            let batch = dataset.build_record_batch(
259                &valid_vertices,
260                &valid_deleted,
261                &valid_versions,
262                &schema,
263            )?;
264            dataset
265                .replace(self.storage.backend(), batch, &schema)
266                .await?;
267        }
268
269        let duration = start.elapsed();
270        let rows_reclaimed = rows_processed as u64 - valid_vertices.len() as u64;
271        metrics::counter!("uni_compaction_rows_reclaimed_total", "type" => "vertex")
272            .increment(rows_reclaimed);
273
274        tracing::Span::current().record("rows_processed", rows_processed);
275        tracing::Span::current().record("duration_ms", duration.as_millis());
276        info!(
277            rows = rows_processed,
278            duration_ms = duration.as_millis(),
279            "Vertex compaction completed"
280        );
281
282        metrics::histogram!("uni_compaction_duration_seconds", "type" => "vertex")
283            .record(duration.as_secs_f64());
284
285        Ok(())
286    }
287
288    fn merge_crdt_values(
289        a: &Value,
290        b: &Value,
291        registry: Option<&uni_plugin::PluginRegistry>,
292    ) -> Result<Value> {
293        if a.is_null() {
294            return Ok(b.clone());
295        }
296        if b.is_null() {
297            return Ok(a.clone());
298        }
299        let mut crdt_a: Crdt = serde_json::from_value(a.clone().into())?;
300        let crdt_b: Crdt = serde_json::from_value(b.clone().into())?;
301        // Operand order: self=existing (a), other=new (b). Preserve — a custom
302        // provider's merge may be non-commutative.
303        match registry {
304            Some(reg) => crdt_a
305                .merge_via_registry(&crdt_b, reg)
306                .map_err(|e| anyhow::anyhow!("{e}"))?,
307            None => crdt_a
308                .try_merge(&crdt_b)
309                .map_err(|e| anyhow::anyhow!("{e}"))?,
310        }
311        Ok(Value::from(serde_json::to_value(crdt_a)?))
312    }
313
314    /// Merge row properties into state based on version comparison.
315    fn merge_row_into_state(
316        row_props: Properties,
317        null_props: Vec<String>,
318        version: u64,
319        current_entry: &mut (Properties, bool),
320        current_version: &mut u64,
321        crdt_props: &HashSet<String>,
322        registry: Option<&uni_plugin::PluginRegistry>,
323    ) -> Result<()> {
324        if version > *current_version {
325            // New version wins for LWW, merge for CRDTs
326            *current_version = version;
327            current_entry.1 = false;
328
329            for (k, v) in row_props {
330                if crdt_props.contains(&k) {
331                    let existing = current_entry.0.entry(k.clone()).or_insert(Value::Null);
332                    *existing = Self::merge_crdt_values(existing, &v, registry)?;
333                } else {
334                    current_entry.0.insert(k, v);
335                }
336            }
337
338            // Remove properties explicitly set to NULL in the newer version
339            for null_prop in &null_props {
340                if !crdt_props.contains(null_prop) {
341                    current_entry.0.remove(null_prop);
342                }
343            }
344        } else if version == *current_version {
345            // Same version: merge all
346            current_entry.1 = false;
347            for (k, v) in row_props {
348                if crdt_props.contains(&k) {
349                    let existing = current_entry.0.entry(k.clone()).or_insert(Value::Null);
350                    *existing = Self::merge_crdt_values(existing, &v, registry)?;
351                } else {
352                    current_entry.0.insert(k, v);
353                }
354            }
355        } else {
356            // Older version: only merge CRDTs
357            if !current_entry.1 {
358                for (k, v) in row_props {
359                    if crdt_props.contains(&k) {
360                        let existing = current_entry.0.entry(k.clone()).or_insert(Value::Null);
361                        *existing = Self::merge_crdt_values(existing, &v, registry)?;
362                    }
363                }
364            }
365        }
366        Ok(())
367    }
368
369    #[instrument(skip(self), fields(delta_count, duration_ms), level = "info")]
370    pub async fn compact_adjacency(
371        &self,
372        edge_type: &str,
373        label: &str,
374        direction: &str,
375    ) -> Result<CompactionInfo> {
376        let start = std::time::Instant::now();
377        let schema = self.storage.schema_manager().schema();
378
379        // 1. Load all L1 Deltas sorted by key
380        let delta_ds = self.storage.delta_dataset(edge_type, direction)?;
381        let deltas = delta_ds
382            .scan_all_backend(self.storage.backend(), &schema)
383            .await?;
384
385        let delta_count = deltas.len();
386        tracing::Span::current().record("delta_count", delta_count);
387
388        if deltas.is_empty() {
389            // Nothing to compact, return info anyway
390            return Ok(CompactionInfo {
391                edge_type: edge_type.to_string(),
392                direction: direction.to_string(),
393            });
394        }
395
396        // Group deltas by src_vid (if fwd) or dst_vid (if bwd)
397        // We'll use a HashMap for now since we loaded all into memory.
398        // Value is list of ops for that vertex.
399        let mut delta_map: HashMap<Vid, Vec<L1Entry>> = HashMap::new();
400        for entry in &deltas {
401            let key = if direction == "fwd" {
402                entry.src_vid
403            } else {
404                entry.dst_vid
405            };
406            delta_map.entry(key).or_default().push(entry.clone());
407        }
408
409        // Sort each VID's ops by version to ensure correct ordering
410        // This guarantees Delete(v=2) beats Insert(v=1) regardless of scan order
411        for ops in delta_map.values_mut() {
412            ops.sort_by_key(|e| e.version);
413        }
414
415        // 2. Open L2 Adjacency stream
416        let adj_ds = self
417            .storage
418            .adjacency_dataset(edge_type, label, direction)?;
419
420        // We need to write a NEW version.
421        // Strategy:
422        // - Read L2 batch by batch.
423        // - For each row (vertex), check if we have deltas.
424        // - Apply deltas.
425        // - Write to new batch.
426        // - Track which vertices from deltas we've processed.
427        // - After L2 stream ends, process remaining "new" vertices from deltas.
428
429        // Output Builders
430        let mut src_vid_builder = UInt64Builder::new();
431        let mut neighbors_builder = ListBuilder::new(UInt64Builder::new());
432        let mut edge_ids_builder = ListBuilder::new(UInt64Builder::new());
433
434        let mut processed_vids = HashSet::new();
435
436        // Try to read from backend (canonical storage)
437        let backend = self.storage.backend();
438        let adj_table_name = adj_ds.table_name();
439        if backend.table_exists(&adj_table_name).await.unwrap_or(false) {
440            let adj_row_count = backend.count_rows(&adj_table_name, None).await?;
441            crate::storage::delta::check_oom_guard(
442                adj_row_count,
443                self.storage.config.max_compaction_rows,
444                &format!("{}_{}", edge_type, label),
445                direction,
446            )?;
447
448            info!(
449                edge_type = %edge_type,
450                label = %label,
451                direction = %direction,
452                adj_row_count,
453                delta_count,
454                estimated_bytes = adj_row_count * 100 + delta_count * ENTRY_SIZE_ESTIMATE,
455                "Starting adjacency compaction"
456            );
457
458            use crate::backend::types::ScanRequest;
459            let batches: Vec<RecordBatch> = backend.scan(ScanRequest::all(&adj_table_name)).await?;
460
461            for batch in batches {
462                let src_col = batch
463                    .column_by_name("src_vid")
464                    .ok_or(anyhow!("Missing src_vid"))?
465                    .as_any()
466                    .downcast_ref::<UInt64Array>()
467                    .ok_or(anyhow!("Invalid src_vid"))?;
468                let neighbors_col = batch
469                    .column_by_name("neighbors")
470                    .ok_or(anyhow!("Missing neighbors"))?
471                    .as_any()
472                    .downcast_ref::<ListArray>()
473                    .ok_or(anyhow!("Invalid neighbors"))?;
474                let edge_ids_col = batch
475                    .column_by_name("edge_ids")
476                    .ok_or(anyhow!("Missing edge_ids"))?
477                    .as_any()
478                    .downcast_ref::<ListArray>()
479                    .ok_or(anyhow!("Invalid edge_ids"))?;
480
481                for i in 0..batch.num_rows() {
482                    let vid = Vid::from(src_col.value(i));
483                    processed_vids.insert(vid);
484
485                    // Reconstruct current adjacency list
486                    let n_list = neighbors_col.value(i);
487                    let n_array = n_list.as_any().downcast_ref::<UInt64Array>().unwrap();
488                    let e_list = edge_ids_col.value(i);
489                    let e_array = e_list.as_any().downcast_ref::<UInt64Array>().unwrap();
490
491                    let mut current_edges: HashMap<Eid, Vid> = HashMap::new();
492                    for j in 0..n_array.len() {
493                        current_edges
494                            .insert(Eid::from(e_array.value(j)), Vid::from(n_array.value(j)));
495                    }
496
497                    if let Some(ops) = delta_map.get(&vid) {
498                        apply_deltas_to_edges(&mut current_edges, ops, direction);
499                    }
500
501                    append_edges_to_builders(
502                        vid,
503                        &current_edges,
504                        &mut src_vid_builder,
505                        &mut neighbors_builder,
506                        &mut edge_ids_builder,
507                    );
508                }
509            }
510        }
511
512        // Process new vertices (in deltas but not in L2)
513        for (vid, ops) in delta_map {
514            if processed_vids.contains(&vid) {
515                continue;
516            }
517
518            let mut current_edges: HashMap<Eid, Vid> = HashMap::new();
519            apply_deltas_to_edges(&mut current_edges, &ops, direction);
520
521            append_edges_to_builders(
522                vid,
523                &current_edges,
524                &mut src_vid_builder,
525                &mut neighbors_builder,
526                &mut edge_ids_builder,
527            );
528        }
529
530        // Final Flush — always replace L2, even when the compacted output is
531        // empty. If every edge for this (edge_type, direction) was deleted the
532        // builders are empty; skipping the replace here would leave the stale
533        // pre-delete L2 rows intact while the tombstone-clear below erases the
534        // Delta L1 deletes, resurrecting the deleted edges on the next read.
535        // Writing the (possibly empty) batch overwrites L2 to match the deltas.
536        {
537            let src_arr = Arc::new(src_vid_builder.finish());
538            let neighbors_arr = Arc::new(neighbors_builder.finish());
539            let edge_ids_arr = Arc::new(edge_ids_builder.finish());
540
541            let schema = adj_ds.get_arrow_schema();
542            let batch = RecordBatch::try_new(schema, vec![src_arr, neighbors_arr, edge_ids_arr])?;
543
544            // Replace the table with compacted data
545            adj_ds.replace(self.storage.backend(), batch).await?;
546        }
547
548        // CRITICAL: Clear Delta L1 after compaction
549        // Topology ops from Delta L1 are now incorporated into L2 adjacency.
550        // Edge properties survive in main_edges (dual-written during flush).
551        // Clearing Delta L1 prevents stale topology data from being read.
552        if !deltas.is_empty() {
553            info!(
554                "Clearing Delta L1 for edge_type={} direction={} after compaction (incorporated {} ops)",
555                edge_type,
556                direction,
557                deltas.len()
558            );
559
560            // Invariant: Every EID in Delta L1 must have a corresponding entry in
561            // main_edges, because Writer::flush_to_l1 performs a dual-write.
562            // Tests that create delta entries directly (for schema/overflow testing)
563            // must not call compact_adjacency without also populating main_edges.
564            #[cfg(debug_assertions)]
565            {
566                use crate::storage::main_edge::MainEdgeDataset;
567
568                let delta_eids: std::collections::HashSet<Eid> =
569                    deltas.iter().map(|e| e.eid).collect();
570
571                for eid in delta_eids {
572                    let main_edge_exists =
573                        MainEdgeDataset::exists_by_eid(self.storage.backend(), eid)
574                            .await
575                            .unwrap_or(false);
576
577                    debug_assert!(
578                        main_edge_exists,
579                        "EID {} from Delta L1 not found in main_edges after compaction. \
580                        This indicates edge properties were not dual-written during flush.",
581                        eid.as_u64()
582                    );
583                }
584            }
585
586            // Clear ONLY the deltas we actually merged into L2 — those at or
587            // below the high-water-mark captured from the rows read at the top
588            // of this function. A concurrent flush that appended new deltas
589            // inside the read→clear window stamped them with a strictly higher
590            // `_version`, so the predicate-delete leaves them intact to be
591            // reprocessed next compaction. This replaces the old unconditional
592            // empty-batch wipe, whose only guard was an instantaneous
593            // `flush_in_progress` check that a flush starting AND finishing in
594            // the window slipped past — silently wiping its rows. (review H11)
595            let clear_hwm = deltas.iter().map(|e| e.version).max().unwrap_or(0);
596            let delta_ds = self.storage.delta_dataset(edge_type, direction)?;
597            delta_ds
598                .delete_up_to_version(self.storage.backend(), clear_hwm)
599                .await?;
600        }
601
602        let duration = start.elapsed();
603        tracing::Span::current().record("duration_ms", duration.as_millis());
604        info!(
605            delta_count,
606            duration_ms = duration.as_millis(),
607            "Adjacency compaction completed"
608        );
609
610        metrics::histogram!("uni_compaction_duration_seconds", "type" => "adjacency")
611            .record(duration.as_secs_f64());
612
613        Ok(CompactionInfo {
614            edge_type: edge_type.to_string(),
615            direction: direction.to_string(),
616        })
617    }
618}
619
620/// Apply delta operations to an edge map, returning the resolved neighbor for the direction.
621fn apply_deltas_to_edges(current_edges: &mut HashMap<Eid, Vid>, ops: &[L1Entry], direction: &str) {
622    for op in ops {
623        match op.op {
624            Op::Insert => {
625                let neighbor = if direction == "fwd" {
626                    op.dst_vid
627                } else {
628                    op.src_vid
629                };
630                current_edges.insert(op.eid, neighbor);
631            }
632            Op::Delete => {
633                current_edges.remove(&op.eid);
634            }
635        }
636    }
637}
638
639/// Write sorted edges from a HashMap into adjacency list builders.
640fn append_edges_to_builders(
641    vid: Vid,
642    current_edges: &HashMap<Eid, Vid>,
643    src_vid_builder: &mut UInt64Builder,
644    neighbors_builder: &mut ListBuilder<UInt64Builder>,
645    edge_ids_builder: &mut ListBuilder<UInt64Builder>,
646) {
647    if current_edges.is_empty() {
648        return;
649    }
650    src_vid_builder.append_value(vid.as_u64());
651
652    let mut sorted_eids: Vec<_> = current_edges.keys().cloned().collect();
653    sorted_eids.sort();
654
655    for eid in sorted_eids {
656        let neighbor = current_edges[&eid];
657        neighbors_builder.values().append_value(neighbor.as_u64());
658        edge_ids_builder.values().append_value(eid.as_u64());
659    }
660    neighbors_builder.append(true);
661    edge_ids_builder.append(true);
662}
663
664/// Information returned by adjacency compaction about what was compacted.
665/// Used to coordinate in-memory CSR re-warm after storage compaction.
666#[derive(Debug, Clone)]
667pub struct CompactionInfo {
668    pub edge_type: String,
669    pub direction: String,
670}
671
672#[cfg(test)]
673mod ws_d_crdt_tests {
674    //! WS-D (P0.4): the compaction merge path must route CRDT merges through
675    //! the plugin registry (previously it called native `try_merge` directly).
676    //! These exercise `Compactor::merge_crdt_values` in isolation: with a
677    //! registered provider the registry path runs; with `None` it falls back
678    //! to native, byte-for-byte. (The full flush→compact integration test that
679    //! proves the stamped `StorageManager` registry reaches `compact_vertices`
680    //! at runtime is a recommended follow-up.)
681
682    use std::sync::Arc;
683    use std::sync::atomic::{AtomicUsize, Ordering};
684
685    use uni_common::Value;
686    use uni_crdt::{Crdt, GCounter};
687    use uni_plugin::traits::crdt::{CrdtKind, CrdtKindProvider, CrdtOp, CrdtState, ScalarValue};
688    use uni_plugin::{
689        Capability, CapabilitySet, FnError, PluginId, PluginRegistrar, PluginRegistry,
690    };
691
692    #[derive(Default)]
693    struct CountingProvider {
694        calls: AtomicUsize,
695    }
696    impl CrdtKindProvider for CountingProvider {
697        fn kind(&self) -> CrdtKind {
698            CrdtKind::new("uni-crdt:g-counter")
699        }
700        fn empty(&self) -> Box<dyn CrdtState> {
701            Box::new(St {
702                inner: Crdt::GCounter(GCounter::new()),
703            })
704        }
705        fn from_persisted(&self, bytes: &[u8]) -> Result<Box<dyn CrdtState>, FnError> {
706            self.calls.fetch_add(1, Ordering::SeqCst);
707            let inner =
708                Crdt::from_msgpack(bytes).map_err(|e| FnError::new(0xA01, format!("{e}")))?;
709            Ok(Box::new(St { inner }))
710        }
711    }
712    struct St {
713        inner: Crdt,
714    }
715    impl CrdtState for St {
716        fn as_any(&self) -> &dyn std::any::Any {
717            self
718        }
719        fn apply(&mut self, _op: &CrdtOp) -> Result<(), FnError> {
720            Ok(())
721        }
722        fn merge(&mut self, other: &dyn CrdtState) -> Result<(), FnError> {
723            let o = other
724                .as_any()
725                .downcast_ref::<St>()
726                .ok_or_else(|| FnError::new(0xA10, "type mismatch"))?;
727            self.inner
728                .try_merge(&o.inner)
729                .map_err(|e| FnError::new(0xA11, format!("{e}")))
730        }
731        fn value(&self) -> Result<ScalarValue, FnError> {
732            Ok(ScalarValue::Utf8(Some(self.inner.type_name().to_owned())))
733        }
734        fn persist(&self) -> Result<Vec<u8>, FnError> {
735            self.inner
736                .to_msgpack()
737                .map_err(|e| FnError::new(0xA12, format!("{e}")))
738        }
739    }
740
741    fn gcounter(replica: &str, by: u64) -> Value {
742        let mut g = GCounter::new();
743        g.increment(replica, by);
744        Value::from(serde_json::to_value(Crdt::GCounter(g)).unwrap())
745    }
746
747    #[test]
748    fn compaction_merge_routes_through_registry() {
749        let registry = PluginRegistry::new();
750        let provider = Arc::new(CountingProvider::default());
751        let caps = CapabilitySet::from_iter_of([Capability::Crdt]);
752        let mut r = PluginRegistrar::new(PluginId::new("test.counting"), &caps, &registry);
753        r.crdt_kind(
754            CrdtKind::new("uni-crdt:g-counter"),
755            Arc::clone(&provider) as Arc<dyn CrdtKindProvider>,
756        )
757        .unwrap();
758        r.commit_to_registry().unwrap();
759
760        let a = gcounter("r1", 5); // existing
761        let b = gcounter("r2", 7); // new
762        let merged = super::Compactor::merge_crdt_values(&a, &b, Some(&registry)).unwrap();
763
764        assert!(
765            provider.calls.load(Ordering::SeqCst) > 0,
766            "compaction merge must route through the registered provider"
767        );
768        let crdt: Crdt = serde_json::from_value(merged.into()).unwrap();
769        match crdt {
770            Crdt::GCounter(g) => assert_eq!(g.value(), 12, "5 + 7 = 12"),
771            other => panic!("expected GCounter, got {other:?}"),
772        }
773    }
774
775    #[test]
776    fn compaction_merge_falls_back_to_native_without_registry() {
777        let a = gcounter("r1", 3);
778        let b = gcounter("r2", 4);
779        let merged = super::Compactor::merge_crdt_values(&a, &b, None).unwrap();
780        let crdt: Crdt = serde_json::from_value(merged.into()).unwrap();
781        match crdt {
782            Crdt::GCounter(g) => assert_eq!(g.value(), 7, "3 + 4 = 7 native fallback"),
783            other => panic!("expected GCounter, got {other:?}"),
784        }
785    }
786}