storage/

tiered.rs

//! Tiered Storage for Buffer
//!
//! Automatic data tiering based on access patterns:
//! - Hot tier: In-memory for frequently accessed data
//! - Warm tier: Local disk cache for recent data
//! - Cold tier: Object storage for infrequently accessed data

use async_trait::async_trait;
use common::{DakeraError, NamespaceId, Result, Vector, VectorId};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};

use crate::traits::VectorStorage;

/// Tiered storage configuration
#[derive(Debug, Clone)]
pub struct TieredStorageConfig {
    /// Hot tier capacity (number of vectors)
    pub hot_tier_capacity: usize,
    /// Hot tier byte budget (approximate, from per-vector size estimates).
    /// When set, the hot tier is bounded by BOTH limits — whichever is
    /// exceeded first triggers LRU demotion to warm.
    pub hot_tier_max_bytes: Option<u64>,
    /// Time before demoting from hot to warm
    pub hot_to_warm_threshold: Duration,
    /// Time before demoting from warm to cold
    pub warm_to_cold_threshold: Duration,
    /// Enable automatic tiering
    pub auto_tier_enabled: bool,
    /// Tier check interval
    pub tier_check_interval: Duration,
}

impl Default for TieredStorageConfig {
    fn default() -> Self {
        Self {
            hot_tier_capacity: 100_000,
            hot_tier_max_bytes: None,
            hot_to_warm_threshold: Duration::from_secs(3600), // 1 hour
            warm_to_cold_threshold: Duration::from_secs(86400), // 24 hours
            auto_tier_enabled: true,
            tier_check_interval: Duration::from_secs(300), // 5 minutes
        }
    }
}

/// Approximate in-memory footprint of a JSON metadata value.
fn approx_json_bytes(value: &serde_json::Value) -> usize {
    match value {
        serde_json::Value::Null => 4,
        serde_json::Value::Bool(_) => 5,
        serde_json::Value::Number(_) => 8,
        serde_json::Value::String(s) => s.len() + 2,
        serde_json::Value::Array(a) => a.iter().map(approx_json_bytes).sum::<usize>() + 2 + a.len(),
        serde_json::Value::Object(o) => {
            o.iter()
                .map(|(k, v)| k.len() + 3 + approx_json_bytes(v))
                .sum::<usize>()
                + 2
        }
    }
}

/// Fixed per-vector overhead: struct fields, map entry, allocator slack.
const VECTOR_FIXED_OVERHEAD: usize = 96;

/// Approximate in-memory footprint of one vector in the hot tier.
fn approx_vector_bytes(v: &Vector) -> usize {
    v.values.len() * std::mem::size_of::<f32>()
        + v.id.len()
        + v.metadata.as_ref().map(approx_json_bytes).unwrap_or(0)
        + VECTOR_FIXED_OVERHEAD
}

/// Saturating subtraction on an atomic counter (avoids wrap-around on
/// accounting drift).
fn sub_saturating(counter: &AtomicU64, value: u64) {
    let _ = counter.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |cur| {
        Some(cur.saturating_sub(value))
    });
}

/// Storage tier for a piece of data
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum StorageTier {
    /// In-memory, fastest access
    Hot,
    /// Local disk, fast access
    Warm,
    /// Object storage, slowest access
    Cold,
}

impl StorageTier {
    pub fn as_str(&self) -> &'static str {
        match self {
            StorageTier::Hot => "hot",
            StorageTier::Warm => "warm",
            StorageTier::Cold => "cold",
        }
    }
}

/// Access tracking for tiering decisions
#[derive(Debug, Clone)]
struct AccessInfo {
    /// Last access timestamp
    last_access: Instant,
    /// Total access count
    access_count: u64,
    /// Current tier
    tier: StorageTier,
    /// Approximate byte footprint (counted toward the hot budget while
    /// tier == Hot)
    approx_bytes: usize,
}

impl Default for AccessInfo {
    fn default() -> Self {
        Self {
            last_access: Instant::now(),
            access_count: 0,
            tier: StorageTier::Hot,
            approx_bytes: 0,
        }
    }
}

/// Statistics for tiered storage
#[derive(Debug, Clone, Default)]
pub struct TieredStorageStats {
    /// Vectors in hot tier
    pub hot_count: u64,
    /// Approximate bytes held by the hot tier
    pub hot_bytes: u64,
    /// Vectors in warm tier
    pub warm_count: u64,
    /// Vectors in cold tier
    pub cold_count: u64,
    /// Hot tier hits
    pub hot_hits: u64,
    /// Warm tier hits
    pub warm_hits: u64,
    /// Cold tier hits
    pub cold_hits: u64,
    /// Promotions to hot tier
    pub promotions_to_hot: u64,
    /// Demotions to warm tier
    pub demotions_to_warm: u64,
    /// Demotions to cold tier
    pub demotions_to_cold: u64,
}

/// Tiered storage manager
pub struct TieredStorage<H, W, C> {
    /// Configuration
    config: TieredStorageConfig,
    /// Hot tier storage (in-memory)
    hot_storage: H,
    /// Warm tier storage (disk cache)
    warm_storage: W,
    /// Cold tier storage (object storage)
    cold_storage: C,
    /// Access tracking per vector
    access_info: RwLock<HashMap<(NamespaceId, VectorId), AccessInfo>>,
    /// Statistics
    stats: TieredStorageStatsInner,
}

struct TieredStorageStatsInner {
    hot_count: AtomicU64,
    hot_bytes: AtomicU64,
    warm_count: AtomicU64,
    cold_count: AtomicU64,
    hot_hits: AtomicU64,
    warm_hits: AtomicU64,
    cold_hits: AtomicU64,
    promotions_to_hot: AtomicU64,
    demotions_to_warm: AtomicU64,
    demotions_to_cold: AtomicU64,
}

