From dd3327b7c3d15ceafd5577318aed409152dde372 Mon Sep 17 00:00:00 2001
From: ad hoc <postma.marin@protonmail.com>
Date: Mon, 13 Feb 2023 15:18:20 +0100
Subject: [PATCH 1/3] implement database hard reset

---
 sqld/src/database/libsql.rs          |  20 +--
 sqld/src/database/write_proxy/mod.rs |  56 +++----
 sqld/src/lib.rs                      | 211 +++++++++++++++------------
 3 files changed, 141 insertions(+), 146 deletions(-)

diff --git a/sqld/src/database/libsql.rs b/sqld/src/database/libsql.rs
index 96ce0c0c..cebb4843 100644
--- a/sqld/src/database/libsql.rs
+++ b/sqld/src/database/libsql.rs
@@ -1,7 +1,5 @@
 use std::path::Path;
 use std::str::FromStr;
-#[cfg(feature = "mwal_backend")]
-use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
 
 use crossbeam::channel::RecvTimeoutError;
@@ -132,16 +130,13 @@ macro_rules! ok_or_exit {
 
 fn open_db(
     path: impl AsRef<Path> + Send + 'static,
-    #[cfg(feature = "mwal_backend")] vwal_methods: Option<
-        Arc<Mutex<sqld_libsql_bindings::mwal::ffi::libsql_wal_methods>>,
-    >,
     wal_hook: impl WalHook + Send + Clone + 'static,
     with_bottomless: bool,
 ) -> anyhow::Result<rusqlite::Connection> {
     let mut retries = 0;
     loop {
         #[cfg(feature = "mwal_backend")]
-        let conn_result = match vwal_methods {
+        let conn_result = match crate::VWAL_METHODS.get().unwrap() {
             Some(ref vwal_methods) => crate::libsql::mwal::open_with_virtual_wal(
                 &path,
                 OpenFlags::SQLITE_OPEN_READ_WRITE
@@ -160,6 +155,7 @@ fn open_db(
                 with_bottomless,
             ),
         };
+
         #[cfg(not(feature = "mwal_backend"))]
         let conn_result = crate::libsql::open_with_regular_wal(
             &path,
@@ -200,23 +196,13 @@ fn open_db(
 impl LibSqlDb {
     pub fn new(
         path: impl AsRef<Path> + Send + 'static,
-        #[cfg(feature = "mwal_backend")] vwal_methods: Option<
-            Arc<Mutex<sqld_libsql_bindings::mwal::ffi::libsql_wal_methods>>,
-        >,
         wal_hook: impl WalHook + Send + Clone + 'static,
         with_bottomless: bool,
     ) -> crate::Result<Self> {
         let (sender, receiver) = crossbeam::channel::unbounded::<Message>();
 
         tokio::task::spawn_blocking(move || {
-            let conn = open_db(
-                path,
-                #[cfg(feature = "mwal_backend")]
-                vwal_methods,
-                wal_hook,
-                with_bottomless,
-            )
-            .unwrap();
+            let conn = open_db(path, wal_hook, with_bottomless).unwrap();
 
             let mut state = ConnectionState::initial();
             let mut timedout = false;
diff --git a/sqld/src/database/write_proxy/mod.rs b/sqld/src/database/write_proxy/mod.rs
index e9fc1433..d8a68d65 100644
--- a/sqld/src/database/write_proxy/mod.rs
+++ b/sqld/src/database/write_proxy/mod.rs
@@ -2,12 +2,11 @@ mod replication;
 
 use std::future::{ready, Ready};
 use std::path::PathBuf;
-#[cfg(feature = "mwal_backend")]
-use std::sync::Arc;
 use std::time::Duration;
 
 use crossbeam::channel::TryRecvError;
 use tokio::sync::Mutex;
+use tokio::task::JoinHandle;
 use tonic::transport::Channel;
 use uuid::Uuid;
 
@@ -27,9 +26,6 @@ use replication::PeriodicDbUpdater;
 pub struct WriteProxyDbFactory {
     write_proxy: ProxyClient<Channel>,
     db_path: PathBuf,
-    #[cfg(feature = "mwal_backend")]
-    vwal_methods:
-        Option<Arc<std::sync::Mutex<sqld_libsql_bindings::mwal::ffi::libsql_wal_methods>>>,
     /// abort handle: abort db update loop on drop
     _abort_handle: crossbeam::channel::Sender<()>,
 }
@@ -42,10 +38,7 @@ impl WriteProxyDbFactory {
         key_path: Option<PathBuf>,
         ca_cert_path: Option<PathBuf>,
         db_path: PathBuf,
-        #[cfg(feature = "mwal_backend")] vwal_methods: Option<
-            Arc<std::sync::Mutex<sqld_libsql_bindings::mwal::ffi::libsql_wal_methods>>,
-        >,
-    ) -> anyhow::Result<Self> {
+    ) -> anyhow::Result<(Self, JoinHandle<anyhow::Result<()>>)> {
         let mut endpoint = Channel::from_shared(addr.to_string())?;
         if tls {
             let cert_pem = std::fs::read_to_string(cert_path.unwrap())?;
@@ -72,22 +65,29 @@ impl WriteProxyDbFactory {
         let mut db_updater =
             PeriodicDbUpdater::new(&db_path, logger, Duration::from_secs(1)).await?;
         let (_abort_handle, receiver) = crossbeam::channel::bounded::<()>(1);
-        tokio::task::spawn_blocking(move || loop {
-            // must abort
-            if let Err(TryRecvError::Disconnected) = receiver.try_recv() {
-                tracing::warn!("periodic updater exiting");
-                break;
+
+        let handle = tokio::task::spawn_blocking(move || -> anyhow::Result<()> {
+            loop {
+                // must abort
+                if let Err(TryRecvError::Disconnected) = receiver.try_recv() {
+                    tracing::warn!("periodic updater exiting");
+                    break Ok(());
+                }
+
+                match db_updater.step() {
+                    Ok(true) => continue,
+                    Ok(false) => return Ok(()),
+                    Err(e) => return Err(e),
+                }
             }
-            db_updater.step();
         });
 
-        Ok(Self {
+        let this = Self {
             write_proxy,
             db_path,
-            #[cfg(feature = "mwal_backend")]
-            vwal_methods,
             _abort_handle,
-        })
+        };
+        Ok((this, handle))
     }
 }
 
@@ -100,8 +100,6 @@ impl DbFactory for WriteProxyDbFactory {
         ready(WriteProxyDatabase::new(
             self.write_proxy.clone(),
             self.db_path.clone(),
-            #[cfg(feature = "mwal_backend")]
-            self.vwal_methods.clone(),
         ))
     }
 }
@@ -114,20 +112,8 @@ pub struct WriteProxyDatabase {
 }
 
 impl WriteProxyDatabase {
-    fn new(
-        write_proxy: ProxyClient<Channel>,
-        path: PathBuf,
-        #[cfg(feature = "mwal_backend")] vwal_methods: Option<
-            Arc<std::sync::Mutex<sqld_libsql_bindings::mwal::ffi::libsql_wal_methods>>,
-        >,
-    ) -> Result<Self> {
-        let read_db = LibSqlDb::new(
-            path,
-            #[cfg(feature = "mwal_backend")]
-            vwal_methods,
-            (),
-            false, // no bottomless replication for replicas
-        )?;
+    fn new(write_proxy: ProxyClient<Channel>, path: PathBuf) -> Result<Self> {
+        let read_db = LibSqlDb::new(path, (), false)?;
         Ok(Self {
             read_db,
             write_proxy,
diff --git a/sqld/src/lib.rs b/sqld/src/lib.rs
index db77a065..12bcae37 100644
--- a/sqld/src/lib.rs
+++ b/sqld/src/lib.rs
@@ -12,8 +12,13 @@ use database::service::DbFactoryService;
 use database::write_proxy::WriteProxyDbFactory;
 use futures::stream::FuturesUnordered;
 use futures::{Future, StreamExt};
+use once_cell::sync::Lazy;
+#[cfg(feature = "mwal_backend")]
+use once_cell::sync::OnceCell;
 use query::{Queries, QueryResult};
 use rpc::run_rpc_server;
+use tokio::sync::Notify;
+use tokio::task::JoinHandle;
 use tower::load::Constant;
 use tower::{Service, ServiceExt};
 use wal_logger::{WalLogger, WalLoggerHook};
@@ -42,6 +47,21 @@ pub enum Backend {
 }
 
 type Result<T> = std::result::Result<T, Error>;
+/// Services handles registry.
+/// All created services must be registered here, so they can be awaited together.
+type Handles = FuturesUnordered<JoinHandle<anyhow::Result<()>>>;
+
+/// Trigger a hard database reset. This cause the database to be wiped, freshly restarted
+/// This is used for replicas that are left in an unrecoverabe state and should restart from a
+/// fresh state.
+///
+/// /!\ use with caution.
+pub(crate) static HARD_RESET: Lazy<Arc<Notify>> = Lazy::new(|| Arc::new(Notify::new()));
+
+#[cfg(feature = "mwal_backend")]
+pub(crate) static VWAL_METHODS: OnceCell<
+    Option<Arc<Mutex<sqld_libsql_bindings::mwal::ffi::libsql_wal_methods>>>,
+> = OnceCell::new();
 
 pub struct Config {
     pub db_path: PathBuf,
@@ -67,7 +87,7 @@ pub struct Config {
     pub create_local_http_tunnel: bool,
 }
 
-async fn run_service<S>(service: S, config: Config) -> anyhow::Result<()>
+async fn run_service<S>(service: S, config: &Config, handles: &mut Handles) -> anyhow::Result<()>
 where
     S: Service<(), Error = Error> + Sync + Send + 'static + Clone,
     S::Response: Service<Queries, Response = Vec<QueryResult>, Error = Error> + Sync + Send,
@@ -84,13 +104,12 @@ where
         server.bind_ws(addr).await?;
     }
 
-    let mut handles = FuturesUnordered::new();
     let factory = PgConnectionFactory::new(service.clone());
     handles.push(tokio::spawn(server.serve(factory)));
 
     if let Some(addr) = config.http_addr {
-        let authorizer =
-            http::auth::parse_auth(config.http_auth).context("failed to parse HTTP auth config")?;
+        let authorizer = http::auth::parse_auth(config.http_auth.clone())
+            .context("failed to parse HTTP auth config")?;
         let handle = tokio::spawn(http::run_http(
             addr,
             authorizer,
@@ -101,9 +120,20 @@ where
         handles.push(handle);
     }
 
-    while let Some(res) = handles.next().await {
-        res??;
-    }
+    Ok(())
+}
+
+/// nukes current DB and start anew
+async fn hard_reset(config: &Config, mut handles: Handles) -> anyhow::Result<()> {
+    tracing::error!("received hard-reset command: reseting replica.");
+
+    tracing::info!("Shutting down all services...");
+    handles.iter_mut().for_each(|h| h.abort());
+    while handles.next().await.is_some() {}
+    tracing::info!("All services have been shut down.");
+
+    let db_path = &config.db_path;
+    tokio::fs::remove_dir_all(db_path).await?;
 
     Ok(())
 }
@@ -132,23 +162,71 @@ impl<F: Future> Future for FutOrNever<F> {
     }
 }
 
+async fn start_primary(config: &Config, handles: &mut Handles, addr: &str) -> anyhow::Result<()> {
+    let (factory, handle) = WriteProxyDbFactory::new(
+        addr,
+        config.writer_rpc_tls,
+        config.writer_rpc_cert.clone(),
+        config.writer_rpc_key.clone(),
+        config.writer_rpc_ca_cert.clone(),
+        config.db_path.clone(),
+    )
+    .await
+    .context("failed to start WriteProxy DB")?;
+
+    handles.push(handle);
+
+    let service = DbFactoryService::new(factory);
+    run_service(service, config, handles).await?;
+
+    Ok(())
+}
+
+async fn start_replica(config: &Config, handles: &mut Handles) -> anyhow::Result<()> {
+    let logger = Arc::new(WalLogger::open(&config.db_path)?);
+    let logger_clone = logger.clone();
+    let path_clone = config.db_path.clone();
+    let enable_bottomless = config.enable_bottomless_replication;
+    let db_factory = move || {
+        let db_path = path_clone.clone();
+        let hook = WalLoggerHook::new(logger.clone());
+        async move { LibSqlDb::new(db_path, hook, enable_bottomless) }
+    };
+    let service = DbFactoryService::new(db_factory.clone());
+    if let Some(ref addr) = config.rpc_server_addr {
+        let handle = tokio::spawn(run_rpc_server(
+            *addr,
+            config.rpc_server_tls,
+            config.rpc_server_cert.clone(),
+            config.rpc_server_key.clone(),
+            config.rpc_server_ca_cert.clone(),
+            db_factory,
+            logger_clone,
+        ));
+
+        handles.push(handle);
+    }
+
+    run_service(service, config, handles).await?;
+
+    Ok(())
+}
+
 pub async fn run_server(config: Config) -> anyhow::Result<()> {
     tracing::trace!("Backend: {:?}", config.backend);
-    #[cfg(feature = "mwal_backend")]
-    if config.backend == Backend::Mwal {
-        std::env::set_var("MVSQLITE_DATA_PLANE", config.mwal_addr.as_ref().unwrap());
-    }
 
     #[cfg(feature = "mwal_backend")]
-    let vwal_methods = config.mwal_addr.as_ref().map(|_| {
-        Arc::new(Mutex::new(
-            sqld_libsql_bindings::mwal::ffi::libsql_wal_methods::new(),
-        ))
-    });
-
-    // create database directory where all the files will go.
-    if !config.db_path.exists() {
-        std::fs::create_dir_all(&config.db_path)?;
+    {
+        if config.backend == Backend::Mwal {
+            std::env::set_var("MVSQLITE_DATA_PLANE", config.mwal_addr.as_ref().unwrap());
+        }
+        VWAL_METHODS
+            .set(config.mwal_addr.as_ref().map(|_| {
+                Arc::new(Mutex::new(
+                    sqld_libsql_bindings::mwal::ffi::libsql_wal_methods::new(),
+                ))
+            }))
+            .map_err(|_| anyhow::anyhow!("wal_methods initialized twice"))?;
     }
 
     if config.enable_bottomless_replication {
@@ -170,84 +248,29 @@ pub async fn run_server(config: Config) -> anyhow::Result<()> {
         println!("HTTP tunnel created: {tunnel}");
     }
 
-    match config.writer_rpc_addr {
-        Some(ref addr) => {
-            let factory = WriteProxyDbFactory::new(
-                addr,
-                config.writer_rpc_tls,
-                config.writer_rpc_cert.clone(),
-                config.writer_rpc_key.clone(),
-                config.writer_rpc_ca_cert.clone(),
-                config.db_path.clone(),
-                #[cfg(feature = "mwal_backend")]
-                vwal_methods,
-            )
-            .await
-            .context("failed to start WriteProxy DB")?;
-            let service = DbFactoryService::new(factory);
-            run_service(service, config).await?;
+    loop {
+        if !config.db_path.exists() {
+            std::fs::create_dir_all(&config.db_path)?;
         }
-        None => {
-            let logger = Arc::new(WalLogger::open(&config.db_path)?);
-            let logger_clone = logger.clone();
-            let path_clone = config.db_path.clone();
-            let db_factory = move || {
-                let db_path = path_clone.clone();
-                #[cfg(feature = "mwal_backend")]
-                let vwal_methods = vwal_methods.clone();
-                let hook = WalLoggerHook::new(logger.clone());
-                async move {
-                    LibSqlDb::new(
-                        db_path,
-                        #[cfg(feature = "mwal_backend")]
-                        vwal_methods,
-                        hook,
-                        config.enable_bottomless_replication,
-                    )
-                }
-            };
-            let service = DbFactoryService::new(db_factory.clone());
-            let rpc_fut: FutOrNever<_> = config
-                .rpc_server_addr
-                .map(|addr| {
-                    tokio::spawn(run_rpc_server(
-                        addr,
-                        config.rpc_server_tls,
-                        config.rpc_server_cert.clone(),
-                        config.rpc_server_key.clone(),
-                        config.rpc_server_ca_cert.clone(),
-                        db_factory,
-                        logger_clone,
-                    ))
-                })
-                .into();
-
-            let svc_fut = run_service(service, config);
+        let mut handles = FuturesUnordered::new();
 
+        match config.writer_rpc_addr {
+            Some(ref addr) => start_primary(&config, &mut handles, addr).await?,
+            None => start_replica(&config, &mut handles).await?,
+        }
+
+        let reset = HARD_RESET.clone();
+        loop {
             tokio::select! {
-                ret = rpc_fut => {
-                    match ret {
-                        Ok(Ok(_)) => {
-                            tracing::info!("Rpc server exited, terminating program.");
-                        }
-                        Err(e) => {
-                            tracing::error!("Rpc server exited with error: {e}");
-                        }
-                        Ok(Err(e)) => {
-                            tracing::error!("Rpc server exited with error: {e}");
-                        }
-                    }
-                }
-                ret = svc_fut => {
-                    match ret {
-                        Ok(_) => tracing::info!("Server exited"),
-                        Err(e) => tracing::error!("Server exited with error: {e}"),
-                    }
-                }
+                _ = reset.notified() => {
+                    hard_reset(&config, handles).await?;
+                    break;
+                },
+                Some(res) = handles.next() => {
+                    res??;
+                },
+                else => return Ok(()),
             }
         }
     }
-
-    let _ = local_tunnel_shutdown.send(());
-    Ok(())
 }

From ce49fb1eac8949e74d7c09dcefcff35a64aae630 Mon Sep 17 00:00:00 2001
From: ad hoc <postma.marin@protonmail.com>
Date: Mon, 13 Feb 2023 15:22:05 +0100
Subject: [PATCH 2/3] enable server handshake

---
 sqld/src/rpc/wal_log.rs | 38 +++++++++++++++++++++-----------------
 sqld/src/wal_logger.rs  | 36 +++++++++++++++++++++++-------------
 2 files changed, 44 insertions(+), 30 deletions(-)

diff --git a/sqld/src/rpc/wal_log.rs b/sqld/src/rpc/wal_log.rs
index 773d3071..9992317c 100644
--- a/sqld/src/rpc/wal_log.rs
+++ b/sqld/src/rpc/wal_log.rs
@@ -6,14 +6,14 @@ pub mod wal_log_rpc {
 use std::net::SocketAddr;
 use std::sync::Arc;
 
-use crate::wal_logger::{FrameId, WalLogger};
-
 use std::collections::HashSet;
 use std::sync::RwLock;
 use tokio_stream::wrappers::ReceiverStream;
 use tonic::Status;
 use wal_log_rpc::wal_log_server::WalLog;
 
+use crate::wal_logger::{FrameId, WalLogger};
+
 use self::wal_log_rpc::{Frame, HelloRequest, HelloResponse, LogOffset};
 
 pub struct WalLogService {
@@ -21,7 +21,7 @@ pub struct WalLogService {
     replicas_with_hello: RwLock<HashSet<SocketAddr>>,
 }
 
-const _NO_HELLO_ERROR_MSG: &str = "NO_HELLO";
+pub const NO_HELLO_ERROR_MSG: &str = "NO_HELLO";
 
 impl WalLogService {
     pub fn new(logger: Arc<WalLogger>) -> Self {
@@ -63,19 +63,17 @@ impl WalLog for WalLogService {
         &self,
         req: tonic::Request<LogOffset>,
     ) -> Result<tonic::Response<Self::LogEntriesStream>, Status> {
-        // TODO: Uncomment once replicas have the ability to handle NO_HELLO error
-        //       Remember to rename _NO_HELLO_ERROR_MSG to NO_HELLO_ERROR_MSG
-        //
-        // let replica_addr = req
-        //     .remote_addr()
-        //     .ok_or(Status::internal("No remote RPC address"))?;
-        // {
-        //     let guard = self.replicas_with_hello.read().unwrap();
-        //     if !guard.contains(&replica_addr) {
-        //         return Err(Status::failed_precondition(NO_HELLO_ERROR_MSG));
-        //     }
-        // }
-        let start_offset = req.into_inner().start_offset;
+        let replica_addr = req
+            .remote_addr()
+            .ok_or(Status::internal("No remote RPC address"))?;
+        {
+            let guard = self.replicas_with_hello.read().unwrap();
+            if !guard.contains(&replica_addr) {
+                return Err(Status::failed_precondition(NO_HELLO_ERROR_MSG));
+            }
+        }
+        // if current_offset is None, then start sending from 0, otherwise return next frame
+        let start_offset = req.into_inner().current_offset.map(|x| x + 1).unwrap_or(0);
         let stream = self.stream_pages(start_offset as _);
         Ok(tonic::Response::new(stream))
     }
@@ -91,6 +89,12 @@ impl WalLog for WalLogService {
             let mut guard = self.replicas_with_hello.write().unwrap();
             guard.insert(replica_addr);
         }
-        Ok(tonic::Response::new(self.logger.hello_response().clone()))
+        let response = HelloResponse {
+            database_id: self.logger.database_id.to_string(),
+            generation_start_index: self.logger.generation.start_index,
+            generation_id: self.logger.generation.id.to_string(),
+        };
+
+        Ok(tonic::Response::new(response))
     }
 }
diff --git a/sqld/src/wal_logger.rs b/sqld/src/wal_logger.rs
index d0a3ce0c..61d394d7 100644
--- a/sqld/src/wal_logger.rs
+++ b/sqld/src/wal_logger.rs
@@ -8,7 +8,7 @@ use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 
 use anyhow::ensure;
-use bytemuck::{cast_ref, try_from_bytes, try_pod_read_unaligned, Pod, Zeroable};
+use bytemuck::{cast_ref, try_from_bytes, Pod, Zeroable};
 use bytes::{BufMut, Bytes, BytesMut};
 use crc::Crc;
 use parking_lot::Mutex;
@@ -19,7 +19,6 @@ use crate::libsql::ffi::{
     PgHdr, Wal,
 };
 use crate::libsql::{ffi::PageHdrIter, wal_hook::WalHook};
-use crate::rpc::wal_log::wal_log_rpc::HelloResponse;
 
 pub const WAL_PAGE_SIZE: i32 = 4096;
 const WAL_MAGIC: u64 = u64::from_le_bytes(*b"SQLDWAL\0");
@@ -171,7 +170,7 @@ struct WalLoggerFileHeader {
 impl WalLoggerFileHeader {
     fn decode(buf: &[u8]) -> anyhow::Result<Self> {
         let this: Self =
-            try_pod_read_unaligned(buf).map_err(|_e| anyhow::anyhow!("invalid WAL log header"))?;
+            *try_from_bytes(buf).map_err(|_e| anyhow::anyhow!("invalid WAL log header"))?;
         ensure!(this.magic == WAL_MAGIC, "invalid WAL log header");
 
         Ok(this)
@@ -208,6 +207,20 @@ impl WalFrameHeader {
     }
 }
 
+pub struct Generation {
+    pub id: Uuid,
+    pub start_index: u64,
+}
+
+impl Generation {
+    fn new(start_index: u64) -> Self {
+        Self {
+            id: Uuid::new_v4(),
+            start_index,
+        }
+    }
+}
+
 pub struct WalLogger {
     /// offset id of the next Frame to write into the log
     next_frame_id: Mutex<FrameId>,
@@ -215,7 +228,8 @@ pub struct WalLogger {
     start_frame_id: FrameId,
     log_file: File,
     current_checksum: AtomicU64,
-    hello_response: HelloResponse,
+    pub database_id: Uuid,
+    pub generation: Generation,
 }
 
 impl WalLogger {
@@ -232,14 +246,16 @@ impl WalLogger {
         let file_end = log_file.metadata()?.len();
         let end_id;
         let current_checksum;
+
         let header = if file_end == 0 {
+            let db_id = Uuid::new_v4();
             let header = WalLoggerFileHeader {
                 version: 1,
                 start_frame_id: 0,
                 magic: WAL_MAGIC,
                 page_size: WAL_PAGE_SIZE,
                 start_checksum: 0,
-                db_id: Uuid::new_v4().as_u128(),
+                db_id: db_id.as_u128(),
             };
 
             let mut header_buf = BytesMut::new();
@@ -264,18 +280,12 @@ impl WalLogger {
             next_frame_id: Mutex::new(end_id),
             start_frame_id: header.start_frame_id,
             log_file,
-            hello_response: HelloResponse {
-                generation_id: uuid::Uuid::new_v4().to_string(),
-                first_frame_id_in_generation: file_end,
-            },
             current_checksum,
+            database_id: Uuid::from_u128(header.db_id),
+            generation: Generation::new(end_id),
         })
     }
 
-    pub fn hello_response(&self) -> &HelloResponse {
-        &self.hello_response
-    }
-
     fn push_page(&self, pages: &[WalPage]) {
         let mut lock = self.next_frame_id.lock();
         let mut current_offset = *lock;

From 690b3ef7270b469124d42c50788a4388bbd9c2b3 Mon Sep 17 00:00:00 2001
From: ad hoc <postma.marin@protonmail.com>
Date: Mon, 13 Feb 2023 15:22:38 +0100
Subject: [PATCH 3/3] enable client handshake

---
 sqld/proto/wal_log.proto                     |   9 +-
 sqld/src/database/write_proxy/replication.rs | 329 +++++++++++++++----
 2 files changed, 264 insertions(+), 74 deletions(-)

diff --git a/sqld/proto/wal_log.proto b/sqld/proto/wal_log.proto
index e814afbe..4ff661a4 100644
--- a/sqld/proto/wal_log.proto
+++ b/sqld/proto/wal_log.proto
@@ -2,14 +2,19 @@ syntax = "proto3";
 package wal_log;
 
 message LogOffset {
-    uint64 startOffset = 1;
+    // offset of the last reveiced frame. If absent, the server should send frame 0.
+    optional uint64 currentOffset = 1;
 }
 
 message HelloRequest { }
 
 message HelloResponse {
+    /// Uuid of the current generation
     string generation_id = 1;
-    uint64 first_frame_id_in_generation = 2;
+    /// First frame index in the current generation
+    uint64 generation_start_index = 2;
+    /// Uuid of the database being replicated
+    string database_id = 3;
 }
 
 message Frame {
diff --git a/sqld/src/database/write_proxy/replication.rs b/sqld/src/database/write_proxy/replication.rs
index ab7c7b4b..592f8c06 100644
--- a/sqld/src/database/write_proxy/replication.rs
+++ b/sqld/src/database/write_proxy/replication.rs
@@ -21,24 +21,33 @@ use std::io::ErrorKind;
 use std::mem::size_of;
 use std::os::unix::prelude::FileExt;
 use std::path::Path;
+use std::str::FromStr;
 use std::time::Duration;
 
-use anyhow::ensure;
+use anyhow::Context;
+use bytemuck::{bytes_of, try_from_bytes, Pod, Zeroable};
+use crossbeam::channel::TryRecvError;
 use futures::StreamExt;
 use rusqlite::ffi::SQLITE_ERROR;
 use rusqlite::OpenFlags;
 use tokio::runtime::Handle;
 use tonic::transport::Channel;
+use tonic::Code;
+use uuid::Uuid;
 
 use crate::libsql::ffi::{types::XWalFrameFn, PgHdr, Wal};
 use crate::libsql::open_with_regular_wal;
 use crate::libsql::wal_hook::WalHook;
 use crate::rpc::wal_log::wal_log_rpc::{wal_log_client::WalLogClient, LogOffset};
+use crate::rpc::wal_log::wal_log_rpc::{HelloRequest, HelloResponse};
+use crate::rpc::wal_log::NO_HELLO_ERROR_MSG;
 use crate::wal_logger::{WalFrame, WAL_PAGE_SIZE};
+use crate::HARD_RESET;
 
 pub struct PeriodicDbUpdater {
     interval: Duration,
     db: rusqlite::Connection,
+    abort_receiver: crossbeam::channel::Receiver<ReplicationError>,
 }
 
 /// The `PeriodicUpdater` role is to periodically trigger a dummy write that will be intercepted by
@@ -49,7 +58,7 @@ impl PeriodicDbUpdater {
         logger: WalLogClient<Channel>,
         interval: Duration,
     ) -> anyhow::Result<Self> {
-        let hook = ReadReplicationHook::new(logger, path).await?;
+        let (hook, abort_receiver) = ReadReplicationHook::new(logger, path).await?;
         let path = path.to_owned();
         let db = tokio::task::spawn_blocking(move || {
             open_with_regular_wal(
@@ -64,23 +73,42 @@ impl PeriodicDbUpdater {
         })
         .await??;
 
-        Ok(Self { interval, db })
+        Ok(Self {
+            interval,
+            db,
+            abort_receiver,
+        })
     }
 
     /// blocking!
-    pub fn step(&mut self) {
-        // dummy write that triggers a call to xFrame
-        let _ = self.db.execute(
-            "create table if not exists __dummy__ (dummy); insert into __dummy__ values (1);",
-            (),
-        );
-        std::thread::sleep(self.interval);
+    /// return whether to continue
+    pub fn step(&mut self) -> anyhow::Result<bool> {
+        match self.abort_receiver.try_recv() {
+            Ok(e) => {
+                // received an error from loop
+                Err(e.into())
+            }
+            Err(TryRecvError::Empty) => {
+                // dummy write that triggers a call to xFrame
+                let _ = self.db.execute(
+                    "create table if not exists __dummy__ (dummy); insert into __dummy__ values (1);",
+                    (),
+                );
+                std::thread::sleep(self.interval);
+
+                Ok(true)
+            }
+
+            Err(TryRecvError::Disconnected) => {
+                // graceful exit
+                Ok(false)
+            }
+        }
     }
 }
 
 struct ReadReplicationHook {
     logger: WalLogClient<Channel>,
-    fetch_frame_index: u64,
     /// Persistent last committed index used for restarts.
     /// The File should contain two little-endian u64:
     /// - The first one is the attempted commit index before the call xFrame
@@ -88,11 +116,16 @@ struct ReadReplicationHook {
     /// After a flight of pages has been successfully written, the two numbers should be the same.
     /// On startup the two number are checked for consistency. If they differ, the database is
     /// considered corrupted, since it is impossible to know what the actually replicated index is.
-    last_applied_index_file: File,
-    last_applied_index: Option<u64>,
+    wal_index_meta_file: File,
+    wal_index_meta: WalIndexMeta,
     /// Buffer for incoming frames
     buffer: VecDeque<WalFrame>,
     rt: Handle,
+    /// A channel to send error back to the polling loop.
+    /// When an error occurs that causes an abort, this handle should be replaced with None, and
+    /// the error sent. This means that if `abort_sender` is None, we should assume a previous
+    /// abort.
+    abort_sender: Option<crossbeam::channel::Sender<ReplicationError>>,
 }
 
 /// Debug assertion. Make sure that all the pages have been applied
@@ -122,6 +155,11 @@ unsafe impl WalHook for ReadReplicationHook {
         sync_flags: c_int,
         orig: XWalFrameFn,
     ) -> c_int {
+        // A fatal error has occured before, don't do anyhting
+        if self.abort_sender.is_none() {
+            return SQLITE_ERROR;
+        }
+
         let rt = self.rt.clone();
         if let Err(e) = rt.block_on(self.fetch_log_entries()) {
             tracing::error!("error fetching log entries: {e}");
@@ -133,24 +171,14 @@ unsafe impl WalHook for ReadReplicationHook {
             assert_ne!(size_after, 0, "commit index points to non commit frame");
 
             tracing::trace!(commit = ?frame_count, size_after = size_after);
-            let attempted_commit_index = self
-                .last_applied_index
-                .map(|x| x + frame_count as u64)
-                .unwrap_or_default();
-            // pre-write index
-            self.last_applied_index_file
-                .write_all_at(&attempted_commit_index.to_le_bytes(), 0)
-                .unwrap();
-
+            self.inc_pre_commit(frame_count as u64)
+                .expect("failed to write pre-commit index");
             let ret = unsafe { orig(wal, WAL_PAGE_SIZE, page_headers, size_after, 1, sync_flags) };
 
             if ret == 0 {
                 debug_assert!(all_applied(page_headers));
-                // persist new commited index
-                self.last_applied_index_file
-                    .write_all_at(&attempted_commit_index.to_le_bytes(), size_of::<u64>() as _)
-                    .unwrap();
-                self.last_applied_index.replace(attempted_commit_index);
+                self.sync_post_commit()
+                    .expect("failed to write post-commit index");
                 // remove commited entries.
                 self.buffer.drain(..frame_count);
                 tracing::trace!("applied frame batch");
@@ -206,39 +234,170 @@ fn free_page_header(h: *const PgHdr) {
     }
 }
 
-impl ReadReplicationHook {
-    async fn new(logger: WalLogClient<Channel>, db_path: &Path) -> anyhow::Result<Self> {
-        let path = db_path.join("client_wal_index");
-        let last_applied_index_file = OpenOptions::new()
-            .create(true)
-            .read(true)
-            .write(true)
-            .open(path)?;
+#[repr(C)]
+#[derive(Debug, Pod, Zeroable, Clone, Copy)]
+struct WalIndexMeta {
+    /// This is the anticipated next frame index to request
+    pre_commit_index: u64,
+    /// After we have written the frames back to the wal, we set this value to the same value as
+    /// pre_commit_index
+    /// On startup we check this value against the pre-commit value to check for consistency
+    post_commit_index: u64,
+    /// Generation Uuid
+    generation_id: u128,
+    /// Uuid of the database this instance is a replica of
+    database_id: u128,
+}
+
+#[derive(Debug, thiserror::Error)]
+enum ReplicationError {
+    #[error("Replica is ahead of primary")]
+    Lagging,
+    #[error("Trying to replicate incompatible databases")]
+    DbIncompatible,
+    #[error("{0}")]
+    Other(#[from] anyhow::Error),
+    #[error("Replication loop exited")]
+    Exit,
+}
 
-        let mut buf = [0; 2 * size_of::<u64>()];
-        let last_applied_index = match last_applied_index_file.read_exact_at(&mut buf, 0) {
+impl WalIndexMeta {
+    fn read(meta_file: &File) -> anyhow::Result<Option<Self>> {
+        let mut buf = [0; size_of::<WalIndexMeta>()];
+        let meta = match meta_file.read_exact_at(&mut buf, 0) {
             Ok(()) => {
-                let pre = u64::from_le_bytes(buf[..size_of::<u64>()].try_into()?);
-                let post = u64::from_le_bytes(buf[size_of::<u64>()..].try_into()?);
-                ensure!(pre == post, "database corrupted");
-                Some(pre)
-            }
-            Err(e) if e.kind() == ErrorKind::UnexpectedEof => {
-                // OK: the file was empty, we don't have a log yet.
-                None
+                meta_file.read_exact_at(&mut buf, 0)?;
+                let meta = *try_from_bytes(&buf)
+                    .map_err(|_| anyhow::anyhow!("invalid index meta file"))?;
+                Some(meta)
             }
+            Err(e) if e.kind() == ErrorKind::UnexpectedEof => None,
             Err(e) => Err(e)?,
         };
 
+        Ok(meta)
+    }
+
+    /// attempts to merge two meta files.
+    fn merge_from_hello(mut self, hello: HelloResponse) -> Result<Self, ReplicationError> {
+        let hello_db_id = Uuid::from_str(&hello.database_id)
+            .context("invalid database id from primary")?
+            .as_u128();
+        let hello_gen_id = Uuid::from_str(&hello.generation_id)
+            .context("invalid generation id from primary")?
+            .as_u128();
+
+        if hello_db_id != self.database_id {
+            return Err(ReplicationError::DbIncompatible);
+        }
+
+        if self.generation_id == hello_gen_id {
+            Ok(self)
+        } else if self.pre_commit_index <= hello.generation_start_index {
+            // Ok: generation changed, but we aren't ahead of primary
+            self.generation_id = hello_gen_id;
+            Ok(self)
+        } else {
+            Err(ReplicationError::Lagging)
+        }
+    }
+
+    fn new_from_hello(hello: HelloResponse) -> anyhow::Result<WalIndexMeta> {
+        let database_id = Uuid::from_str(&hello.database_id)
+            .context("invalid database id from primary")?
+            .as_u128();
+        let generation_id = Uuid::from_str(&hello.generation_id)
+            .context("invalid generation id from primary")?
+            .as_u128();
+
         Ok(Self {
+            pre_commit_index: 0,
+            post_commit_index: 0,
+            generation_id,
+            database_id,
+        })
+    }
+}
+
+impl ReadReplicationHook {
+    async fn new(
+        logger: WalLogClient<Channel>,
+        db_path: &Path,
+    ) -> anyhow::Result<(Self, crossbeam::channel::Receiver<ReplicationError>)> {
+        let path = db_path.join("client_wal_index");
+        let index_meta_file = OpenOptions::new()
+            .create(true)
+            .read(true)
+            .write(true)
+            .open(path)?;
+
+        let (sender, receiver) = crossbeam::channel::bounded(1);
+        let mut abort_sender = Some(sender);
+        let meta = WalIndexMeta::read(&index_meta_file)?;
+        let wal_index_meta =
+            Self::perform_handshake(meta, logger.clone(), &mut abort_sender).await?;
+        let this = Self {
             logger,
-            // ask for the frame right after the one we last applied
-            fetch_frame_index: last_applied_index.map(|x| x + 1).unwrap_or_default(),
-            last_applied_index_file,
-            last_applied_index,
+            wal_index_meta_file: index_meta_file,
+            wal_index_meta,
             buffer: Default::default(),
+            abort_sender,
             rt: Handle::current(),
-        })
+        };
+
+        Ok((this, receiver))
+    }
+
+    async fn perform_handshake(
+        meta: Option<WalIndexMeta>,
+        mut logger: WalLogClient<Channel>,
+        abort_sender: &mut Option<crossbeam::channel::Sender<ReplicationError>>,
+    ) -> anyhow::Result<WalIndexMeta> {
+        let hello = logger.hello(HelloRequest {}).await?.into_inner();
+        tracing::debug!(?hello);
+        match meta {
+            Some(meta) => match meta.merge_from_hello(hello) {
+                Ok(meta) => Ok(meta),
+                Err(e @ ReplicationError::Lagging) => {
+                    tracing::error!("Replica ahead of primary: hard-reseting replica");
+                    HARD_RESET.notify_waiters();
+
+                    // We don't send an error here because we don't want the program to exit: the
+                    // hard reset flow will take care of cleaning behind us and restart fresh.
+                    abort_sender.take();
+
+                    Err(e.into())
+                }
+                Err(e @ ReplicationError::DbIncompatible) => {
+                    if let Some(sender) = abort_sender.take() {
+                        let _ = sender.send(e);
+                    }
+
+                    Err(ReplicationError::Exit.into())
+                }
+                // non-fatal error
+                Err(e) => Err(e.into()),
+            },
+            None => Ok(WalIndexMeta::new_from_hello(hello)?),
+        }
+    }
+
+    fn flush_meta(&self) -> anyhow::Result<()> {
+        self.wal_index_meta_file
+            .write_all_at(bytes_of(&self.wal_index_meta), 0)?;
+
+        Ok(())
+    }
+
+    /// Increment the pre-commit index by n, and flush the meta file
+    fn inc_pre_commit(&mut self, n: u64) -> anyhow::Result<()> {
+        self.wal_index_meta.pre_commit_index += n;
+        self.flush_meta()
+    }
+
+    fn sync_post_commit(&mut self) -> anyhow::Result<()> {
+        self.wal_index_meta.post_commit_index = self.wal_index_meta.pre_commit_index;
+        self.flush_meta()
     }
 
     /// Returns the next page headers list the log truncate count, and the commit frame for the
@@ -270,29 +429,55 @@ impl ReadReplicationHook {
 
     /// Asks the writer for new log frames to apply.
     async fn fetch_log_entries(&mut self) -> anyhow::Result<()> {
-        // try to fetch next page.
-        let start_offset = self.fetch_frame_index;
-        let req = LogOffset { start_offset };
-
-        let mut stream = self.logger.log_entries(req).await?.into_inner();
-        while let Some(raw_frame) = stream.next().await {
-            let raw_frame = raw_frame?;
-            let frame = WalFrame::decode(raw_frame.data)?;
-            debug_assert_eq!(
-                frame.header.frame_id, self.fetch_frame_index,
-                "out of order log frame"
-            );
-            self.fetch_frame_index = frame.header.frame_id + 1;
-            self.buffer.push_back(frame);
-        }
+        let current_offset = self.current_frame_index();
+        let req = LogOffset { current_offset };
 
-        tracing::trace!(
-            start_offset,
-            self.fetch_frame_index,
-            "received {} frames",
-            self.fetch_frame_index - start_offset,
-        );
+        match self.logger.log_entries(req).await {
+            Ok(stream) => {
+                let mut stream = stream.into_inner();
+                let mut frame_count = 0;
+                while let Some(raw_frame) = stream.next().await {
+                    let raw_frame = raw_frame?;
+                    let frame = WalFrame::decode(raw_frame.data)?;
+                    debug_assert_eq!(
+                        frame.header.frame_id, self.wal_index_meta.pre_commit_index,
+                        "out of order log frame"
+                    );
+                    frame_count += 1;
+                    self.buffer.push_back(frame);
+                }
 
-        Ok(())
+                tracing::debug!(current_offset, frame_count,);
+
+                Ok(())
+            }
+            Err(s) if s.code() == Code::FailedPrecondition && s.message() == NO_HELLO_ERROR_MSG => {
+                tracing::info!("Primary restarted, perfoming hanshake again");
+                self.wal_index_meta = Self::perform_handshake(
+                    Some(self.wal_index_meta),
+                    self.logger.clone(),
+                    &mut self.abort_sender,
+                )
+                .await?;
+
+                Ok(())
+            }
+            Err(e) => Err(e)?,
+        }
+    }
+
+    /// Return the current frame index. None if we haven't received any frame yet
+    fn current_frame_index(&self) -> Option<u64> {
+        // This is how we detect a fresh start.
+        // While we could _potentially_ already have requested for index 0, and commited *only*
+        // frame 0, it's ok to reapply it.
+        if self.buffer.is_empty() && self.wal_index_meta.pre_commit_index == 0 {
+            None
+        } else {
+            // the next frame we want to fetch is the one that after the last we have commiter +
+            // the ones we have buffered
+            let index = self.wal_index_meta.pre_commit_index + self.buffer.len() as u64;
+            Some(index)
+        }
     }
 }
