quatzal-storage 0.1.0

// SPDX-License-Identifier: Apache-2.0
//! Cross-shard router: owns one glommio executor OS thread per shard and dispatches
//! requests to the shard that owns a given key by hash. `UACE-FR-1.7`. No shared mutable
//! state ever crosses a shard boundary — everything below this line is message-passing.

use std::collections::BTreeMap;
use std::hash::{Hash, Hasher};
use std::path::PathBuf;
use std::sync::mpsc::{self, Sender, SyncSender};

use glommio::{ExecutorJoinHandle, LocalExecutorBuilder, Placement};
use quatzal_schema::{Row, UaceError, UaceResult, Value};

use crate::shard::{Shard, ShardConfig, ShardStats};

enum ShardRequest {
    Put(Row, SyncSender<UaceResult<()>>),
    UpdateScalars(Vec<u8>, BTreeMap<String, Value>, SyncSender<UaceResult<()>>),
    Get(Vec<u8>, SyncSender<UaceResult<Option<Row>>>),
    Delete(Vec<u8>, SyncSender<UaceResult<()>>),
    Stats(SyncSender<ShardStats>),
    Scan(SyncSender<UaceResult<Vec<Row>>>),
    PutBatch(Vec<Row>, SyncSender<UaceResult<()>>),
    GetBatch(Vec<Vec<u8>>, SyncSender<UaceResult<Vec<Option<Row>>>>),
}

pub struct Engine {
    senders: Vec<Sender<ShardRequest>>,
    joined: bool,
    raw_handles: Vec<ExecutorJoinHandle<()>>,
}

impl Engine {
    /// Open (or create) an engine rooted at `dir`, sharded across `num_shards` glommio
    /// executor threads (default: available parallelism).
    pub fn open(dir: impl Into<PathBuf>, num_shards: Option<usize>) -> UaceResult<Self> {
        let dir = dir.into();
        std::fs::create_dir_all(&dir)
            .map_err(|e| UaceError::Io(format!("create data dir {dir:?}: {e}")))?;
        let available = std::thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(1);
        let num_shards = num_shards.unwrap_or(available).max(1);

        let mut senders = Vec::with_capacity(num_shards);
        let mut raw_handles = Vec::with_capacity(num_shards);

        for shard_id in 0..num_shards {
            let (tx, rx) = mpsc::channel::<ShardRequest>();
            let shard_dir = dir.join(format!("shard-{shard_id:04}"));
            let placement = if shard_id < available {
                Placement::Fixed(shard_id)
            } else {
                Placement::Unbound
            };

            let handle = LocalExecutorBuilder::new(placement)
                .name(&format!("uace-shard-{shard_id}"))
                .spawn(move || async move {
                    let cfg = ShardConfig::new(shard_dir);
                    let mut shard = Shard::open(cfg).await.expect("shard open failed");
                    while let Ok(req) = rx.recv() {
                        // A shard has exactly one long-running task and runs compaction
                        // inline (see CLAUDE.md); blocking on `rx.recv()` between requests
                        // is therefore safe — there is nothing else on this executor for it
                        // to starve.
                        match req {
                            ShardRequest::Put(row, reply) => {
                                let r = shard.put(row).await;
                                let _ = reply.send(r);
                            }
                            ShardRequest::UpdateScalars(key, scalars, reply) => {
                                let r = shard.update_scalars(&key, scalars).await;
                                let _ = reply.send(r);
                            }
                            ShardRequest::Get(key, reply) => {
                                let r = shard.get(&key);
                                let _ = reply.send(r);
                            }
                            ShardRequest::Delete(key, reply) => {
                                let r = shard.delete(&key).await;
                                let _ = reply.send(r);
                            }
                            ShardRequest::Stats(reply) => {
                                let _ = reply.send(shard.stats());
                            }
                            ShardRequest::Scan(reply) => {
                                let _ = reply.send(shard.scan());
                            }
                            ShardRequest::PutBatch(rows, reply) => {
                                let r = shard.put_batch(rows).await;
                                let _ = reply.send(r);
                            }
                            ShardRequest::GetBatch(keys, reply) => {
                                let _ = reply.send(shard.get_batch(&keys));
                            }
                        }
                    }
                })
                .map_err(|e| UaceError::Io(format!("spawn shard executor {shard_id}: {e}")))?;

            senders.push(tx);
            raw_handles.push(handle);
        }

        Ok(Engine {
            senders,
            joined: false,
            raw_handles,
        })
    }

    fn shard_index(&self, key: &[u8]) -> usize {
        let mut hasher = std::collections::hash_map::DefaultHasher::new();
        key.hash(&mut hasher);
        (hasher.finish() as usize) % self.senders.len()
    }

    pub fn num_shards(&self) -> usize {
        self.senders.len()
    }

    pub fn put(&self, row: Row) -> UaceResult<()> {
        let idx = self.shard_index(&row.key);
        let (tx, rx) = mpsc::sync_channel(1);
        self.senders[idx]
            .send(ShardRequest::Put(row, tx))
            .map_err(|_| UaceError::ShardClosed)?;
        rx.recv().map_err(|_| UaceError::ShardClosed)?
    }

    pub fn update_scalars(&self, key: &[u8], scalars: BTreeMap<String, Value>) -> UaceResult<()> {
        let idx = self.shard_index(key);
        let (tx, rx) = mpsc::sync_channel(1);
        self.senders[idx]
            .send(ShardRequest::UpdateScalars(key.to_vec(), scalars, tx))
            .map_err(|_| UaceError::ShardClosed)?;
        rx.recv().map_err(|_| UaceError::ShardClosed)?
    }

    pub fn get(&self, key: &[u8]) -> UaceResult<Option<Row>> {
        let idx = self.shard_index(key);
        let (tx, rx) = mpsc::sync_channel(1);
        self.senders[idx]
            .send(ShardRequest::Get(key.to_vec(), tx))
            .map_err(|_| UaceError::ShardClosed)?;
        rx.recv().map_err(|_| UaceError::ShardClosed)?
    }

    pub fn delete(&self, key: &[u8]) -> UaceResult<()> {
        let idx = self.shard_index(key);
        let (tx, rx) = mpsc::sync_channel(1);
        self.senders[idx]
            .send(ShardRequest::Delete(key.to_vec(), tx))
            .map_err(|_| UaceError::ShardClosed)?;
        rx.recv().map_err(|_| UaceError::ShardClosed)?
    }

    /// `UACE-FR-1.9`: group-committed batch insert. Rows are partitioned by the same key
    /// hash single `put` uses, then each shard applies its slice under **one** WAL fsync.
    /// This is the fix for the constraint `UACE-NFR-33` measured end-to-end: per-row fsync
    /// capped the product surface at ~540 rows/sec while the vector tier was 68x faster.
    /// Shards are addressed sequentially here (each is a channel round trip), but the
    /// fsyncs themselves happen concurrently across shard threads.
    pub fn put_batch(&self, rows: Vec<Row>) -> UaceResult<()> {
        if rows.is_empty() {
            return Ok(());
        }
        let mut buckets: Vec<Vec<Row>> = (0..self.senders.len()).map(|_| Vec::new()).collect();
        for row in rows {
            let idx = self.shard_index(&row.key);
            buckets[idx].push(row);
        }
        // Dispatch every shard's slice first, then collect -- otherwise each shard's fsync
        // would serialize behind the previous shard's.
        let mut pending = Vec::new();
        for (idx, bucket) in buckets.into_iter().enumerate() {
            if bucket.is_empty() {
                continue;
            }
            let (tx, rx) = mpsc::sync_channel(1);
            self.senders[idx]
                .send(ShardRequest::PutBatch(bucket, tx))
                .map_err(|_| UaceError::ShardClosed)?;
            pending.push(rx);
        }
        for rx in pending {
            rx.recv().map_err(|_| UaceError::ShardClosed)??;
        }
        Ok(())
    }

    /// `UACE-FR-1.10`: many point reads in one round trip per shard, results returned
    /// positionally (`None` for absent keys). `UACE-NFR-33` traced part of query latency
    /// to provenance lookups doing k *sequential* channel round trips per search.
    pub fn get_batch(&self, keys: &[Vec<u8>]) -> UaceResult<Vec<Option<Row>>> {
        if keys.is_empty() {
            return Ok(Vec::new());
        }
        // Remember each key's original position so the caller gets answers in its own order.
        let mut buckets: Vec<(Vec<Vec<u8>>, Vec<usize>)> = (0..self.senders.len())
            .map(|_| (Vec::new(), Vec::new()))
            .collect();
        for (position, key) in keys.iter().enumerate() {
            let idx = self.shard_index(key);
            buckets[idx].0.push(key.clone());
            buckets[idx].1.push(position);
        }

        let mut pending = Vec::new();
        for (idx, (shard_keys, positions)) in buckets.into_iter().enumerate() {
            if shard_keys.is_empty() {
                continue;
            }
            let (tx, rx) = mpsc::sync_channel(1);
            self.senders[idx]
                .send(ShardRequest::GetBatch(shard_keys, tx))
                .map_err(|_| UaceError::ShardClosed)?;
            pending.push((rx, positions));
        }

        let mut out: Vec<Option<Row>> = (0..keys.len()).map(|_| None).collect();
        for (rx, positions) in pending {
            let rows = rx.recv().map_err(|_| UaceError::ShardClosed)??;
            for (position, row) in positions.into_iter().zip(rows) {
                out[position] = row;
            }
        }
        Ok(out)
    }

    /// `UACE-FR-1.8`: every live row across all shards -- newest version per key, no
    /// tombstones. A full-scan primitive (whole result set in memory), deliberately not
    /// an iterator API yet; revisit when result sets stop fitting in memory.
    pub fn scan(&self) -> UaceResult<Vec<Row>> {
        let mut rows = Vec::new();
        for sender in &self.senders {
            let (tx, rx) = mpsc::sync_channel(1);
            sender
                .send(ShardRequest::Scan(tx))
                .map_err(|_| UaceError::ShardClosed)?;
            rows.extend(rx.recv().map_err(|_| UaceError::ShardClosed)??);
        }
        Ok(rows)
    }

    pub fn stats(&self) -> Vec<ShardStats> {
        self.senders
            .iter()
            .map(|s| {
                let (tx, rx) = mpsc::sync_channel(1);
                if s.send(ShardRequest::Stats(tx)).is_err() {
                    return ShardStats::default();
                }
                rx.recv().unwrap_or_default()
            })
            .collect()
    }
}

impl Drop for Engine {
    fn drop(&mut self) {
        if self.joined {
            return;
        }
        self.joined = true;
        // Dropping senders closes every shard's channel, which ends its request loop and
        // lets its executor thread finish; only then do we join, so this never deadlocks.
        self.senders.clear();
        for handle in self.raw_handles.drain(..) {
            let _ = handle.join();
        }
    }
}