Merge #323

323: Replica connection wait for replication r=MarinPostma a=MarinPostma

What this PR do is quite simple: it makes every replica database connection to wait for the replica instance to catch up with the last write that they performed.

When a replica detects a write, this write is proxied to the primary, that executes it, and then starts replicating the state. The replica receives a response that contains a logical timestamp issued after the write was performed. The replica's connection keep this timestamp and on subsequent operation on the _local_ connection, waits until the replicator has applied changes up to the stored timestamp.

Note that each connection holds a different timestamp, so a connection may proceed with a read, while another is still waiting.

The basic guarantee that it gives is read your write: a connection is always guaranteed to witness it's write. In fact the guarantee is much stronger because of the underlying consistency of sqlite, the real-time consistency model of sqld is now as strong as `sequential consistency`.


Co-authored-by: ad hoc <[email protected]>

Author: bors[bot]

sqld/proto/proxy.proto

@@ -89,6 +89,8 @@ message ExecuteResults {
     }
     /// State after executing the queries
     State state = 2;
+    /// Primary frame_no after executing the request.
+    uint64 current_frame_no = 3;
 }
 
 message Program {

sqld/src/database/write_proxy.rs

@@ -1,14 +1,16 @@
 use std::path::PathBuf;
 use std::sync::Arc;
 
-use tokio::sync::Mutex;
+use parking_lot::Mutex as PMutex;
+use tokio::sync::{watch, Mutex};
 use tonic::transport::Channel;
 use uuid::Uuid;
 
 use crate::auth::{Authenticated, Authorized};
 use crate::error::Error;
 use crate::query::{QueryResponse, QueryResult};
 use crate::query_analysis::State;
+use crate::replication::FrameNo;
 use crate::rpc::proxy::rpc::proxy_client::ProxyClient;
 use crate::rpc::proxy::rpc::query_result::RowResult;
 use crate::rpc::proxy::rpc::DisconnectMessage;
@@ -23,6 +25,7 @@ pub struct WriteProxyDbFactory {
     client: ProxyClient<Channel>,
     db_path: PathBuf,
     stats: Stats,
+    applied_frame_no_receiver: watch::Receiver<FrameNo>,
 }
 
 impl WriteProxyDbFactory {
@@ -31,12 +34,14 @@ impl WriteProxyDbFactory {
         channel: Channel,
         uri: tonic::transport::Uri,
         stats: Stats,
+        applied_frame_no_receiver: watch::Receiver<FrameNo>,
     ) -> Self {
         let client = ProxyClient::with_origin(channel, uri);
         Self {
             client,
             db_path,
             stats,
+            applied_frame_no_receiver,
         }
     }
 }
@@ -48,6 +53,7 @@ impl DbFactory for WriteProxyDbFactory {
             self.client.clone(),
             self.db_path.clone(),
             self.stats.clone(),
+            self.applied_frame_no_receiver.clone(),
         )?;
         Ok(Arc::new(db))
     }
@@ -58,16 +64,29 @@ pub struct WriteProxyDatabase {
     write_proxy: ProxyClient<Channel>,
     state: Mutex<State>,
     client_id: Uuid,
+    /// FrameNo of the last write performed by this connection on the primary.
+    /// any subsequent read on this connection must wait for the replicator to catch up with this
+    /// frame_no
+    last_write_frame_no: PMutex<FrameNo>,
+    /// Notifier from the repliator of the currently applied frameno
+    applied_frame_no_receiver: watch::Receiver<FrameNo>,
 }
 
 impl WriteProxyDatabase {
-    fn new(write_proxy: ProxyClient<Channel>, path: PathBuf, stats: Stats) -> Result<Self> {
+    fn new(
+        write_proxy: ProxyClient<Channel>,
+        path: PathBuf,
+        stats: Stats,
+        applied_frame_no_receiver: watch::Receiver<FrameNo>,
+    ) -> Result<Self> {
         let read_db = LibSqlDb::new(path, (), false, stats)?;
         Ok(Self {
             read_db,
             write_proxy,
             state: Mutex::new(State::Init),
             client_id: Uuid::new_v4(),
+            last_write_frame_no: PMutex::new(FrameNo::MAX),
+            applied_frame_no_receiver,
         })
     }
 
@@ -104,6 +123,8 @@ impl WriteProxyDatabase {
                     })
                     .collect();
 
+                self.update_last_write_frame_no(execute_result.current_frame_no);
+
                 Ok((results, *state))
             }
             Err(e) => {
@@ -114,6 +135,35 @@ impl WriteProxyDatabase {
             }
         }
     }
