photoncloud-monorepo/flaredb/crates/flaredb-client/src/region_cache.rs

use flaredb_proto::pdpb::{Region, Store};
use std::sync::Arc;
use tokio::sync::RwLock;

#[derive(Clone)]
pub struct RegionCache {
    // Map: StartKey -> RegionInfo
    // For MVP, strict match or prefix?
    // Regions are: [start, end).
    // We need to find a region where start <= key < end.
    // BTreeMap is better for range search. But Rust BTreeMap is sync.
    // Using Vec + Scan for MVP is fine (few regions).
    regions: Arc<RwLock<Vec<(Region, Store)>>>,
}

impl RegionCache {
    pub fn new() -> Self {
        Self {
            regions: Arc::new(RwLock::new(Vec::new())),
        }
    }

    pub async fn update(&self, region: Region, leader: Store) {
        let mut cache = self.regions.write().await;
        // Remove overlaps? Or just append and filter?
        // Ideally replace.
        cache.retain(|(r, _)| r.id != region.id);
        cache.push((region, leader));
    }

    pub async fn get_store_addr(&self, key: &[u8]) -> Option<String> {
        let cache = self.regions.read().await;
        for (region, store) in cache.iter() {
            let start_ok = region.start_key.is_empty() || key >= &region.start_key[..];
            let end_ok = region.end_key.is_empty() || key < &region.end_key[..];
            if start_ok && end_ok {
                return Some(store.addr.clone());
            }
        }
        None
    }
}

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