value_log/
scanner.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
// Copyright (c) 2024-present, fjall-rs
// This source code is licensed under both the Apache 2.0 and MIT License
// (found in the LICENSE-* files in the repository)

use crate::{id::SegmentId, ValueHandle};
use std::{collections::BTreeMap, sync::MutexGuard};

#[derive(Debug, Default)]
pub struct SegmentCounter {
    pub size: u64,
    pub item_count: u64,
}

pub type SizeMap = BTreeMap<SegmentId, SegmentCounter>;

/// Scans a value log, building a size map for the GC report
pub struct Scanner<'a, I: Iterator<Item = std::io::Result<(ValueHandle, u32)>>> {
    iter: I,

    #[allow(unused)]
    lock_guard: MutexGuard<'a, ()>,

    size_map: SizeMap,
}

impl<'a, I: Iterator<Item = std::io::Result<(ValueHandle, u32)>>> Scanner<'a, I> {
    pub fn new(iter: I, lock_guard: MutexGuard<'a, ()>, ids: &[SegmentId]) -> Self {
        let mut size_map = BTreeMap::default();

        for &id in ids {
            size_map.insert(id, SegmentCounter::default());
        }

        Self {
            iter,
            lock_guard,
            size_map,
        }
    }

    pub fn finish(self) -> SizeMap {
        self.size_map
    }

    pub fn scan(&mut self) -> crate::Result<()> {
        for vhandle in self.iter.by_ref() {
            let (vhandle, size) = vhandle.map_err(|_| {
                crate::Error::Io(std::io::Error::new(
                    std::io::ErrorKind::Other,
                    "Index returned error",
                ))
            })?;
            let size = u64::from(size);

            self.size_map
                .entry(vhandle.segment_id)
                .and_modify(|x| {
                    x.item_count += 1;
                    x.size += size;
                })
                .or_insert_with(|| SegmentCounter {
                    size,
                    item_count: 1,
                });
        }

        Ok(())
    }
}