+
+    fn update_last_write_frame_no(&self, new_frame_no: FrameNo) {
+        let mut last_frame_no = self.last_write_frame_no.lock();
+        if *last_frame_no == FrameNo::MAX || new_frame_no > *last_frame_no {
+            *last_frame_no = new_frame_no
+        }
+    }
+
+    /// wait for the replicator to have caught up with our current write frame_no
+    async fn wait_replication_sync(&self) -> Result<()> {
+        let current_frame_no = *self.last_write_frame_no.lock();
+
+        if current_frame_no == FrameNo::MAX {
+            return Ok(());
+        }
+
+        let mut receiver = self.applied_frame_no_receiver.clone();
+        let mut last_applied = *receiver.borrow_and_update();
+
+        while last_applied < current_frame_no {
+            receiver
+                .changed()
+                .await
+                .map_err(|_| Error::ReplicatorExited)?;
+            last_applied = *receiver.borrow_and_update();
+        }
+
+        Ok(())
+    }
 }
 
 #[async_trait::async_trait]
@@ -125,6 +175,7 @@ impl Database for WriteProxyDatabase {
     ) -> Result<(Vec<Option<QueryResult>>, State)> {
         let mut state = self.state.lock().await;
         if *state == State::Init && pgm.is_read_only() {
+            self.wait_replication_sync().await?;
             // We know that this program won't perform any writes. We attempt to run it on the
             // replica. If it leaves an open transaction, then this program is an interactive
             // transaction, so we rollback the replica, and execute again on the primary.

sqld/src/error.rs

@@ -25,6 +25,8 @@ pub enum Error {
     InvalidBatchStep(usize),
     #[error("Not authorized to execute query: {0}")]
     NotAuthorized(String),
+    #[error("The replicator exited, instance cannot make any progress.")]
+    ReplicatorExited,
 }
 
 impl From<tokio::sync::oneshot::error::RecvError> for Error {

sqld/src/lib.rs

@@ -201,10 +201,17 @@ async fn start_replica(
 ) -> anyhow::Result<()> {
     let (channel, uri) = configure_rpc(config)?;
     let replicator = Replicator::new(config.db_path.clone(), channel.clone(), uri.clone());
+    let applied_frame_no_receiver = replicator.current_frame_no_notifier.subscribe();
 
     join_set.spawn(replicator.run());
 
-    let factory = WriteProxyDbFactory::new(config.db_path.clone(), channel, uri, stats.clone());
+    let factory = WriteProxyDbFactory::new(
+        config.db_path.clone(),
+        channel,
+        uri,
+        stats.clone(),
+        applied_frame_no_receiver,
+    );
     run_service(
         Arc::new(factory),
         config,

sqld/src/replication/replica/replicator.rs

@@ -3,6 +3,7 @@ use std::time::Duration;
 
 use anyhow::bail;
 use futures::StreamExt;
+use tokio::sync::watch;
 use tonic::transport::Channel;
 
 use crate::replication::frame::Frame;
@@ -27,16 +28,19 @@ pub struct Replicator {
     db_path: PathBuf,
     injector: Option<FrameInjectorHandle>,
     current_frame_no: FrameNo,
+    pub current_frame_no_notifier: watch::Sender<FrameNo>,
 }
 
 impl Replicator {
     pub fn new(db_path: PathBuf, channel: Channel, uri: tonic::transport::Uri) -> Self {
         let client = Client::with_origin(channel, uri);
+        let (applied_frame_notifier, _) = watch::channel(FrameNo::MAX);
         Self {
             client,
             db_path,
             injector: None,
             current_frame_no: FrameNo::MAX,
+            current_frame_no_notifier: applied_frame_notifier,
         }
     }
 
@@ -72,7 +76,7 @@ impl Replicator {
                     }
                     let (injector, last_applied_frame_no) =
                         FrameInjectorHandle::new(self.db_path.clone(), hello).await?;
-                    self.current_frame_no = last_applied_frame_no;
+                    self.update_current_frame_no(last_applied_frame_no);
                     self.injector.replace(injector);
                     return Ok(());
                 }
@@ -88,6 +92,11 @@ impl Replicator {
         bail!("couldn't connect to primary after {HANDSHAKE_MAX_RETRIES} tries.");
     }
 
+    fn update_current_frame_no(&mut self, new_frame_no: FrameNo) {
+        self.current_frame_no = new_frame_no;
+        self.current_frame_no_notifier.send_replace(new_frame_no);
+    }
+
     async fn replicate(&mut self) -> anyhow::Result<()> {
         let offset = LogOffset {
             // if current == FrameNo::Max then it means that we're starting fresh
@@ -142,13 +151,15 @@ impl Replicator {
     }
 
     async fn flush_txn(&mut self, frames: Vec<Frame>) -> anyhow::Result<()> {
-        self.current_frame_no = self
+        let new_frame_no = self
             .injector
             .as_mut()
             .expect("invalid state")
             .apply_frames(Frames::Vec(frames))
             .await?;
 
+        self.update_current_frame_no(new_frame_no);
+
         Ok(())
     }
 

sqld/src/rpc/mod.rs

@@ -26,7 +26,7 @@ pub async fn run_rpc_server(
     logger: Arc<ReplicationLogger>,
     idle_shutdown_layer: Option<IdleShutdownLayer>,
 ) -> anyhow::Result<()> {
-    let proxy_service = ProxyService::new(factory);
+    let proxy_service = ProxyService::new(factory, logger.new_frame_notifier.subscribe());
     let logger_service = ReplicationLogService::new(logger);
 
     tracing::info!("serving write proxy server at {addr}");

sqld/src/rpc/proxy.rs

@@ -3,11 +3,13 @@ use std::str::FromStr;
 use std::sync::Arc;
 
 use async_lock::{RwLock, RwLockUpgradableReadGuard};
+use tokio::sync::watch;
 use uuid::Uuid;
 
 use crate::auth::{Authenticated, Authorized};
 use crate::database::factory::DbFactory;
 use crate::database::{Database, Program};
+use crate::replication::FrameNo;
 
 use self::rpc::execute_results::State;
 use self::rpc::proxy_server::Proxy;
@@ -272,13 +274,15 @@ pub mod rpc {
 pub struct ProxyService {
     clients: RwLock<HashMap<Uuid, Arc<dyn Database>>>,
     factory: Arc<dyn DbFactory>,
+    new_frame_notifier: watch::Receiver<FrameNo>,
 }
 
 impl ProxyService {
-    pub fn new(factory: Arc<dyn DbFactory>) -> Self {
+    pub fn new(factory: Arc<dyn DbFactory>, new_frame_notifier: watch::Receiver<FrameNo>) -> Self {
         Self {
             clients: Default::default(),
             factory,
+            new_frame_notifier,
         }
     }
 }
@@ -323,10 +327,12 @@ impl Proxy for ProxyService {
             .await
             .map_err(|e| tonic::Status::new(tonic::Code::PermissionDenied, e.to_string()))?;
         let results = results.into_iter().map(|r| r.into()).collect();
+        let current_frame_no = *self.new_frame_notifier.borrow();
 
         Ok(tonic::Response::new(ExecuteResults {
             results,
             state: State::from(state).into(),
+            current_frame_no,
         }))
     }