Simplify
This commit is contained in:
parent
1461b41a36
commit
7f3ec69cf6
7 changed files with 236 additions and 158 deletions
194
src/db.rs
194
src/db.rs
|
|
@ -1,7 +1,6 @@
|
|||
use rusqlite::{params, Connection, OpenFlags};
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
|
||||
use crate::error::AppError;
|
||||
|
||||
|
|
@ -151,174 +150,73 @@ pub fn all_records(conn: &Connection) -> Result<Vec<Record>, AppError> {
|
|||
Ok(records)
|
||||
}
|
||||
|
||||
// --- Write commands ---
|
||||
|
||||
pub enum WriteCmd {
|
||||
Put {
|
||||
key: String,
|
||||
volumes: Vec<String>,
|
||||
size: Option<i64>,
|
||||
reply: oneshot::Sender<Result<(), AppError>>,
|
||||
},
|
||||
Delete {
|
||||
key: String,
|
||||
reply: oneshot::Sender<Result<(), AppError>>,
|
||||
},
|
||||
BulkPut {
|
||||
records: Vec<(String, Vec<String>, Option<i64>)>,
|
||||
reply: oneshot::Sender<Result<(), AppError>>,
|
||||
},
|
||||
}
|
||||
|
||||
fn execute_cmd(
|
||||
conn: &Connection,
|
||||
cmd: WriteCmd,
|
||||
) -> (Result<(), AppError>, oneshot::Sender<Result<(), AppError>>) {
|
||||
match cmd {
|
||||
WriteCmd::Put {
|
||||
key,
|
||||
volumes,
|
||||
size,
|
||||
reply,
|
||||
} => {
|
||||
let volumes_json = encode_volumes(&volumes);
|
||||
let result = conn
|
||||
.prepare_cached(
|
||||
"INSERT INTO kv (key, volumes, size) VALUES (?1, ?2, ?3)
|
||||
ON CONFLICT(key) DO UPDATE SET volumes = ?2, size = ?3",
|
||||
)
|
||||
.and_then(|mut s| s.execute(params![key, volumes_json, size]))
|
||||
.map(|_| ())
|
||||
.map_err(AppError::from);
|
||||
(result, reply)
|
||||
}
|
||||
WriteCmd::Delete { key, reply } => {
|
||||
let result = conn
|
||||
.prepare_cached("DELETE FROM kv WHERE key = ?1")
|
||||
.and_then(|mut s| s.execute(params![key]))
|
||||
.map(|_| ())
|
||||
.map_err(AppError::from);
|
||||
(result, reply)
|
||||
}
|
||||
WriteCmd::BulkPut { records, reply } => {
|
||||
let result = (|| -> Result<(), AppError> {
|
||||
let mut stmt = conn.prepare_cached(
|
||||
"INSERT INTO kv (key, volumes, size) VALUES (?1, ?2, ?3)
|
||||
ON CONFLICT(key) DO UPDATE SET volumes = ?2, size = ?3",
|
||||
)?;
|
||||
for (key, volumes, size) in &records {
|
||||
let volumes_json = encode_volumes(volumes);
|
||||
stmt.execute(params![key, volumes_json, size])?;
|
||||
}
|
||||
Ok(())
|
||||
})();
|
||||
(result, reply)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- WriterHandle ---
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct WriterHandle {
|
||||
tx: mpsc::Sender<WriteCmd>,
|
||||
conn: Arc<Mutex<Connection>>,
|
||||
}
|
||||
|
||||
impl WriterHandle {
|
||||
pub fn new(path: &str) -> Self {
|
||||
let conn = open_readwrite(path);
|
||||
create_tables(&conn);
|
||||
Self {
|
||||
conn: Arc::new(Mutex::new(conn)),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn put(
|
||||
&self,
|
||||
key: String,
|
||||
volumes: Vec<String>,
|
||||
size: Option<i64>,
|
||||
) -> Result<(), AppError> {
|
||||
let (reply_tx, reply_rx) = oneshot::channel();
|
||||
self.tx
|
||||
.send(WriteCmd::Put {
|
||||
key,
|
||||
volumes,
|
||||
size,
|
||||
reply: reply_tx,
|
||||
})
|
||||
.await
|
||||
.map_err(|_| AppError::WriterDead)?;
|
||||
reply_rx.await.map_err(|_| AppError::WriterDroppedReply)?
|
||||
let conn = self.conn.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = conn.lock().unwrap();
|
||||
let volumes_json = encode_volumes(&volumes);
|
||||
conn.prepare_cached(
|
||||
"INSERT INTO kv (key, volumes, size) VALUES (?1, ?2, ?3)
|
||||
ON CONFLICT(key) DO UPDATE SET volumes = ?2, size = ?3",
|
||||
)?
|
||||
.execute(params![key, volumes_json, size])?;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
pub async fn delete(&self, key: String) -> Result<(), AppError> {
|
||||
let (reply_tx, reply_rx) = oneshot::channel();
|
||||
self.tx
|
||||
.send(WriteCmd::Delete {
|
||||
key,
|
||||
reply: reply_tx,
|
||||
})
|
||||
.await
|
||||
.map_err(|_| AppError::WriterDead)?;
|
||||
reply_rx.await.map_err(|_| AppError::WriterDroppedReply)?
|
||||
let conn = self.conn.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = conn.lock().unwrap();
|
||||
conn.prepare_cached("DELETE FROM kv WHERE key = ?1")?
|
||||
.execute(params![key])?;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
pub async fn bulk_put(
|
||||
&self,
|
||||
records: Vec<(String, Vec<String>, Option<i64>)>,
|
||||
) -> Result<(), AppError> {
|
||||
let (reply_tx, reply_rx) = oneshot::channel();
|
||||
self.tx
|
||||
.send(WriteCmd::BulkPut {
|
||||
records,
|
||||
reply: reply_tx,
|
||||
})
|
||||
.await
|
||||
.map_err(|_| AppError::WriterDead)?;
|
||||
reply_rx.await.map_err(|_| AppError::WriterDroppedReply)?
|
||||
let conn = self.conn.clone();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let conn = conn.lock().unwrap();
|
||||
let mut stmt = conn.prepare_cached(
|
||||
"INSERT INTO kv (key, volumes, size) VALUES (?1, ?2, ?3)
|
||||
ON CONFLICT(key) DO UPDATE SET volumes = ?2, size = ?3",
|
||||
)?;
|
||||
for (key, volumes, size) in &records {
|
||||
let volumes_json = encode_volumes(volumes);
|
||||
stmt.execute(params![key, volumes_json, size])?;
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
// --- spawn_writer ---
|
||||
|
||||
pub fn spawn_writer(path: String) -> (WriterHandle, oneshot::Receiver<()>) {
|
||||
let (tx, mut rx) = mpsc::channel::<WriteCmd>(4096);
|
||||
let (ready_tx, ready_rx) = oneshot::channel();
|
||||
|
||||
std::thread::spawn(move || {
|
||||
let conn = open_readwrite(&path);
|
||||
create_tables(&conn);
|
||||
let _ = ready_tx.send(());
|
||||
|
||||
loop {
|
||||
let Some(first) = rx.blocking_recv() else {
|
||||
break;
|
||||
};
|
||||
|
||||
let mut batch = vec![first];
|
||||
while batch.len() < 512 {
|
||||
match rx.try_recv() {
|
||||
Ok(cmd) => batch.push(cmd),
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
|
||||
let _ = conn.execute_batch("BEGIN");
|
||||
let mut replies: Vec<(Result<(), AppError>, oneshot::Sender<Result<(), AppError>>)> =
|
||||
Vec::with_capacity(batch.len());
|
||||
|
||||
for (i, cmd) in batch.into_iter().enumerate() {
|
||||
let sp = format!("sp{i}");
|
||||
let _ = conn.execute(&format!("SAVEPOINT {sp}"), []);
|
||||
let (result, reply) = execute_cmd(&conn, cmd);
|
||||
if result.is_ok() {
|
||||
let _ = conn.execute(&format!("RELEASE {sp}"), []);
|
||||
} else {
|
||||
let _ = conn.execute(&format!("ROLLBACK TO {sp}"), []);
|
||||
let _ = conn.execute(&format!("RELEASE {sp}"), []);
|
||||
}
|
||||
replies.push((result, reply));
|
||||
}
|
||||
|
||||
let _ = conn.execute_batch("COMMIT");
|
||||
for (result, reply) in replies {
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
(WriterHandle { tx }, ready_rx)
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue