hypercore-kb / opencog-kb / basal-atomspace / rust-atomspace / src / shard.rs

//! Sharded AtomSpace for billion-scale storage
//!
//! Distributes atoms across multiple storage backends by hash.

use crate::atom::{Atom, Hash, Value};
use crate::query::{query, Pattern, QueryResult};
use crate::storage::AtomStore;
use rayon::prelude::*;
use std::path::Path;

#[cfg(feature = "rocksdb-storage")]
use crate::storage::RocksStore;

/// Sharded AtomSpace - distributes across N shards by hash
pub struct ShardedAtomSpace<S: AtomStore + Send + Sync> {
    shards: Vec<S>,
    num_shards: usize,
}

impl<S: AtomStore + Send + Sync> ShardedAtomSpace<S> {
    /// Create from pre-built shards
    pub fn from_shards(shards: Vec<S>) -> Self {
        let num_shards = shards.len();
        Self { shards, num_shards }
    }

    /// Get shard index for a hash
    fn shard_idx(&self, hash: &Hash) -> usize {
        // Use first byte of hash for distribution
        (hash[0] as usize) % self.num_shards
    }

    /// Get shard for a hash
    fn shard(&self, hash: &Hash) -> &S {
        &self.shards[self.shard_idx(hash)]
    }

    /// Add an atom
    pub fn add(&self, atom: &Atom) -> Hash {
        let hash = atom.hash();
        self.shard(&hash).add(atom);
        hash
    }

    /// Get an atom by hash
    pub fn get(&self, hash: &Hash) -> Option<Atom> {
        self.shard(hash).get(hash)
    }

    /// Get a node by type and name
    pub fn get_node(&self, typ: &str, name: &str) -> Option<(Hash, Atom)> {
        // Need to check all shards since we don't know the hash
        // In practice, could add a global name index
        for shard in &self.shards {
            if let Some(result) = shard.get_node(typ, name) {
                return Some(result);
            }
        }
        None
    }

    /// Query across all shards in parallel
    pub fn query(&self, pattern: &Pattern) -> Vec<(Hash, Atom)> {
        self.shards
            .par_iter()
            .flat_map(|shard| query(shard, pattern).collect::<Vec<_>>())
            .collect()
    }

    /// Set a value on an atom
    pub fn set_value(&self, hash: &Hash, key: &str, value: &Value) {
        self.shard(hash).set_value(hash, key, value);
    }

    /// Get a value from an atom
    pub fn get_value(&self, hash: &Hash, key: &str) -> Option<Value> {
        self.shard(hash).get_value(hash, key)
    }

    /// Get incoming links (may need to check all shards)
    pub fn get_incoming(&self, target: &Hash) -> Vec<Hash> {
        // Links could be in any shard, need to check all
        self.shards
            .par_iter()
            .flat_map(|shard| shard.get_incoming(target))
            .collect()
    }

    /// Total atom count
    pub fn len(&self) -> usize {
        self.shards.iter().map(|s| s.len()).sum()
    }

    /// Is empty?
    pub fn is_empty(&self) -> bool {
        self.shards.iter().all(|s| s.is_empty())
    }
}

#[cfg(feature = "rocksdb-storage")]
impl ShardedAtomSpace<RocksStore> {
    /// Open a sharded atomspace with RocksDB backends
    pub fn open(base_path: impl AsRef<Path>, num_shards: usize) -> Result<Self, rocksdb::Error> {
        let base = base_path.as_ref();
        let shards: Result<Vec<_>, _> = (0..num_shards)
            .map(|i| {
                let path = base.join(format!("shard_{:04}", i));
                std::fs::create_dir_all(&path).ok();
                RocksStore::open(&path)
            })
            .collect();

        Ok(Self {
            shards: shards?,
            num_shards,
        })
    }
}

/// Builder for creating a sharded atomspace
pub struct ShardedAtomSpaceBuilder {
    num_shards: usize,
    base_path: Option<std::path::PathBuf>,
}

impl Default for ShardedAtomSpaceBuilder {
    fn default() -> Self {
        Self::new()
    }
}

impl ShardedAtomSpaceBuilder {
    pub fn new() -> Self {
        Self {
            num_shards: num_cpus::get(), // Default: one shard per CPU
            base_path: None,
        }
    }

    pub fn num_shards(mut self, n: usize) -> Self {
        self.num_shards = n;
        self
    }

    pub fn path(mut self, path: impl AsRef<Path>) -> Self {
        self.base_path = Some(path.as_ref().to_path_buf());
        self
    }

    #[cfg(feature = "rocksdb-storage")]
    pub fn build_rocks(self) -> Result<ShardedAtomSpace<RocksStore>, rocksdb::Error> {
        let path = self.base_path.unwrap_or_else(|| std::path::PathBuf::from("./atomspace"));
        ShardedAtomSpace::open(path, self.num_shards)
    }

    pub fn build_memory(self) -> ShardedAtomSpace<crate::storage::MemoryStore> {
        let shards = (0..self.num_shards)
            .map(|_| crate::storage::MemoryStore::new())
            .collect();
        ShardedAtomSpace::from_shards(shards)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::storage::MemoryStore;

    #[test]
    fn test_sharded_atomspace() {
        let shards: Vec<MemoryStore> = (0..4).map(|_| MemoryStore::new()).collect();
        let as_ = ShardedAtomSpace::from_shards(shards);

        // Add many atoms
        let mut hashes = Vec::new();
        for i in 0..1000 {
            let atom = Atom::node("ConceptNode", format!("node_{}", i));
            hashes.push(as_.add(&atom));
        }

        // Verify distribution
        let shard_counts: Vec<usize> = as_.shards.iter().map(|s| s.len()).collect();
        println!("Shard distribution: {:?}", shard_counts);

        // Should be roughly even (within 2x)
        let max = *shard_counts.iter().max().unwrap();
        let min = *shard_counts.iter().min().unwrap();
        assert!(max < min * 3, "Shards too uneven: {:?}", shard_counts);

        // Verify retrieval
        for hash in &hashes {
            assert!(as_.get(hash).is_some());
        }

        // Test query
        let pattern = Pattern::by_type("ConceptNode");
        let results = as_.query(&pattern);
        assert_eq!(results.len(), 1000);
    }
}