sqlite/libsql
0
0
Fork 0
mirror of https://github.com/tursodatabase/libsql.git synced 2025-05-25 01:40:30 +00:00
Code Activity

abstract replicator injector and introduce SqliteInjector

Browse Source
This commit is contained in:
ad hoc
2024-08-08 10:22:19 +02:00
parent 807794891f
commit 7c4ea18c75
8 changed files with 418 additions and 338 deletions
bottomless/src
replicator.rs
libsql-replication/src
injector
error.rsmod.rs
sqlite_injector
headers.rsinjector_wal.rsmod.rs
replicator.rs
libsql/src/replication
mod.rs

9
bottomless/src/replicator.rs

@ -17,6 +17,7 @@ use aws_sdk_s3::primitives::ByteStream;
use aws_sdk_s3::{Client, Config};
use bytes::{Buf, Bytes};
use chrono::{DateTime, NaiveDateTime, TimeZone, Utc};
use libsql_replication::injector::Injector as _;
use libsql_sys::{Cipher, EncryptionConfig};
use std::ops::Deref;
use std::path::{Path, PathBuf};
@ -1449,12 +1450,12 @@ impl Replicator {
db_path: &Path,
) -> Result<bool> {
let encryption_config = self.encryption_config.clone();
let mut injector = libsql_replication::injector::Injector::new(
db_path,
let mut injector = libsql_replication::injector::SqliteInjector::new(
db_path.to_path_buf(),
4096,
libsql_sys::connection::NO_AUTOCHECKPOINT,
encryption_config,
)?;
).await?;
let prefix = format!("{}-{}/", self.db_name, generation);
let mut page_buf = {
let mut v = Vec::with_capacity(page_size);
@ -1552,7 +1553,7 @@ impl Replicator {
},
page_buf.as_slice(),
);
injector.inject_frame(frame_to_inject)?;
injector.inject_frame(frame_to_inject).await?;
applied_wal_frame = true;
}
}

2
libsql-replication/src/injector/error.rs

@ -1,3 +1,5 @@
pub type Result<T, E=Error> = std::result::Result<T, E>;
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("IO error: {0}")]

298
libsql-replication/src/injector/mod.rs

