Files
ChuSQL/chusql-storage/src/bin/chusql-storage.rs
T

827 lines
28 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
use std::collections::{BTreeSet, HashMap};
use std::io::{BufRead, BufReader, Write};
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::thread;
use chusql_storage::catalog::Catalog;
use chusql_storage::config::{self, Config, Loaded};
use chusql_storage::heap::HeapTable;
use chusql_storage::log;
use chusql_storage::protocol::{
ColumnStatWire, IndexWire, Request, Response, Row, SchemaColumn, TableSchemaWire,
};
use chusql_storage::wal::{Wal, WalOp};
use chusql_storage::{log_debug, log_error, log_info, log_warn};
use interprocess::local_socket::{prelude::*, GenericNamespaced, ListenerOptions};
use interprocess::TryClone;
// 命名管道服务:按行读写 JSON,Windows 上映射成 \\.\pipe\<name>。
// 启动时读配置并把每项来源打进日志;表句柄按名字缓存复用。
struct Server {
cfg: Config,
tables: Mutex<HashMap<String, Arc<Mutex<HeapTable>>>>,
write_lock: Mutex<()>,
catalog: Mutex<Catalog>,
wal: Wal,
}
impl Server {
/// 新建服务器状态
fn new(cfg: Config) -> std::io::Result<Self> {
std::fs::create_dir_all(&cfg.data_dir)?;
let catalog = Catalog::load(cfg.data_dir.join("catalog.json"))?;
let wal = Wal::new(cfg.data_dir.join("wal.log"));
Ok(Server {
cfg,
tables: Mutex::new(HashMap::new()),
write_lock: Mutex::new(()),
catalog: Mutex::new(catalog),
wal,
})
}
/// 表名 -> 数据文件
fn table_path(&self, table: &str) -> PathBuf {
self.cfg.data_dir.join(format!("{}.db", table))
}
/// (表名, 列名) 对应哪个索引文件
fn index_path(&self, table: &str, column: &str) -> PathBuf {
if column == "id" {
self.cfg.data_dir.join(format!("{}.idx", table))
} else {
self.cfg
.data_dir
.join(format!("{}.{}.idx", table, column))
}
}
/// 数据字典文件
fn catalog_path(&self) -> PathBuf {
self.cfg.data_dir.join("catalog.json")
}
/// 表文件存在且非空
fn table_exists(&self, table: &str) -> bool {
if self.tables.lock().unwrap().contains_key(table) {
return true;
}
if let Ok(c) = Catalog::load(self.catalog_path()) {
if c.describe(table).is_some() {
return true;
}
}
std::fs::metadata(self.table_path(table))
.map(|m| m.len() > 0)
.unwrap_or(false)
}
/// 一张表该开哪些索引
fn index_columns(&self, table: &str) -> Vec<String> {
let mut cols = vec!["id".to_string()];
let c = self.catalog.lock().unwrap();
if let Some(s) = c.describe(table) {
for i in &s.indexes {
if !cols.contains(&i.column) {
cols.push(i.column.clone());
}
}
}
cols
}
/// 拿表句柄(缓存复用)
fn table_handle(&self, table: &str) -> std::io::Result<Arc<Mutex<HeapTable>>> {
if let Some(h) = self.tables.lock().unwrap().get(table) {
return Ok(h.clone());
}
let columns = self.index_columns(table);
let paths: Vec<(String, PathBuf)> = columns
.iter()
.map(|c| (c.clone(), self.index_path(table, c)))
.collect();
let mut map = self.tables.lock().unwrap();
if let Some(h) = map.get(table) {
return Ok(h.clone());
}
std::fs::create_dir_all(&self.cfg.data_dir)?;
let t = HeapTable::open_with_indexes(
self.table_path(table),
&paths,
self.cfg.page_size,
self.cfg.btree_order,
self.cfg.pool_size,
)?;
let h = Arc::new(Mutex::new(t));
map.insert(table.to_string(), h.clone());
Ok(h)
}
/// 在表上跑闭包
fn with_table<R>(
&self,
table: &str,
f: impl FnOnce(&mut HeapTable) -> std::io::Result<R>,
) -> std::io::Result<R> {
let h = self.table_handle(table)?;
let mut guard = h.lock().unwrap();
f(&mut guard)
}
/// 表不存在就报错
fn with_existing_table<R>(
&self,
table: &str,
f: impl FnOnce(&mut HeapTable) -> std::io::Result<R>,
) -> std::io::Result<R> {
let cached = self.tables.lock().unwrap().contains_key(table);
if !cached && !self.table_exists(table) {
return Err(std::io::Error::new(
std::io::ErrorKind::NotFound,
format!("unknown table: {}", table),
));
}
self.with_table(table, f)
}
/// 把一条 WAL 操作落到磁盘
fn apply_op(&self, op: &WalOp) -> std::io::Result<()> {
match op {
WalOp::Insert { table, row } => {
self.with_table(table, |t| t.insert_row(row))?;
let mut c = self.catalog.lock().unwrap();
if c.record_insert(table, row) {
c.save(self.catalog_path())?;
}
Ok(())
}
WalOp::InsertBatch { table, rows } => {
self.with_table(table, |t| t.insert_rows(rows))?;
let mut c = self.catalog.lock().unwrap();
let mut changed = false;
for r in rows {
if c.record_insert(table, r) {
changed = true;
}
}
if changed {
c.save(self.catalog_path())?;
}
Ok(())
}
WalOp::DeleteKeys { table, keys } => {
let deleted = self.with_table(table, |t| t.delete_by_keys(keys))?;
let mut c = self.catalog.lock().unwrap();
c.record_delete(table, &deleted);
Ok(())
}
WalOp::ReplaceAll { table, rows } => {
self.with_table(table, |t| t.replace_all(rows))?;
let mut c = self.catalog.lock().unwrap();
c.rebuild_stats(table, rows);
c.save(self.catalog_path())
}
WalOp::DropColumn { table, column, rows } => {
self.with_table(table, |t| {
t.detach_index(column);
t.replace_all(rows)
})?;
let mut c = self.catalog.lock().unwrap();
c.remove_column(table, column)?;
c.rebuild_stats(table, rows);
c.save(self.catalog_path())?;
drop(c);
let _ = std::fs::remove_file(self.index_path(table, column));
Ok(())
}
}
}
/// 先写 WAL 再改数据
fn write_via_wal(&self, op: &WalOp) -> std::io::Result<()> {
let _guard = self.write_lock.lock().unwrap();
std::fs::create_dir_all(&self.cfg.data_dir)?;
let wal = &self.wal;
wal.truncate()?;
wal.append(op)?;
let result = self.apply_op(op);
if result.is_ok() {
if let Err(e) = self.flush_op_tables(op) {
return Err(e);
}
}
wal.clear()?;
result
}
/// 把涉及的表写回磁盘
fn flush_op_tables(&self, op: &WalOp) -> std::io::Result<()> {
let table = match op {
WalOp::Insert { table, .. }
| WalOp::InsertBatch { table, .. }
| WalOp::DeleteKeys { table, .. }
| WalOp::ReplaceAll { table, .. }
| WalOp::DropColumn { table, .. } => table,
};
self.with_table(table, |t| t.flush())
}
/// 启动时重放残留 WAL
fn recover_wal(&self) -> std::io::Result<()> {
std::fs::create_dir_all(&self.cfg.data_dir)?;
let wal = &self.wal;
match wal.read()? {
None => Ok(()),
Some(op) => {
log_warn!(core, "replaying incomplete WAL entry");
self.apply_op(&op)?;
wal.clear()?;
log_info!(core, "WAL replay complete");
Ok(())
}
}
}
/// 启动时重建列统计
fn rebuild_stats(&self) -> std::io::Result<()> {
let names: Vec<String> = self.catalog.lock().unwrap().table_names();
let mut touched = false;
for table in names {
if !self.table_path(&table).exists() {
continue;
}
let rows = self.with_table(&table, |t| t.scan())?;
let mut c = self.catalog.lock().unwrap();
c.rebuild_stats(&table, &rows);
touched = true;
}
if touched {
let c = self.catalog.lock().unwrap();
c.save(self.catalog_path())?;
}
Ok(())
}
/// 给一列建索引
fn create_index(&self, table: &str, column: &str) -> std::io::Result<()> {
if column == "id" {
return Err(std::io::Error::new(
std::io::ErrorKind::AlreadyExists,
"id is the built-in index and cannot be created twice",
));
}
{
let c = self.catalog.lock().unwrap();
let Some(schema) = c.describe(table) else {
return Err(std::io::Error::new(
std::io::ErrorKind::NotFound,
format!("unknown table: {}", table),
));
};
if !schema.columns.iter().any(|c| c.name == column) {
return Err(std::io::Error::new(
std::io::ErrorKind::NotFound,
format!("unknown column: {}", column),
));
}
if schema.indexes.iter().any(|i| i.column == column) {
return Err(std::io::Error::new(
std::io::ErrorKind::AlreadyExists,
format!("index on column \"{}\" already exists", column),
));
}
}
let path = self.index_path(table, column);
let (page_size, order, pool) = (self.cfg.page_size, self.cfg.btree_order, self.cfg.pool_size);
self.with_table(table, |t| {
t.build_index(column, &path, page_size, order, pool)
})?;
let mut c = self.catalog.lock().unwrap();
c.add_index(table, column)?;
c.save(self.catalog_path())
}
/// 去掉一列的索引(连文件一起删)
fn drop_index(&self, table: &str, column: &str) -> std::io::Result<()> {
if column == "id" {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"id is the built-in index and cannot be dropped",
));
}
{
let c = self.catalog.lock().unwrap();
if c.describe(table).is_none() {
return Err(std::io::Error::new(
std::io::ErrorKind::NotFound,
format!("unknown table: {}", table),
));
}
}
self.with_table(table, |t| {
t.detach_index(column);
Ok(())
})?;
let mut c = self.catalog.lock().unwrap();
c.remove_index(table, column)?;
c.save(self.catalog_path())?;
drop(c);
let _ = std::fs::remove_file(self.index_path(table, column));
Ok(())
}
/// 删一列:定义、索引、行里的字段
fn drop_column(&self, table: &str, column: &str) -> std::io::Result<()> {
if column == "id" {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidInput,
"the built-in id column cannot be dropped",
));
}
{
let c = self.catalog.lock().unwrap();
let Some(schema) = c.describe(table) else {
return Err(std::io::Error::new(
std::io::ErrorKind::NotFound,
format!("unknown table: {}", table),
));
};
if !schema.columns.iter().any(|c| c.name == column) {
return Err(std::io::Error::new(
std::io::ErrorKind::NotFound,
format!("unknown column: {}", column),
));
}
}
let rows: Vec<Row> = self
.with_existing_table(table, |t| t.scan())?
.into_iter()
.map(|row| {
row.into_iter()
.filter(|(k, _)| k != column)
.collect::<Row>()
})
.collect();
let op = WalOp::DropColumn {
table: table.to_string(),
column: column.to_string(),
rows,
};
self.write_via_wal(&op)
}
/// 一条请求翻成一条响应
fn handle_request(&self, req: Request) -> Response {
match req {
Request::Ping => {
log_debug!(pipe, "ping");
Response::Pong
}
Request::Scan { table } => {
log_debug!(pipe, "scan table={}", table);
let result = self.with_existing_table(&table, |t| t.scan());
match result {
Ok(rows) => Response::Rows { rows },
Err(e) => {
log_warn!(pipe, "scan {}: {}", table, e);
Response::Error {
message: format!("scan {}: {}", table, e),
}
}
}
}
Request::Insert { table, row } => {
log_debug!(pipe, "insert table={}", table);
let tname = table.clone();
let op = WalOp::Insert { table, row };
match self.write_via_wal(&op) {
Ok(()) => Response::Ok,
Err(e) => {
log_warn!(pipe, "insert {}: {}", tname, e);
Response::Error {
message: format!("insert {}: {}", tname, e),
}
}
}
}
Request::InsertBatch { table, rows } => {
log_debug!(pipe, "insert_batch table={} rows={}", table, rows.len());
let tname = table.clone();
let op = WalOp::InsertBatch { table, rows };
match self.write_via_wal(&op) {
Ok(()) => Response::Ok,
Err(e) => {
log_warn!(pipe, "insert_batch {}: {}", tname, e);
Response::Error {
message: format!("insert_batch {}: {}", tname, e),
}
}
}
}
Request::DeleteKeys { table, keys } => {
log_debug!(pipe, "delete_keys table={} keys={}", table, keys.len());
let tname = table.clone();
let op = WalOp::DeleteKeys { table, keys };
match self.write_via_wal(&op) {
Ok(()) => Response::Ok,
Err(e) => {
log_warn!(pipe, "delete_keys {}: {}", tname, e);
Response::Error {
message: format!("delete_keys {}: {}", tname, e),
}
}
}
}
Request::LookupByIndex { table, column, key } => {
log_debug!(pipe, "lookup_by_index table={} column={} key={}", table, column, key);
let indexed =
match self.with_existing_table(&table, |t| Ok(t.has_index(&column))) {
Ok(v) => v,
Err(e) => {
log_warn!(pipe, "lookup_by_index {}: {}", table, e);
return Response::Error {
message: format!("lookup_by_index {}: {}", table, e),
};
}
};
if !indexed {
return Response::NoIndex;
}
match self.with_table(&table, |t| t.get_by_column_key(&column, key)) {
Ok(Some(row)) => Response::Rows { rows: vec![row] },
Ok(None) => Response::Rows { rows: vec![] },
Err(e) => {
log_warn!(pipe, "lookup_by_index {}: {}", table, e);
Response::Error {
message: format!("lookup_by_index {}: {}", table, e),
}
}
}
}
Request::ReplaceAll { table, rows } => {
log_debug!(pipe, "replace_all table={} rows={}", table, rows.len());
let tname = table.clone();
let op = WalOp::ReplaceAll { table, rows };
match self.write_via_wal(&op) {
Ok(()) => Response::Ok,
Err(e) => {
log_warn!(pipe, "replace_all {}: {}", tname, e);
Response::Error {
message: format!("replace_all {}: {}", tname, e),
}
}
}
}
Request::ListTables => {
log_debug!(pipe, "list_tables");
match self.list_tables() {
Ok(tables) => Response::Tables { tables },
Err(e) => {
log_warn!(pipe, "list_tables: {}", e);
Response::Error {
message: format!("list_tables: {}", e),
}
}
}
}
Request::ListCatalog => {
log_debug!(pipe, "list_catalog");
let c = self.catalog.lock().unwrap();
let schemas: Vec<TableSchemaWire> = c
.all_tables()
.into_iter()
.map(|(name, ts)| TableSchemaWire {
table: name.to_string(),
columns: ts.columns.clone(),
row_count: ts.row_count,
indexes: index_wires(ts),
stats: stat_wires(ts),
})
.collect();
Response::Catalog { schemas }
}
Request::DescribeTable { table } => {
log_debug!(pipe, "describe_table table={}", table);
match self.describe_table(&table) {
Ok((cols, count, indexes, stats)) => Response::Schema {
table,
columns: cols,
row_count: count,
indexes,
stats,
},
Err(e) => {
log_warn!(pipe, "describe_table {}: {}", table, e);
Response::Error {
message: format!("describe_table {}: {}", table, e),
}
}
}
}
Request::CreateTable { table, columns } => {
log_debug!(pipe, "create_table table={} cols={}", table, columns.len());
{
let c = self.catalog.lock().unwrap();
if c.describe(&table).is_some() {
return Response::Error {
message: format!("create_table {}: table already exists", table),
};
}
}
if let Err(e) = self.with_table(&table, |_t| Ok(())) {
return Response::Error {
message: format!("create_table {}: {}", table, e),
};
}
let mut c = self.catalog.lock().unwrap();
if let Err(e) = c.create_table(&table, columns) {
return Response::Error {
message: format!("create_table {}: {}", table, e),
};
}
if let Err(e) = c.save(self.catalog_path()) {
return Response::Error {
message: format!("create_table {}: catalog save: {}", table, e),
};
}
Response::Ok
}
Request::DropTable { table } => {
log_debug!(pipe, "drop_table table={}", table);
let dropped_indexes: Vec<String> = {
let mut c = self.catalog.lock().unwrap();
let idx = c
.describe(&table)
.map(|s| s.indexes.iter().map(|i| i.column.clone()).collect())
.unwrap_or_default();
if let Err(e) = c.drop_table(&table) {
return Response::Error {
message: format!("drop_table {}: {}", table, e),
};
}
if let Err(e) = c.save(self.catalog_path()) {
return Response::Error {
message: format!("drop_table {}: catalog save: {}", table, e),
};
}
idx
};
self.tables.lock().unwrap().remove(&table);
let _ = std::fs::remove_file(self.table_path(&table));
let _ = std::fs::remove_file(self.index_path(&table, "id"));
for col in dropped_indexes {
let _ = std::fs::remove_file(self.index_path(&table, &col));
}
Response::Ok
}
Request::CreateIndex { table, column } => {
log_debug!(pipe, "create_index table={} column={}", table, column);
match self.create_index(&table, &column) {
Ok(()) => Response::Ok,
Err(e) => {
log_warn!(pipe, "create_index {}({}): {}", table, column, e);
Response::Error {
message: format!("create_index {}({}): {}", table, column, e),
}
}
}
}
Request::DropIndex { table, column } => {
log_debug!(pipe, "drop_index table={} column={}", table, column);
match self.drop_index(&table, &column) {
Ok(()) => Response::Ok,
Err(e) => {
log_warn!(pipe, "drop_index {}({}): {}", table, column, e);
Response::Error {
message: format!("drop_index {}({}): {}", table, column, e),
}
}
}
}
Request::DropColumn { table, column } => {
log_debug!(pipe, "drop_column table={} column={}", table, column);
match self.drop_column(&table, &column) {
Ok(()) => Response::Ok,
Err(e) => {
log_warn!(pipe, "drop_column {}({}): {}", table, column, e);
Response::Error {
message: format!("drop_column {}({}): {}", table, column, e),
}
}
}
}
}
}
/// 表名:catalog 与目录取并集
fn list_tables(&self) -> std::io::Result<Vec<String>> {
let mut names = BTreeSet::new();
{
let c = self.catalog.lock().unwrap();
for t in c.table_names() {
names.insert(t);
}
}
let dir = &self.cfg.data_dir;
if dir.exists() {
for entry in std::fs::read_dir(dir)? {
let entry = entry?;
let name = entry.file_name().to_string_lossy().to_string();
if let Some(stem) = name.strip_suffix(".db") {
names.insert(stem.to_string());
}
}
}
Ok(names.into_iter().collect())
}
/// 查一张表的 schema
fn describe_table(
&self,
table: &str,
) -> std::io::Result<(Vec<SchemaColumn>, u64, Vec<IndexWire>, Vec<ColumnStatWire>)> {
let c = self.catalog.lock().unwrap();
if let Some(schema) = c.describe(table) {
return Ok((
schema.columns.clone(),
schema.row_count,
index_wires(schema),
stat_wires(schema),
));
}
drop(c);
if std::fs::metadata(self.table_path(table)).is_ok() {
return Ok((Vec::new(), 0, Vec::new(), Vec::new()));
}
Err(std::io::Error::new(
std::io::ErrorKind::NotFound,
format!("unknown table: {}", table),
))
}
}
/// 一张表的索引定义转线上格式
fn index_wires(schema: &chusql_storage::catalog::TableSchema) -> Vec<IndexWire> {
let mut out = vec![IndexWire {
column: "id".to_string(),
}];
for i in &schema.indexes {
if i.column != "id" {
out.push(IndexWire {
column: i.column.clone(),
});
}
}
out
}
/// 一张表的列统计 -> 线上格式
fn stat_wires(schema: &chusql_storage::catalog::TableSchema) -> Vec<ColumnStatWire> {
schema
.stats
.iter()
.map(|(name, s)| ColumnStatWire {
name: name.clone(),
distinct: s.distinct,
capped: s.capped,
})
.collect()
}
impl Drop for Server {
fn drop(&mut self) {
if let Ok(c) = self.catalog.lock() {
let _ = c.save(self.catalog_path());
}
}
}
/// 处理一条连接
fn handle(server: &Server, conn: LocalSocketStream) -> std::io::Result<()> {
let mut writer = conn.try_clone()?;
let reader = BufReader::new(conn);
for line in reader.lines() {
let line = line?;
if line.trim().is_empty() {
continue;
}
let resp = match serde_json::from_str::<Request>(&line) {
Ok(req) => server.handle_request(req),
Err(e) => {
log_warn!(pipe, "bad request line: {}", e);
Response::Error {
message: format!("parse error: {}", e),
}
}
};
writeln!(writer, "{}", serde_json::to_string(&resp).unwrap())?;
writer.flush()?;
}
log_debug!(pipe, "connection closed");
Ok(())
}
/// 把配置来源打进日志
fn log_config(loaded: &Loaded) {
match &loaded.config_path {
Some(p) => log_info!(core, "config file: {}", p.display()),
None => log_info!(core, "config file: none"),
}
for (name, value, origin) in &loaded.origins {
log_info!(core, " {:<18} {} [{}]", name, value, origin.describe());
}
}
/// 进程入口
fn main() -> std::io::Result<()> {
let loaded = match config::load() {
Ok(l) => l,
Err(e) => {
log_error!(core, "config error: {}", e);
return Err(std::io::Error::new(std::io::ErrorKind::InvalidInput, e));
}
};
log::init();
log::set_level(loaded.config.log_level);
log_info!(core, "chusql-storage {} starting", env!("CARGO_PKG_VERSION"));
log_config(&loaded);
let server = Arc::new(Server::new(loaded.config)?);
if let Err(e) = server.recover_wal() {
log_error!(core, "WAL recovery failed: {}", e);
return Err(e);
}
if let Err(e) = server.rebuild_stats() {
log_warn!(core, "stats rebuild skipped: {}", e);
}
let ns_name = server
.cfg
.pipe_name
.as_str()
.to_ns_name::<GenericNamespaced>()
.map_err(|e| std::io::Error::other(e.to_string()))?;
let listener = ListenerOptions::new().name(ns_name).create_sync()?;
log_info!(
core,
"listening on pipe {} (\\\\.\\pipe\\{})",
server.cfg.pipe_name,
server.cfg.pipe_name
);
for conn in listener.incoming() {
let conn = conn?;
let server = Arc::clone(&server);
thread::spawn(move || {
log_debug!(pipe, "connection accepted");
if let Err(e) = handle(&server, conn) {
log_error!(pipe, "connection error: {}", e);
}
});
}
Ok(())
}