engramdb_io/

batch.rs

//! 批式 badge 读取：排序去重 + 预取计划 + 多线程 gather。

use std::collections::HashMap;
use std::fs::File;
use std::path::Path;

use crate::backend::{default_backend, IoBackend};
use engramdb_core::layout::Layout;

/// 预取计划：按分片分组的 badge 块列表（每 shard 内部升序，供顺序预读/合并）。
#[derive(Debug, Default, Clone)]
pub struct PrefetchPlan {
    /// shard_id -> 去重升序的 badge 块号
    pub shard_badges: HashMap<u64, Vec<u64>>,
    pub n_badges: usize,
}

impl PrefetchPlan {
    pub fn build(keys: &[u64], layout: &Layout) -> Self {
        let mut shard_badges: HashMap<u64, Vec<u64>> = HashMap::new();
        for &k in keys {
            let (shard, badge, _) = layout.locate(k);
            shard_badges.entry(shard).or_default().push(badge);
        }
        for v in shard_badges.values_mut() {
            v.sort_unstable();
            v.dedup();
        }
        let n_badges = shard_badges.values().map(|v| v.len()).sum();
        Self {
            shard_badges,
            n_badges,
        }
    }

    pub fn badges(&self, shard: u64) -> &[u64] {
        self.shard_badges
            .get(&shard)
            .map(|v| v.as_slice())
            .unwrap_or(&[])
    }

    /// 追加一个（shard, badge）到计划（流水线增量累积；重复由 settle 处理）。
    pub fn entry(&mut self, shard: u64, badge: u64) {
        self.shard_badges.entry(shard).or_default().push(badge);
        self.n_badges += 1;
    }

    /// 结算：shard 内排序 + 去重 + 重算 n_badges（计划被消费前调用）。
    pub fn settle(&mut self) {
        let mut n = 0usize;
        for v in self.shard_badges.values_mut() {
            v.sort_unstable();
            v.dedup();
            n += v.len();
        }
        self.n_badges = n;
    }
}

/// 读路径（pread 池化）：一次 `gather_plan` = 按 plan 拉取全部 badge 并组装 `[n,width]`。
pub struct BadgeGather<'a> {
    pub layout: &'a Layout,
    files: Vec<File>,
    backend: Box<dyn IoBackend>,
}

impl<'a> BadgeGather<'a> {
    pub fn open(dir: &Path, layout: &'a Layout) -> std::io::Result<Self> {
        Self::open_with_backend(dir, layout, default_backend())
    }

    /// 可注入后端（测试/基准：Preadv；未来 Linux：io_uring）。
    pub fn open_with_backend(
        dir: &Path,
        layout: &'a Layout,
        backend: Box<dyn IoBackend>,
    ) -> std::io::Result<Self> {
        let n = layout.shards as usize;
        let mut files = Vec::with_capacity(n);
        for i in 0..n {
            let s = dir.join(format!("shard_{:03}.bin", i));
            let b = dir.join(format!("badge_{:03}.bin", i));
            let p = if s.exists() { s } else { b };
            files.push(File::open(p)?);
        }
        Ok(Self {
            layout,
            files,
            backend,
        })
    }

    pub fn into_files(self) -> Vec<File> {
        self.files
    }

    /// 多线程 gather：`keys`/`out` 按 chunk 分片并行（各线程独立 badge 缓冲区）。
    pub fn gather_parallel(
        &self,
        keys: &[u64],
        out: &mut [u8],
        threads: usize,
    ) -> std::io::Result<()> {
        if threads <= 1 || keys.len() <= 1024 {
            return self.gather_naive(keys, out);
        }
        let w = self.layout.width as usize;
        let chunk_keys = keys.len().div_ceil(threads);
        std::thread::scope(|s| {
            for (kc, oc) in keys.chunks(chunk_keys).zip(out.chunks_mut(chunk_keys * w)) {
                s.spawn(|| {
                    let _ = self.gather_naive(kc, oc);
                });
            }
        });
        Ok(())
    }

    /// 朴素单线程：逐 key 定位 badge 读取（每个线程独立缓冲区）。
    pub fn gather_naive(&self, keys: &[u64], out: &mut [u8]) -> std::io::Result<()> {
        let w = self.layout.width as usize;
        let rb = self.layout.row_bytes as usize;
        let mut badge_buf = vec![0u8; self.layout.badge_bytes() as usize];
        let mut last: Option<(u64, u64)> = None;
        for (i, &k) in keys.iter().enumerate() {
            let (shard, badge, in_badge) = self.layout.locate(k);
            if last != Some((shard, badge)) {
                let off = badge * self.layout.badge_bytes();
                self.backend
                    .read_exact_at(&self.files[shard as usize], &mut badge_buf, off)?;
                last = Some((shard, badge));
            }
            let src = in_badge as usize * rb;
            out[i * w..(i + 1) * w].copy_from_slice(&badge_buf[src..src + w]);
        }
        Ok(())
    }

