tidaldb/tidal-server/src/state.rs

use std::collections::HashMap;
use std::sync::Arc;

use tidaldb::query::retrieve::Retrieve;
use tidaldb::query::search::Search;
use tidaldb::schema::{EntityId, Schema};
use tidaldb::replication::RegionId;
use tidaldb::testing::cluster::SimulatedCluster;
use tidaldb::TidalDb;

use crate::config::ClusterLayout;
use crate::error::{Result, ServerError};

#[derive(Clone)]
pub struct ServerState {
    mode: Mode,
}

#[derive(Clone)]
enum Mode {
    Standalone(Arc<TidalDb>),
    Cluster(ClusterState),
}

#[derive(Clone)]
struct ClusterState {
    cluster: Arc<SimulatedCluster>,
    regions: RegionDirectory,
}

#[derive(Clone)]
struct RegionDirectory {
    #[allow(dead_code)]
    default_region: RegionId,
    name_to_id: HashMap<String, RegionId>,
    id_to_name: HashMap<u16, String>,
}

#[derive(Debug, serde::Serialize)]
pub struct ClusterHealth {
    pub leader: String,
    pub relay_log_len: u64,
    pub regions: Vec<RegionStatus>,
}

#[derive(Debug, serde::Serialize)]
pub struct RegionStatus {
    pub name: String,
    pub applied_events: u64,
    pub lag_events: i64,
    pub partitioned: bool,
}

impl ServerState {
    pub fn standalone(db: TidalDb) -> Self {
        Self {
            mode: Mode::Standalone(Arc::new(db)),
        }
    }

    pub fn cluster(schema: Schema, layout: ClusterLayout) -> Result<Self> {
        use tidaldb::testing::cluster::ClusterConfig;

        let mut regions = Vec::new();
        for (idx, name) in layout.regions.iter().enumerate() {
            regions.push((RegionId(idx as u16), name.clone()));
        }

        let leader_id = regions
            .iter()
            .find(|(_, name)| name.eq_ignore_ascii_case(&layout.leader))
            .map(|(id, _)| *id)
            .ok_or_else(|| {
                ServerError::ClusterConfig(format!(
                    "leader '{}' not found in region list",
                    layout.leader
                ))
            })?;

        let config = ClusterConfig {
            regions: regions.iter().map(|(id, _)| *id).collect(),
            leader_region: leader_id,
            schema,
        };
        let cluster = Arc::new(SimulatedCluster::build(config));
        let mut name_to_id = HashMap::new();
        let mut id_to_name = HashMap::new();
        for (id, name) in &regions {
            name_to_id.insert(name.to_lowercase(), *id);
            id_to_name.insert(id.0, name.clone());
        }

        let directory = RegionDirectory {
            default_region: leader_id,
            name_to_id,
            id_to_name,
        };

        Ok(Self {
            mode: Mode::Cluster(ClusterState {
                cluster,
                regions: directory,
            }),
        })
    }

    pub fn is_cluster(&self) -> bool {
        matches!(self.mode, Mode::Cluster(_))
    }

    pub fn write_item(
        &self,
        entity_id: EntityId,
        metadata: &HashMap<String, String>,
    ) -> Result<()> {
        match &self.mode {
            Mode::Standalone(db) => db
                .write_item_with_metadata(entity_id, metadata)
                .map_err(ServerError::from),
            Mode::Cluster(cluster) => cluster
                .cluster
                .write_item_with_metadata(entity_id, metadata)
                .map_err(ServerError::from),
        }
    }

    pub fn write_embedding(&self, entity_id: EntityId, embedding: &[f32]) -> Result<()> {
        match &self.mode {
            Mode::Standalone(db) => db
                .write_item_embedding(entity_id, embedding)
                .map_err(ServerError::from),
            Mode::Cluster(cluster) => cluster.cluster.write_item_embedding(entity_id, embedding).map_err(ServerError::from),
        }
    }

    pub fn signal(
        &self,
        signal_name: &str,
        entity_id: EntityId,
        weight: f64,
        user_id: Option<u64>,
        creator_id: Option<u64>,
    ) -> Result<()> {
        match &self.mode {
            Mode::Standalone(db) => {
                if user_id.is_some() || creator_id.is_some() {
                    db.signal_with_context(
                        signal_name,
                        entity_id,
                        weight,
                        tidaldb::schema::Timestamp::now(),
                        user_id,
                        creator_id,
                    )
                    .map_err(ServerError::from)
                } else {
                    db.signal(
                        signal_name,
                        entity_id,
                        weight,
                        tidaldb::schema::Timestamp::now(),
                    )
                    .map_err(ServerError::from)
                }
            }
            Mode::Cluster(cluster) => {
                if user_id.is_some() || creator_id.is_some() {
                    return Err(ServerError::BadRequest(
                        "cluster mode currently supports only global signals (no user_id/creator_id)".into(),
                    ));
                }
                cluster.cluster.write_signal(signal_name, entity_id, weight);
                Ok(())
            }
        }
    }

