bottomless/src/replicator.rs
Rama main · réplica de lectura del repositorio original.
use crate::backup::WalCopier;
use crate::completion_progress::{CompletionProgress, SavepointTracker};
use crate::read::BatchReader;
use crate::uuid_utils::decode_unix_timestamp;
use crate::wal::WalFileReader;
use anyhow::{anyhow, bail};
use arc_swap::ArcSwapOption;
use async_compression::tokio::write::{GzipEncoder, ZstdEncoder};
use aws_config::BehaviorVersion;
use aws_sdk_s3::config::{Credentials, Region, SharedCredentialsProvider};
use aws_sdk_s3::error::SdkError;
use aws_sdk_s3::operation::get_object::builders::GetObjectFluentBuilder;
use aws_sdk_s3::operation::get_object::GetObjectError;
use aws_sdk_s3::operation::list_objects::builders::ListObjectsFluentBuilder;
use aws_sdk_s3::operation::list_objects::ListObjectsOutput;
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_replication::rpc::replication::Frame as RpcFrame;
use libsql_sys::{Cipher, EncryptionConfig};
use metrics::{counter, gauge, histogram};
use std::ops::Deref;
use std::path::{Path, PathBuf};
use std::str::FromStr;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::Arc;
use tokio::fs::{File, OpenOptions};
use tokio::io::AsyncWriteExt;
use tokio::sync::watch::{channel, Receiver, Sender};
use tokio::sync::{Mutex, Semaphore};
use tokio::task::JoinHandle;
use tokio::task::JoinSet;
use tokio::time::Duration;
use tokio::time::{timeout_at, Instant};
use uuid::{NoContext, Uuid};
/// Maximum number of generations that can participate in database restore procedure.
/// This effectively means that at least one in [MAX_RESTORE_STACK_DEPTH] number of
/// consecutive generations has to have a snapshot included.
const MAX_RESTORE_STACK_DEPTH: usize = 100;
pub type Result<T> = anyhow::Result<T>;
#[derive(Debug)]
pub struct Replicator {
pub client: Client,
/// Frame number, incremented whenever a new frame is written from SQLite.
next_frame_no: Arc<AtomicU32>,
/// Last frame which has been requested to be sent to S3.
/// Always: [last_sent_frame_no] <= [next_frame_no].
last_sent_frame_no: Arc<AtomicU32>,
/// Last frame which has been confirmed as stored locally outside of WAL file.
/// Always: [last_committed_frame_no] <= [last_sent_frame_no].
last_committed_frame_no: Receiver<Result<u32>>,
flush_trigger: Option<Sender<()>>,
shutdown_trigger: Option<tokio::sync::watch::Sender<()>>,
snapshot_waiter: Receiver<Result<Option<Uuid>>>,
snapshot_notifier: Arc<Sender<Result<Option<Uuid>>>>,
pub page_size: usize,
generation: Arc<ArcSwapOption<Uuid>>,
verify_crc: bool,
pub bucket: String,
pub db_path: String,
pub db_name: String,
use_compression: CompressionKind,
encryption_config: Option<EncryptionConfig>,
max_frames_per_batch: usize,
s3_max_parallelism: usize,
join_set: JoinSet<()>,
upload_progress: Arc<Mutex<CompletionProgress>>,
last_uploaded_frame_no: Receiver<u32>,
skip_snapshot: bool,
skip_shutdown_upload: bool,
}
#[derive(Debug)]
pub struct FetchedResults {
pub pages: Vec<(i32, Bytes)>,
pub next_marker: Option<String>,
}
#[derive(Debug)]
pub enum RestoreAction {
SnapshotMainDbFile,
ReuseGeneration(Uuid),
}
#[derive(Clone, Debug)]
pub struct Options {
pub create_bucket_if_not_exists: bool,
/// If `true` when restoring, frames checksums will be verified prior their pages being flushed
/// into the main database file.
pub verify_crc: bool,
/// Kind of compression algorithm used on the WAL frames to be sent to S3.
pub use_compression: CompressionKind,
pub encryption_config: Option<EncryptionConfig>,
pub aws_endpoint: Option<String>,
pub access_key_id: Option<String>,
pub secret_access_key: Option<String>,
pub session_token: Option<String>,
pub region: Option<String>,
pub db_id: Option<String>,
/// Bucket directory name where all S3 objects are backed up. General schema is:
/// - `{db-name}-{uuid-v7}` subdirectories:
/// - `.meta` file with database page size and initial WAL checksum.
/// - Series of files `{first-frame-no}-{last-frame-no}.{compression-kind}` containing
/// the batches of frames from which the restore will be made.
pub bucket_name: String,
/// Max number of WAL frames per S3 object.
pub max_frames_per_batch: usize,
/// Max time before next frame of batched frames should be synced. This works in the case
/// when we don't explicitly run into `max_frames_per_batch` threshold and the corresponding
/// checkpoint never commits.
pub max_batch_interval: Duration,
/// Maximum number of S3 file upload requests that may happen in parallel.
pub s3_max_parallelism: usize,
/// Max number of retries for S3 operations
pub s3_max_retries: u32,
/// Skip snapshot upload per checkpoint.
pub skip_snapshot: bool,
/// Skip uploading snapshots on shutdown
pub skip_shutdown_upload: bool,
}
impl Options {
pub async fn client_config(&self) -> Result<Config> {
let mut loader = aws_config::SdkConfig::builder();
if let Some(endpoint) = self.aws_endpoint.as_deref() {
loader = loader.endpoint_url(endpoint);
}
let region = self
.region
.clone()
.ok_or(anyhow!("LIBSQL_BOTTOMLESS_AWS_DEFAULT_REGION was not set"))?;
let access_key_id = self
.access_key_id
.clone()
.ok_or(anyhow!("LIBSQL_BOTTOMLESS_AWS_ACCESS_KEY_ID was not set"))?;
let secret_access_key = self.secret_access_key.clone().ok_or(anyhow!(
"LIBSQL_BOTTOMLESS_AWS_SECRET_ACCESS_KEY was not set"
))?;
let session_token: Option<String> = self.session_token.clone();
let conf = loader
.behavior_version(BehaviorVersion::latest())
.region(Region::new(region))
.credentials_provider(SharedCredentialsProvider::new(Credentials::new(
access_key_id,
secret_access_key,
session_token,
None,
"Static",
)))
.retry_config(
aws_sdk_s3::config::retry::RetryConfig::standard()
.with_max_attempts(self.s3_max_retries),
)
.build();
let s3_config = aws_sdk_s3::config::Builder::from(&conf)
.force_path_style(true)
.build();
Ok(s3_config)
}
pub fn from_env() -> Result<Self> {
fn env_var(key: &str) -> Result<String> {
match std::env::var(key) {
Ok(res) => {
let res = res.trim().to_string();
if res.is_empty() {
bail!("{} environment variable is empty", key)
} else {
Ok(res)
}
}
Err(_) => bail!("{} environment variable not set", key),
}
}
fn env_var_or<S: ToString>(key: &str, default_value: S) -> String {
match env_var(key) {
Ok(res) => res,
Err(_) => default_value.to_string(),
}
}
let db_id = env_var("LIBSQL_BOTTOMLESS_DATABASE_ID").ok();
let aws_endpoint = env_var("LIBSQL_BOTTOMLESS_ENDPOINT").ok();
let bucket_name = env_var_or("LIBSQL_BOTTOMLESS_BUCKET", "bottomless");
let max_batch_interval = Duration::from_secs(
env_var_or("LIBSQL_BOTTOMLESS_BATCH_INTERVAL_SECS", 15).parse::<u64>()?,
);
let access_key_id = env_var("LIBSQL_BOTTOMLESS_AWS_ACCESS_KEY_ID").ok();
let secret_access_key = env_var("LIBSQL_BOTTOMLESS_AWS_SECRET_ACCESS_KEY").ok();
let session_token = env_var("LIBSQL_BOTTOMLESS_AWS_SESSION_TOKEN").ok();
let region = env_var("LIBSQL_BOTTOMLESS_AWS_DEFAULT_REGION").ok();
let max_frames_per_batch =
env_var_or("LIBSQL_BOTTOMLESS_BATCH_MAX_FRAMES", 10000).parse::<usize>()?;
let s3_max_parallelism =
env_var_or("LIBSQL_BOTTOMLESS_S3_PARALLEL_MAX", 32).parse::<usize>()?;
let use_compression =
CompressionKind::parse(&env_var_or("LIBSQL_BOTTOMLESS_COMPRESSION", "zstd"))
.map_err(|e| anyhow!("unknown compression kind: {}", e))?;
let encryption_cipher = env_var("LIBSQL_BOTTOMLESS_ENCRYPTION_CIPHER").ok();
let encryption_key = env_var("LIBSQL_BOTTOMLESS_ENCRYPTION_KEY")
.map(Bytes::from)
.ok();
let verify_crc = match env_var_or("LIBSQL_BOTTOMLESS_VERIFY_CRC", true)
.to_lowercase()
.as_ref()
{
"yes" | "true" | "1" | "y" | "t" => true,
"no" | "false" | "0" | "n" | "f" => false,
other => bail!(
"Invalid LIBSQL_BOTTOMLESS_VERIFY_CRC environment variable: {}",
other
),
};
let skip_snapshot = match env_var_or("LIBSQL_BOTTOMLESS_SKIP_SNAPSHOT", false)
.to_lowercase()
.as_ref()
{
"yes" | "true" | "1" | "y" | "t" => true,
"no" | "false" | "0" | "n" | "f" => false,
other => bail!(
"Invalid LIBSQL_BOTTOMLESS_SKIP_SNAPSHOT environment variable: {}",
other
),
};
let s3_max_retries = env_var_or("LIBSQL_BOTTOMLESS_S3_MAX_RETRIES", 10).parse::<u32>()?;
let cipher = match encryption_cipher {
Some(cipher) => Cipher::from_str(&cipher)?,
None => Cipher::default(),
};
let encryption_config = match encryption_key {
Some(key) => Some(EncryptionConfig::new(cipher, key)),
None => None,
};
let skip_shutdown_upload =
env_var_or("LIBSQL_BOTTOMLESS_SKIP_SHUTDOWN_UPLOAD", false).parse::<bool>()?;
Ok(Options {
db_id,
create_bucket_if_not_exists: true,
verify_crc,
use_compression,
encryption_config,
max_batch_interval,
max_frames_per_batch,
s3_max_parallelism,
aws_endpoint,
access_key_id,
secret_access_key,
session_token,
region,
bucket_name,
s3_max_retries,
skip_snapshot,
skip_shutdown_upload,
})
}
}
impl Replicator {
pub const UNSET_PAGE_SIZE: usize = usize::MAX;
pub async fn new<S: Into<String>>(db_path: S) -> Result<Self> {
Self::with_options(db_path, Options::from_env()?).await
}
fn set_local_last_frame_no(db_name: &str, last_frame_no: u32) {
let db_name = db_name.to_string();
gauge!("bottomless_local_last_frame_no", last_frame_no as f64, "db_name" => db_name);
}
fn increment_local_ready_frame_ranges(db_name: &str, ready_ranges: u32) {
let db_name = db_name.to_string();
counter!("bottomless_local_ready_frame_ranges", ready_ranges as u64, "db_name" => db_name);
}
fn record_local_flush_time(db_name: &str, duration: Duration) {
let db_name = db_name.to_string();
histogram!("bottomless_local_flush_time", duration.as_secs_f64(), "db_name" => db_name);
}
fn record_s3_write_time(db_name: &str, duration: Duration) {
let db_name = db_name.to_string();
histogram!("bottomless_s3_write_time", duration.as_secs_f64(), "db_name" => db_name);
}
fn increment_s3_processed_frame_ranges(db_name: &str, processed: u64) {
let db_name = db_name.to_string();
counter!("bottomless_s3_processed_frame_ranges", processed, "db_name" => db_name);
}
fn record_snapshot_upload_time(db_name: &str, duration: Duration) {
let db_name = db_name.to_string();
histogram!("bottomless_snapshot_upload_time", duration.as_secs_f64(), "db_name" => db_name);
}
fn record_restore_upload_files_time(db_name: &str, duration: Duration) {
let db_name = db_name.to_string();
histogram!("bottomless_restore_upload_files_time", duration.as_secs_f64(), "db_name" => db_name);
}
fn record_restore_time(db_name: &str, duration: Duration) {
let db_name = db_name.to_string();
histogram!("bottomless_restore_time", duration.as_secs_f64(), "db_name" => db_name);
}
fn set_s3_processing_frame_no(db_name: &str, frame_no: u32) {
let db_name = db_name.to_string();
gauge!("bottomless_s3_processing_frame_no", frame_no as f64, "db_name" => db_name);
}
fn set_s3_queue_size(db_name: &str, size: usize) {
let db_name = db_name.to_string();
gauge!("bottomless_s3_queue_size", size as f64, "db_name" => db_name);
}
pub async fn with_options<S: Into<String>>(db_path: S, options: Options) -> Result<Self> {
let config = options.client_config().await?;
let client = Client::from_conf(config);
let bucket = options.bucket_name.clone();
let generation = Arc::new(ArcSwapOption::default());
match client.head_bucket().bucket(&bucket).send().await {
Ok(_) => tracing::info!("Bucket {} exists and is accessible", bucket),
Err(SdkError::ServiceError(err)) if err.err().is_not_found() => {
if options.create_bucket_if_not_exists {
tracing::info!("Bucket {} not found, recreating", bucket);
client.create_bucket().bucket(&bucket).send().await?;
} else {
tracing::error!("Bucket {} does not exist", bucket);
return Err(SdkError::ServiceError(err).into());
}
}
Err(e) => {
tracing::error!("Bucket checking error: {}", e);
return Err(e.into());
}
}
let db_path = db_path.into();
let db_name = if let Some(db_id) = options.db_id.clone() {
db_id
} else {
bail!("database id was not set")
};
tracing::debug!("Database path: '{}', name: '{}'", db_path, db_name);
let skip_shutdown_upload = options.skip_shutdown_upload;
if skip_shutdown_upload {
tracing::warn!("skipping upload on shutdown");
}
let (flush_trigger, mut flush_trigger_rx) = channel(());
let (last_committed_frame_no_sender, last_committed_frame_no) = channel(Ok(0));
let next_frame_no = Arc::new(AtomicU32::new(1));
let last_sent_frame_no = Arc::new(AtomicU32::new(0));
let mut join_set = JoinSet::new();
let (shutdown_trigger, shutdown_watch) = tokio::sync::watch::channel(());
let (frames_outbox, mut frames_inbox) = tokio::sync::mpsc::unbounded_channel();
let _local_backup = {
let mut copier = WalCopier::new(
bucket.clone(),
db_name.clone().into(),
generation.clone(),
&db_path,
options.max_frames_per_batch,
options.use_compression,
frames_outbox,
);
let next_frame_no = next_frame_no.clone();
let last_sent_frame_no = last_sent_frame_no.clone();
let batch_interval = options.max_batch_interval;
let db_name = db_name.clone();
join_set.spawn(async move {
loop {
let timeout = Instant::now() + batch_interval;
let trigger = match timeout_at(timeout, flush_trigger_rx.changed()).await {
Ok(Ok(())) => true,
Ok(Err(_)) => {
return;
}
Err(_) => {
true // timeout reached
}
};
if trigger {
let next_frame = next_frame_no.load(Ordering::Acquire);
let last_sent_frame =
last_sent_frame_no.swap(next_frame - 1, Ordering::Acquire);
let frames = (last_sent_frame + 1)..next_frame;
if !frames.is_empty() {
let start_time = Instant::now();
let res = copier.flush(frames).await;
Self::record_local_flush_time(&db_name, start_time.elapsed());
if let Ok((last_frame_no, ready_ranges)) = res {
Self::set_local_last_frame_no(&db_name, last_frame_no);
Self::increment_local_ready_frame_ranges(&db_name, ready_ranges);
}
if last_committed_frame_no_sender
.send(res.map(|r| r.0))
.is_err()
{
// Replicator was probably dropped and therefore corresponding
// receiver has been closed
return;
}
}
}
}
})
};
let (upload_progress, last_uploaded_frame_no) = CompletionProgress::new(0);
let upload_progress = Arc::new(Mutex::new(upload_progress));
let _s3_upload = {
let client = client.clone();
let bucket = options.bucket_name.clone();
let max_parallelism = options.s3_max_parallelism;
let upload_progress = upload_progress.clone();
let db_name = db_name.clone();
let shutdown_watch = Arc::new(shutdown_watch);
join_set.spawn(async move {
let sem = Arc::new(tokio::sync::Semaphore::new(max_parallelism));
let mut join_set = JoinSet::new();
while let Some(req) = frames_inbox.recv().await {
tracing::trace!("Received S3 upload request: {}", req.path);
let start = Instant::now();
let sem = sem.clone();
let permit = sem.acquire_owned().await.unwrap();
let client = client.clone();
let bucket = bucket.clone();
let upload_progress = upload_progress.clone();
if let Some(ref frames) = req.frames {
Self::set_s3_processing_frame_no(&db_name, *frames.end());
}
Self::set_s3_queue_size(&db_name, frames_inbox.len());
let db_name = db_name.clone();
let shutdown_watch = shutdown_watch.clone();
join_set.spawn(async move {
let fpath = format!("{}/{}", &bucket, &req.path);
loop {
let start_time = Instant::now();
let body = ByteStream::from_path(&fpath).await.unwrap();
let response = client
.put_object()
.bucket(&bucket)
.key(&req.path)
.body(body)
.send()
.await;
Self::record_s3_write_time(&db_name, start_time.elapsed());
if response.is_ok() {
break;
}
tracing::error!("Failed to send {} to S3: {}, will retry after 1 second", fpath, response.err().unwrap());
if shutdown_watch.has_changed().is_err() {
tracing::error!("stop retry for failed S3 frames upload because shutdown was requested");
return
}
tokio::time::sleep(Duration::from_millis(1000)).await;
}
tokio::fs::remove_file(&fpath).await.unwrap();
let elapsed = Instant::now() - start;
tracing::debug!("Uploaded to S3: {} in {:?}", fpath, elapsed);
if let Some(frames) = req.frames {
let mut up = upload_progress.lock().await;
up.update(*frames.start(), *frames.end());
Self::increment_s3_processed_frame_ranges(&db_name, 1);
}
drop(permit);
});
}
while join_set.join_next().await.is_some() {}
})
};
let (snapshot_notifier, snapshot_waiter) = channel(Ok(None));
Ok(Self {
client,
bucket,
page_size: Self::UNSET_PAGE_SIZE,
generation,
next_frame_no,
last_sent_frame_no,
flush_trigger: Some(flush_trigger),
shutdown_trigger: Some(shutdown_trigger),
last_committed_frame_no,
verify_crc: options.verify_crc,
db_path,
db_name,
snapshot_waiter,
snapshot_notifier: Arc::new(snapshot_notifier),
use_compression: options.use_compression,
encryption_config: options.encryption_config,
max_frames_per_batch: options.max_frames_per_batch,
s3_max_parallelism: options.s3_max_parallelism,
skip_snapshot: options.skip_snapshot,
join_set,
upload_progress,
last_uploaded_frame_no,
skip_shutdown_upload,
})
}
/// Checks if there exists any backup of given database
pub async fn has_backup_of(db_name: impl AsRef<str>, options: &Options) -> Result<bool> {
let prefix = match &options.db_id {
Some(db_id) => format!("{db_id}-"),
None => format!("ns-:{}-", db_name.as_ref()),
};
let config = options.client_config().await?;
let client = Client::from_conf(config);
let bucket = options.bucket_name.clone();
match client.head_bucket().bucket(&bucket).send().await {
Ok(_) => tracing::trace!("Bucket {bucket} exists and is accessible"),
Err(e) => {
tracing::trace!("Bucket checking error: {e}");
return Err(e.into());
}
}
let mut last_frame = 0;
let list_objects = client.list_objects().bucket(&bucket).prefix(&prefix);
let response = list_objects.send().await?;
let _ = Self::try_get_last_frame_no(response, &mut last_frame);
tracing::trace!("Last frame of {prefix}: {last_frame}");
Ok(last_frame > 0)
}
pub async fn shutdown_gracefully(&mut self) -> Result<()> {
if !self.skip_shutdown_upload {
tracing::info!("bottomless replicator: shutting down...");
// 1. wait for all committed WAL frames to be committed locally
let last_frame_no = self.last_known_frame();
// force flush in order to not wait for periodic wake up of local back up process
if let Some(tx) = &self.flush_trigger {
let _ = tx.send(());
}
self.wait_until_committed(last_frame_no).await?;
tracing::info!(
"bottomless replicator: local backup replicated frames until {}",
last_frame_no
);
// 2. wait for snapshot upload to S3 to finish
self.wait_until_snapshotted().await?;
tracing::info!("bottomless replicator: snapshot succesfully uploaded to S3");
// 3. drop flush trigger, which will cause WAL upload loop to close. Since this action will
// close the channel used by wait_until_committed, it must happen after wait_until_committed
// has finished. If trigger won't be dropped, tasks from join_set will never finish.
self.flush_trigger.take();
// 4. drop shutdown trigger which will notify S3 upload process to stop all retry attempts
// and finish upload process
self.shutdown_trigger.take();
while let Some(t) = self.join_set.join_next().await {
// one of the tasks we're waiting for is upload of local WAL segment from pt.1 to S3
// this should ensure that all WAL frames are one S3
t?;
}
} else {
tracing::warn!("skipping snapshot upload during shutdown");
}
tracing::info!("bottomless replicator: shutdown complete");
Ok(())
}
pub fn next_frame_no(&self) -> u32 {
self.next_frame_no.load(Ordering::Acquire)
}
pub fn last_known_frame(&self) -> u32 {
self.next_frame_no() - 1
}
pub fn last_sent_frame_no(&self) -> u32 {
self.last_sent_frame_no.load(Ordering::Acquire)
}
pub fn compression_kind(&self) -> CompressionKind {
self.use_compression
}
pub async fn is_snapshotted(&mut self) -> bool {
if let Ok(generation) = self.generation() {
if !self.main_db_exists_and_not_empty().await {
tracing::debug!("Not snapshotting, the main db file does not exist or is empty");
let _ = self.snapshot_notifier.send(Ok(Some(generation)));
return false;
}
tracing::debug!("waiting for generation snapshot {} to complete", generation);
let current = self.snapshot_waiter.borrow();
match &*current {
Ok(Some(gen)) => *gen == generation,
_ => false,
}
} else {
false
}
}
/// FIXME: I'm pretty sure that this function is buggy. First of all we don't check the output
/// of that function in the wal, second of all, we assume that an error means that we have
/// snapshotted. This whole stuff is a mess.
pub async fn wait_until_snapshotted(&mut self) -> Result<bool> {
if let Ok(generation) = self.generation() {
if !self.main_db_exists_and_not_empty().await {
tracing::debug!("Not snapshotting, the main db file does not exist or is empty");
let _ = self.snapshot_notifier.send(Ok(Some(generation)));
return Ok(false);
}
tracing::debug!("waiting for generation snapshot {} to complete", generation);
let res = self
.snapshot_waiter
.wait_for(|result| match result {
Ok(Some(gen)) => *gen == generation,
Ok(None) => false,
Err(_) => true,
})
.await?;
tracing::debug!("done waiting");
match res.deref() {
Ok(_) => Ok(true),
Err(e) => Err(anyhow!("Failed snapshot generation {}: {}", generation, e)),
}
} else {
Ok(false)
}
}
async fn reset_wal_tracker(&self) {
let mut lock = self.upload_progress.lock().await;
lock.reset();
}
pub fn savepoint(&self) -> SavepointTracker {
tracing::debug!(
"calling for backup savepoint for `{}` on generation `{}`, frame no.: {}",
self.db_name,
match self.generation() {
Ok(gen) => gen.to_string(),
_ => "".to_string(),
},
self.next_frame_no() - 1
);
if let Some(tx) = &self.flush_trigger {
let _ = tx.send(());
}
SavepointTracker::new(
self.generation.clone(),
self.snapshot_waiter.clone(),
self.next_frame_no.clone(),
self.last_uploaded_frame_no.clone(),
self.db_path.clone(),
)
}
/// Waits until the commit for a given frame_no or higher was given.
pub async fn wait_until_committed(&mut self, frame_no: u32) -> Result<u32> {
let res = self
.last_committed_frame_no
.wait_for(|result| match result {
Ok(last_committed) => *last_committed >= frame_no,
Err(_) => true,
})
.await?;
match res.deref() {
Ok(last_committed) => {
tracing::trace!(
"Confirmed commit of frame no. {} (waited for >= {})",
last_committed,
frame_no
);
Ok(*last_committed)
}
Err(e) => Err(anyhow!("Failed to flush frames: {}", e)),
}
}
/// Returns number of frames waiting to be replicated.
pub fn pending_frames(&self) -> u32 {
self.next_frame_no() - self.last_sent_frame_no() - 1
}
// The database can use different page size - as soon as it's known,
// it should be communicated to the replicator via this call.
// NOTICE: in practice, WAL journaling mode does not allow changing page sizes,
// so verifying that it hasn't changed is a panic check. Perhaps in the future
// it will be useful, if WAL ever allows changing the page size.
pub fn set_page_size(&mut self, page_size: usize) -> Result<()> {
if self.page_size != page_size {
tracing::trace!("Setting page size to: {}", page_size);
}
if self.page_size != Self::UNSET_PAGE_SIZE && self.page_size != page_size {
return Err(anyhow::anyhow!(
"Cannot set page size to {}, it was already set to {}",
page_size,
self.page_size
));
}
self.page_size = page_size;
Ok(())
}
// Gets an object from the current bucket
fn get_object(&self, key: String) -> GetObjectFluentBuilder {
self.client.get_object().bucket(&self.bucket).key(key)
}
// Lists objects from the current bucket
fn list_objects(&self) -> ListObjectsFluentBuilder {
self.client.list_objects().bucket(&self.bucket)
}
fn reset_frames(&mut self, frame_no: u32) {
let last_sent = self.last_sent_frame_no();
self.next_frame_no.store(frame_no + 1, Ordering::Release);
self.last_sent_frame_no
.store(last_sent.min(frame_no), Ordering::Release);
}
// Generates a new generation UUID v7, which contains a timestamp and is binary-sortable.
// This timestamp goes back in time - that allows us to list newest generations
// first in the S3-compatible bucket, under the assumption that fetching newest generations
// is the most common operation.
// NOTICE: at the time of writing, uuid v7 is an unstable feature of the uuid crate
fn generate_generation() -> Uuid {
let ts = uuid::timestamp::Timestamp::now(uuid::NoContext);
Self::generation_from_timestamp(ts)
}
fn generation_from_timestamp(ts: uuid::Timestamp) -> Uuid {
let (seconds, nanos) = ts.to_unix();
let (seconds, nanos) = (253370761200 - seconds, 999999999 - nanos);
let synthetic_ts = uuid::Timestamp::from_unix(uuid::NoContext, seconds, nanos);
crate::uuid_utils::new_v7(synthetic_ts)
}
pub fn generation_to_timestamp(generation: &Uuid) -> Option<uuid::Timestamp> {
let ts = decode_unix_timestamp(generation);
let (seconds, nanos) = ts.to_unix();
let (seconds, nanos) = (253370761200 - seconds, 999999999 - nanos);
Some(uuid::Timestamp::from_unix(NoContext, seconds, nanos))
}
// Starts a new generation for this replicator instance
pub async fn new_generation(&mut self) -> Option<Uuid> {
self.reset_wal_tracker().await;
let curr = Self::generate_generation();
let prev = self.set_generation(curr);
if let Some(prev) = prev {
if prev != curr {
// try to store dependency between previous and current generation
tracing::trace!("New generation {} (parent: {})", curr, prev);
self.store_dependency(prev, curr)
}
}
prev
}
// Sets a generation for this replicator instance. This function
// should be called if a generation number from S3-compatible storage
// is reused in this session.
pub fn set_generation(&mut self, generation: Uuid) -> Option<Uuid> {
let prev_generation = self.generation.swap(Some(Arc::new(generation)));
self.reset_frames(0);
if let Some(prev) = prev_generation.as_deref() {
tracing::debug!("Generation changed from {} -> {}", prev, generation);
Some(*prev)
} else {
tracing::debug!("Generation set {}", generation);
None
}
}
pub fn generation(&self) -> Result<Uuid> {
let guard = self.generation.load();
guard
.as_deref()
.cloned()
.ok_or(anyhow!("Replicator generation was not initialized"))
}
/// Request to store dependency between current generation and its predecessor on S3 object.
/// This works asynchronously on best-effort rules, as putting object to S3 introduces an
/// extra undesired latency and this method may be called during SQLite checkpoint.
fn store_dependency(&self, prev: Uuid, curr: Uuid) {
let key = format!("{}-{}/.dep", self.db_name, curr);
let request =
self.client
.put_object()
.bucket(&self.bucket)
.key(key)
.body(ByteStream::from(Bytes::copy_from_slice(
prev.into_bytes().as_slice(),
)));
tokio::spawn(async move {
if let Err(e) = request.send().await {
tracing::error!(
"Failed to store dependency between generations {} -> {}: {}",
prev,
curr,
e
);
} else {
tracing::trace!(
"Stored dependency between parent ({}) and child ({})",
prev,
curr
);
}
});
}
pub async fn get_dependency(&self, generation: &Uuid) -> Result<Option<Uuid>> {
let key = format!("{}-{}/.dep", self.db_name, generation);
let resp = self
.client
.get_object()
.bucket(&self.bucket)
.key(key)
.send()
.await;
match resp {
Ok(out) => {
let bytes = out.body.collect().await?.into_bytes();
let prev_generation = Uuid::from_bytes(bytes.as_ref().try_into()?);
Ok(Some(prev_generation))
}
Err(SdkError::ServiceError(se)) => match se.into_err() {
GetObjectError::NoSuchKey(_) => Ok(None),
e => Err(e.into()),
},
Err(e) => Err(e.into()),
}
}
// Returns the current last valid frame in the replicated log
pub fn peek_last_valid_frame(&self) -> u32 {
self.next_frame_no().saturating_sub(1)
}
// Sets the last valid frame in the replicated log.
pub fn register_last_valid_frame(&mut self, frame: u32) {
let last_valid_frame = self.peek_last_valid_frame();
if frame != last_valid_frame {
// If frame >= last_valid_frame, it comes from a transaction large enough
// that it got split to multiple xFrames calls. In this case, we just
// update the last_valid_frame to this one, all good.
if last_valid_frame != 0 && frame < last_valid_frame {
tracing::error!(
"[BUG] Local max valid frame is {}, while replicator thinks it's {}",
frame,
last_valid_frame
);
}
self.reset_frames(frame);
}
}
/// Submit next `frame_count` of frames to be replicated.
pub fn submit_frames(&mut self, frame_count: u32) {
let prev = self.next_frame_no.fetch_add(frame_count, Ordering::SeqCst);
let last_sent = self.last_sent_frame_no();
let most_recent = prev + frame_count - 1;
if most_recent - last_sent >= self.max_frames_per_batch as u32 {
self.request_flush();
}
}
pub fn request_flush(&self) {
if let Some(tx) = self.flush_trigger.as_ref() {
tracing::trace!("Requesting flush");
let _ = tx.send(());
} else {
tracing::warn!("Cannot request flush - replicator is closing");
}
}
// Drops uncommitted frames newer than given last valid frame
pub fn rollback_to_frame(&mut self, last_valid_frame: u32) {
// NOTICE: O(size), can be optimized to O(removed) if ever needed
self.reset_frames(last_valid_frame);
tracing::debug!("Rolled back to {}", last_valid_frame);
}
// Opens a raw libSQL connection that doesn't checkpoint.
// Useful for reading metadata from the database file.
fn open_db(&self) -> Result<libsql_sys::Connection<libsql_sys::wal::Sqlite3Wal>> {
use libsql_sys::connection::OpenFlags;
use libsql_sys::wal::Sqlite3WalManager;
let flags = OpenFlags::SQLITE_OPEN_READ_WRITE
| OpenFlags::SQLITE_OPEN_CREATE
| OpenFlags::SQLITE_OPEN_URI
| OpenFlags::SQLITE_OPEN_NO_MUTEX;
let conn = libsql_sys::Connection::open(
&self.db_path,
flags,
Sqlite3WalManager::new(),
libsql_sys::connection::NO_AUTOCHECKPOINT, // no checkpointing
self.encryption_config.clone(),
)?;
Ok(conn)
}
// Tries to read the local change counter from the given database file
fn read_change_counter(&self) -> Result<[u8; 4]> {
if !<str as AsRef<Path>>::as_ref(self.db_path.as_str()).try_exists()? {
return Ok([0; 4]);
}
let conn = self.open_db()?;
let change_counter = conn.db_change_counter().map_err(|rc| {
anyhow::anyhow!(
"Failed to read local change counter from `{}`: {rc}",
self.db_path,
)
})?;
tracing::trace!("Local change counter: {change_counter}");
// TODO: we shouldn't leak the connection here but for some reason when this connection get
// dropped it seems to checkpoint the database
if std::env::var("LIBSQL_BOTTOMLESS_DISABLE_INIT_CHECKPOINTING").is_ok()
|| std::env::var("LIBSQL_DISABLE_INIT_CHECKPOINTING").is_ok()
{
std::mem::forget(conn);
}
Ok(change_counter.to_be_bytes())
}
// Tries to read the local page size from the given database file
async fn read_page_size(&self) -> Result<usize> {
let conn = self.open_db()?;
let page_size = conn.query_row("PRAGMA page_size", (), |r| r.get::<usize, usize>(0))?;
tracing::trace!("Local page size: {page_size}");
Ok(page_size)
}
// Returns the compressed database file path and its change counter, extracted
// from the header of page1 at offset 24..27 (as per SQLite documentation).
pub async fn maybe_compress_main_db_file(
db_path: &Path,
compression: CompressionKind,
) -> Result<ByteStream> {
if !tokio::fs::try_exists(db_path).await? {
bail!("database file was not found at `{}`", db_path.display())
}
match compression {
CompressionKind::None => Ok(ByteStream::from_path(db_path).await?),
CompressionKind::Gzip => {
let mut reader = File::open(db_path).await?;
let gzip_path = Self::db_compressed_path(db_path, "gz");
let compressed_file = OpenOptions::new()
.create(true)
.write(true)
.read(true)
.truncate(true)
.open(&gzip_path)
.await?;
let mut writer = GzipEncoder::new(compressed_file);
let size = tokio::io::copy(&mut reader, &mut writer).await?;
writer.shutdown().await?;
tracing::debug!(
"Compressed database file ({} bytes) into `{}`",
size,
gzip_path.display()
);
Ok(ByteStream::from_path(gzip_path).await?)
}
CompressionKind::Zstd => {
let mut reader = File::open(db_path).await?;
let zstd_path = Self::db_compressed_path(db_path, "zstd");
let compressed_file = OpenOptions::new()
.create(true)
.write(true)
.read(true)
.truncate(true)
.open(&zstd_path)
.await?;
let mut writer = ZstdEncoder::new(compressed_file);
let size = tokio::io::copy(&mut reader, &mut writer).await?;
writer.shutdown().await?;
tracing::debug!(
"Compressed database file ({} bytes) into `{}`",
size,
zstd_path.display()
);
Ok(ByteStream::from_path(zstd_path).await?)
}
}
}
fn db_compressed_path(db_path: &Path, suffix: &'static str) -> PathBuf {
let mut compressed_path: PathBuf = db_path.to_path_buf();
compressed_path.pop();
compressed_path.join(format!("db.{suffix}"))
}
fn restore_db_path(&self) -> PathBuf {
let mut gzip_path = PathBuf::from(&self.db_path);
gzip_path.pop();
gzip_path.join("data.tmp")
}
// Replicates local WAL pages to S3, if local WAL is present.
// This function is called under the assumption that if local WAL
// file is present, it was already detected to be newer than its
// remote counterpart.
pub async fn maybe_replicate_wal(&mut self) -> Result<()> {
let wal = match WalFileReader::open(&format!("{}-wal", &self.db_path)).await {
Ok(Some(file)) => file,
_ => {
tracing::info!("Local WAL not present - not replicating");
return Ok(());
}
};
self.store_metadata(wal.page_size(), wal.checksum()).await?;
let frame_count = wal.frame_count().await;
tracing::trace!("Local WAL pages: {}", frame_count);
self.submit_frames(frame_count);
self.request_flush();
let last_written_frame = self.wait_until_committed(frame_count - 1).await?;
tracing::info!("Backed up WAL frames up to {}", last_written_frame);
let pending_frames = self.pending_frames();
if pending_frames != 0 {
tracing::warn!(
"Uncommitted WAL entries: {} frames in total",
pending_frames
);
}
tracing::info!("Local WAL replicated");
Ok(())
}
// Check if the local database file exists and contains data
async fn main_db_exists_and_not_empty(&self) -> bool {
let file = match File::open(&self.db_path).await {
Ok(file) => file,
Err(_) => return false,
};
match file.metadata().await {
Ok(metadata) => metadata.len() > 0,
Err(_) => false,
}
}
pub fn skip_snapshot_for_current_generation(&self) {
let generation = self.generation.load().as_deref().cloned();
let _ = self.snapshot_notifier.send(Ok(generation));
}
// Sends the main database file to S3 - if -wal file is present, it's replicated
// too - it means that the local file was detected to be newer than its remote
// counterpart.
pub async fn snapshot_main_db_file(&mut self, force: bool) -> Result<Option<JoinHandle<()>>> {
if self.skip_snapshot && !force {
tracing::trace!("database snapshot skipped");
let _ = self.snapshot_notifier.send(Ok(self.generation().ok()));
return Ok(None);
}
tracing::debug!("snapshotting db file");
if !self.main_db_exists_and_not_empty().await {
let generation = self.generation()?;
tracing::debug!(
"Not snapshotting {}, the main db file does not exist or is empty",
generation
);
let _ = self.snapshot_notifier.send(Ok(Some(generation)));
return Ok(None);
}
let generation = self.generation()?;
let start_ts = Instant::now();
let client = self.client.clone();
let change_counter = self.read_change_counter()?;
let snapshot_req = client.put_object().bucket(self.bucket.clone()).key(format!(
"{}-{}/db.{}",
self.db_name, generation, self.use_compression
));
/* FIXME: we can't rely on the change counter in WAL mode:
** "In WAL mode, changes to the database are detected using the wal-index and
** so the change counter is not needed. Hence, the change counter might not be
** incremented on each transaction in WAL mode."
** Instead, we need to consult WAL checksums.
*/
let change_counter_key = format!("{}-{}/.changecounter", self.db_name, generation);
let change_counter_req = self
.client
.put_object()
.bucket(&self.bucket)
.key(change_counter_key)
.body(ByteStream::from(Bytes::copy_from_slice(
change_counter.as_ref(),
)));
let snapshot_notifier = self.snapshot_notifier.clone();
let compression = self.use_compression;
let db_path = PathBuf::from(self.db_path.clone());
let db_name = self.db_name.clone();
let handle = tokio::spawn(async move {
tracing::trace!("Start snapshotting generation {}", generation);
let start = Instant::now();
let body = match Self::maybe_compress_main_db_file(&db_path, compression).await {
Ok(file) => file,
Err(e) => {
tracing::error!(
"Failed to compress db file (generation `{}`, path: `{}`): {:?}",
generation,
db_path.display(),
e
);
let _ = snapshot_notifier.send(Err(e));
return;
}
};
let mut result = snapshot_req.body(body).send().await;
if let Err(e) = result {
tracing::error!(
"Failed to upload snapshot for generation {}: {:?}",
generation,
e
);
let _ = snapshot_notifier.send(Err(e.into()));
return;
}
result = change_counter_req.send().await;
if let Err(e) = result {
tracing::error!(
"Failed to upload change counter for generation {}: {:?}",
generation,
e
);
let _ = snapshot_notifier.send(Err(e.into()));
return;
}
let _ = snapshot_notifier.send(Ok(Some(generation)));
let elapsed = Instant::now() - start;
tracing::info!("Snapshot upload finished (took {:?})", elapsed);
Self::record_snapshot_upload_time(&db_name, elapsed);
// cleanup gzip/zstd database snapshot if exists
for suffix in &["gz", "zstd"] {
let _ = tokio::fs::remove_file(Self::db_compressed_path(&db_path, suffix)).await;
}
});
let elapsed = Instant::now() - start_ts;
tracing::debug!("Scheduled DB snapshot {} (took {:?})", generation, elapsed);
Ok(Some(handle))
}
// Returns newest replicated generation, or None, if one is not found.
// FIXME: assumes that this bucket stores *only* generations for databases,
// it should be more robust and continue looking if the first item does not
// match the <db-name>-<generation-uuid>/ pattern.
#[tracing::instrument(skip(self))]
pub async fn latest_generation_before(
&self,
timestamp: Option<&NaiveDateTime>,
) -> Option<Uuid> {
tracing::debug!("last generation before");
let mut next_marker: Option<String> = None;
let prefix = format!("{}-", self.db_name);
let threshold = timestamp.map(|ts| ts.and_utc().timestamp() as u64);
loop {
let mut request = self.list_objects().prefix(prefix.clone());
if threshold.is_none() {
request = request.max_keys(1);
}
if let Some(marker) = next_marker.take() {
request = request.marker(marker);
}
let response = request.send().await.ok()?;
let objs = response.contents();
if objs.is_empty() {
tracing::debug!("no objects found in bucket");
break;
}
let mut last_key = None;
let mut last_gen = None;
for obj in objs {
let key = obj.key();
last_key = key;
if let Some(key) = last_key {
let key = match key.find('/') {
Some(index) => &key[self.db_name.len() + 1..index],
None => key,
};
if Some(key) != last_gen {
last_gen = Some(key);
if let Ok(generation) = Uuid::parse_str(key) {
match threshold.as_ref() {
None => return Some(generation),
Some(threshold) => match Self::generation_to_timestamp(&generation)
{
None => {
tracing::warn!(
"Generation {} is not valid UUID v7",
generation
);
}
Some(ts) => {
let (unix_seconds, _) = ts.to_unix();
if tracing::enabled!(tracing::Level::DEBUG) {
let ts = Utc
.timestamp_millis_opt((unix_seconds * 1000) as i64)
.unwrap()
.to_rfc3339();
tracing::debug!(
"Generation candidate: {} - timestamp: {}",
generation,
ts
);
}
if &unix_seconds <= threshold {
return Some(generation);
}
}
},
}
}
}
}
}
next_marker = last_key.map(String::from);
}
None
}
// Tries to fetch the remote database change counter from given generation
pub async fn get_remote_change_counter(&self, generation: &Uuid) -> Result<[u8; 4]> {
let mut remote_change_counter = [0u8; 4];
if let Ok(response) = self
.get_object(format!("{}-{}/.changecounter", self.db_name, generation))
.send()
.await
{
response
.body
.collect()
.await?
.copy_to_slice(&mut remote_change_counter)
}
Ok(remote_change_counter)
}
// Returns the number of pages stored in the local WAL file, or 0, if there aren't any.
async fn get_local_wal_page_count(&mut self) -> u32 {
match WalFileReader::open(&format!("{}-wal", &self.db_path)).await {
Ok(None) => 0,
Ok(Some(wal)) => {
let page_size = wal.page_size();
if self.set_page_size(page_size as usize).is_err() {
return 0;
}
wal.frame_count().await
}
Err(_) => 0,
}
}
// Parses the frame and page number from given key.
// Format: <db-name>-<generation>/<first-frame-no>-<last-frame-no>-<timestamp>.<compression-kind>
pub fn parse_frame_range(key: &str) -> Option<(u32, u32, u64, CompressionKind)> {
let frame_delim = key.rfind('/')?;
let frame_suffix = &key[(frame_delim + 1)..];
let timestamp_delim = frame_suffix.rfind('-')?;
let last_frame_delim = frame_suffix[..timestamp_delim].rfind('-')?;
let compression_delim = frame_suffix.rfind('.')?;
let first_frame_no = frame_suffix[0..last_frame_delim].parse::<u32>().ok()?;
let last_frame_no = frame_suffix[(last_frame_delim + 1)..timestamp_delim]
.parse::<u32>()
.ok()?;
let timestamp = frame_suffix[(timestamp_delim + 1)..compression_delim]
.parse::<u64>()
.ok()?;
let compression_kind =
CompressionKind::parse(&frame_suffix[(compression_delim + 1)..]).ok()?;
Some((first_frame_no, last_frame_no, timestamp, compression_kind))
}
/// Restores the database state from given remote generation
/// On success, returns the RestoreAction, and whether the database was recovered from backup.
#[tracing::instrument(skip(self))]
async fn restore_from(
&mut self,
generation: Uuid,
timestamp: Option<NaiveDateTime>,
) -> Result<(RestoreAction, bool)> {
tracing::debug!("restoring from");
if let Some(tombstoned_at) = self.is_tombstoned(generation).await? {
bail!(
"Couldn't restore from generation {}. Database '{}' has been tombstoned at {}.",
generation,
self.db_name,
tombstoned_at
);
}
let start_ts = Instant::now();
// first check if there are any remaining files that we didn't manage to upload
// on time in the last run
self.upload_remaining_files(&generation).await?;
Self::record_restore_upload_files_time(&self.db_name, start_ts.elapsed());
tracing::debug!("done uploading remaining files");
let last_frame = self.get_last_consistent_frame(&generation).await?;
tracing::debug!("Last consistent remote frame in generation {generation}: {last_frame}.");
if let Some(action) = self.compare_with_local(generation, last_frame).await? {
return Ok((action, false));
}
// at this point we know, we should do a full restore
let restore_path = self.restore_db_path();
let _ = tokio::fs::remove_file(&restore_path).await; // remove previous (failed) restoration
match self
.full_restore(&restore_path, generation, timestamp, last_frame)
.await
{
Ok(result) => {
tokio::fs::rename(&restore_path, &self.db_path).await?;
let _ = self.remove_wal_files().await; // best effort, WAL files may not exists
let elapsed = Instant::now() - start_ts;
tracing::info!("Finished database restoration in {:?}", elapsed);
Self::record_restore_time(&self.db_name, elapsed);
Ok(result)
}
Err(e) => {
tracing::error!("failed to restore the database: {}. Rollback", e);
let _ = tokio::fs::remove_file(restore_path).await;
Err(e)
}
}
}
async fn full_restore(
&mut self,
restore_path: &Path,
generation: Uuid,
timestamp: Option<NaiveDateTime>,
last_frame: u32,
) -> Result<(RestoreAction, bool)> {
tracing::debug!("Restoring database to `{}`", restore_path.display());
let mut db = OpenOptions::new()
.create(true)
.read(true)
.write(true)
.open(restore_path)
.await?;
let mut restore_stack = Vec::new();
// If the db file is not present, the database could have been empty
let mut current = Some(generation);
while let Some(curr) = current.take() {
// stash current generation - we'll use it to replay WAL across generations since the
// last snapshot
restore_stack.push(curr);
let restored = self.restore_from_snapshot(&curr, &mut db).await?;
if restored {
break;
} else {
if restore_stack.len() > MAX_RESTORE_STACK_DEPTH {
bail!("Restoration failed: maximum number of generations to restore from was reached.");
}
tracing::debug!("No snapshot found on the generation {}", curr);
// there was no snapshot to restore from, it means that we either:
// 1. Have only WAL to restore from - case when we're at the initial generation
// of the database.
// 2. Snapshot never existed - in that case try to reach for parent generation
// of the current one and read snapshot from there.
current = self.get_dependency(&curr).await?;
if let Some(prev) = ¤t {
tracing::debug!("Rolling restore back from generation {} to {}", curr, prev);
}
}
}
tracing::trace!(
"Restoring database from {} generations",
restore_stack.len()
);
let mut applied_wal_frame = false;
while let Some(gen) = restore_stack.pop() {
if let Some((page_size, checksum)) = self.get_metadata(&gen).await? {
self.set_page_size(page_size as usize)?;
let last_frame = if restore_stack.is_empty() {
// we're at the last generation to restore from, it may still being written to
// so we constraint the restore to a frame checked at the beginning of the
// restore procedure
Some(last_frame)
} else {
None
};
self.restore_wal(
&gen,
page_size as usize,
last_frame,
checksum,
timestamp,
restore_path,
)
.await?;
applied_wal_frame = true;
} else {
tracing::info!(".meta object not found, skipping WAL restore.");
};
}
db.shutdown().await?;
{
let mut guard = self.upload_progress.lock().await;
guard.update(1, last_frame);
}
if applied_wal_frame {
tracing::info!("WAL file has been applied onto database file in generation {}. Requesting snapshot.", generation);
Ok::<_, anyhow::Error>((RestoreAction::SnapshotMainDbFile, true))
} else {
tracing::info!("Reusing generation {}.", generation);
// since WAL was not applied, we can reuse the latest generation
Ok::<_, anyhow::Error>((RestoreAction::ReuseGeneration(generation), true))
}
}
/// Compares S3 generation backup state against current local database file to determine
/// if we are up to date (returned restore action) or should we perform restoration.
async fn compare_with_local(
&mut self,
generation: Uuid,
last_consistent_frame: u32,
) -> Result<Option<RestoreAction>> {
// Check if the database needs to be restored by inspecting the database
// change counter and the WAL size.
let local_counter = self.read_change_counter().unwrap_or([0u8; 4]);
if local_counter != [0u8; 4] && local_counter != [0, 0, 0, 1] {
// if a non-empty database file exists always treat it as new and more up to date,
// skipping the restoration process and calling for a new generation to be made
return Ok(Some(RestoreAction::SnapshotMainDbFile));
}
let remote_counter = self.get_remote_change_counter(&generation).await?;
tracing::debug!("Counters: l={:?}, r={:?}", local_counter, remote_counter);
let wal_pages = self.get_local_wal_page_count().await;
// We impersonate as a given generation, since we're comparing against local backup at that
// generation. This is used later in [Self::new_generation] to create a dependency between
// this generation and a new one.
self.generation.store(Some(Arc::new(generation)));
match local_counter.cmp(&remote_counter) {
std::cmp::Ordering::Equal => {
tracing::debug!(
"Consistent: {}; wal pages: {}",
last_consistent_frame,
wal_pages
);
match wal_pages.cmp(&last_consistent_frame) {
std::cmp::Ordering::Equal => {
if local_counter == [0u8; 4] && wal_pages == 0 {
if self.get_dependency(&generation).await?.is_some() {
// empty generation and empty local state, but we have a dependency
// to previous generation: restore required
return Ok(None);
}
}
tracing::info!(
"Remote generation is up-to-date, reusing it in this session"
);
self.reset_frames(wal_pages + 1);
Ok(Some(RestoreAction::ReuseGeneration(generation)))
}
std::cmp::Ordering::Greater => {
tracing::info!("Local change counter matches the remote one, but local WAL contains newer data from generation {}, which needs to be replicated.", generation);
Ok(Some(RestoreAction::SnapshotMainDbFile))
}
std::cmp::Ordering::Less => Ok(None),
}
}
std::cmp::Ordering::Greater => {
tracing::info!("Local change counter is larger than its remote counterpart - a new snapshot needs to be replicated (generation: {})", generation);
Ok(Some(RestoreAction::SnapshotMainDbFile))
}
std::cmp::Ordering::Less => Ok(None),
}
}
async fn restore_from_snapshot(&mut self, generation: &Uuid, db: &mut File) -> Result<bool> {
let algos_to_try = match self.use_compression {
CompressionKind::None => &[
CompressionKind::None,
CompressionKind::Zstd,
CompressionKind::Gzip,
],
CompressionKind::Gzip => &[
CompressionKind::Gzip,
CompressionKind::Zstd,
CompressionKind::None,
],
CompressionKind::Zstd => &[
CompressionKind::Zstd,
CompressionKind::Gzip,
CompressionKind::None,
],
};
for algo in algos_to_try {
let main_db_path = match algo {
CompressionKind::None => format!("{}-{}/db.raw", self.db_name, generation),
CompressionKind::Gzip => format!("{}-{}/db.gz", self.db_name, generation),
CompressionKind::Zstd => format!("{}-{}/db.zstd", self.db_name, generation),
};
if let Ok(db_file) = self.get_object(main_db_path).send().await {
let mut body_reader = db_file.body.into_async_read();
let db_size = match algo {
CompressionKind::None => tokio::io::copy(&mut body_reader, db).await?,
CompressionKind::Gzip => {
let mut decompress_reader =
async_compression::tokio::bufread::GzipDecoder::new(
tokio::io::BufReader::new(body_reader),
);
tokio::io::copy(&mut decompress_reader, db).await?
}
CompressionKind::Zstd => {
let mut decompress_reader =
async_compression::tokio::bufread::ZstdDecoder::new(
tokio::io::BufReader::new(body_reader),
);
tokio::io::copy(&mut decompress_reader, db).await?
}
};
db.flush().await?;
let page_size = self.read_page_size().await?;
self.set_page_size(page_size)?;
tracing::info!("Restored the main database file ({} bytes)", db_size);
return Ok(true);
}
}
Ok(false)
}
async fn restore_wal(
&self,
generation: &Uuid,
page_size: usize,
last_consistent_frame: Option<u32>,
mut checksum: (u32, u32),
utc_time: Option<NaiveDateTime>,
db_path: &Path,
) -> Result<bool> {
let encryption_config = self.encryption_config.clone();
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);
v.spare_capacity_mut();
unsafe { v.set_len(page_size) };
v
};
let mut next_marker = None;
let mut applied_wal_frame = false;
let mut last_injected_frame_no = 0;
'restore_wal: loop {
let mut list_request = self.list_objects().prefix(&prefix);
if let Some(marker) = next_marker {
list_request = list_request.marker(marker);
}
let response = list_request.send().await?;
let objs = response.contents();
if objs.is_empty() {
tracing::debug!("No objects found in generation {}", generation);
break;
}
for obj in objs {
let key = obj
.key()
.ok_or_else(|| anyhow::anyhow!("Failed to get key for an object"))?;
tracing::debug!("Loading {}", key);
let (first_frame_no, last_frame_no, timestamp, compression_kind) =
match Self::parse_frame_range(key) {
Some(result) => result,
None => {
if !key.ends_with(".gz")
&& !key.ends_with(".zstd")
&& !key.ends_with(".db")
&& !key.ends_with(".meta")
&& !key.ends_with(".dep")
&& !key.ends_with(".changecounter")
{
tracing::warn!("Failed to parse frame/page from key {}", key);
}
continue;
}
};
if first_frame_no != last_injected_frame_no + 1 {
tracing::warn!("Missing series of consecutive frames. Last applied frame: {}, next found: {}. Stopping the restoration process",
last_injected_frame_no, first_frame_no);
break;
}
if let Some(frame) = last_consistent_frame {
if last_frame_no > frame {
tracing::warn!("Remote log contains frame {} larger than last consistent frame ({}), stopping the restoration process",
last_frame_no, frame);
break;
}
}
if let Some(threshold) = utc_time.as_ref() {
match DateTime::from_timestamp(timestamp as i64, 0).map(|t| t.naive_utc()) {
Some(timestamp) => {
if ×tamp > threshold {
tracing::info!("Frame batch {} has timestamp more recent than expected {}. Stopping recovery.", key, timestamp);
break 'restore_wal; // reached end of restoration timestamp
}
}
_ => {
tracing::trace!("Couldn't parse requested frame batch {} timestamp. Stopping recovery.", key);
break 'restore_wal;
}
}
}
let frame = self.get_object(key.into()).send().await?;
let mut reader = BatchReader::new(
first_frame_no,
frame.body.into_async_read(),
self.page_size,
compression_kind,
);
while let Some(frame) = reader.next_frame_header().await? {
last_injected_frame_no = reader.next_frame_no();
reader.next_page(&mut page_buf).await?;
if self.verify_crc {
checksum = frame.verify(checksum, &page_buf)?;
}
let (crc1, crc2) = frame.crc();
let checksum = (crc1 as u64) << 32 | crc2 as u64;
let frame_to_inject = libsql_replication::frame::Frame::from_parts(
&libsql_replication::frame::FrameHeader {
frame_no: (last_injected_frame_no as u64).into(),
checksum: checksum.into(),
page_no: frame.pgno().into(),
size_after: frame.size_after().into(),
},
page_buf.as_slice(),
);
let frame = RpcFrame {
data: frame_to_inject.bytes(),
timestamp: None,
durable_frame_no: None,
};
injector.inject_frame(frame).await?;
applied_wal_frame = true;
}
}
next_marker = response
.is_truncated()
// This previously was not optional but when upgrading to the s3 sdk to 1.0 we must
// check for this situation, defaulting to true to stop the search seems safe. From
// https://docs.aws.amazon.com/AmazonS3/latest/API/API_ListObjects.html#API_ListObjects_ResponseSyntax
// it looks like this value is always set.
.unwrap_or(true)
.then(|| objs.last().map(|elem| elem.key().unwrap().to_string()))
.flatten();
if next_marker.is_none() {
tracing::trace!("Restored DB from S3 backup using generation {}", generation);
break;
}
}
Ok(applied_wal_frame)
}
async fn remove_wal_files(&self) -> Result<()> {
tracing::debug!("Overwriting any existing WAL file: {}-wal", &self.db_path);
tokio::fs::remove_file(&format!("{}-wal", &self.db_path)).await?;
tokio::fs::remove_file(&format!("{}-shm", &self.db_path)).await?;
Ok(())
}
pub async fn copy(&mut self, generation: Option<Uuid>, to_dir: String) -> Result<()> {
let generation = self
.choose_generation(generation, None)
.await
.ok_or(anyhow!("generation not found"))?;
if let Some(tombstoned_at) = self.is_tombstoned(generation).await? {
tracing::warn!("generation was tombstoned at {tombstoned_at}");
}
std::fs::create_dir_all(PathBuf::from(&to_dir))?;
let prefix = format!("{}-{}/", self.db_name, generation);
let mut marker: Option<String> = None;
tracing::info!(
"ready to copy S3 content from directory {} to local directory {} (parallelism: {})",
&prefix,
&to_dir,
self.s3_max_parallelism
);
loop {
let mut list_request = self.list_objects().prefix(&prefix);
if let Some(marker) = marker.take() {
list_request = list_request.marker(marker);
}
let semaphore = Arc::new(Semaphore::new(self.s3_max_parallelism));
let mut group = JoinSet::new();
let list_response = list_request.send().await?;
for entry in list_response.contents() {
let key = String::from(entry.key().unwrap());
marker = Some(key.clone());
let request = self
.client
.get_object()
.bucket(&self.bucket)
.key(key.clone());
let to_dir = to_dir.clone();
let entry_size = entry.size().unwrap_or(0);
let semaphore = semaphore.clone();
group.spawn(async move {
let acquired = semaphore.acquire().await.unwrap();
if let Ok(response) = request.send().await {
tracing::debug!(
"start copy of entry {} (size {} bytes)",
&key,
entry_size,
);
let entry_name = key.split("/").last().unwrap();
let mut entry_path = PathBuf::from(&to_dir);
entry_path.push(entry_name);
let mut entry_file = OpenOptions::new()
.create(true)
.write(true)
.read(true)
.truncate(true)
.open(entry_path)
.await
.unwrap();
let mut body_reader = response.body.into_async_read();
tokio::io::copy(&mut body_reader, &mut entry_file).await.unwrap();
tracing::debug!("finish copy of entry {}", &key);
}
drop(acquired);
});
}
while let Some(_) = group.join_next().await {}
if !marker.is_some() {
break;
}
}
Ok(())
}
async fn choose_generation(
&mut self,
generation: Option<Uuid>,
timestamp: Option<NaiveDateTime>,
) -> Option<Uuid> {
match generation {
Some(gen) => Some(gen),
None => match self.latest_generation_before(timestamp.as_ref()).await {
Some(gen) => Some(gen),
None => None,
},
}
}
/// Restores the database state from newest remote generation
/// On success, returns the RestoreAction, and whether the database was recovered from backup.
pub async fn restore(
&mut self,
generation: Option<Uuid>,
timestamp: Option<NaiveDateTime>,
) -> Result<(RestoreAction, bool)> {
tracing::debug!("restoring with {generation:?} at {timestamp:?}");
let generation = match self.choose_generation(generation, timestamp).await {
Some(generation) => generation,
None => {
tracing::debug!("No generation found, nothing to restore");
return Ok((RestoreAction::SnapshotMainDbFile, false));
}
};
let (action, recovered) = self.restore_from(generation, timestamp).await?;
let _ = self.snapshot_notifier.send(Ok(Some(generation)));
tracing::info!(
"Restoring from generation {generation}: action={action:?}, recovered={recovered}"
);
Ok((action, recovered))
}
#[tracing::instrument(skip(self))]
pub async fn get_last_consistent_frame(&self, generation: &Uuid) -> Result<u32> {
tracing::debug!("get last consistent frame");
let prefix = format!("{}-{}/", self.db_name, generation);
let mut marker: Option<String> = None;
let mut last_frame = 0;
while {
let mut list_objects = self.list_objects().prefix(&prefix);
if let Some(marker) = marker.take() {
list_objects = list_objects.marker(marker);
}
let response = list_objects.send().await?;
marker = Self::try_get_last_frame_no(response, &mut last_frame);
marker.is_some()
} {}
Ok(last_frame)
}
fn try_get_last_frame_no(response: ListObjectsOutput, frame_no: &mut u32) -> Option<String> {
let objs = response.contents();
let mut last_key = None;
for obj in objs {
last_key = Some(obj.key()?);
if let Some(key) = last_key {
if let Some((_, last_frame_no, _, _)) = Self::parse_frame_range(key) {
*frame_no = last_frame_no;
}
}
}
last_key.map(String::from)
}
async fn upload_remaining_files(&self, generation: &Uuid) -> Result<()> {
let prefix = format!("{}-{}", self.db_name, generation);
let dir = format!("{}/{}-{}", self.bucket, self.db_name, generation);
if tokio::fs::try_exists(&dir).await? {
let mut files = tokio::fs::read_dir(&dir).await?;
let sem = Arc::new(tokio::sync::Semaphore::new(self.s3_max_parallelism));
while let Some(file) = files.next_entry().await? {
let fpath = file.path();
if let Some(key) = Self::fpath_to_key(&fpath, &prefix) {
tracing::trace!("Requesting upload of the remaining backup file: {}", key);
let permit = sem.clone().acquire_owned().await?;
let bucket = self.bucket.clone();
let key = key.to_string();
let client = self.client.clone();
tokio::spawn(async move {
let body = ByteStream::from_path(&fpath).await.unwrap();
if let Err(e) = client
.put_object()
.bucket(bucket)
.key(key.clone())
.body(body)
.send()
.await
{
tracing::error!("Failed to send {} to S3: {}", key, e);
} else {
tokio::fs::remove_file(&fpath).await.unwrap();
tracing::trace!("Uploaded to S3: {}", key);
}
drop(permit);
});
}
}
// wait for all started upload tasks to finish
let _ = sem.acquire_many(self.s3_max_parallelism as u32).await?;
if let Err(e) = tokio::fs::remove_dir(&dir).await {
tracing::warn!("Couldn't remove backed up directory {}: {}", dir, e);
}
}
Ok(())
}
fn fpath_to_key<'a>(fpath: &'a Path, dir: &str) -> Option<&'a str> {
let str = fpath.to_str()?;
if str.ends_with(".db")
| str.ends_with(".gz")
| str.ends_with(".zstd")
| str.ends_with(".raw")
| str.ends_with(".meta")
| str.ends_with(".dep")
| str.ends_with(".changecounter")
{
let idx = str.rfind(dir)?;
return Some(&str[idx..]);
}
None
}
async fn store_metadata(&self, page_size: u32, checksum: (u32, u32)) -> Result<()> {
let generation = self.generation()?;
let key = format!("{}-{}/.meta", self.db_name, generation);
tracing::debug!(
"Storing metadata at '{}': page size - {}, crc - {},{}",
key,
page_size,
checksum.0,
checksum.1,
);
let mut body = Vec::with_capacity(12);
body.extend_from_slice(page_size.to_be_bytes().as_slice());
body.extend_from_slice(checksum.0.to_be_bytes().as_slice());
body.extend_from_slice(checksum.1.to_be_bytes().as_slice());
let _ = self
.client
.put_object()
.bucket(self.bucket.clone())
.key(key)
.body(ByteStream::from(body))
.send()
.await?;
Ok(())
}
pub async fn get_metadata(&self, generation: &Uuid) -> Result<Option<(u32, (u32, u32))>> {
let key = format!("{}-{}/.meta", self.db_name, generation);
if let Ok(obj) = self
.client
.get_object()
.bucket(&self.bucket)
.key(key)
.send()
.await
{
let mut data = obj.body.collect().await?;
let page_size = data.get_u32();
let checksum = (data.get_u32(), data.get_u32());
Ok(Some((page_size, checksum)))
} else {
Ok(None)
}
}
/// Marks current replicator database as deleted, invalidating all generations.
pub async fn delete_all(&self, older_than: Option<NaiveDateTime>) -> Result<DeleteAll> {
tracing::info!(
"Called for tombstoning of all contents of the '{}' database",
self.db_name
);
let key = format!("{}.tombstone", self.db_name);
let threshold = older_than.unwrap_or(NaiveDateTime::MAX);
self.client
.put_object()
.bucket(&self.bucket)
.key(key)
.body(ByteStream::from(
threshold.and_utc().timestamp().to_be_bytes().to_vec(),
))
.send()
.await?;
let delete_task = DeleteAll::new(
self.client.clone(),
self.bucket.clone(),
self.db_name.clone(),
threshold,
);
Ok(delete_task)
}
pub async fn is_tombstoned(&self, generation: Uuid) -> Result<Option<NaiveDateTime>> {
if let Some(tombstone) = self.get_tombstone().await? {
if let Some(timestamp) = Self::generation_to_timestamp(&generation) {
if (tombstone.and_utc().timestamp() as u64) >= timestamp.to_unix().0 {
return Ok(Some(tombstone));
};
}
}
Ok(None)
}
/// Checks if current replicator database has been marked as deleted.
pub async fn get_tombstone(&self) -> Result<Option<NaiveDateTime>> {
let key = format!("{}.tombstone", self.db_name);
let resp = self
.client
.get_object()
.bucket(&self.bucket)
.key(key)
.send()
.await;
match resp {
Ok(out) => {
let mut buf = [0u8; 8];
out.body.collect().await?.copy_to_slice(&mut buf);
let timestamp = i64::from_be_bytes(buf);
let tombstone = DateTime::from_timestamp(timestamp, 0).map(|t| t.naive_utc());
Ok(tombstone)
}
Err(SdkError::ServiceError(se)) => match se.into_err() {
GetObjectError::NoSuchKey(_) => Ok(None),
e => Err(e.into()),
},
Err(e) => Err(e.into()),
}
}
}
/// This structure is returned by [Replicator::delete_all] after tombstoning (soft deletion) has
/// been confirmed. It may be called using [DeleteAll::commit] to trigger a follow up procedure that
/// performs hard deletion of corresponding S3 objects.
#[derive(Debug)]
pub struct DeleteAll {
client: Client,
bucket: String,
db_name: String,
threshold: NaiveDateTime,
}
impl DeleteAll {
fn new(client: Client, bucket: String, db_name: String, threshold: NaiveDateTime) -> Self {
DeleteAll {
client,
bucket,
db_name,
threshold,
}
}
pub fn threshold(&self) -> &NaiveDateTime {
&self.threshold
}
/// Performs hard deletion of all bottomless generations older than timestamp provided in
/// current request.
pub async fn commit(self) -> Result<u32> {
let mut next_marker = None;
let mut removed_count = 0;
loop {
let mut list_request = self
.client
.list_objects()
.bucket(&self.bucket)
.set_delimiter(Some("/".to_string()))
.prefix(&self.db_name);
if let Some(marker) = next_marker {
list_request = list_request.marker(marker)
}
let response = list_request.send().await?;
let prefixes = response.common_prefixes();
if prefixes.is_empty() {
tracing::debug!("no generations found to delete");
return Ok(0);
}
for prefix in prefixes {
if let Some(prefix) = &prefix.prefix {
let prefix = &prefix[self.db_name.len() + 1..prefix.len() - 1];
let uuid = Uuid::try_parse(prefix)?;
if let Some(datetime) = Replicator::generation_to_timestamp(&uuid) {
if datetime.to_unix().0 >= self.threshold.and_utc().timestamp() as u64 {
continue;
}
tracing::debug!("Removing generation {}", uuid);
self.remove(uuid).await?;
removed_count += 1;
}
}
}
next_marker = response.next_marker().map(|s| s.to_owned());
if next_marker.is_none() {
break;
}
}
tracing::debug!("Removed {} generations", removed_count);
self.remove_tombstone().await?;
Ok(removed_count)
}
pub async fn remove_tombstone(&self) -> Result<()> {
let key = format!("{}.tombstone", self.db_name);
self.client
.delete_object()
.bucket(&self.bucket)
.key(key)
.send()
.await?;
Ok(())
}
async fn remove(&self, generation: Uuid) -> Result<()> {
let mut removed = 0;
let mut next_marker = None;
loop {
let mut list_request = self
.client
.list_objects()
.bucket(&self.bucket)
.prefix(format!("{}-{}/", &self.db_name, generation));
if let Some(marker) = next_marker {
list_request = list_request.marker(marker)
}
let response = list_request.send().await?;
let objs = response.contents();
if objs.is_empty() {
return Ok(());
}
for obj in objs {
if let Some(key) = obj.key() {
tracing::trace!("Removing {}", key);
self.client
.delete_object()
.bucket(&self.bucket)
.key(key)
.send()
.await?;
removed += 1;
}
}
next_marker = response.next_marker().map(|s| s.to_owned());
if next_marker.is_none() {
tracing::trace!("Removed {} snapshot generations", removed);
return Ok(());
}
}
}
}
pub struct Context {
pub replicator: Replicator,
pub runtime: tokio::runtime::Runtime,
}
#[derive(Debug, Clone, Copy, Default, Ord, PartialOrd, Eq, PartialEq)]
pub enum CompressionKind {
#[default]
None,
Gzip,
Zstd,
}
impl CompressionKind {
pub fn parse(kind: &str) -> std::result::Result<Self, &str> {
match kind {
"gz" | "gzip" => Ok(CompressionKind::Gzip),
"raw" | "" => Ok(CompressionKind::None),
"zstd" => Ok(CompressionKind::Zstd),
other => Err(other),
}
}
}
impl std::fmt::Display for CompressionKind {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
CompressionKind::None => write!(f, "raw"),
CompressionKind::Gzip => write!(f, "gz"),
CompressionKind::Zstd => write!(f, "zstd"),
}
}
}