    /// 页对齐专用读路径：按 shard 分线程，shard 内按键升序聚页（4KiB 对齐，每页只读一次），
    /// 行跨页边界时补读该行尾部。各线程只写自己的采集缓冲，主线程回填 out（无共享 &mut）。
    pub fn gather_pp(&self, keys: &[u64], out: &mut [u8], threads: usize) -> std::io::Result<()> {
        const PAGE: u64 = 4096;
        let w = self.layout.width as usize;
        let rb = self.layout.row_bytes as usize;
        let mut groups: HashMap<u64, Vec<(u64, usize)>> = HashMap::new();
        for (i, &k) in keys.iter().enumerate() {
            let (shard, _, _) = self.layout.locate(k);
            groups.entry(shard).or_default().push((k, i));
        }
        let mut tasks: Vec<(u64, Vec<(u64, usize)>)> = groups.into_iter().collect();
        tasks.sort_unstable_by_key(|&(s, _)| s);

        // 各任务独立产出 (idxs 升序, rows 扁平)
        let nt = threads.max(1).min(tasks.len());
        let chunk = tasks.len().div_ceil(nt);
        let mut results: Vec<(Vec<usize>, Vec<u8>)> = Vec::new();

        std::thread::scope(|s| {
            let mut handles = Vec::new();
            let mut task_iter = tasks.into_iter();
            while task_iter.len() > 0 {
                let t: Vec<(u64, Vec<(u64, usize)>)> = task_iter.by_ref().take(chunk).collect();
                handles.push(s.spawn(move || {
                    let mut out_rows: Vec<u8> = Vec::new();
                    let mut out_idxs: Vec<usize> = Vec::new();
                    for (shard, mut pairs) in t {
                        pairs.sort_unstable();
                        let f = &self.files[shard as usize];
                        let mut last_page: Option<u64> = None;
                        let mut page = vec![0u8; (PAGE + 2 * (rb as u64)) as usize];
                        let mut prev_key: Option<u64> = None;
                        for (k, oi) in pairs {
                            let (_, _, in_b) = self.layout.locate(k);
                            // gather_pp groups by shard and reads from that
                            // shard's file, so the byte offset must be local
                            // to the shard rather than the global rowid.
                            let local_row = k % self.layout.rows_per_shard;
                            let byte_off = local_row * rb as u64;
                            let page_id = byte_off & !(PAGE - 1);
                            if last_page != Some(page_id) {
                                let want = (PAGE + rb as u64) as usize;
                                let n = self
                                    .backend
                                    .read_at(f, &mut page[..want], page_id)
                                    .unwrap_or(0);
                                let _ = n;
                                last_page = Some(page_id);
                            }
                            let in_page = (byte_off - page_id) as usize;
                            if in_page + rb <= PAGE as usize {
                                out_rows.extend_from_slice(&page[in_page..in_page + rb]);
                            } else {
                                let mut tmp = vec![0u8; rb];
                                let _ = self.backend.read_exact_at(f, &mut tmp, byte_off);
                                out_rows.extend_from_slice(&tmp);
                            }
                            out_idxs.push(oi);
                            let _ = (in_b, prev_key);
                            prev_key = Some(k);
                        }
                    }
                    (out_idxs, out_rows)
                }));
            }
            for h in handles {
                if let Ok(r) = h.join() {
                    results.push(r);
                }
            }
        });

        for (idxs, rows) in results {
            for (j, &oi) in idxs.iter().enumerate() {
                let slice = &rows[j * w..(j + 1) * w];
                out[oi * w..(oi + 1) * w].copy_from_slice(slice);
            }
        }
        Ok(())
    }