    pub fn retrieve(
        &self,
        region_name: Option<&str>,
        query: &Retrieve,
    ) -> Result<tidaldb::query::retrieve::Results> {
        match &self.mode {
            Mode::Standalone(db) => db.retrieve(query).map_err(ServerError::from),
            Mode::Cluster(cluster) => {
                let region = cluster.resolve_region(region_name)?;
                cluster.cluster.retrieve(region, query).map_err(ServerError::from)
            }
        }
    }

    pub fn search(
        &self,
        region_name: Option<&str>,
        query: &Search,
    ) -> Result<tidaldb::query::search::SearchResults> {
        match &self.mode {
            Mode::Standalone(db) => db.search(query).map_err(ServerError::from),
            Mode::Cluster(cluster) => {
                let region = cluster.resolve_region(region_name)?;
                cluster.cluster.search(region, query).map_err(ServerError::from)
            }
        }
    }

    pub fn item_count(&self, region_name: Option<&str>) -> Result<u64> {
        match &self.mode {
            Mode::Standalone(db) => Ok(db.item_count()),
            Mode::Cluster(cluster) => {
                let region = cluster.resolve_region(region_name)?;
                Ok(cluster.cluster.item_count(region))
            }
        }
    }

    pub fn cluster_status(&self) -> Option<ClusterHealth> {
        match &self.mode {
            Mode::Cluster(cluster) => {
                let leader_id = cluster.cluster.leader_region();
                let leader_name = cluster
                    .regions
                    .id_to_name
                    .get(&leader_id.0)
                    .cloned()
                    .unwrap_or_else(|| format!("region-{}", leader_id.0));
                let leader_seqno = cluster.cluster.leader_seqno();
                let statuses = cluster
                    .regions
                    .id_to_name
                    .iter()
                    .map(|(id, name)| {
                        let rid = RegionId(*id);
                        let applied = cluster.cluster.applied_count(rid);
                        let lag = leader_seqno as i64 - applied as i64;
                        RegionStatus {
                            name: name.clone(),
                            applied_events: applied,
                            lag_events: lag.max(0),
                            partitioned: cluster.cluster.is_partitioned(rid),
                        }
                    })
                    .collect();
                Some(ClusterHealth {
                    leader: leader_name,
                    relay_log_len: cluster.cluster.relay_log_len(),
                    regions: statuses,
                })
            }
            _ => None,
        }
    }

    pub fn promote_leader(&self, region_name: &str) -> Result<()> {
        match &self.mode {
            Mode::Cluster(cluster) => {
                let region = cluster.regions.lookup(region_name).ok_or_else(|| {
                    ServerError::BadRequest(format!("unknown region '{region_name}'"))
                })?;
                cluster.cluster.promote_leader(region);
                Ok(())
            }
            _ => Err(ServerError::BadRequest(
                "leader promotion only supported in cluster mode".into(),
            )),
        }
    }

    pub fn partition_region(&self, region_name: &str) -> Result<()> {
        match &self.mode {
            Mode::Cluster(cluster) => {
                let region = cluster.regions.lookup(region_name).ok_or_else(|| {
                    ServerError::BadRequest(format!("unknown region '{region_name}'"))
                })?;
                cluster.cluster.partition_region(region);
                Ok(())
            }
            _ => Err(ServerError::BadRequest(
                "partitions only supported in cluster mode".into(),
            )),
        }
    }

    pub fn heal_region(&self, region_name: &str) -> Result<()> {
        match &self.mode {
            Mode::Cluster(cluster) => {
                let region = cluster.regions.lookup(region_name).ok_or_else(|| {
                    ServerError::BadRequest(format!("unknown region '{region_name}'"))
                })?;
                cluster.cluster.heal_region(region);
                Ok(())
            }
            _ => Err(ServerError::BadRequest(
                "partitions only supported in cluster mode".into(),
            )),
        }
    }
}

impl ClusterState {
    fn resolve_region(&self, name: Option<&str>) -> Result<RegionId> {
        if let Some(name) = name {
            self.regions
                .lookup(name)
                .ok_or_else(|| ServerError::BadRequest(format!("unknown region '{name}'")))
        } else {
            Ok(self.cluster.leader_region())
        }
    }
}

impl RegionDirectory {
    fn lookup(&self, name: &str) -> Option<RegionId> {
        self.name_to_id.get(&name.trim().to_lowercase()).copied()
    }
}