Merge pull request #661 from tursodatabase/improve-replication-error-…

…handling

improve replication error handling

libsql-replication/src/rpc.rs

@@ -29,6 +29,8 @@ pub mod replication {
 
     pub const NO_HELLO_ERROR_MSG: &str = "NO_HELLO";
     pub const NEED_SNAPSHOT_ERROR_MSG: &str = "NEED_SNAPSHOT";
+    /// A tonic error code to signify that a namespace doesn't exist.
+    pub const NAMESPACE_DOESNT_EXIST: &str = "NAMESPACE_DOESNT_EXIST";
 
     pub const SESSION_TOKEN_KEY: &str = "x-session-token";
     pub const NAMESPACE_METADATA_KEY: &str = "x-namespace-bin";

libsql-server/src/namespace/mod.rs

@@ -36,7 +36,6 @@ use crate::database::{Database, PrimaryDatabase, ReplicaDatabase};
 use crate::error::{Error, LoadDumpError};
 use crate::metrics::NAMESPACE_LOAD_LATENCY;
 use crate::replication::primary::logger::{ReplicationLoggerHookCtx, REPLICATION_METHODS};
-use crate::replication::replicator_client::NamespaceDoesntExist;
 use crate::replication::{FrameNo, NamespacedSnapshotCallback, ReplicationLogger};
 use crate::stats::Stats;
 use crate::{
@@ -720,15 +719,12 @@ impl Namespace<ReplicaDatabase> {
             use libsql_replication::replicator::Error;
             loop {
                 match replicator.run().await {
-                    Error::Fatal(e) => {
-                        if let Ok(err) = e.downcast::<NamespaceDoesntExist>() {
-                                tracing::error!("namespace {namespace} doesn't exist, destroying...");
-                                (reset)(ResetOp::Destroy(namespace.clone()));
-                                Err(err)?;
-                        } else {
-                            unreachable!("unexpected fatal replication error")
-                        }
-                    },
+                    err @ Error::Fatal(_) => Err(err)?,
+                    err @ Error::NamespaceDoesntExist => {
+                        tracing::error!("namespace {namespace} doesn't exist, destroying...");
+                        (reset)(ResetOp::Destroy(namespace.clone()));
+                        Err(err)?;
+                    }
                     Error::Meta(err) => {
                         use libsql_replication::meta::Error;
                         match err {

libsql-server/src/replication/replicator_client.rs

@@ -7,18 +7,16 @@ use libsql_replication::meta::WalIndexMeta;
 use libsql_replication::replicator::{map_frame_err, Error, ReplicatorClient};
 use libsql_replication::rpc::replication::replication_log_client::ReplicationLogClient;
 use libsql_replication::rpc::replication::{
-    verify_session_token, HelloRequest, LogOffset, NAMESPACE_METADATA_KEY, NEED_SNAPSHOT_ERROR_MSG,
-    NO_HELLO_ERROR_MSG, SESSION_TOKEN_KEY,
+    verify_session_token, HelloRequest, LogOffset, NAMESPACE_METADATA_KEY, SESSION_TOKEN_KEY,
 };
 use tokio::sync::watch;
 use tokio_stream::{Stream, StreamExt};
 use tonic::metadata::{AsciiMetadataValue, BinaryMetadataValue};
 use tonic::transport::Channel;
-use tonic::{Code, Request};
+use tonic::Request;
 
 use crate::namespace::NamespaceName;
 use crate::replication::FrameNo;
-use crate::rpc::NAMESPACE_DOESNT_EXIST;
 
 pub struct Client {
     client: ReplicationLogClient<Channel>,
@@ -78,37 +76,23 @@ impl Client {
     }
 }
 
-#[derive(Debug, thiserror::Error)]
-#[error("namespace doesn't exist")]
-pub struct NamespaceDoesntExist;
-
 #[async_trait::async_trait]
 impl ReplicatorClient for Client {
     type FrameStream = Pin<Box<dyn Stream<Item = Result<Frame, Error>> + Send + 'static>>;
 
     async fn handshake(&mut self) -> Result<(), Error> {
         tracing::info!("Attempting to perform handshake with primary.");
         let req = self.make_request(HelloRequest::new());
-        match self.client.hello(req).await {
-            Ok(resp) => {
-                let hello = resp.into_inner();
-                verify_session_token(&hello.session_token).map_err(Error::Client)?;
-                self.primary_replication_index = hello.current_replication_index;
-                self.session_token.replace(hello.session_token.clone());
-                self.meta.init_from_hello(hello)?;
-                self.current_frame_no_notifier
-                    .send_replace(self.meta.current_frame_no());
-
-                Ok(())
-            }
-            Err(e)
-                if e.code() == Code::FailedPrecondition
-                    && e.message() == NAMESPACE_DOESNT_EXIST =>
-            {
-                Err(Error::Fatal(NamespaceDoesntExist.into()))?
-            }
-            Err(e) => Err(Error::Client(e.into()))?,
-        }
+        let resp = self.client.hello(req).await?;
+        let hello = resp.into_inner();
+        verify_session_token(&hello.session_token).map_err(Error::Client)?;
+        self.primary_replication_index = hello.current_replication_index;
+        self.session_token.replace(hello.session_token.clone());
+        self.meta.init_from_hello(hello)?;
+        self.current_frame_no_notifier
+            .send_replace(self.meta.current_frame_no());
+
+        Ok(())
     }
 
     async fn next_frames(&mut self) -> Result<Self::FrameStream, Error> {
@@ -119,14 +103,7 @@ impl ReplicatorClient for Client {
         let stream = self
             .client
             .log_entries(req)
-            .await
-            .map_err(|e| {
-                if e.code() == Code::FailedPrecondition && e.message() == NEED_SNAPSHOT_ERROR_MSG {
-                    Error::NeedSnapshot
-                } else {
-                    Error::Client(e.into())
-                }
-            })?
+            .await?
             .into_inner()
             .map(map_frame_err);
 
@@ -141,8 +118,7 @@ impl ReplicatorClient for Client {
         let stream = self
             .client
             .snapshot(req)
-            .await
-            .map_err(map_status)?
+            .await?
             .into_inner()
             .map(map_frame_err);
         Ok(Box::pin(stream))
@@ -162,11 +138,3 @@ impl ReplicatorClient for Client {
         self.meta.current_frame_no()
     }
 }
-
-fn map_status(status: tonic::Status) -> Error {
-    if status.code() == Code::FailedPrecondition && status.message() == NO_HELLO_ERROR_MSG {
-        Error::NoHandshake
-    } else {
-        Error::Client(status.into())
-    }
-}

libsql-server/src/rpc/mod.rs

@@ -26,9 +26,6 @@ pub mod replication_log;
 pub mod replication_log_proxy;
 pub mod streaming_exec;
 
-/// A tonic error code to signify that a namespace doesn't exist.
-pub const NAMESPACE_DOESNT_EXIST: &str = "NAMESPACE_DOESNT_EXIST";
-
 pub async fn run_rpc_server<A: crate::net::Accept>(
     proxy_service: ProxyService,
     acceptor: A,

libsql-server/src/rpc/proxy.rs

@@ -11,6 +11,7 @@ use libsql_replication::rpc::proxy::{
     describe_result, Ack, DescribeRequest, DescribeResult, Description, DisconnectMessage, ExecReq,
     ExecResp, ExecuteResults, QueryResult, ResultRows, Row,
 };
+use libsql_replication::rpc::replication::NAMESPACE_DOESNT_EXIST;
 use rusqlite::types::ValueRef;
 use uuid::Uuid;
 
@@ -24,8 +25,6 @@ use crate::query_result_builder::{
 use crate::replication::FrameNo;
 use crate::rpc::streaming_exec::make_proxy_stream;
 
-use super::NAMESPACE_DOESNT_EXIST;
-
 pub mod rpc {
     use std::sync::Arc;
 

libsql-server/src/rpc/replication_log.rs

@@ -8,8 +8,8 @@ use futures::stream::BoxStream;
 pub use libsql_replication::rpc::replication as rpc;
 use libsql_replication::rpc::replication::replication_log_server::ReplicationLog;
 use libsql_replication::rpc::replication::{
-    Frame, Frames, HelloRequest, HelloResponse, LogOffset, NEED_SNAPSHOT_ERROR_MSG,
-    NO_HELLO_ERROR_MSG, SESSION_TOKEN_KEY,
+    Frame, Frames, HelloRequest, HelloResponse, LogOffset, NAMESPACE_DOESNT_EXIST,
+    NEED_SNAPSHOT_ERROR_MSG, NO_HELLO_ERROR_MSG, SESSION_TOKEN_KEY,
 };
 use tokio_stream::StreamExt;
 use tonic::transport::server::TcpConnectInfo;
@@ -22,7 +22,7 @@ use crate::replication::primary::frame_stream::FrameStream;
 use crate::replication::LogReadError;
 use crate::utils::services::idle_shutdown::IdleShutdownKicker;
 
-use super::{extract_namespace, NAMESPACE_DOESNT_EXIST};
+use super::extract_namespace;
 
 pub struct ReplicationLogService {
     namespaces: NamespaceStore<PrimaryNamespaceMaker>,

libsql-server/tests/cluster/replication.rs

@@ -115,6 +115,12 @@ fn apply_partial_snapshot() {
         notify.notify_waiters();
 
         let client = Client::new();
+
+        // wait for replica to start up
+        while client.get("http://replica:8080/").await.is_err() {
+            tokio::time::sleep(Duration::from_millis(100)).await;
+        }
+
         loop {
             let resp = client
                 .get("http://replica:9090/v1/namespaces/default/stats")

libsql/src/replication/remote_client.rs

@@ -8,11 +8,10 @@ use libsql_replication::frame::{Frame, FrameHeader, FrameNo};
 use libsql_replication::meta::WalIndexMeta;
 use libsql_replication::replicator::{map_frame_err, Error, ReplicatorClient};
 use libsql_replication::rpc::replication::{
-    verify_session_token, HelloRequest, LogOffset, NEED_SNAPSHOT_ERROR_MSG, SESSION_TOKEN_KEY,
+    verify_session_token, HelloRequest, LogOffset, SESSION_TOKEN_KEY,
 };
 use tokio_stream::Stream;
 use tonic::metadata::AsciiMetadataValue;
-use tonic::Code;
 
 /// A remote replicator client, that pulls frames over RPC
 pub struct RemoteClient {
@@ -64,31 +63,27 @@ impl ReplicatorClient for RemoteClient {
     async fn handshake(&mut self) -> Result<(), Error> {
         tracing::info!("Attempting to perform handshake with primary.");
         let req = self.make_request(HelloRequest::new());
-        match self.remote.replication.hello(req).await {
-            Ok(resp) => {
-                let hello = resp.into_inner();
-                verify_session_token(&hello.session_token).map_err(Error::Client)?;
-                self.session_token = Some(hello.session_token.clone());
-                if self.dirty {
-                    self.meta.reset();
-                    self.dirty = false;
-                }
-                if let Err(e) = self.meta.init_from_hello(hello) {
-                    // set the meta as dirty. The caller should catch the error and clean the db
-                    // file. On the next call to replicate, the db will be replicated from the new
-                    // log.
-                    if let libsql_replication::meta::Error::LogIncompatible = e {
-                        self.dirty = true;
-                    }
-
-                    Err(e)?;
-                }
-                self.meta.flush().await?;
-
-                Ok(())
+        let resp = self.remote.replication.hello(req).await?;
+        let hello = resp.into_inner();
+        verify_session_token(&hello.session_token).map_err(Error::Client)?;
+        self.session_token = Some(hello.session_token.clone());
+        if self.dirty {
+            self.meta.reset();
+            self.dirty = false;
+        }
+        if let Err(e) = self.meta.init_from_hello(hello) {
+            // set the meta as dirty. The caller should catch the error and clean the db
+            // file. On the next call to replicate, the db will be replicated from the new
+            // log.
+            if let libsql_replication::meta::Error::LogIncompatible = e {
+                self.dirty = true;
             }
-            Err(e) => Err(Error::Client(e.into()))?,
+
+            Err(e)?;
         }
+        self.meta.flush().await?;
+
+        Ok(())
     }
 
     /// Return a stream of frames to apply to the database
@@ -100,14 +95,7 @@ impl ReplicatorClient for RemoteClient {
             .remote
             .replication
             .batch_log_entries(req)
-            .await
-            .map_err(|e| {
-                if e.code() == Code::FailedPrecondition && e.message() == NEED_SNAPSHOT_ERROR_MSG {
-                    Error::NeedSnapshot
-                } else {
-                    Error::Client(e.into())
-                }
-            })?
+            .await?
             .into_inner()
             .frames;
 
@@ -136,8 +124,7 @@ impl ReplicatorClient for RemoteClient {
             .remote
             .replication
             .snapshot(req)
-            .await
-            .map_err(|e| Error::Client(e.into()))?
+            .await?
             .into_inner()
             .map(map_frame_err)
             .peekable();