remdb 0.3.1

嵌入式内存数据库
Documentation
use crate::try_lock;

use crate::sql::query_parser::IndexType as SqlIndexType;
use crate::types::{IndexType, RemDbError, Result};
use alloc::sync::Arc;
use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::collections::{HashMap, VecDeque};
use std::sync::Mutex;
use std::thread::{self, JoinHandle};
use std::time::Duration;

#[cfg(feature = "log")]
use crate::log::error;

/// 索引构建状态
enum IndexBuildState {
    /// 等待执行
    Pending,
    /// 正在执行
    Running,
    /// 执行完成
    Completed,
    /// 执行失败
    Failed(String),
}

/// 索引构建任务ID
type IndexBuildTaskId = u64;

/// 索引构建参数
pub struct IndexBuildParams {
    /// 索引类型
    pub index_type: IndexType,
    /// 向量索引类型(如果是向量索引)
    pub vector_index_type: Option<crate::types::VectorIndexType>,
    /// HNSW M参数
    pub hnsw_m: Option<u8>,
    /// HNSW ef_construction参数
    pub hnsw_ef_construction: Option<u32>,
    /// HNSW ef_search参数
    pub hnsw_ef_search: Option<u32>,
    /// IVF nlist参数
    pub ivf_nlist: Option<u32>,
    /// IVF nprobe参数
    pub ivf_nprobe: Option<u32>,
    /// 构建模式(在线/离线)
    pub online: bool,
    /// 存储位置
    pub storage: String,
    /// 压缩类型
    pub compression: String,
}

impl Default for IndexBuildParams {
    fn default() -> Self {
        Self {
            index_type: IndexType::BTree,
            vector_index_type: None,
            hnsw_m: None,
            hnsw_ef_construction: None,
            hnsw_ef_search: None,
            ivf_nlist: None,
            ivf_nprobe: None,
            online: true,
            storage: "MEMORY".to_string(),
            compression: "NONE".to_string(),
        }
    }
}

/// 索引构建任务
struct IndexBuildTask {
    /// 任务ID
    id: IndexBuildTaskId,
    /// 表名
    table_name: String,
    /// 字段名
    column_name: Vec<String>,
    /// 索引类型
    sql_index_type: SqlIndexType,
    /// 索引构建参数
    params: IndexBuildParams,
    /// 取消标志
    canceled: Arc<AtomicBool>,
}

/// 索引构建状态信息
pub struct IndexBuildStatus {
    /// 任务ID
    pub id: IndexBuildTaskId,
    /// 表名
    pub table_name: String,
    /// 字段名
    pub column_name: String,
    /// 索引类型
    pub index_type: String,
    /// 构建状态
    state: IndexBuildState,
    /// 构建进度(0-100)
    pub progress: AtomicUsize,
    /// 已处理行数
    pub processed_rows: AtomicUsize,
    /// 总行数
    pub total_rows: AtomicUsize,
    /// 已运行时间(毫秒)
    pub elapsed_time: AtomicUsize,
    /// 错误信息(如果失败)
    pub error: Mutex<Option<String>>,
}

impl IndexBuildStatus {
    /// 创建新的索引构建状态
    fn new(
        id: IndexBuildTaskId,
        table_name: String,
        column_name: String,
        index_type: String,
    ) -> Self {
        Self {
            id,
            table_name,
            column_name,
            index_type,
            state: IndexBuildState::Pending,
            progress: AtomicUsize::new(0),
            processed_rows: AtomicUsize::new(0),
            total_rows: AtomicUsize::new(0),
            elapsed_time: AtomicUsize::new(0),
            error: Mutex::new(None),
        }
    }

    /// 更新状态为运行中
    fn set_running(&mut self, total_rows: usize) {
        self.state = IndexBuildState::Running;
        self.total_rows.store(total_rows, Ordering::SeqCst);
    }

    /// 更新状态为已完成
    fn set_completed(&mut self) {
        self.state = IndexBuildState::Completed;
        self.progress.store(100, Ordering::SeqCst);
    }

    /// 更新状态为失败
    fn set_failed(&mut self, error: String) {
        self.state = IndexBuildState::Failed(error.clone());
        *try_lock!(self.error) = Some(error);
    }

    /// 检查是否已取消
    fn is_canceled(&self) -> bool {
        matches!(self.state, IndexBuildState::Failed(_))
    }

    /// 获取状态字符串
    pub fn get_state_str(&self) -> &'static str {
        match self.state {
            IndexBuildState::Pending => "PENDING",
            IndexBuildState::Running => "RUNNING",
            IndexBuildState::Completed => "COMPLETED",
            IndexBuildState::Failed(_) => "FAILED",
        }
    }
}

/// 索引构建线程池
pub struct IndexBuildThreadPool {
    /// 线程数量
    thread_count: usize,
    /// 工作线程
    workers: Vec<JoinHandle<()>>,
    /// 任务队列
    task_queue: Arc<Mutex<VecDeque<IndexBuildTask>>>,
    /// 任务ID计数器
    next_task_id: AtomicUsize,
    /// 索引构建状态映射
    build_status: Mutex<HashMap<IndexBuildTaskId, Arc<Mutex<IndexBuildStatus>>>>,
    /// 线程池停止标志
    stop_flag: Arc<AtomicBool>,
}

impl IndexBuildThreadPool {
    /// 创建新的索引构建线程池
    pub fn new(thread_count: usize) -> Self {
        let stop_flag = Arc::new(AtomicBool::new(false));
        let task_queue = Arc::new(Mutex::new(VecDeque::new()));

        let mut workers = Vec::with_capacity(thread_count);
        for _ in 0..thread_count {
            let task_queue_clone = task_queue.clone();
            let stop = stop_flag.clone();
            let handle = thread::Builder::new()
                .stack_size(8 * 1024 * 1024)
                .spawn(move || Self::worker_loop(task_queue_clone, stop))
                .expect("Failed to spawn builder thread");
            workers.push(handle);
        }

        Self {
            thread_count,
            workers,
            task_queue,
            next_task_id: AtomicUsize::new(0),
            build_status: Mutex::new(HashMap::new()),
            stop_flag,
        }
    }

    /// 工作线程主循环
    fn worker_loop(task_queue: Arc<Mutex<VecDeque<IndexBuildTask>>>, stop: Arc<AtomicBool>) {
        while !stop.load(Ordering::SeqCst) {
            // 尝试从任务队列中获取任务
            let task = {
                let mut queue = try_lock!(task_queue);
                queue.pop_front()
            };

            if let Some(task) = task {
                // 执行索引构建任务
                Self::execute_index_build(task);
            } else {
                // 任务队列为空,等待一段时间后继续检查
                thread::sleep(Duration::from_millis(100));
            }
        }
    }

    /// 执行索引构建任务
    fn execute_index_build(task: IndexBuildTask) {
        // 1. 查找目标表
        let db = unsafe { crate::get_global_db() };
        if db.is_none() {
            #[cfg(feature = "log")]
            error!("Database not initialized");
            return;
        }
        let db = db.unwrap();

        // 查找表ID
        let mut table_id = None;
        for (id, table_opt) in db.tables.iter().enumerate() {
            if let Some(table) = table_opt {
                if table.def.name == task.table_name {
                    table_id = Some(id);
                    break;
                }
            }
        }

        if table_id.is_none() {
            #[cfg(feature = "log")]
            error!("Table {} not found", task.table_name);
            return;
        }
        let table_id = table_id.unwrap();

        // 获取表引用
        let _table = match db.get_table(table_id) {
            Ok(table) => table,
            Err(e) => {
                #[cfg(feature = "log")]
                error!("Error getting table: {:?}", e);
                return;
            }
        };

        // Index building is not supported yet
        #[cfg(feature = "log")]
        error!("Index building not supported yet");
        return;
    }

    /// 提交索引构建任务
    pub fn submit_task(
        &self,
        table_name: String,
        column_name: Vec<String>,
        sql_index_type: SqlIndexType,
        params: IndexBuildParams,
    ) -> IndexBuildTaskId {
        let task_id = self.next_task_id.fetch_add(1, Ordering::SeqCst) as u64;

        // 创建索引构建状态
        let status = Arc::new(Mutex::new(IndexBuildStatus::new(
            task_id,
            table_name.clone(),
            column_name.join(", "),
            sql_index_type.to_string(),
        )));

        // 存储状态
        try_lock!(self.build_status).insert(task_id, status);

        // 创建并提交任务
        let task = IndexBuildTask {
            id: task_id,
            table_name,
            column_name,
            sql_index_type,
            params,
            canceled: Arc::new(AtomicBool::new(false)),
        };

        // 将任务添加到队列
        try_lock!(self.task_queue).push_back(task);

        task_id
    }

    /// 获取索引构建状态
    pub fn get_build_status(
        &self,
        task_id: Option<IndexBuildTaskId>,
    ) -> Vec<Arc<Mutex<IndexBuildStatus>>> {
        let status_map = try_lock!(self.build_status);

        match task_id {
            Some(id) => {
                // 获取指定任务的状态
                if let Some(status) = status_map.get(&id) {
                    vec![status.clone()]
                } else {
                    vec![]
                }
            }
            None => {
                // 获取所有任务的状态
                status_map.values().cloned().collect()
            }
        }
    }

    /// 停止线程池
    pub fn stop(&mut self) {
        self.stop_flag.store(true, Ordering::SeqCst);
        for worker in self.workers.drain(..) {
            worker.join().expect("Failed to join builder thread");
        }
    }
}

/// 全局索引构建线程池实例
pub static mut INDEX_BUILD_THREAD_POOL: Option<Arc<IndexBuildThreadPool>> = None;

/// 初始化索引构建线程池
pub fn init_index_build_thread_pool(thread_count: usize) {
    unsafe {
        INDEX_BUILD_THREAD_POOL = Some(Arc::new(IndexBuildThreadPool::new(thread_count)));
    }
}

/// 获取索引构建线程池
pub fn get_index_build_thread_pool() -> Result<Arc<IndexBuildThreadPool>> {
    unsafe {
        INDEX_BUILD_THREAD_POOL
            .as_ref()
            .ok_or(RemDbError::UnsupportedOperation)
            .cloned()
    }
}