impl Default for TieredStorageStatsInner {
    fn default() -> Self {
        Self {
            hot_count: AtomicU64::new(0),
            hot_bytes: AtomicU64::new(0),
            warm_count: AtomicU64::new(0),
            cold_count: AtomicU64::new(0),
            hot_hits: AtomicU64::new(0),
            warm_hits: AtomicU64::new(0),
            cold_hits: AtomicU64::new(0),
            promotions_to_hot: AtomicU64::new(0),
            demotions_to_warm: AtomicU64::new(0),
            demotions_to_cold: AtomicU64::new(0),
        }
    }
}

impl<H, W, C> TieredStorage<H, W, C>
where
    H: VectorStorage,
    W: VectorStorage,
    C: VectorStorage,
{
    /// Create a new tiered storage
    pub fn new(
        config: TieredStorageConfig,
        hot_storage: H,
        warm_storage: W,
        cold_storage: C,
    ) -> Self {
        Self {
            config,
            hot_storage,
            warm_storage,
            cold_storage,
            access_info: RwLock::new(HashMap::new()),
            stats: TieredStorageStatsInner::default(),
        }
    }

    /// Get the tiered storage configuration
    pub fn config(&self) -> &TieredStorageConfig {
        &self.config
    }

    /// Record access to a vector
    fn record_access(&self, namespace: &NamespaceId, id: &VectorId, tier: StorageTier) {
        let key = (namespace.clone(), id.clone());
        let mut access_map = self.access_info.write();

        let info = access_map.entry(key).or_default();
        info.last_access = Instant::now();
        info.access_count += 1;
        // Keep the hot byte budget consistent if the observed tier disagrees
        // with the tracked one (e.g. served from warm while still marked hot).
        if info.tier == StorageTier::Hot && tier != StorageTier::Hot {
            sub_saturating(&self.stats.hot_bytes, info.approx_bytes as u64);
        } else if info.tier != StorageTier::Hot && tier == StorageTier::Hot {
            self.stats
                .hot_bytes
                .fetch_add(info.approx_bytes as u64, Ordering::Relaxed);
        }
        info.tier = tier;

        // Update hit counters
        match tier {
            StorageTier::Hot => self.stats.hot_hits.fetch_add(1, Ordering::Relaxed),
            StorageTier::Warm => self.stats.warm_hits.fetch_add(1, Ordering::Relaxed),
            StorageTier::Cold => self.stats.cold_hits.fetch_add(1, Ordering::Relaxed),
        };
    }

    /// Get the current tier for a vector
    fn get_tier(&self, namespace: &NamespaceId, id: &VectorId) -> Option<StorageTier> {
        let access_map = self.access_info.read();
        access_map
            .get(&(namespace.clone(), id.clone()))
            .map(|info| info.tier)
    }

    /// Promote a vector to a higher tier
    pub async fn promote(&self, namespace: &NamespaceId, id: &VectorId) -> Result<bool> {
        let current_tier = self.get_tier(namespace, id);

        match current_tier {
            Some(StorageTier::Warm) => {
                // Promote warm -> hot
                let vectors = self
                    .warm_storage
                    .get(namespace, std::slice::from_ref(id))
                    .await?;
                if !vectors.is_empty() {
                    let est = vectors.first().map(approx_vector_bytes).unwrap_or(0);
                    self.hot_storage.upsert(namespace, vectors).await?;
                    self.warm_storage
                        .delete(namespace, std::slice::from_ref(id))
                        .await?;

                    self.set_approx_bytes(namespace, id, est);
                    self.update_tier(namespace, id, StorageTier::Hot);
                    self.stats.promotions_to_hot.fetch_add(1, Ordering::Relaxed);
                    self.stats.hot_count.fetch_add(1, Ordering::Relaxed);
                    self.stats.warm_count.fetch_sub(1, Ordering::Relaxed);

                    return Ok(true);
                }
            }
            Some(StorageTier::Cold) => {
                // Promote cold -> warm (or directly to hot if frequently accessed)
                let vectors = self
                    .cold_storage
                    .get(namespace, std::slice::from_ref(id))
                    .await?;
                if !vectors.is_empty() {
                    // Check if should go directly to hot based on access frequency
                    let should_be_hot = {
                        let access_map = self.access_info.read();
                        access_map
                            .get(&(namespace.clone(), id.clone()))
                            .map(|info| info.access_count > 10)
                            .unwrap_or(false)
                    };

                    if should_be_hot {
                        let est = vectors.first().map(approx_vector_bytes).unwrap_or(0);
                        self.hot_storage.upsert(namespace, vectors).await?;
                        self.set_approx_bytes(namespace, id, est);
                        self.update_tier(namespace, id, StorageTier::Hot);
                        self.stats.promotions_to_hot.fetch_add(1, Ordering::Relaxed);
                        self.stats.hot_count.fetch_add(1, Ordering::Relaxed);
                    } else {
                        self.warm_storage.upsert(namespace, vectors).await?;
                        self.update_tier(namespace, id, StorageTier::Warm);
                        self.stats.warm_count.fetch_add(1, Ordering::Relaxed);
                    }
                    // Cold tier is the durable source of truth — never delete on promotion.
                    // The tier map tracks which tier is "active" for reads; cold remains
                    // as the persistent backup in case warm/hot are lost on restart.

                    return Ok(true);
                }
            }
            _ => {}
        }

        Ok(false)
    }

    /// Demote a vector to a lower tier
    pub async fn demote(&self, namespace: &NamespaceId, id: &VectorId) -> Result<bool> {
        let current_tier = self.get_tier(namespace, id);

        match current_tier {
            Some(StorageTier::Hot) => {
                // Demote hot -> warm
                let vectors = self
                    .hot_storage
                    .get(namespace, std::slice::from_ref(id))
                    .await?;
                if !vectors.is_empty() {
                    self.warm_storage.upsert(namespace, vectors).await?;
                    self.hot_storage
                        .delete(namespace, std::slice::from_ref(id))
                        .await?;

                    self.update_tier(namespace, id, StorageTier::Warm);
                    self.stats.demotions_to_warm.fetch_add(1, Ordering::Relaxed);
                    self.stats.hot_count.fetch_sub(1, Ordering::Relaxed);
                    self.stats.warm_count.fetch_add(1, Ordering::Relaxed);

                    return Ok(true);
                }
            }
            Some(StorageTier::Warm) => {
                // Demote warm -> cold
                let vectors = self
                    .warm_storage
                    .get(namespace, std::slice::from_ref(id))
                    .await?;
                if !vectors.is_empty() {
                    self.cold_storage.upsert(namespace, vectors).await?;
                    self.warm_storage
                        .delete(namespace, std::slice::from_ref(id))
                        .await?;

                    self.update_tier(namespace, id, StorageTier::Cold);
                    self.stats.demotions_to_cold.fetch_add(1, Ordering::Relaxed);
                    self.stats.warm_count.fetch_sub(1, Ordering::Relaxed);
                    self.stats.cold_count.fetch_add(1, Ordering::Relaxed);

                    return Ok(true);
                }
            }
            _ => {}
        }

        Ok(false)
    }

    /// Update tier tracking, keeping the hot-tier byte budget in sync with
    /// Hot ↔ non-Hot transitions.
    fn update_tier(&self, namespace: &NamespaceId, id: &VectorId, tier: StorageTier) {
        let mut access_map = self.access_info.write();
        let key = (namespace.clone(), id.clone());
        let info = access_map.entry(key).or_default();
        let was_hot = info.tier == StorageTier::Hot;
        let now_hot = tier == StorageTier::Hot;
        if was_hot && !now_hot {
            sub_saturating(&self.stats.hot_bytes, info.approx_bytes as u64);
        } else if !was_hot && now_hot {
            self.stats
                .hot_bytes
                .fetch_add(info.approx_bytes as u64, Ordering::Relaxed);
        }
        info.tier = tier;
    }

    /// Refresh the byte estimate for a tracked vector (e.g. on promotion,
    /// when the vector payload is available again).
    fn set_approx_bytes(&self, namespace: &NamespaceId, id: &VectorId, bytes: usize) {
        let mut access_map = self.access_info.write();
        if let Some(info) = access_map.get_mut(&(namespace.clone(), id.clone())) {
            if info.tier == StorageTier::Hot {
                sub_saturating(&self.stats.hot_bytes, info.approx_bytes as u64);
                self.stats
                    .hot_bytes
                    .fetch_add(bytes as u64, Ordering::Relaxed);
            }
            info.approx_bytes = bytes;
        }
    }

    fn is_permanent_storage_error(error: &DakeraError) -> bool {
        match error {
            DakeraError::NamespaceNotFound(_) => true,
            DakeraError::Storage(msg) => {
                msg.contains("PermissionDenied") || msg.contains("permission denied")
            }
            _ => false,
        }
    }

    /// Run automatic tiering based on access patterns
    pub async fn run_auto_tiering(&self) -> Result<TieringResult> {
        if !self.config.auto_tier_enabled {
            // Time-based tiering is opt-out, but the hot capacity bound is a
            // hard limit — enforce it regardless.
            let demoted_to_warm = self.enforce_hot_capacity().await;
            return Ok(TieringResult {
                demoted_to_warm,
                ..Default::default()
            });
        }

        let now = Instant::now();
        let mut to_demote_to_warm = Vec::new();
        let mut to_demote_to_cold = Vec::new();

        // Collect vectors to demote
        {
            let access_map = self.access_info.read();
            for ((namespace, id), info) in access_map.iter() {
                let elapsed = now.duration_since(info.last_access);

                match info.tier {
                    StorageTier::Hot if elapsed > self.config.hot_to_warm_threshold => {
                        to_demote_to_warm.push((namespace.clone(), id.clone()));
                    }
                    StorageTier::Warm if elapsed > self.config.warm_to_cold_threshold => {
                        to_demote_to_cold.push((namespace.clone(), id.clone()));
                    }
                    _ => {}
                }
            }
        }

        // Execute demotions — errors on individual vectors are handled gracefully
        // to prevent a single deleted namespace from blocking all tiering.
        let mut demoted_to_warm = 0;
        let mut demoted_to_cold = 0;
        let mut stale_entries = Vec::new();

        for (namespace, id) in to_demote_to_warm {
            match self.demote(&namespace, &id).await {
                Ok(true) => demoted_to_warm += 1,
                Ok(false) => {}
                Err(e) => {
                    if Self::is_permanent_storage_error(&e) {
                        stale_entries.push((namespace, id));
                    } else {
                        tracing::debug!(error = %e, "Transient demotion error, will retry next cycle");
                    }
                }
            }
        }

        for (namespace, id) in to_demote_to_cold {
            match self.demote(&namespace, &id).await {
                Ok(true) => demoted_to_cold += 1,
                Ok(false) => {}
                Err(e) => {
                    if Self::is_permanent_storage_error(&e) {
                        stale_entries.push((namespace, id));
                    } else {
                        tracing::debug!(error = %e, "Transient demotion error, will retry next cycle");
                    }
                }
            }
        }

        if !stale_entries.is_empty() {
            let removed = stale_entries.len();
            let mut access_map = self.access_info.write();
            for (ns, id) in &stale_entries {
                if let Some(info) = access_map.remove(&(ns.clone(), id.clone())) {
                    match info.tier {
                        StorageTier::Hot => {
                            sub_saturating(&self.stats.hot_bytes, info.approx_bytes as u64);
                            sub_saturating(&self.stats.hot_count, 1);
                        }
                        StorageTier::Warm => sub_saturating(&self.stats.warm_count, 1),
                        StorageTier::Cold => sub_saturating(&self.stats.cold_count, 1),
                    }
                }
            }
            tracing::warn!(
                removed = removed,
                "Purged stale entries from tiering queue (deleted namespace or permanent error)"
            );
        }

        // Capacity enforcement after time-based demotion — the hot tier must
        // respect both the vector-count cap and the byte budget.
        let demoted_by_capacity = self.enforce_hot_capacity().await;

        Ok(TieringResult {
            demoted_to_warm: demoted_to_warm + demoted_by_capacity,
            demoted_to_cold,
            promoted_to_hot: 0,
            promoted_to_warm: 0,
        })
    }

    /// Demote least-recently-accessed hot vectors to warm until the hot tier
    /// is within both `hot_tier_capacity` (count) and `hot_tier_max_bytes`
    /// (approximate bytes). Returns the number of vectors demoted.
    ///
    /// This is the enforcement half of the L1 budget: without it,
    /// `hot_tier_capacity` is advisory only and the hot tier grows unbounded
    /// until the memory-pressure guard trips (DAK-7227 / DAK-7337).
    pub async fn enforce_hot_capacity(&self) -> u64 {
        let cap_count = self.config.hot_tier_capacity;
        let cap_bytes = self.config.hot_tier_max_bytes;

        let mut hot_entries: Vec<((NamespaceId, VectorId), Instant, u64)> = {
            let access_map = self.access_info.read();
            access_map
                .iter()
                .filter(|(_, info)| info.tier == StorageTier::Hot)
                .map(|(key, info)| (key.clone(), info.last_access, info.approx_bytes as u64))
                .collect()
        };

        let mut cur_count = hot_entries.len();
        let mut cur_bytes: u64 = hot_entries.iter().map(|(_, _, b)| *b).sum();
        let over = |count: usize, bytes: u64| {
            count > cap_count || cap_bytes.map(|cap| bytes > cap).unwrap_or(false)
        };

        if !over(cur_count, cur_bytes) {
            return 0;
        }

        // Oldest access first — LRU demotion order.
        hot_entries.sort_by_key(|(_, last_access, _)| *last_access);

        let mut demoted = 0u64;
        for ((ns, id), _, bytes) in hot_entries {
            if !over(cur_count, cur_bytes) {
                break;
            }
            match self.demote(&ns, &id).await {
                Ok(true) => {
                    demoted += 1;
                    cur_count -= 1;
                    cur_bytes = cur_bytes.saturating_sub(bytes);
                }
                Ok(false) => {
                    // Vector no longer present in hot storage — tracking is
                    // stale for this entry; skip it in the local accounting so
                    // the loop still terminates.
                    cur_count = cur_count.saturating_sub(1);
                    cur_bytes = cur_bytes.saturating_sub(bytes);
                }
                Err(e) => {
                    tracing::debug!(
                        error = %e,
                        namespace = %ns,
                        "Capacity demotion error, will retry next cycle"
                    );
                    cur_count = cur_count.saturating_sub(1);
                    cur_bytes = cur_bytes.saturating_sub(bytes);
                }
            }
        }

        if demoted > 0 {
            tracing::info!(
                demoted,
                cap_count,
                cap_bytes = cap_bytes.unwrap_or(0),
                "Hot tier capacity enforcement demoted LRU vectors to warm"
            );
        }
        demoted
    }

    /// Get storage statistics
    pub fn stats(&self) -> TieredStorageStats {
        TieredStorageStats {
            hot_count: self.stats.hot_count.load(Ordering::Relaxed),
            hot_bytes: self.stats.hot_bytes.load(Ordering::Relaxed),
            warm_count: self.stats.warm_count.load(Ordering::Relaxed),
            cold_count: self.stats.cold_count.load(Ordering::Relaxed),
            hot_hits: self.stats.hot_hits.load(Ordering::Relaxed),
            warm_hits: self.stats.warm_hits.load(Ordering::Relaxed),
            cold_hits: self.stats.cold_hits.load(Ordering::Relaxed),
            promotions_to_hot: self.stats.promotions_to_hot.load(Ordering::Relaxed),
            demotions_to_warm: self.stats.demotions_to_warm.load(Ordering::Relaxed),
            demotions_to_cold: self.stats.demotions_to_cold.load(Ordering::Relaxed),
        }
    }

    /// Get tier distribution by namespace
    pub fn tier_distribution(&self, namespace: &NamespaceId) -> TierDistribution {
        let access_map = self.access_info.read();
        let mut hot = 0u64;
        let mut warm = 0u64;
        let mut cold = 0u64;

        for ((ns, _), info) in access_map.iter() {
            if ns == namespace {
                match info.tier {
                    StorageTier::Hot => hot += 1,
                    StorageTier::Warm => warm += 1,
                    StorageTier::Cold => cold += 1,
                }
            }
        }

        TierDistribution { hot, warm, cold }
    }
}

