#![allow(clippy::upper_case_acronyms)]
use crate::helpers::{deserialize, set_path_access};
use crate::migration::Migration;
use crate::query::rows::RowOwned;
use crate::store::state_machine::sqlite::TypeConfigSqlite;
use crate::store::state_machine::sqlite::param::Param;
use crate::store::state_machine::sqlite::snapshot_builder::SQLiteSnapshotBuilder;
use crate::store::state_machine::sqlite::writer::WriterRequest::MetadataRead;
use crate::store::state_machine::sqlite::writer::{
self, MetaPersistRequest, SqlBatch, SqlTransaction, WriterRequest,
};
use crate::store::{StorageResult, logs};
use crate::{Error, Node, NodeId};
use openraft::storage::RaftStateMachine;
use openraft::{
EntryPayload, LogId, OptionalSend, Snapshot, SnapshotId, SnapshotMeta, StorageError,
StorageIOError, StoredMembership,
};
use rusqlite::functions::FunctionFlags;
use rusqlite::{OpenFlags, OptionalExtension, ToSql};
use serde::{Deserialize, Serialize};
use std::borrow::Cow;
use std::clone::Clone;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
use tokio::sync::oneshot;
use tokio::{fs, task, time};
use tracing::{debug, error, info, warn};
use uuid::Uuid;
type Entry = openraft::Entry<TypeConfigSqlite>;
type SnapshotData = tokio::fs::File;
pub type SqlitePool = deadpool::unmanaged::Pool<rusqlite::Connection>;
pub type Params = Vec<Param>;
pub struct PathDb(pub String);
pub struct PathBackups(pub String);
pub struct PathSnapshots(pub String);
pub struct PathLockFile(pub String);
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum QueryWrite {
Execute(Query),
ExecuteReturning(Query),
Transaction(Vec<Query>),
Batch(Cow<'static, str>),
Migration(Vec<Migration>),
#[cfg(feature = "backup")]
Backup((NodeId, i64)),
RTT,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Query {
pub sql: Cow<'static, str>,
pub params: Params,
}
#[derive(Debug, Serialize, Deserialize)]
pub enum Response {
Empty,
Execute(ResponseExecute),
ExecuteReturning(ResponseExecuteReturning),
Transaction(Result<Vec<Result<usize, Error>>, Error>),
Batch(ResponseBatch),
Migrate(Result<(), Error>),
Backup(Result<(), Error>),
RTT,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct ResponseExecute {
pub result: Result<usize, Error>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct ResponseExecuteReturning {
pub result: Result<Vec<Result<RowOwned, Error>>, Error>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct ResponseBatch {
pub result: Result<Vec<Result<usize, Error>>, Error>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoredSnapshot {
pub meta: SnapshotMeta<NodeId, Node>,
pub path: String,
}
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct StateMachineData {
pub last_applied_log_id: Option<LogId<NodeId>>,
pub last_membership: StoredMembership<NodeId, Node>,
pub last_snapshot_id: Option<String>,
}
#[derive(Debug, Clone)]
pub struct StateMachineSqlite {
this_node: NodeId,
path_snapshots: String,
#[cfg(feature = "backup")]
path_backups: String,
path_lock_file: String,
#[cfg(feature = "s3")]
s3_config: Option<Arc<crate::s3::S3Config>>,
pub read_pool: SqlitePool,
pub(crate) write_tx: flume::Sender<WriterRequest>,
}
impl StateMachineSqlite {
#[allow(clippy::too_many_arguments)]
pub(crate) async fn new(
data_dir: &str,
filename_db: &str,
this_node: NodeId,
log_statements: bool,
prepared_statement_cache_capacity: usize,
read_pool_size: usize,
#[cfg(feature = "s3")] s3_config: Option<Arc<crate::s3::S3Config>>,
do_reset_metadata: bool,
#[cfg(feature = "backup")] local_backup_keep_days: u16,
) -> Result<StateMachineSqlite, StorageError<NodeId>> {
let mut db_exists = Self::db_exists(data_dir, filename_db).await;
debug!("db_exists in state_machine::new(): {db_exists}");
let (
PathDb(path_db),
PathBackups(path_backups),
PathSnapshots(path_snapshots),
PathLockFile(path_lock_file),
) = Self::build_folders(data_dir, true).await;
Self::check_set_lock_file(&path_lock_file, &path_db, &mut db_exists).await;
let conn = Self::connect(
path_db.to_string(),
filename_db.to_string(),
false,
prepared_statement_cache_capacity,
)
.await
.map_err(|err| StorageError::IO {
source: StorageIOError::write(&err),
})?;
let write_tx = writer::spawn_writer(
conn,
this_node,
path_lock_file.clone(),
log_statements,
do_reset_metadata,
#[cfg(feature = "backup")]
local_backup_keep_days,
);
let read_pool = Self::connect_read_pool(
path_db.as_ref(),
filename_db,
prepared_statement_cache_capacity,
read_pool_size,
)
.await
.map_err(|err| StorageError::IO {
source: StorageIOError::read(&err),
})?;
let mut slf = Self {
this_node,
path_snapshots,
#[cfg(feature = "backup")]
path_backups,
path_lock_file,
#[cfg(feature = "s3")]
s3_config,
read_pool,
write_tx,
};
if !db_exists && let Some(snapshot) = slf.read_current_snapshot().await? {
slf.update_state_machine_(snapshot.path).await?;
}
Ok(slf)
}
async fn db_exists(data_dir: &str, filename_db: &str) -> bool {
let path_db = Self::path_db(data_dir);
let path_db_full = format!("{path_db}/{filename_db}");
fs::File::open(&path_db_full).await.is_ok()
}
pub fn path_base(data_dir: &str) -> String {
format!("{data_dir}/state_machine")
}
fn path_db(data_dir: &str) -> String {
format!("{}/db", Self::path_base(data_dir))
}
pub async fn build_folders(
data_dir: &str,
create: bool,
) -> (PathDb, PathBackups, PathSnapshots, PathLockFile) {
let path_base = Self::path_base(data_dir);
let path_db = Self::path_db(data_dir);
let path_backups = format!("{path_base}/backups");
let path_snapshots = format!("{path_base}/snapshots");
let path_lock_file = format!("{path_base}/lock");
if create {
let _ = fs::create_dir_all(&path_db).await;
set_path_access(&path_base, 0o700)
.await
.expect("Cannot set access rights for path_base");
set_path_access(&path_db, 0o700)
.await
.expect("Cannot set access rights for path_db");
fs::create_dir_all(&path_backups)
.await
.expect("create state machine folder backups");
set_path_access(&path_backups, 0o700)
.await
.expect("Cannot set access rights for path_backups");
fs::create_dir_all(&path_snapshots)
.await
.expect("create state machine folder snapshots");
set_path_access(&path_snapshots, 0o700)
.await
.expect("Cannot set access rights for path_snapshots");
}
(
PathDb(path_db),
PathBackups(path_backups),
PathSnapshots(path_snapshots),
PathLockFile(path_lock_file),
)
}
async fn check_set_lock_file(path_lock_file: &str, path_db: &str, db_exists: &mut bool) {
let is_locked = fs::File::open(path_lock_file).await.is_ok();
if is_locked {
#[cfg(feature = "auto-heal")]
{
warn!(
"Lock file already exists: {path_lock_file}\n\
Node did not shut down gracefully - auto-rebuilding State Machine"
);
let _ = fs::remove_dir_all(path_db).await;
if let Err(err) = fs::create_dir_all(path_db).await {
panic!("Cannot re-create DB folder {path_db}: {err}");
}
*db_exists = false;
}
#[cfg(not(feature = "auto-heal"))]
panic!(
"Lock file already exists: {}\n\
Node did not shut down gracefully - needs manual interaction",
path_lock_file
);
} else if let Err(err) = fs::File::create(path_lock_file).await {
panic!("Error creating lock file {path_lock_file}: {err}");
}
}
pub(crate) fn remove_lock_file(path: &str) {
let _ = std::fs::remove_file(path);
}
pub async fn connect(
path: String,
filename_db: String,
read_only: bool,
prepared_statement_cache_capacity: usize,
) -> Result<rusqlite::Connection, Error> {
task::spawn_blocking(move || {
let path_full = format!("{path}/{filename_db}");
let conn = rusqlite::Connection::open(path_full)?;
Self::apply_pragmas(&conn, read_only, prepared_statement_cache_capacity)?;
if !read_only {
Self::overwrite_non_det_fns(&conn);
}
Ok(conn)
})
.await?
}
async fn connect_read_pool(
path: &str,
filename_db: &str,
prepared_statement_cache_capacity: usize,
pool_size: usize,
) -> Result<SqlitePool, Error> {
let path_full = format!("{path}/{filename_db}");
let mut conns = Vec::with_capacity(pool_size);
for _ in 0..pool_size {
let mut conn = Self::connect(
path.to_string(),
filename_db.to_string(),
true,
prepared_statement_cache_capacity,
)
.await;
while conn.is_err() {
time::sleep(Duration::from_millis(10)).await;
conn = Self::connect(
path.to_string(),
filename_db.to_string(),
true,
prepared_statement_cache_capacity,
)
.await;
}
conns.push(conn?);
}
let pool = deadpool::unmanaged::Pool::from(conns);
let conn = pool.get().await?;
task::spawn_blocking(move || {
let _ = conn.query_row("SELECT 1", (), |row| {
let res: i64 = row.get(0)?;
Ok(res)
})?;
Ok::<(), Error>(())
})
.await?;
Ok(pool)
}
fn apply_pragmas(
conn: &rusqlite::Connection,
read_only: bool,
prepared_statement_cache_capacity: usize,
) -> Result<(), rusqlite::Error> {
conn.pragma_update(None, "journal_mode", "WAL")?;
conn.pragma_update(None, "synchronous", "OFF")?;
conn.pragma_update(None, "page_size", 4096)?;
conn.pragma_update(None, "journal_size_limit", 16384)?;
conn.pragma_update(None, "wal_autocheckpoint", 4_000)?;
conn.pragma_update(None, "foreign_keys", "ON")?;
conn.pragma_update(None, "auto_vacuum", "INCREMENTAL")?;
conn.pragma_update(None, "optimize", "0x10002")?;
if read_only {
conn.pragma_update(None, "query_only", true)?;
} else {
}
conn.set_prepared_statement_cache_capacity(prepared_statement_cache_capacity);
Ok(())
}
fn overwrite_non_det_fns(conn: &rusqlite::Connection) {
conn.create_scalar_function(
"date",
-1,
FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
|_| -> rusqlite::Result<String> {
panic!(
"forbidden usage of `date()` - non-deterministic functions must never be \
used for writing connections in a Raft cluster"
)
},
);
conn.create_scalar_function(
"datetime",
-1,
FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
|_| -> rusqlite::Result<String> {
panic!(
"forbidden usage of `datetime()` - non-deterministic functions must never be \
used for writing connections in a Raft cluster"
)
},
);
conn.create_scalar_function(
"julianday",
-1,
FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
|_| -> rusqlite::Result<String> {
panic!(
"forbidden usage of `julianday()` - non-deterministic functions must never be \
used for writing connections in a Raft cluster"
);
},
);
conn.create_scalar_function(
"now",
-1,
FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
|_| -> rusqlite::Result<String> {
panic!(
"forbidden usage of `now()` - non-deterministic functions must never be \
used for writing connections in a Raft cluster"
)
},
);
conn.create_scalar_function(
"random",
-1,
FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
|_| -> rusqlite::Result<String> {
panic!(
"forbidden usage of `random()` - non-deterministic functions must never be \
used for writing connections in a Raft cluster"
)
},
);
conn.create_scalar_function(
"randomblob",
-1,
FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
|_| -> rusqlite::Result<String> {
panic!(
"forbidden usage of `randomblob()` - non-deterministic functions must never be \
used for writing connections in a Raft cluster"
)
},
);
conn.create_scalar_function(
"strftime",
-1,
FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
|_| -> rusqlite::Result<String> {
panic!(
"forbidden usage of `strftime()` - non-deterministic functions must never be \
used for writing connections in a Raft cluster"
)
},
);
conn.create_scalar_function(
"time",
-1,
FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
|_| -> rusqlite::Result<String> {
panic!(
"forbidden usage of `time()` - non-deterministic functions must never be \
used for writing connections in a Raft cluster"
)
},
);
conn.create_scalar_function(
"timediff",
-1,
FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
|_| -> rusqlite::Result<String> {
panic!(
"forbidden usage of `timediff()` - non-deterministic functions must never be \
used for writing connections in a Raft cluster"
)
},
);
conn.create_scalar_function(
"unixepoch",
-1,
FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC,
|_| -> rusqlite::Result<String> {
panic!(
"forbidden usage of `unixepoch()` - non-deterministic functions must never be \
used for writing connections in a Raft cluster"
)
},
);
}
async fn update_state_machine_(
&mut self,
snapshot_path: String,
) -> Result<(), StorageError<NodeId>> {
let (tx, rx) = oneshot::channel();
self.write_tx
.send_async(WriterRequest::SnapshotApply((snapshot_path, tx)))
.await
.expect("SQLite Writer rx to always be listening");
rx.await.expect("Snapshot installation to succeed");
Ok(())
}
async fn read_current_snapshot(&mut self) -> StorageResult<Option<StoredSnapshot>> {
let mut list = tokio::fs::read_dir(&self.path_snapshots)
.await
.map_err(|err| StorageError::IO {
source: StorageIOError::read(&err),
})?;
let mut snapshot_id: Option<Uuid> = None;
while let Ok(Some(entry)) = list.next_entry().await {
let file_name = entry.file_name();
let name = file_name.to_str().unwrap_or("UNKNOWN");
let id = match Uuid::parse_str(name) {
Ok(uuid) => uuid,
Err(_) => {
debug!("Non-UUID in snapshots folder");
continue;
}
};
let meta = entry.metadata().await.map_err(|err| StorageError::IO {
source: StorageIOError::read(&err),
})?;
if meta.is_dir() {
warn!("Invalid folder in snapshots dir: {}", name);
continue;
}
if let Some(curr) = &snapshot_id {
if &id > curr {
snapshot_id = Some(id);
}
} else {
snapshot_id = Some(id);
}
}
if snapshot_id.is_none() {
return Ok(None);
}
let id = snapshot_id.unwrap();
let path_snapshot = format!("{}/{}", self.path_snapshots, id);
let db_path = self.path_snapshots.clone();
let filename_db = id.to_string();
let conn = Self::connect(db_path, filename_db, false, 2)
.await
.map_err(|err| StorageError::IO {
source: StorageIOError::write(&err),
})?;
let path_dbg = path_snapshot.clone();
let metadata = task::spawn_blocking(move || {
let mut stmt = conn
.prepare("SELECT data FROM _metadata WHERE key = 'meta'")
.map_err(|err| {
Error::Sqlite(
format!(
"Error preparing metadata read stmt in read from snapshot: {}",
err
)
.into(),
)
})?;
let mut metadata = stmt
.query_row((), |row| {
let meta_bytes: Vec<u8> = row.get(0)?;
let metadata: StateMachineData =
deserialize(&meta_bytes).expect("Metadata to deserialize ok");
Ok(metadata)
})
.map_err(|err| {
Error::Sqlite(
format!(
"Error reading metadata from Snapshot '{}': {}",
path_dbg, err
)
.into(),
)
})?;
Ok::<StateMachineData, Error>(metadata)
})
.await
.map_err(|err| StorageError::IO {
source: StorageIOError::write(&err),
})?;
let metadata = metadata.map_err(|err| StorageError::IO {
source: StorageIOError::write(&err),
})?;
let snapshot_id = id.to_string();
assert_eq!(
Some(snapshot_id.as_str()),
metadata.last_snapshot_id.as_deref()
);
let meta = SnapshotMeta {
last_log_id: metadata.last_applied_log_id,
last_membership: metadata.last_membership,
snapshot_id,
};
let snapshot = StoredSnapshot {
meta,
path: path_snapshot,
};
Ok(Some(snapshot))
}
}
impl RaftStateMachine<TypeConfigSqlite> for StateMachineSqlite {
type SnapshotBuilder = SQLiteSnapshotBuilder;
async fn applied_state(
&mut self,
) -> Result<(Option<LogId<NodeId>>, StoredMembership<NodeId, Node>), StorageError<NodeId>> {
let (ack, rx) = oneshot::channel();
self.write_tx
.send_async(WriterRequest::MetadataRead(ack))
.await
.map_err(|err| StorageError::IO {
source: StorageIOError::read(&err),
})?;
let data = rx.await.expect("To always get Metadata from DB");
debug!("applied_state: {:?}", data);
Ok((data.last_applied_log_id, data.last_membership))
}
async fn apply<I>(&mut self, entries: I) -> Result<Vec<Response>, StorageError<NodeId>>
where
I: IntoIterator<Item = Entry> + OptionalSend,
I::IntoIter: OptionalSend,
{
let entries = entries.into_iter();
let (bound_lower, bound_upper) = entries.size_hint();
let entries_len = bound_upper
.expect("We always expect an upper bound to entries in apply()")
- bound_lower
+ 1;
let mut replies = Vec::with_capacity(entries_len);
for entry in entries {
let last_applied_log_id = Some(entry.log_id);
let resp = match entry.payload {
EntryPayload::Blank => Response::Empty,
EntryPayload::Normal(QueryWrite::Execute(Query { sql, params })) => {
let (tx, rx) = oneshot::channel();
let query = writer::Query::Execute(writer::SqlExecute {
sql,
params,
last_applied_log_id,
tx,
});
self.write_tx
.send_async(WriterRequest::Query(query))
.await
.expect("sql writer to always be listening");
let result = rx.await.expect("to always get a response from sql writer");
Response::Execute(ResponseExecute { result })
}
EntryPayload::Normal(QueryWrite::ExecuteReturning(Query { sql, params })) => {
let (tx, rx) = oneshot::channel();
let query = writer::Query::ExecuteReturning(writer::SqlExecuteReturning {
sql,
params,
last_applied_log_id,
tx,
});
self.write_tx
.send_async(WriterRequest::Query(query))
.await
.expect("sql writer to always be listening");
let result = rx.await.expect("to always get a response from sql writer");
Response::ExecuteReturning(ResponseExecuteReturning { result })
}
EntryPayload::Normal(QueryWrite::Transaction(queries)) => {
let (tx, rx) = oneshot::channel();
let req = WriterRequest::Query(writer::Query::Transaction(SqlTransaction {
queries,
last_applied_log_id,
tx,
}));
self.write_tx
.send_async(req)
.await
.expect("sql writer to always be listening");
let result = rx.await.expect("to always get a response from sql writer");
Response::Transaction(result)
}
EntryPayload::Normal(QueryWrite::Batch(sql)) => {
let (tx, rx) = oneshot::channel();
let req = WriterRequest::Query(writer::Query::Batch(SqlBatch {
sql,
last_applied_log_id,
tx,
}));
self.write_tx
.send_async(req)
.await
.expect("sql writer to always be listening");
let result = rx.await.expect("to always get a response from sql writer");
Response::Batch(ResponseBatch { result })
}
#[cfg(feature = "backup")]
EntryPayload::Normal(QueryWrite::Backup((node_id, ts))) => {
let (ack, rx) = oneshot::channel();
let req = WriterRequest::Backup(writer::BackupRequest {
node_id,
target_folder: self.path_backups.clone(),
ts,
#[cfg(feature = "s3")]
s3_config: self.s3_config.clone(),
last_applied_log_id,
ack,
});
self.write_tx
.send_async(req)
.await
.expect("sql writer to always be listening");
let result = rx.await.expect("to always get a response from sql writer");
Response::Backup(result)
}
EntryPayload::Normal(QueryWrite::Migration(migrations)) => {
let (tx, rx) = oneshot::channel();
let req = WriterRequest::Migrate(writer::Migrate {
migrations,
last_applied_log_id,
tx,
});
self.write_tx
.send_async(req)
.await
.expect("sql writer to always be listening");
let result = rx.await.expect("to always get a response from sql writer");
Response::Migrate(result)
}
EntryPayload::Normal(QueryWrite::RTT) => {
let (ack, rx) = oneshot::channel();
let req = WriterRequest::RTT(writer::RTTRequest {
last_applied_log_id,
ack,
});
self.write_tx
.send_async(req)
.await
.expect("sql writer to always be listening");
rx.await.expect("to always get a response from sql writer");
Response::RTT
}
EntryPayload::Membership(mem) => {
let (ack, rx) = oneshot::channel();
let req = WriterRequest::MetadataMembership(writer::MetaMembershipRequest {
last_membership: StoredMembership::new(Some(entry.log_id), mem),
last_applied_log_id,
ack,
});
self.write_tx
.send_async(req)
.await
.expect("sql writer to always be listening");
rx.await.expect("to always get a response from sql writer");
Response::Empty
}
};
replies.push(resp);
}
Ok(replies)
}
#[tracing::instrument(level = "trace", skip(self))]
async fn get_snapshot_builder(&mut self) -> Self::SnapshotBuilder {
SQLiteSnapshotBuilder {
#[cfg(feature = "backup")]
path_backups: self.path_backups.clone(),
path_snapshots: self.path_snapshots.clone(),
write_tx: self.write_tx.clone(),
}
}
#[tracing::instrument(level = "trace", skip(self))]
async fn begin_receiving_snapshot(&mut self) -> Result<Box<fs::File>, StorageError<NodeId>> {
let path = format!("{}/temp", self.path_snapshots);
let _ = fs::remove_file(&path).await;
match fs::File::create(path).await {
Ok(file) => Ok(Box::new(file)),
Err(err) => Err(StorageError::IO {
source: StorageIOError::write(&err),
}),
}
}
#[tracing::instrument(level = "trace", skip(self, _snapshot))]
async fn install_snapshot(
&mut self,
meta: &SnapshotMeta<NodeId, Node>,
_snapshot: Box<SnapshotData>,
) -> Result<(), StorageError<NodeId>> {
let src = format!("{}/temp", self.path_snapshots);
let dest = format!("{}/{}", self.path_snapshots, meta.snapshot_id);
fs::copy(&src, &dest)
.await
.map_err(|err| StorageError::IO {
source: StorageIOError::write(&err),
})?;
fs::remove_file(src).await.map_err(|err| StorageError::IO {
source: StorageIOError::write(&err),
})?;
self.update_state_machine_(dest).await?;
Ok(())
}
#[tracing::instrument(level = "trace", skip(self))]
async fn get_current_snapshot(
&mut self,
) -> Result<Option<Snapshot<TypeConfigSqlite>>, StorageError<NodeId>> {
match self.read_current_snapshot().await? {
None => Ok(None),
Some(snap) => {
let file = fs::File::open(&snap.path)
.await
.map_err(|err| StorageError::IO {
source: StorageIOError::read(&err),
})?;
Ok(Some(Snapshot {
meta: snap.meta,
snapshot: Box::new(file),
}))
}
}
}
}