@ -1,299 +1,27 @@
use std::path::Path;
use std::sync::Arc;
use std::{collections::VecDeque, path::PathBuf};
use std::future::Future;
use parking_lot::Mutex;
use rusqlite::OpenFlags;
pub use sqlite_injector::SqliteInjector;
use crate::frame::{Frame, FrameNo};
use error::Result;
pub use error::Error;
use self::injector_wal::{
InjectorWal, InjectorWalManager, LIBSQL_INJECT_FATAL, LIBSQL_INJECT_OK, LIBSQL_INJECT_OK_TXN,
};
mod error;
mod headers;
mod injector_wal;
mod sqlite_injector;
#[derive(Debug)]
pub enum InjectError {}
pub trait Injector {
/// Inject a singular frame.
fn inject_frame(
&mut self,
frame: Frame,
) -> impl Future<Output = Result<Option<FrameNo>>> + Send;
pub type FrameBuffer = Arc<Mutex<VecDeque<Frame>>>;
pub struct Injector {
/// The injector is in a transaction state
is_txn: bool,
/// Buffer for holding current transaction frames
buffer: FrameBuffer,
/// Maximum capacity of the frame buffer
capacity: usize,
/// Injector connection
// connection must be dropped before the hook context
connection: Arc<Mutex<libsql_sys::Connection<InjectorWal>>>,
biggest_uncommitted_seen: FrameNo,
// Connection config items used to recreate the injection connection
path: PathBuf,
encryption_config: Option<libsql_sys::EncryptionConfig>,
auto_checkpoint: u32,
}
/// Methods from this trait are called before and after performing a frame injection.
/// This trait trait is used to record the last committed frame_no to the log.
/// The implementer can persist the pre and post commit frame no, and compare them in the event of
/// a crash; if the pre and post commit frame_no don't match, then the log may be corrupted.
impl Injector {
pub fn new(
path: impl AsRef<Path>,
capacity: usize,
auto_checkpoint: u32,
encryption_config: Option<libsql_sys::EncryptionConfig>,
) -> Result<Self, Error> {
let path = path.as_ref().to_path_buf();
let buffer = FrameBuffer::default();
let wal_manager = InjectorWalManager::new(buffer.clone());
let connection = libsql_sys::Connection::open(
&path,
OpenFlags::SQLITE_OPEN_READ_WRITE
| OpenFlags::SQLITE_OPEN_CREATE
| OpenFlags::SQLITE_OPEN_URI
| OpenFlags::SQLITE_OPEN_NO_MUTEX,
wal_manager,
auto_checkpoint,
encryption_config.clone(),
)?;
Ok(Self {
is_txn: false,
buffer,
capacity,
connection: Arc::new(Mutex::new(connection)),
biggest_uncommitted_seen: 0,
path,
encryption_config,
auto_checkpoint,
})
}
/// Inject a frame into the log. If this was a commit frame, returns Ok(Some(FrameNo)).
pub fn inject_frame(&mut self, frame: Frame) -> Result<Option<FrameNo>, Error> {
let frame_close_txn = frame.header().size_after.get() != 0;
self.buffer.lock().push_back(frame);
if frame_close_txn || self.buffer.lock().len() >= self.capacity {
return self.flush();
}
Ok(None)
}
pub fn rollback(&mut self) {
let conn = self.connection.lock();
let mut rollback = conn.prepare_cached("ROLLBACK").unwrap();
let _ = rollback.execute(());
self.is_txn = false;
}
/// Discard any uncommintted frames.
fn rollback(&mut self) -> impl Future<Output = ()> + Send;
/// Flush the buffer to libsql WAL.
/// Trigger a dummy write, and flush the cache to trigger a call to xFrame. The buffer's frame
/// are then injected into the wal.
pub fn flush(&mut self) -> Result<Option<FrameNo>, Error> {
match self.try_flush() {
Err(e) => {
// something went wrong, rollback the connection to make sure we can retry in a
// clean state
self.biggest_uncommitted_seen = 0;
self.rollback();
Err(e)
}
Ok(ret) => Ok(ret),
}
}
fn try_flush(&mut self) -> Result<Option<FrameNo>, Error> {
if !self.is_txn {
self.begin_txn()?;
}
let lock = self.buffer.lock();
// the frames in the buffer are either monotonically increasing (log) or decreasing
// (snapshot). Either way, we want to find the biggest frameno we're about to commit, and
// that is either the front or the back of the buffer
let last_frame_no = match lock.back().zip(lock.front()) {
Some((b, f)) => f.header().frame_no.get().max(b.header().frame_no.get()),
None => {
tracing::trace!("nothing to inject");
return Ok(None);
}
};
self.biggest_uncommitted_seen = self.biggest_uncommitted_seen.max(last_frame_no);
drop(lock);
let connection = self.connection.lock();
// use prepare cached to avoid parsing the same statement over and over again.
let mut stmt =
connection.prepare_cached("INSERT INTO libsql_temp_injection VALUES (42)")?;
// We execute the statement, and then force a call to xframe if necesacary. If the execute
// succeeds, then xframe wasn't called, in this case, we call cache_flush, and then process
// the error.
// It is unexpected that execute flushes, but it is possible, so we handle that case.
match stmt.execute(()).and_then(|_| connection.cache_flush()) {
Ok(_) => panic!("replication hook was not called"),
Err(e) => {
if let Some(e) = e.sqlite_error() {
if e.extended_code == LIBSQL_INJECT_OK {
// refresh schema
connection.pragma_update(None, "writable_schema", "reset")?;
let mut rollback = connection.prepare_cached("ROLLBACK")?;
let _ = rollback.execute(());
self.is_txn = false;
assert!(self.buffer.lock().is_empty());
let commit_frame_no = self.biggest_uncommitted_seen;
self.biggest_uncommitted_seen = 0;
return Ok(Some(commit_frame_no));
} else if e.extended_code == LIBSQL_INJECT_OK_TXN {
self.is_txn = true;
assert!(self.buffer.lock().is_empty());
return Ok(None);
} else if e.extended_code == LIBSQL_INJECT_FATAL {
return Err(Error::FatalInjectError);
}
}
Err(Error::FatalInjectError)
}
}
}
fn begin_txn(&mut self) -> Result<(), Error> {
let mut conn = self.connection.lock();
{
let wal_manager = InjectorWalManager::new(self.buffer.clone());
let new_conn = libsql_sys::Connection::open(
&self.path,
OpenFlags::SQLITE_OPEN_READ_WRITE
| OpenFlags::SQLITE_OPEN_CREATE
| OpenFlags::SQLITE_OPEN_URI
| OpenFlags::SQLITE_OPEN_NO_MUTEX,
wal_manager,
self.auto_checkpoint,
self.encryption_config.clone(),
)?;
let _ = std::mem::replace(&mut *conn, new_conn);
}
conn.pragma_update(None, "writable_schema", "true")?;
let mut stmt = conn.prepare_cached("BEGIN IMMEDIATE")?;
stmt.execute(())?;
// we create a dummy table. This table MUST not be persisted, otherwise the replica schema
// would differ with the primary's.
let mut stmt =
conn.prepare_cached("CREATE TABLE IF NOT EXISTS libsql_temp_injection (x)")?;
stmt.execute(())?;
Ok(())
}
pub fn clear_buffer(&mut self) {
self.buffer.lock().clear()
}
#[cfg(test)]
pub fn is_txn(&self) -> bool {
self.is_txn
}
}
#[cfg(test)]
mod test {
use crate::frame::FrameBorrowed;
use std::mem::size_of;
use super::*;
/// this this is generated by creating a table test, inserting 5 rows into it, and then
/// truncating the wal file of it's header.
const WAL: &[u8] = include_bytes!("../../assets/test/test_wallog");
fn wal_log() -> impl Iterator<Item = Frame> {
WAL.chunks(size_of::<FrameBorrowed>())
.map(|b| Frame::try_from(b).unwrap())
}
#[test]
fn test_simple_inject_frames() {
let temp = tempfile::tempdir().unwrap();
let mut injector = Injector::new(temp.path().join("data"), 10, 10000, None).unwrap();
let log = wal_log();
for frame in log {
injector.inject_frame(frame).unwrap();
}
let conn = rusqlite::Connection::open(temp.path().join("data")).unwrap();
conn.query_row("SELECT COUNT(*) FROM test", (), |row| {
assert_eq!(row.get::<_, usize>(0).unwrap(), 5);
Ok(())
})
.unwrap();
}
#[test]
fn test_inject_frames_split_txn() {
let temp = tempfile::tempdir().unwrap();
// inject one frame at a time
let mut injector = Injector::new(temp.path().join("data"), 1, 10000, None).unwrap();
let log = wal_log();
for frame in log {
injector.inject_frame(frame).unwrap();
}
let conn = rusqlite::Connection::open(temp.path().join("data")).unwrap();
conn.query_row("SELECT COUNT(*) FROM test", (), |row| {
assert_eq!(row.get::<_, usize>(0).unwrap(), 5);
Ok(())
})
.unwrap();
}
#[test]
fn test_inject_partial_txn_isolated() {
let temp = tempfile::tempdir().unwrap();
// inject one frame at a time
let mut injector = Injector::new(temp.path().join("data"), 10, 1000, None).unwrap();
let mut frames = wal_log();
assert!(injector
.inject_frame(frames.next().unwrap())
.unwrap()
.is_none());
let conn = rusqlite::Connection::open(temp.path().join("data")).unwrap();
assert!(conn
.query_row("SELECT COUNT(*) FROM test", (), |_| Ok(()))
.is_err());
while injector
.inject_frame(frames.next().unwrap())
.unwrap()
.is_none()
{}
// reset schema
conn.pragma_update(None, "writable_schema", "reset")
.unwrap();
conn.query_row("SELECT COUNT(*) FROM test", (), |_| Ok(()))
.unwrap();
}
fn flush(&mut self) -> impl Future<Output = Result<Option<FrameNo>>> + Send;
}

0
libsql-replication/src/injector/headers.rs → libsql-replication/src/injector/sqlite_injector/headers.rs

0
libsql-replication/src/injector/injector_wal.rs → libsql-replication/src/injector/sqlite_injector/injector_wal.rs

345
libsql-replication/src/injector/sqlite_injector/mod.rs Normal file

@ -0,0 +1,345 @@
use std::path::Path;
use std::sync::Arc;
use std::{collections::VecDeque, path::PathBuf};
use parking_lot::Mutex;
use rusqlite::OpenFlags;
use tokio::task::spawn_blocking;
use crate::frame::{Frame, FrameNo};
use self::injector_wal::{
InjectorWal, InjectorWalManager, LIBSQL_INJECT_FATAL, LIBSQL_INJECT_OK, LIBSQL_INJECT_OK_TXN,
};
use super::error::Result;
use super::{Error, Injector};
mod headers;
mod injector_wal;
pub type FrameBuffer = Arc<Mutex<VecDeque<Frame>>>;
pub struct SqliteInjector {
pub(in super::super) inner: Arc<Mutex<SqliteInjectorInner>>,
}
impl Injector for SqliteInjector {
async fn inject_frame(
&mut self,
frame: Frame,
) -> Result<Option<FrameNo>> {
let inner = self.inner.clone();
spawn_blocking(move || {
inner.lock().inject_frame(frame)
}).await.unwrap()
}
async fn rollback(&mut self) {
let inner = self.inner.clone();
spawn_blocking(move || {
inner.lock().rollback()
}).await.unwrap();
}
async fn flush(&mut self) -> Result<Option<FrameNo>> {
let inner = self.inner.clone();
spawn_blocking(move || {
inner.lock().flush()
}).await.unwrap()
}
}
impl SqliteInjector {
pub async fn new(
path: PathBuf,
capacity: usize,
auto_checkpoint: u32,
encryption_config: Option<libsql_sys::EncryptionConfig>,
) ->super::Result<Self> {
let inner = spawn_blocking(move || {
SqliteInjectorInner::new(path, capacity, auto_checkpoint, encryption_config)
}).await.unwrap()?;
Ok(Self {
inner: Arc::new(Mutex::new(inner))
})
}
}
pub(in super::super) struct SqliteInjectorInner {
/// The injector is in a transaction state
is_txn: bool,
/// Buffer for holding current transaction frames
buffer: FrameBuffer,
/// Maximum capacity of the frame buffer
capacity: usize,
/// Injector connection
// connection must be dropped before the hook context
connection: Arc<Mutex<libsql_sys::Connection<InjectorWal>>>,
biggest_uncommitted_seen: FrameNo,
// Connection config items used to recreate the injection connection
path: PathBuf,
encryption_config: Option<libsql_sys::EncryptionConfig>,
auto_checkpoint: u32,
}
/// Methods from this trait are called before and after performing a frame injection.
/// This trait trait is used to record the last committed frame_no to the log.
/// The implementer can persist the pre and post commit frame no, and compare them in the event of
/// a crash; if the pre and post commit frame_no don't match, then the log may be corrupted.
impl SqliteInjectorInner {
fn new(
path: impl AsRef<Path>,
capacity: usize,
auto_checkpoint: u32,
encryption_config: Option<libsql_sys::EncryptionConfig>,
) -> Result<Self, Error> {
let path = path.as_ref().to_path_buf();
let buffer = FrameBuffer::default();
let wal_manager = InjectorWalManager::new(buffer.clone());
let connection = libsql_sys::Connection::open(
&path,
OpenFlags::SQLITE_OPEN_READ_WRITE
| OpenFlags::SQLITE_OPEN_CREATE
| OpenFlags::SQLITE_OPEN_URI
| OpenFlags::SQLITE_OPEN_NO_MUTEX,
wal_manager,
auto_checkpoint,
encryption_config.clone(),
)?;
Ok(Self {
is_txn: false,
buffer,
capacity,
connection: Arc::new(Mutex::new(connection)),
biggest_uncommitted_seen: 0,
path,
encryption_config,
auto_checkpoint,
})
}
/// Inject a frame into the log. If this was a commit frame, returns Ok(Some(FrameNo)).
pub fn inject_frame(&mut self, frame: Frame) -> Result<Option<FrameNo>, Error> {
let frame_close_txn = frame.header().size_after.get() != 0;
self.buffer.lock().push_back(frame);
if frame_close_txn || self.buffer.lock().len() >= self.capacity {
return self.flush();
}
Ok(None)
}
pub fn rollback(&mut self) {
self.clear_buffer();
let conn = self.connection.lock();
let mut rollback = conn.prepare_cached("ROLLBACK").unwrap();
let _ = rollback.execute(());
self.is_txn = false;
}
/// Flush the buffer to libsql WAL.
/// Trigger a dummy write, and flush the cache to trigger a call to xFrame. The buffer's frame
/// are then injected into the wal.
pub fn flush(&mut self) -> Result<Option<FrameNo>, Error> {
match self.try_flush() {
Err(e) => {
// something went wrong, rollback the connection to make sure we can retry in a
// clean state
self.biggest_uncommitted_seen = 0;
self.rollback();
Err(e)
}
Ok(ret) => Ok(ret),
}
}
fn try_flush(&mut self) -> Result<Option<FrameNo>, Error> {
if !self.is_txn {
self.begin_txn()?;
}
let lock = self.buffer.lock();
// the frames in the buffer are either monotonically increasing (log) or decreasing
// (snapshot). Either way, we want to find the biggest frameno we're about to commit, and
// that is either the front or the back of the buffer
let last_frame_no = match lock.back().zip(lock.front()) {
Some((b, f)) => f.header().frame_no.get().max(b.header().frame_no.get()),
None => {
tracing::trace!("nothing to inject");
return Ok(None);
}
};
self.biggest_uncommitted_seen = self.biggest_uncommitted_seen.max(last_frame_no);
drop(lock);
let connection = self.connection.lock();
// use prepare cached to avoid parsing the same statement over and over again.
let mut stmt =
connection.prepare_cached("INSERT INTO libsql_temp_injection VALUES (42)")?;
// We execute the statement, and then force a call to xframe if necesacary. If the execute
// succeeds, then xframe wasn't called, in this case, we call cache_flush, and then process
// the error.
// It is unexpected that execute flushes, but it is possible, so we handle that case.
match stmt.execute(()).and_then(|_| connection.cache_flush()) {
Ok(_) => panic!("replication hook was not called"),
Err(e) => {
if let Some(e) = e.sqlite_error() {
if e.extended_code == LIBSQL_INJECT_OK {
// refresh schema
connection.pragma_update(None, "writable_schema", "reset")?;
let mut rollback = connection.prepare_cached("ROLLBACK")?;
let _ = rollback.execute(());
self.is_txn = false;
assert!(self.buffer.lock().is_empty());
let commit_frame_no = self.biggest_uncommitted_seen;
self.biggest_uncommitted_seen = 0;
return Ok(Some(commit_frame_no));
} else if e.extended_code == LIBSQL_INJECT_OK_TXN {
self.is_txn = true;
assert!(self.buffer.lock().is_empty());
return Ok(None);
} else if e.extended_code == LIBSQL_INJECT_FATAL {
return Err(Error::FatalInjectError);
}
}
Err(Error::FatalInjectError)
}
}
}
fn begin_txn(&mut self) -> Result<(), Error> {
let mut conn = self.connection.lock();
{
let wal_manager = InjectorWalManager::new(self.buffer.clone());
let new_conn = libsql_sys::Connection::open(
&self.path,
OpenFlags::SQLITE_OPEN_READ_WRITE
| OpenFlags::SQLITE_OPEN_CREATE
| OpenFlags::SQLITE_OPEN_URI
| OpenFlags::SQLITE_OPEN_NO_MUTEX,
wal_manager,
self.auto_checkpoint,
self.encryption_config.clone(),
)?;
let _ = std::mem::replace(&mut *conn, new_conn);
}
conn.pragma_update(None, "writable_schema", "true")?;
let mut stmt = conn.prepare_cached("BEGIN IMMEDIATE")?;
stmt.execute(())?;
// we create a dummy table. This table MUST not be persisted, otherwise the replica schema
// would differ with the primary's.
let mut stmt =
conn.prepare_cached("CREATE TABLE IF NOT EXISTS libsql_temp_injection (x)")?;
stmt.execute(())?;
Ok(())
}
pub fn clear_buffer(&mut self) {
self.buffer.lock().clear()
}
#[cfg(test)]
pub fn is_txn(&self) -> bool {
self.is_txn
}
}
#[cfg(test)]
mod test {
use crate::frame::FrameBorrowed;
use std::mem::size_of;
use super::*;
/// this this is generated by creating a table test, inserting 5 rows into it, and then
/// truncating the wal file of it's header.
const WAL: &[u8] = include_bytes!("../../../assets/test/test_wallog");
fn wal_log() -> impl Iterator<Item = Frame> {
WAL.chunks(size_of::<FrameBorrowed>())
.map(|b| Frame::try_from(b).unwrap())
}
#[test]
fn test_simple_inject_frames() {
let temp = tempfile::tempdir().unwrap();
let mut injector = SqliteInjectorInner::new(temp.path().join("data"), 10, 10000, None).unwrap();
let log = wal_log();
for frame in log {
injector.inject_frame(frame).unwrap();
}
let conn = rusqlite::Connection::open(temp.path().join("data")).unwrap();
conn.query_row("SELECT COUNT(*) FROM test", (), |row| {
assert_eq!(row.get::<_, usize>(0).unwrap(), 5);
Ok(())
})
.unwrap();
}
#[test]
fn test_inject_frames_split_txn() {
let temp = tempfile::tempdir().unwrap();
// inject one frame at a time
let mut injector = SqliteInjectorInner::new(temp.path().join("data"), 1, 10000, None).unwrap();
let log = wal_log();
for frame in log {
injector.inject_frame(frame).unwrap();
}
let conn = rusqlite::Connection::open(temp.path().join("data")).unwrap();
conn.query_row("SELECT COUNT(*) FROM test", (), |row| {
assert_eq!(row.get::<_, usize>(0).unwrap(), 5);
Ok(())
})
.unwrap();
}
#[test]
fn test_inject_partial_txn_isolated() {
let temp = tempfile::tempdir().unwrap();
// inject one frame at a time
let mut injector = SqliteInjectorInner::new(temp.path().join("data"), 10, 1000, None).unwrap();
let mut frames = wal_log();
assert!(injector
.inject_frame(frames.next().unwrap())
.unwrap()
.is_none());
let conn = rusqlite::Connection::open(temp.path().join("data")).unwrap();
assert!(conn
.query_row("SELECT COUNT(*) FROM test", (), |_| Ok(()))
.is_err());
while injector
.inject_frame(frames.next().unwrap())
.unwrap()
.is_none()
{}
// reset schema
conn.pragma_update(None, "writable_schema", "reset")
.unwrap();
conn.query_row("SELECT COUNT(*) FROM test", (), |_| Ok(()))
.unwrap();
}
}

95
libsql-replication/src/replicator.rs

@ -1,14 +1,11 @@
use std::path::PathBuf;
use std::sync::Arc;
use parking_lot::Mutex;
use tokio::task::spawn_blocking;
use tokio::time::Duration;
use tokio_stream::{Stream, StreamExt};
use tonic::{Code, Status};
use crate::frame::{Frame, FrameNo};
use crate::injector::Injector;
use crate::injector::{Injector, SqliteInjector};
use crate::rpc::replication::{
Frame as RpcFrame, NAMESPACE_DOESNT_EXIST, NEED_SNAPSHOT_ERROR_MSG, NO_HELLO_ERROR_MSG,
};
@ -137,9 +134,9 @@ where
/// The `Replicator`'s duty is to download frames from the primary, and pass them to the injector at
/// transaction boundaries.
pub struct Replicator<C> {
pub struct Replicator<C, I> {
client: C,
injector: Arc<Mutex<Injector>>,
injector: I,
state: ReplicatorState,
frames_synced: usize,
}
@ -154,33 +151,42 @@ enum ReplicatorState {
Exit,
}
impl<C: ReplicatorClient> Replicator<C> {
impl<C> Replicator<C, SqliteInjector>
where
C: ReplicatorClient,
{
/// Creates a replicator for the db file pointed at by `db_path`
pub async fn new(
pub async fn new_sqlite(
client: C,
db_path: PathBuf,
auto_checkpoint: u32,
encryption_config: Option<libsql_sys::EncryptionConfig>,
) -> Result<Self, Error> {
let injector = {
let db_path = db_path.clone();
spawn_blocking(move || {
Injector::new(
db_path,
INJECTOR_BUFFER_CAPACITY,
auto_checkpoint,
encryption_config,
)
})
.await??
};
let injector = SqliteInjector::new(
db_path.clone(),
INJECTOR_BUFFER_CAPACITY,
auto_checkpoint,
encryption_config,
)
.await?;
Ok(Self {
Ok(Self::new(client, injector))
}
}
impl<C, I> Replicator<C, I>
where
C: ReplicatorClient,
I: Injector,
{
pub fn new(client: C, injector: I) -> Self {
Self {
client,
injector: Arc::new(Mutex::new(injector)),
injector,
state: ReplicatorState::NeedHandshake,
frames_synced: 0,
})
}
}
/// for a handshake on next call to replicate.
@ -250,7 +256,7 @@ impl<C: ReplicatorClient> Replicator<C> {
// in case of error we rollback the current injector transaction, and start over.
if ret.is_err() {
self.client.rollback();
self.injector.lock().rollback();
self.injector.rollback().await;
}
self.state = match ret {
@ -293,7 +299,8 @@ impl<C: ReplicatorClient> Replicator<C> {
}
async fn load_snapshot(&mut self) -> Result<(), Error> {
self.injector.lock().clear_buffer();
self.client.rollback();
self.injector.rollback().await;
loop {
match self.client.snapshot().await {
Ok(mut stream) => {
@ -315,26 +322,22 @@ impl<C: ReplicatorClient> Replicator<C> {
async fn inject_frame(&mut self, frame: Frame) -> Result<(), Error> {
self.frames_synced += 1;
let injector = self.injector.clone();
match spawn_blocking(move || injector.lock().inject_frame(frame)).await? {
Ok(Some(commit_fno)) => {
match self.injector.inject_frame(frame).await? {
Some(commit_fno) => {
self.client.commit_frame_no(commit_fno).await?;
}
Ok(None) => (),
Err(e) => Err(e)?,
None => (),
}
Ok(())
}
pub async fn flush(&mut self) -> Result<(), Error> {
let injector = self.injector.clone();
match spawn_blocking(move || injector.lock().flush()).await? {
Ok(Some(commit_fno)) => {
match self.injector.flush().await? {
Some(commit_fno) => {
self.client.commit_frame_no(commit_fno).await?;
}
Ok(None) => (),
Err(e) => Err(e)?,
None => (),
}
Ok(())
@ -395,7 +398,7 @@ mod test {
fn rollback(&mut self) {}
}
let mut replicator = Replicator::new(Client, tmp.path().to_path_buf(), 10000, None)
let mut replicator = Replicator::new_sqlite(Client, tmp.path().to_path_buf(), 10000, None)
.await
.unwrap();
@ -438,7 +441,7 @@ mod test {
fn rollback(&mut self) {}
}
let mut replicator = Replicator::new(Client, tmp.path().to_path_buf(), 10000, None)
let mut replicator = Replicator::new_sqlite(Client, tmp.path().to_path_buf(), 10000, None)
.await
.unwrap();
// we assume that we already received the handshake and the handshake is not valid anymore
@ -482,7 +485,7 @@ mod test {
fn rollback(&mut self) {}
}
let mut replicator = Replicator::new(Client, tmp.path().to_path_buf(), 10000, None)
let mut replicator = Replicator::new_sqlite(Client, tmp.path().to_path_buf(), 10000, None)
.await
.unwrap();
// we assume that we already received the handshake and the handshake is not valid anymore
@ -526,7 +529,7 @@ mod test {
fn rollback(&mut self) {}
}
let mut replicator = Replicator::new(Client, tmp.path().to_path_buf(), 10000, None)
let mut replicator = Replicator::new_sqlite(Client, tmp.path().to_path_buf(), 10000, None)
.await
.unwrap();
// we assume that we already received the handshake and the handshake is not valid anymore
@ -568,7 +571,7 @@ mod test {
fn rollback(&mut self) {}
}
let mut replicator = Replicator::new(Client, tmp.path().to_path_buf(), 10000, None)
let mut replicator = Replicator::new_sqlite(Client, tmp.path().to_path_buf(), 10000, None)
.await
.unwrap();
// we assume that we already received the handshake and the handshake is not valid anymore
@ -610,7 +613,7 @@ mod test {
fn rollback(&mut self) {}
}
let mut replicator = Replicator::new(Client, tmp.path().to_path_buf(), 10000, None)
let mut replicator = Replicator::new_sqlite(Client, tmp.path().to_path_buf(), 10000, None)
.await
.unwrap();
replicator.state = ReplicatorState::NeedSnapshot;
@ -653,7 +656,7 @@ mod test {
fn rollback(&mut self) {}
}
let mut replicator = Replicator::new(Client, tmp.path().to_path_buf(), 10000, None)
let mut replicator = Replicator::new_sqlite(Client, tmp.path().to_path_buf(), 10000, None)
.await
.unwrap();
// we assume that we already received the handshake and the handshake is not valid anymore
@ -696,7 +699,7 @@ mod test {
fn rollback(&mut self) {}
}
let mut replicator = Replicator::new(Client, tmp.path().to_path_buf(), 10000, None)
let mut replicator = Replicator::new_sqlite(Client, tmp.path().to_path_buf(), 10000, None)
.await
.unwrap();
replicator.state = ReplicatorState::NeedHandshake;
@ -784,7 +787,7 @@ mod test {
committed_frame_no: None,
};
let mut replicator = Replicator::new(client, tmp.path().to_path_buf(), 10000, None)
let mut replicator = Replicator::new_sqlite(client, tmp.path().to_path_buf(), 10000, None)
.await
.unwrap();
@ -795,7 +798,7 @@ mod test {
replicator.try_replicate_step().await.unwrap_err(),
Error::Client(_)
));
assert!(!replicator.injector.lock().is_txn());
assert!(!replicator.injector.inner.lock().is_txn());
assert!(replicator.client_mut().committed_frame_no.is_none());
assert_eq!(replicator.state, ReplicatorState::NeedHandshake);
@ -805,7 +808,7 @@ mod test {
replicator.client_mut().should_error = false;
replicator.try_replicate_step().await.unwrap();
assert!(!replicator.injector.lock().is_txn());
assert!(!replicator.injector.inner.lock().is_txn());
assert_eq!(replicator.state, ReplicatorState::Exit);
assert_eq!(replicator.client_mut().committed_frame_no, Some(6));
}

7
libsql/src/replication/mod.rs

@ -6,6 +6,7 @@ use std::sync::Arc;
use std::time::Duration;
pub use libsql_replication::frame::{Frame, FrameNo};
use libsql_replication::injector::SqliteInjector;
use libsql_replication::replicator::{Either, Replicator};
pub use libsql_replication::snapshot::SnapshotFile;
@ -129,7 +130,7 @@ impl Writer {
#[derive(Clone)]
pub(crate) struct EmbeddedReplicator {
replicator: Arc<Mutex<Replicator<Either<RemoteClient, LocalClient>>>>,
replicator: Arc<Mutex<Replicator<Either<RemoteClient, LocalClient>, SqliteInjector>>>,
bg_abort: Option<Arc<DropAbort>>,
last_frames_synced: Arc<AtomicUsize>,
}
@ -149,7 +150,7 @@ impl EmbeddedReplicator {
perodic_sync: Option<Duration>,
) -> Result<Self> {
let replicator = Arc::new(Mutex::new(
Replicator::new(
Replicator::new_sqlite(
Either::Left(client),
db_path,
auto_checkpoint,
@ -193,7 +194,7 @@ impl EmbeddedReplicator {
encryption_config: Option<EncryptionConfig>,
) -> Result<Self> {
let replicator = Arc::new(Mutex::new(
Replicator::new(
Replicator::new_sqlite(
Either::Right(client),
db_path,
auto_checkpoint,