/// Result of automatic tiering
#[derive(Debug, Clone, Default)]
pub struct TieringResult {
    /// Vectors demoted to warm tier
    pub demoted_to_warm: u64,
    /// Vectors demoted to cold tier
    pub demoted_to_cold: u64,
    /// Vectors promoted to hot tier
    pub promoted_to_hot: u64,
    /// Vectors promoted to warm tier
    pub promoted_to_warm: u64,
}

/// Tier distribution for a namespace
#[derive(Debug, Clone)]
pub struct TierDistribution {
    pub hot: u64,
    pub warm: u64,
    pub cold: u64,
}

#[async_trait]
impl<H, W, C> VectorStorage for TieredStorage<H, W, C>
where
    H: VectorStorage,
    W: VectorStorage,
    C: VectorStorage + Clone + Send + Sync + 'static,
{
    async fn upsert(&self, namespace: &NamespaceId, vectors: Vec<Vector>) -> Result<usize> {
        // Clone vectors for the background cold-tier flush before consuming them.
        let cold_vectors = vectors.clone();

        // Write hot tier first — lowest latency path, immediately available for reads.
        let count = self.hot_storage.upsert(namespace, vectors).await?;

        // Track access info + hot byte budget in a single pass (before the cold
        // flush consumes cold_vectors). Overwrites of vectors already in the hot
        // tier replace their byte estimate instead of double-counting (both
        // bytes and hot_count).
        {
            let mut access_map = self.access_info.write();
            let mut newly_hot = 0u64;
            for v in &cold_vectors {
                let est = approx_vector_bytes(v);
                let key = (namespace.clone(), v.id.clone());
                match access_map.get_mut(&key) {
                    Some(info) => {
                        if info.tier == StorageTier::Hot {
                            sub_saturating(&self.stats.hot_bytes, info.approx_bytes as u64);
                        } else {
                            newly_hot += 1;
                        }
                        info.tier = StorageTier::Hot;
                        info.last_access = Instant::now();
                        info.access_count += 1;
                        info.approx_bytes = est;
                    }
                    None => {
                        newly_hot += 1;
                        access_map.insert(
                            key,
                            AccessInfo {
                                last_access: Instant::now(),
                                access_count: 1,
                                tier: StorageTier::Hot,
                                approx_bytes: est,
                            },
                        );
                    }
                }
                self.stats
                    .hot_bytes
                    .fetch_add(est as u64, Ordering::Relaxed);
                self.stats.hot_hits.fetch_add(1, Ordering::Relaxed);
            }
            self.stats.hot_count.fetch_add(newly_hot, Ordering::Relaxed);
        }

        // Flush to cold tier (S3) in the background — failure is logged but not fatal.
        // The hot tier (RocksDB/in-memory) is the primary durable read path.
        let cold = self.cold_storage.clone();
        let cold_ns = namespace.clone();
        tokio::spawn(async move {
            if let Err(e) = cold.ensure_namespace(&cold_ns).await {
                tracing::error!(
                    error = %e,
                    namespace = %cold_ns,
                    "Cold tier namespace ensure failed (S3 flush aborted)"
                );
                return;
            }
            let retry_vectors = cold_vectors.clone();
            match cold.upsert(&cold_ns, cold_vectors).await {
                Ok(_) => {}
                Err(DakeraError::DimensionMismatch { expected, actual }) => {
                    tracing::warn!(
                        namespace = %cold_ns,
                        cold_dim = expected,
                        hot_dim = actual,
                        "Cold tier dimension mismatch — resetting stale cold namespace"
                    );
                    if let Err(e) = cold.delete_namespace(&cold_ns).await {
                        tracing::error!(error = %e, namespace = %cold_ns,
                            "Failed to delete stale cold namespace");
                        return;
                    }
                    if let Err(e) = cold.ensure_namespace(&cold_ns).await {
                        tracing::error!(error = %e, namespace = %cold_ns,
                            "Failed to recreate cold namespace after dimension fix");
                        return;
                    }
                    if let Err(e) = cold.upsert(&cold_ns, retry_vectors).await {
                        tracing::error!(
                            error = %e,
                            namespace = %cold_ns,
                            "Cold tier S3 flush failed after dimension fix — data is durable in hot tier"
                        );
                    }
                }
                Err(e) => {
                    tracing::error!(
                        error = %e,
                        namespace = %cold_ns,
                        "Cold tier S3 flush failed — data is durable in hot tier"
                    );
                }
            }
        });

        Ok(count)
    }

    async fn get(&self, namespace: &NamespaceId, ids: &[VectorId]) -> Result<Vec<Vector>> {
        let mut results = Vec::with_capacity(ids.len());
        let mut remaining_ids: Vec<VectorId> = ids.to_vec();

        // Try hot tier first (NamespaceNotFound is normal after restart when hot cache is empty)
        let hot_results = match self.hot_storage.get(namespace, &remaining_ids).await {
            Ok(v) => v,
            Err(DakeraError::NamespaceNotFound(_)) => vec![],
            Err(e) => return Err(e),
        };
        for v in &hot_results {
            self.record_access(namespace, &v.id, StorageTier::Hot);
        }

        // Remove found IDs
        let found_ids: std::collections::HashSet<_> = hot_results.iter().map(|v| &v.id).collect();
        remaining_ids.retain(|id| !found_ids.contains(id));
        results.extend(hot_results);

        if remaining_ids.is_empty() {
            return Ok(results);
        }

        // Try warm tier (NamespaceNotFound is normal after restart when warm cache is empty)
        let warm_results = match self.warm_storage.get(namespace, &remaining_ids).await {
            Ok(v) => v,
            Err(common::DakeraError::NamespaceNotFound(_)) => vec![],
            Err(e) => return Err(e),
        };
        for v in &warm_results {
            self.record_access(namespace, &v.id, StorageTier::Warm);
        }

        let found_ids: std::collections::HashSet<_> = warm_results.iter().map(|v| &v.id).collect();
        remaining_ids.retain(|id| !found_ids.contains(id));
        results.extend(warm_results);

        if remaining_ids.is_empty() {
            return Ok(results);
        }

        // Try cold tier
        let cold_results = match self.cold_storage.get(namespace, &remaining_ids).await {
            Ok(v) => v,
            Err(DakeraError::NamespaceNotFound(_)) => vec![],
            Err(e) => return Err(e),
        };
        for v in &cold_results {
            self.record_access(namespace, &v.id, StorageTier::Cold);
        }
        results.extend(cold_results);

        Ok(results)
    }

    async fn get_all(&self, namespace: &NamespaceId) -> Result<Vec<Vector>> {
        let mut seen = std::collections::HashSet::new();
        let mut results = Vec::new();

        // Helper: treat NamespaceNotFound as empty — normal for hot/warm after restart.
        let tier_get_all = |res: common::Result<Vec<Vector>>| -> common::Result<Vec<Vector>> {
            match res {
                Ok(v) => Ok(v),
                Err(common::DakeraError::NamespaceNotFound(_)) => Ok(vec![]),
                Err(e) => Err(e),
            }
        };

        // Gather from all tiers, preferring hot over warm over cold.
        // Deduplicate by vector ID since write-through means a vector
        // can exist in both hot and cold simultaneously.
        for v in tier_get_all(self.hot_storage.get_all(namespace).await)? {
            if seen.insert(v.id.clone()) {
                results.push(v);
            }
        }
        for v in tier_get_all(self.warm_storage.get_all(namespace).await)? {
            if seen.insert(v.id.clone()) {
                results.push(v);
            }
        }
        for v in tier_get_all(self.cold_storage.get_all(namespace).await)? {
            if seen.insert(v.id.clone()) {
                results.push(v);
            }
        }

        Ok(results)
    }

    async fn delete(&self, namespace: &NamespaceId, ids: &[VectorId]) -> Result<usize> {
        let mut deleted = 0;

        // Delete from all tiers. Tolerate NamespaceNotFound from individual tiers
        // since data may only reside in a subset of tiers (e.g. cold but not hot).
        match self.hot_storage.delete(namespace, ids).await {
            Ok(n) => deleted += n,
            Err(DakeraError::NamespaceNotFound(_)) => {}
            Err(e) => return Err(e),
        }
        match self.warm_storage.delete(namespace, ids).await {
            Ok(n) => deleted += n,
            Err(DakeraError::NamespaceNotFound(_)) => {}
            Err(e) => return Err(e),
        }
        match self.cold_storage.delete(namespace, ids).await {
            Ok(n) => deleted += n,
            Err(DakeraError::NamespaceNotFound(_)) => {}
            Err(e) => return Err(e),
        }

        // Remove from tracking, releasing hot budget and tier counts.
        {
            let mut access_map = self.access_info.write();
            for id in ids {
                if let Some(info) = access_map.remove(&(namespace.clone(), id.clone())) {
                    match info.tier {
                        StorageTier::Hot => {
                            sub_saturating(&self.stats.hot_bytes, info.approx_bytes as u64);
                            sub_saturating(&self.stats.hot_count, 1);
                        }
                        StorageTier::Warm => sub_saturating(&self.stats.warm_count, 1),
                        StorageTier::Cold => sub_saturating(&self.stats.cold_count, 1),
                    }
                }
            }
        }

        Ok(deleted)
    }

    async fn namespace_exists(&self, namespace: &NamespaceId) -> Result<bool> {
        // Check any tier
        Ok(self.hot_storage.namespace_exists(namespace).await?
            || self.warm_storage.namespace_exists(namespace).await?
            || self.cold_storage.namespace_exists(namespace).await?)
    }

    async fn ensure_namespace(&self, namespace: &NamespaceId) -> Result<()> {
        // Ensure in all tiers
        self.hot_storage.ensure_namespace(namespace).await?;
        self.warm_storage.ensure_namespace(namespace).await?;
        self.cold_storage.ensure_namespace(namespace).await?;
        Ok(())
    }

    async fn count(&self, namespace: &NamespaceId) -> Result<usize> {
        // With write-through, cold tier is the source of truth for total count.
        // Hot/warm are caches that hold subsets of the same data.
        // Use cold count as the baseline, then add any vectors that are ONLY
        // in hot or warm (shouldn't happen with write-through, but safe).
        let cold = self.cold_storage.count(namespace).await?;
        if cold > 0 {
            return Ok(cold);
        }
        // Fallback: if cold is empty, count from hot + warm (non-tiered data)
        let hot = self.hot_storage.count(namespace).await?;
        let warm = self.warm_storage.count(namespace).await?;
        Ok(hot + warm)
    }

    async fn dimension(&self, namespace: &NamespaceId) -> Result<Option<usize>> {
        // Check hot first, then warm, then cold
        if let Some(dim) = self.hot_storage.dimension(namespace).await? {
            return Ok(Some(dim));
        }
        if let Some(dim) = self.warm_storage.dimension(namespace).await? {
            return Ok(Some(dim));
        }
        self.cold_storage.dimension(namespace).await
    }

    async fn list_namespaces(&self) -> Result<Vec<NamespaceId>> {
        let mut namespaces = std::collections::HashSet::new();

        namespaces.extend(self.hot_storage.list_namespaces().await?);
        namespaces.extend(self.warm_storage.list_namespaces().await?);
        namespaces.extend(self.cold_storage.list_namespaces().await?);

        Ok(namespaces.into_iter().collect())
    }

    async fn delete_namespace(&self, namespace: &NamespaceId) -> Result<bool> {
        // Delete from all tiers
        let hot_deleted = self.hot_storage.delete_namespace(namespace).await?;
        let warm_deleted = self.warm_storage.delete_namespace(namespace).await?;
        let cold_deleted = self.cold_storage.delete_namespace(namespace).await?;

        // Remove from access tracking
        {
            let mut access_map = self.access_info.write();
            access_map.retain(|(ns, _), _| ns != namespace);
        }

        Ok(hot_deleted || warm_deleted || cold_deleted)
    }

    async fn cleanup_expired(&self, namespace: &NamespaceId) -> Result<usize> {
        // Cleanup from all tiers
        let mut total = 0;
        total += self.hot_storage.cleanup_expired(namespace).await?;
        total += self.warm_storage.cleanup_expired(namespace).await?;
        total += self.cold_storage.cleanup_expired(namespace).await?;
        Ok(total)
    }

    async fn cleanup_all_expired(&self) -> Result<usize> {
        // Cleanup from all tiers
        let mut total = 0;
        total += self.hot_storage.cleanup_all_expired().await?;
        total += self.warm_storage.cleanup_all_expired().await?;
        total += self.cold_storage.cleanup_all_expired().await?;
        Ok(total)
    }
}

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

    fn create_test_vector(id: &str, dim: usize) -> Vector {
        Vector {
            id: id.to_string(),
            values: vec![1.0; dim],
            metadata: None,
            ttl_seconds: None,
            expires_at: None,
        }
    }

    #[tokio::test]
    async fn test_tiered_storage_basic() {
        let config = TieredStorageConfig::default();
        let storage = TieredStorage::new(
            config,
            InMemoryStorage::new(),
            InMemoryStorage::new(),
            InMemoryStorage::new(),
        );

        let namespace = "test".to_string();
        storage.ensure_namespace(&namespace).await.unwrap();

        // Upsert goes to hot tier
        let vectors = vec![create_test_vector("v1", 4)];
        let count = storage.upsert(&namespace, vectors).await.unwrap();
        assert_eq!(count, 1);

        // Get should find in hot tier
        let results = storage.get(&namespace, &["v1".to_string()]).await.unwrap();
        assert_eq!(results.len(), 1);

        let stats = storage.stats();
        assert_eq!(stats.hot_hits, 2); // One from upsert record_access, one from get
    }

    #[tokio::test]
    async fn test_tiered_storage_promotion_demotion() {
        let config = TieredStorageConfig::default();
        let storage = TieredStorage::new(
            config,
            InMemoryStorage::new(),
            InMemoryStorage::new(),
            InMemoryStorage::new(),
        );

        let namespace = "test".to_string();
        storage.ensure_namespace(&namespace).await.unwrap();

        // Add to hot tier
        storage
            .upsert(&namespace, vec![create_test_vector("v1", 4)])
            .await
            .unwrap();

        // Verify in hot
        assert_eq!(
            storage.get_tier(&namespace, &"v1".to_string()),
            Some(StorageTier::Hot)
        );

        // Demote to warm
        let demoted = storage.demote(&namespace, &"v1".to_string()).await.unwrap();
        assert!(demoted);
        assert_eq!(
            storage.get_tier(&namespace, &"v1".to_string()),
            Some(StorageTier::Warm)
        );

        // Demote to cold
        let demoted = storage.demote(&namespace, &"v1".to_string()).await.unwrap();
        assert!(demoted);
        assert_eq!(
            storage.get_tier(&namespace, &"v1".to_string()),
            Some(StorageTier::Cold)
        );

        // Still accessible
        let results = storage.get(&namespace, &["v1".to_string()]).await.unwrap();
        assert_eq!(results.len(), 1);

        let stats = storage.stats();
        assert_eq!(stats.demotions_to_warm, 1);
        assert_eq!(stats.demotions_to_cold, 1);
    }

    #[tokio::test]
    async fn test_tiered_storage_multi_tier_get() {
        let config = TieredStorageConfig::default();
        let storage = TieredStorage::new(
            config,
            InMemoryStorage::new(),
            InMemoryStorage::new(),
            InMemoryStorage::new(),
        );

        let namespace = "test".to_string();
        storage.ensure_namespace(&namespace).await.unwrap();

        // Add vectors and demote some
        for i in 0..3 {
            storage
                .upsert(&namespace, vec![create_test_vector(&format!("v{}", i), 4)])
                .await
                .unwrap();
        }

        // v0 stays hot, v1 goes warm, v2 goes cold
        storage.demote(&namespace, &"v1".to_string()).await.unwrap();
        storage.demote(&namespace, &"v2".to_string()).await.unwrap();
        storage.demote(&namespace, &"v2".to_string()).await.unwrap();

        // Get all at once
        let ids: Vec<_> = (0..3).map(|i| format!("v{}", i)).collect();
        let results = storage.get(&namespace, &ids).await.unwrap();
        assert_eq!(results.len(), 3);
    }

    #[tokio::test]
    async fn test_tier_distribution() {
        let config = TieredStorageConfig::default();
        let storage = TieredStorage::new(
            config,
            InMemoryStorage::new(),
            InMemoryStorage::new(),
            InMemoryStorage::new(),
        );

        let namespace = "test".to_string();
        storage.ensure_namespace(&namespace).await.unwrap();

        // Add 5 vectors
        for i in 0..5 {
            storage
                .upsert(&namespace, vec![create_test_vector(&format!("v{}", i), 4)])
                .await
                .unwrap();
        }

        // Demote 2 to warm, 1 to cold
        storage.demote(&namespace, &"v3".to_string()).await.unwrap();
        storage.demote(&namespace, &"v4".to_string()).await.unwrap();
        storage.demote(&namespace, &"v4".to_string()).await.unwrap();

        let dist = storage.tier_distribution(&namespace);
        assert_eq!(dist.hot, 3);
        assert_eq!(dist.warm, 1);
        assert_eq!(dist.cold, 1);
    }

    #[tokio::test]
    async fn test_auto_tiering_purges_stale_namespace_entries() {
        let config = TieredStorageConfig {
            auto_tier_enabled: true,
            hot_to_warm_threshold: Duration::from_millis(0),
            warm_to_cold_threshold: Duration::from_millis(0),
            ..Default::default()
        };
        let storage = TieredStorage::new(
            config,
            InMemoryStorage::new(),
            InMemoryStorage::new(),
            InMemoryStorage::new(),
        );

        let namespace = "deleted_bench_ns".to_string();
        storage.ensure_namespace(&namespace).await.unwrap();
        storage
            .upsert(&namespace, vec![create_test_vector("v1", 4)])
            .await
            .unwrap();

        // Delete the namespace (simulates bench cleanup)
        storage.delete_namespace(&namespace).await.unwrap();

        // Re-insert stale entries into access_info (simulates race / leftover)
        {
            let mut access_map = storage.access_info.write();
            access_map.insert(
                (namespace.clone(), "v1".to_string()),
                AccessInfo {
                    last_access: Instant::now() - Duration::from_secs(7200),
                    access_count: 1,
                    tier: StorageTier::Hot,
                    approx_bytes: 0,
                },
            );
        }

        // Auto-tiering should succeed (not bail) and purge the stale entry
        let result = storage.run_auto_tiering().await.unwrap();
        assert_eq!(result.demoted_to_warm, 0);

        let access_map = storage.access_info.read();
        assert!(
            access_map.is_empty(),
            "Stale entries should have been purged from tiering queue"
        );
    }

    #[test]
    fn test_is_permanent_storage_error() {
        assert!(TieredStorage::<
            InMemoryStorage,
            InMemoryStorage,
            InMemoryStorage,
        >::is_permanent_storage_error(
            &DakeraError::NamespaceNotFound("test".into())
        ));
        assert!(TieredStorage::<
            InMemoryStorage,
            InMemoryStorage,
            InMemoryStorage,
        >::is_permanent_storage_error(&DakeraError::Storage(
            "PermissionDenied (permanent) at write".into()
        )));
        assert!(!TieredStorage::<
            InMemoryStorage,
            InMemoryStorage,
            InMemoryStorage,
        >::is_permanent_storage_error(
            &DakeraError::Storage("connection timeout".into())
        ));
        assert!(!TieredStorage::<
            InMemoryStorage,
            InMemoryStorage,
            InMemoryStorage,
        >::is_permanent_storage_error(
            &DakeraError::DimensionMismatch {
                expected: 384,
                actual: 1024,
            }
        ));
    }

    #[test]
    fn test_approx_vector_bytes_scales_with_payload() {
        let small = approx_vector_bytes(&create_test_vector("v", 4));
        let big = approx_vector_bytes(&create_test_vector("v", 1024));
        // Embedding values dominate: (1024 - 4) * 4 bytes difference.
        assert_eq!(big - small, (1024 - 4) * std::mem::size_of::<f32>());

        let mut with_meta = create_test_vector("v", 4);
        with_meta.metadata = Some(serde_json::json!({ "content": "x".repeat(100) }));
        assert!(approx_vector_bytes(&with_meta) > small + 100);
    }

    #[tokio::test]
    async fn test_hot_bytes_accounting() {
        let storage = TieredStorage::new(
            TieredStorageConfig::default(),
            InMemoryStorage::new(),
            InMemoryStorage::new(),
            InMemoryStorage::new(),
        );
        let namespace = "test".to_string();
        storage.ensure_namespace(&namespace).await.unwrap();

        storage
            .upsert(
                &namespace,
                vec![create_test_vector("v1", 4), create_test_vector("v2", 4)],
            )
            .await
            .unwrap();
        let after_insert = storage.stats().hot_bytes;
        assert!(after_insert > 0);

        // Overwriting an existing hot vector must not double-count bytes or count.
        storage
            .upsert(&namespace, vec![create_test_vector("v1", 4)])
            .await
            .unwrap();
        assert_eq!(storage.stats().hot_bytes, after_insert);
        assert_eq!(storage.stats().hot_count, 2);

        // Delete releases the budget.
        storage
            .delete(&namespace, &["v1".to_string()])
            .await
            .unwrap();
        assert!(storage.stats().hot_bytes < after_insert);
        assert_eq!(storage.stats().hot_count, 1);

        storage
            .delete(&namespace, &["v2".to_string()])
            .await
            .unwrap();
        assert_eq!(storage.stats().hot_bytes, 0);
        assert_eq!(storage.stats().hot_count, 0);
    }

    #[tokio::test]
    async fn test_capacity_enforcement_by_count() {
        let config = TieredStorageConfig {
            hot_tier_capacity: 2,
            // Capacity is a hard bound — enforced even with time-based tiering off.
            auto_tier_enabled: false,
            ..Default::default()
        };
        let storage = TieredStorage::new(
            config,
            InMemoryStorage::new(),
            InMemoryStorage::new(),
            InMemoryStorage::new(),
        );
        let namespace = "test".to_string();
        storage.ensure_namespace(&namespace).await.unwrap();

        for (i, id) in ["v1", "v2", "v3", "v4"].iter().enumerate() {
            storage
                .upsert(&namespace, vec![create_test_vector(id, 4)])
                .await
                .unwrap();
            // Stagger last_access so LRU order is deterministic.
            if i < 3 {
                tokio::time::sleep(Duration::from_millis(5)).await;
            }
        }

        let result = storage.run_auto_tiering().await.unwrap();
        assert_eq!(result.demoted_to_warm, 2);
        let stats = storage.stats();
        assert_eq!(stats.hot_count, 2);
        assert_eq!(stats.warm_count, 2);

        // Oldest two demoted, newest two stayed hot; all four still readable.
        assert_eq!(
            storage.get_tier(&namespace, &"v1".to_string()),
            Some(StorageTier::Warm)
        );
        assert_eq!(
            storage.get_tier(&namespace, &"v2".to_string()),
            Some(StorageTier::Warm)
        );
        assert_eq!(
            storage.get_tier(&namespace, &"v4".to_string()),
            Some(StorageTier::Hot)
        );
        let all = storage
            .get(
                &namespace,
                &[
                    "v1".to_string(),
                    "v2".to_string(),
                    "v3".to_string(),
                    "v4".to_string(),
                ],
            )
            .await
            .unwrap();
        assert_eq!(all.len(), 4);
    }

    #[tokio::test]
    async fn test_capacity_enforcement_by_bytes() {
        let per_vec = approx_vector_bytes(&create_test_vector("v1", 128)) as u64;
        let budget = per_vec * 2 + 8; // room for ~2 vectors
        let config = TieredStorageConfig {
            hot_tier_capacity: usize::MAX,
            hot_tier_max_bytes: Some(budget),
            auto_tier_enabled: false,
            ..Default::default()
        };
        let storage = TieredStorage::new(
            config,
            InMemoryStorage::new(),
            InMemoryStorage::new(),
            InMemoryStorage::new(),
        );
        let namespace = "test".to_string();
        storage.ensure_namespace(&namespace).await.unwrap();

        for (i, id) in ["v1", "v2", "v3", "v4"].iter().enumerate() {
            storage
                .upsert(&namespace, vec![create_test_vector(id, 128)])
                .await
                .unwrap();
            if i < 3 {
                tokio::time::sleep(Duration::from_millis(5)).await;
            }
        }
        assert!(storage.stats().hot_bytes > budget);

        let result = storage.run_auto_tiering().await.unwrap();
        assert_eq!(result.demoted_to_warm, 2);
        let stats = storage.stats();
        assert!(stats.hot_bytes <= budget);
        assert_eq!(stats.hot_count, 2);

        // Nothing lost — demoted vectors are served from warm.
        let all = storage
            .get(
                &namespace,
                &[
                    "v1".to_string(),
                    "v2".to_string(),
                    "v3".to_string(),
                    "v4".to_string(),
                ],
            )
            .await
            .unwrap();
        assert_eq!(all.len(), 4);
    }
}