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),
}
type IndexBuildTaskId = u64;
pub struct IndexBuildParams {
pub index_type: IndexType,
pub vector_index_type: Option<crate::types::VectorIndexType>,
pub hnsw_m: Option<u8>,
pub hnsw_ef_construction: Option<u32>,
pub hnsw_ef_search: Option<u32>,
pub ivf_nlist: Option<u32>,
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: IndexBuildTaskId,
table_name: String,
column_name: Vec<String>,
sql_index_type: SqlIndexType,
params: IndexBuildParams,
canceled: Arc<AtomicBool>,
}
pub struct IndexBuildStatus {
pub id: IndexBuildTaskId,
pub table_name: String,
pub column_name: String,
pub index_type: String,
state: IndexBuildState,
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>>>,
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) {
let db = unsafe { crate::get_global_db() };
if db.is_none() {
#[cfg(feature = "log")]
error!("Database not initialized");
return;
}
let db = db.unwrap();
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;
}
};
#[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()
}
}