    /// 有序批式读：对同一 badge 的 keys 组内合并读（按 plan 的排序），
    /// 期望预取服务器将 plan 先落地 —— 本实现直接做"计划->读取"。
    pub fn gather_planned(&self, keys: &[u64], out: &mut [u8]) -> std::io::Result<()> {
        // 简单实现：按 (shard,badge) 分组，逐组读一次，再按 key 顺序回填
        let plan = PrefetchPlan::build(keys, self.layout);
        let rb = self.layout.row_bytes as usize;
        let w = self.layout.width as usize;
        let mut cache: HashMap<(u64, u64), Vec<u8>> = HashMap::new();
        let mut groups: HashMap<(u64, u64), Vec<usize>> = HashMap::new();
        for (i, &k) in keys.iter().enumerate() {
            let (s, b, _) = self.layout.locate(k);
            groups.entry((s, b)).or_default().push(i);
        }
        for (&(s, b), idxs) in &groups {
            let buf = cache.entry((s, b)).or_insert_with(|| {
                let mut buf = vec![0u8; self.layout.badge_bytes() as usize];
                let off = b * self.layout.badge_bytes();
                let _ = self
                    .backend
                    .read_exact_at(&self.files[s as usize], &mut buf, off);
                buf
            });
            for &i in idxs {
                let (_, _, in_badge) = self.layout.locate(keys[i]);
                let src = in_badge as usize * rb;
                out[i * w..(i + 1) * w].copy_from_slice(&buf[src..src + w]);
            }
        }
        let _ = plan;
        Ok(())
    }
    /// 按计划消费：`plan` 决定需要读的 badge 集合，本函数按键组读取一次完整 badge
    /// 并回填 `out`（与 `keys` 顺序一致）。与 `gather_planned` 的差异：计划预先生成
    /// （由 StreamingPlanner/外部预取服务器），这里只做"计划 -> 数据"。
    pub fn gather_plan(
        &self,
        keys: &[u64],
        plan: &PrefetchPlan,
        out: &mut [u8],
    ) -> std::io::Result<()> {
        let rb = self.layout.row_bytes as usize;
        let w = self.layout.width as usize;
        let mut groups: HashMap<(u64, u64), Vec<usize>> = HashMap::new();
        for (i, &k) in keys.iter().enumerate() {
            let (s, b, _) = self.layout.locate(k);
            if plan.badges(s).contains(&b) {
                groups.entry((s, b)).or_default().push(i);
            } else {
                // 计划外的 key：退化为直接行读（防御；正常流式路径不会出现）
                let mut tmp = vec![0u8; rb];
                self.backend
                    .read_exact_at(&self.files[s as usize], &mut tmp, k * rb as u64)?;
                out[i * w..(i + 1) * w].copy_from_slice(&tmp);
            }
        }
        for (&(s, b), idxs) in &groups {
            let mut buf = vec![0u8; self.layout.badge_bytes() as usize];
            let off = b * self.layout.badge_bytes();
            self.backend
                .read_exact_at(&self.files[s as usize], &mut buf, off)?;
            for &i in idxs {
                let (_, _, in_badge) = self.layout.locate(keys[i]);
                let src = in_badge as usize * rb;
                out[i * w..(i + 1) * w].copy_from_slice(&buf[src..src + w]);
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use engramdb_core::layout::Layout;

    #[test]
    fn plan_dedup_sort() {
        let layout = Layout::new(1, 10_000, 160, 1); // 单分片
        let keys = vec![9999, 0, 5, 9999, 250, 250];
        let p = PrefetchPlan::build(&keys, &layout);
        assert_eq!(p.badges(0), &[0, 10, 399]);
    }

    #[test]
    fn gather_plan_end_to_end() {
        use std::io::Write;
        let dir = std::env::temp_dir().join("engramdb-gather-plan-test");
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        // 单分片 100 行 × 8 字节（u64 值 = rowid）；文件填满完整 badge（512 行 × 8B = 4096B）
        let layout = Layout::new(1, 100, 8, 1);
        let mut f = std::fs::File::create(dir.join("shard_000.bin")).unwrap();
        for i in 0..512u64 {
            f.write_all(&i.to_le_bytes()).unwrap();
        }
        drop(f);
        let bg = BadgeGather::open(&dir, &layout).unwrap();
        let keys = vec![3u64, 77, 3, 99, 60];
        let plan = PrefetchPlan::build(&keys, &layout);
        let mut out = vec![0u8; keys.len() * 8];
        bg.gather_plan(&keys, &plan, &mut out).unwrap();
        for (j, &want) in keys.iter().enumerate() {
            let got = u64::from_le_bytes(out[j * 8..(j + 1) * 8].try_into().unwrap());
            assert_eq!(got, want, "rowid {want} at out[{j}]");
        }
        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn gather_pp_multishard_uses_local_row_offsets() {
        use std::io::Write;
        let dir = std::env::temp_dir().join("engramdb-gather-pp-multishard-test");
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        let shards = 3u64;
        let rows_per_shard = 32u64;
        let width = 4u64;
        let layout = Layout::new(shards, rows_per_shard, width, 1);
        for s in 0..shards {
            let mut f = std::fs::File::create(dir.join(format!("shard_{s:03}.bin"))).unwrap();
            for local in 0..rows_per_shard {
                let val = (s * 1_000_000 + local) as u32;
                f.write_all(&val.to_le_bytes()).unwrap();
            }
        }
        let bg = BadgeGather::open(&dir, &layout).unwrap();
        let keys = vec![0u64, 1, 31, 32, 63, 64, 95];
        let mut out = vec![0u8; keys.len() * width as usize];
        bg.gather_pp(&keys, &mut out, 4).unwrap();
        for (j, &k) in keys.iter().enumerate() {
            let shard = k / rows_per_shard;
            let local = k % rows_per_shard;
            let want = (shard * 1_000_000 + local) as u32;
            let got = u32::from_le_bytes(out[j * 4..(j + 1) * 4].try_into().unwrap());
            assert_eq!(got, want, "gather_pp rowid {k} at out[{j}]");
        }
        let _ = std::fs::remove_dir_all(&dir);
    }
}