Refactor replicas for better statuses, break every other driver =)

Author: farost

Cargo.lock

@@ -19,7 +19,7 @@
 
 use std::{ffi::c_char, ptr::addr_of_mut};
 
-use typedb_driver::{ReplicaType, ServerReplica};
+use typedb_driver::{ReplicaRole, ServerReplica};
 
 use crate::common::{
     iterator::{iterator_next, CIterator},
@@ -62,8 +62,8 @@ pub extern "C" fn server_replica_get_address(replica_info: *const ServerReplica)
 
 /// Returns whether this is the primary replica of the raft cluster or any of the supporting types.
 #[no_mangle]
-pub extern "C" fn server_replica_get_type(replica_info: *const ServerReplica) -> ReplicaType {
-    borrow(replica_info).replica_type()
+pub extern "C" fn server_replica_get_type(replica_info: *const ServerReplica) -> ReplicaRole {
+    borrow(replica_info).role().unwrap() // TODO: Return optional!
 }
 
 /// Checks whether this is the primary replica of the raft cluster.
@@ -75,5 +75,5 @@ pub extern "C" fn server_replica_is_primary(replica_info: *const ServerReplica)
 /// Returns the raft protocol ‘term’ of this replica.
 #[no_mangle]
 pub extern "C" fn server_replica_get_term(replica_info: *const ServerReplica) -> i64 {
-    borrow(replica_info).term() as i64
+    borrow(replica_info).term().unwrap() as i64 // TODO: Return optional!
 }

c/src/server/server_replica.rs

@@ -30,7 +30,7 @@ def typedb_protocol():
     git_repository(
         name = "typedb_protocol",
         remote = "https://github.com/typedb/typedb-protocol",
-        commit = "55ed1423598425fa5ad56c7b7ca417a05d05dac3",  # sync-marker: do not remove this comment, this is used for sync-dependencies by @typedb_protocol
+        commit = "7653c86e6ca7f19ade1d055793c1082b685cac80",  # sync-marker: do not remove this comment, this is used for sync-dependencies by @typedb_protocol
     )
 
 def typedb_behaviour():

dependencies/typedb/repositories.bzl

@@ -60,7 +60,7 @@
 
 	[dependencies.typedb-protocol]
 		features = []
-		rev = "55ed1423598425fa5ad56c7b7ca417a05d05dac3"
+		rev = "7653c86e6ca7f19ade1d055793c1082b685cac80"
 		git = "https://github.com/typedb/typedb-protocol"
 		default-features = false
 

rust/Cargo.toml

@@ -137,10 +137,12 @@ error_messages! { ConnectionError
         4: "Unable to connect to TypeDB server(s), received network or protocol error: \n{error}",
     ServerConnectionIsClosed =
         5: "The connection has been closed and no further operation is allowed.",
+    ServerConnectionIsClosedUnexpectedly =
+        6: "The connection has been closed unexpectedly and no further operation is allowed.",
     TransactionIsClosed =
-        6: "The transaction is closed and no further operation is allowed.",
+        7: "The transaction is closed and no further operation is allowed.",
     TransactionIsClosedWithErrors { errors: String } =
-        7: "The transaction is closed because of the error(s):\n{errors}",
+        8: "The transaction is closed because of the error(s):\n{errors}",
     MissingResponseField { field: &'static str } =
         9: "Missing field in message received from server: '{field}'. This is either a version compatibility issue or a bug.",
     UnknownRequestId { request_id: RequestID } =
@@ -331,6 +333,8 @@ impl Error {
     fn from_message(message: &str) -> Self {
         if is_rst_stream(message) || is_tcp_connect_error(message) {
             Self::Connection(ConnectionError::ServerConnectionFailedNetworking { error: message.to_owned() })
+        } else if is_reading_body_from_connection_error(message) {
+            Self::Connection(ConnectionError::ServerConnectionIsClosedUnexpectedly)
         } else {
             println!("Error Other from message: {message}");
             Self::Other(message.to_owned())
@@ -455,6 +459,11 @@ fn is_rst_stream(message: &str) -> bool {
     message.contains("Received Rst Stream")
 }
 
+fn is_reading_body_from_connection_error(message: &str) -> bool {
+    // This error can be returned when the server crashes
+    message.contains("error reading a body from connection")
+}
+
 fn is_tcp_connect_error(message: &str) -> bool {
     // No TCP connection
     message.contains("tcp connect error")

rust/src/common/error.rs

@@ -325,12 +325,7 @@ impl TryIntoProto<authentication::token::create::Req> for Credentials {
 
 impl TryFromProto<connection::open::Res> for Response {
     fn try_from_proto(proto: connection::open::Res) -> Result<Self> {
-        let mut servers = Vec::new();
-        for server_proto in
-            proto.servers_all.ok_or(ConnectionError::MissingResponseField { field: "servers_all" })?.servers
-        {
-            servers.push(ServerReplica::try_from_proto(server_proto)?);
-        }
+        let servers = proto.servers_all.ok_or(ConnectionError::MissingResponseField { field: "servers_all" })?.servers.into_iter().map(|server_proto| ServerReplica::try_from_proto(server_proto)).try_collect()?;
         Ok(Self::ConnectionOpen {
             connection_id: Uuid::from_slice(
                 proto
@@ -357,8 +352,7 @@ impl TryFromProto<server_manager::all::Res> for Response {
 impl TryFromProto<server_manager::get::Res> for Response {
     fn try_from_proto(proto: server_manager::get::Res) -> Result<Self> {
         let server_manager::get::Res { server } = proto;
-        let server =
-            ServerReplica::try_from_proto(server.ok_or(ConnectionError::MissingResponseField { field: "server" })?)?;
+        let server = ServerReplica::try_from_proto(server.ok_or(ConnectionError::MissingResponseField { field: "server" })?)?;
         Ok(Self::ServersGet { server })
     }
 }

rust/src/connection/network/proto/message.rs

@@ -26,39 +26,41 @@ use super::TryFromProto;
 use crate::{
     common::Result,
     connection::{
-        server_replica::{ReplicaStatus, ReplicaType, ServerReplica},
+        server_replica::{ReplicaStatus, ReplicaRole, ServerReplica},
         server_version::ServerVersion,
     },
     error::ConnectionError,
 };
 
 impl TryFromProto<ServerProto> for ServerReplica {
-    fn try_from_proto(proto: ServerProto) -> Result<Self> {
-        let address = proto.address.parse()?;
-        let replica_status = match proto.replica_status {
-            Some(replica_status) => ReplicaStatus::try_from_proto(replica_status)?,
-            None => ReplicaStatus::default(),
-        };
-        Ok(Self::from_private(address, replica_status))
+    fn try_from_proto(proto: ServerProto) -> Result<ServerReplica> {
+        let replica_status = proto.replica_status.map(|status| ReplicaStatus::try_from_proto(status)).transpose()?;
+        match proto.address {
+            Some(address) => Ok(ServerReplica::available_from_private(address.parse()?, replica_status)),
+            None => Ok(ServerReplica::Unavailable { replica_status }),
+        }
     }
 }
 
 impl TryFromProto<ReplicaStatusProto> for ReplicaStatus {
     fn try_from_proto(proto: ReplicaStatusProto) -> Result<Self> {
         Ok(Self {
             id: proto.replica_id,
-            replica_type: ReplicaType::try_from_proto(proto.replica_type)?,
+            role: Option::<ReplicaRole>::try_from_proto(proto.replica_type)?,
             term: proto.term,
         })
     }
 }
 
-impl TryFromProto<i32> for ReplicaType {
-    fn try_from_proto(replica_type: i32) -> Result<Self> {
+impl TryFromProto<Option<i32>> for Option<ReplicaRole> {
+    fn try_from_proto(replica_type: Option<i32>) -> Result<Option<ReplicaRole>> {
+        let Some(replica_type) = replica_type else {
+            return Ok(None);
+        };
         match replica_type {
-            0 => Ok(Self::Primary),
-            1 => Ok(Self::Candidate),
-            2 => Ok(Self::Secondary),
+            0 => Ok(Some(ReplicaRole::Primary)),
+            1 => Ok(Some(ReplicaRole::Candidate)),
+            2 => Ok(Some(ReplicaRole::Follower)),
             _ => Err(ConnectionError::UnexpectedReplicaType { replica_type }.into()),
         }
     }

rust/src/connection/network/proto/server.rs

@@ -41,10 +41,11 @@ use crate::{
     error::{ConnectionError, InternalError},
     Credentials, DriverOptions, Error, Result,
 };
+use crate::connection::server_replica::{AvailableServerReplica, Replica};
 
 pub(crate) struct ServerManager {
     configured_addresses: Addresses,
-    replicas: RwLock<HashSet<ServerReplica>>,
+    replicas: RwLock<HashSet<AvailableServerReplica>>,
     replica_connections: RwLock<HashMap<Address, ServerConnection>>,
     address_translation: RwLock<AddressTranslation>,
 
@@ -87,11 +88,10 @@ impl ServerManager {
         )
         .await?;
         let address_translation = addresses.address_translation();
-
         println!("INIT REPLICA CONNECTIONS: {:?}", source_connections);
         let server_manager = Self {
             configured_addresses: addresses,
-            replicas: RwLock::new(replicas),
+            replicas: RwLock::new(Self::filter_unavailable_replicas(replicas)),
             replica_connections: RwLock::new(source_connections),
             address_translation: RwLock::new(address_translation),
             background_runtime,
@@ -207,10 +207,6 @@ impl ServerManager {
         self.replica_connections.write().expect("Expected server connections write access")
     }
 
-    fn read_replicas(&self) -> RwLockReadGuard<'_, HashSet<ServerReplica>> {
-        self.replicas.read().expect("Expected a read replica lock")
-    }
-
     fn read_address_translation(&self) -> RwLockReadGuard<'_, AddressTranslation> {
         self.address_translation.read().expect("Expected address translation read access")
     }
@@ -219,16 +215,6 @@ impl ServerManager {
         self.read_replica_connections().values().map(ServerConnection::force_close).try_collect().map_err(Into::into)
     }
 
-    fn replicas(&self) -> HashSet<ServerReplica> {
-        let res = self.read_replicas().iter().cloned().collect();
-        println!("Current replicas: {res:?}");
-        res
-    }
-
-    pub(crate) fn primary_replica(&self) -> Option<ServerReplica> {
-        self.read_replicas().iter().filter(|replica| replica.is_primary()).max_by_key(|replica| replica.term()).cloned()
-    }
-
     pub(crate) fn username(&self) -> Result<String> {
         match self.read_replica_connections().iter().next() {
             Some((_, replica_connection)) => Ok(replica_connection.username().to_string()),
@@ -271,9 +257,12 @@ impl ServerManager {
         P: Future<Output = Result<R>>,
     {
         println!("STRONG");
-        let mut primary_replica = match self.primary_replica() {
+        let mut primary_replica = match self.read_primary_replica() {
             Some(replica) => replica,
-            None => self.seek_primary_replica_in(self.replicas()).await?,
+            None => {
+                let replicas: HashSet<_> = self.read_replicas().iter().cloned().collect();
+                self.seek_primary_replica_in(replicas).await?
+            },
         };
 
         let retries = self.driver_options.primary_failover_retries;
@@ -283,7 +272,8 @@ impl ServerManager {
             let private_address = primary_replica.private_address().clone();
             match self.execute_on(primary_replica.address(), &private_address, &task).await {
                 Err(Error::Connection(connection_error)) => {
-                    let replicas_without_old_primary = self.replicas().into_iter().filter(|replica| {
+                    let replicas: HashSet<_> = self.read_replicas().iter().cloned().collect();
+                    let replicas_without_old_primary = replicas.into_iter().filter(|replica| {
                         println!("REPLICAS: Filter out? {private_address:?} vs mine {:?}", replica.private_address());
                         replica.private_address() != &private_address
                     });
@@ -326,12 +316,12 @@ impl ServerManager {
         F: Fn(ServerConnection) -> P,
         P: Future<Output = Result<R>>,
     {
-        let replicas = self.replicas();
+        let replicas: HashSet<_> = self.read_replicas().iter().cloned().collect();
         self.execute_on_any(replicas, task).await
     }
 
     #[cfg_attr(feature = "sync", maybe_async::must_be_sync)]
-    async fn execute_on_any<F, P, R>(&self, replicas: impl IntoIterator<Item = ServerReplica>, task: F) -> Result<R>
+    async fn execute_on_any<F, P, R>(&self, replicas: impl IntoIterator<Item = AvailableServerReplica>, task: F) -> Result<R>
     where
         F: Fn(ServerConnection) -> P,
         P: Future<Output = Result<R>>,
@@ -379,8 +369,8 @@ impl ServerManager {
     #[cfg_attr(feature = "sync", maybe_async::must_be_sync)]
     async fn seek_primary_replica_in(
         &self,
-        source_replicas: impl IntoIterator<Item = ServerReplica>,
-    ) -> Result<ServerReplica> {
+        source_replicas: impl IntoIterator<Item = AvailableServerReplica>,
+    ) -> Result<AvailableServerReplica> {
         // TODO: Add retries with sleeps in between? Maybe replica_discovery_attempts should work / be used differently
         println!("SEEK");
         self.execute_on_any(source_replicas, |replica_connection| async {
@@ -390,13 +380,13 @@ impl ServerManager {
     }
 
     #[cfg_attr(feature = "sync", maybe_async::must_be_sync)]
-    async fn seek_primary_replica(&self, replica_connection: ServerConnection) -> Result<ServerReplica> {
+    async fn seek_primary_replica(&self, replica_connection: ServerConnection) -> Result<AvailableServerReplica> {
         let address_translation = self.read_address_translation().clone();
         println!("SEEK PRIMARY REPLICA");
         let replicas = Self::fetch_replicas_from_connection(&replica_connection, &address_translation).await?;
         println!("WRITE NEW PRIMARY");
-        *self.replicas.write().expect("Expected replicas write lock") = replicas;
-        if let Some(replica) = self.primary_replica() {
+        *self.replicas.write().expect("Expected replicas write lock") = Self::filter_unavailable_replicas(replicas);
+        if let Some(replica) = self.read_primary_replica() {
             self.update_replica_connections().await?;
             Ok(replica)
         } else {
@@ -438,7 +428,7 @@ impl ServerManager {
                         return Ok((source_connections, translated_replicas));
                     } else {
                         if let Some(target_replica) =
-                            translated_replicas.into_iter().find(|replica| replica.address() == address)
+                            translated_replicas.into_iter().find(|replica| replica.address() == Some(address))
                         {
                             let source_connections = HashMap::from([(address.clone(), replica_connection)]);
                             return Ok((source_connections, HashSet::from([target_replica])));
@@ -478,7 +468,7 @@ impl ServerManager {
                 async move { Self::fetch_replicas_from_connection(&replica_connection, &address_translation).await }
             })
             .await?;
-        *self.replicas.write().expect("Expected replicas write lock") = replicas.clone();
+        *self.replicas.write().expect("Expected replicas write lock") = Self::filter_unavailable_replicas(replicas.clone());
         Ok(replicas)
     }
 
@@ -497,6 +487,31 @@ impl ServerManager {
         replicas.into_iter().map(|replica| replica.translated(address_translation)).collect()
     }
 
+    fn filter_unavailable_replicas(replicas: impl IntoIterator<Item = ServerReplica>) -> HashSet<AvailableServerReplica> {
+        replicas.into_iter().filter_map(|replica| match replica {
+            ServerReplica::Available(available) => Some(available),
+            ServerReplica::Unavailable { .. } => None,
+        }).collect()
+    }
+
+    #[cfg_attr(feature = "sync", maybe_async::must_be_sync)]
+    pub(crate) async fn fetch_primary_replica(&self) -> Result<Option<ServerReplica>> {
+        let replicas = self.fetch_replicas().await?;
+        Ok(Self::filter_primary_replica(replicas.iter()))
+    }
+
+    fn read_replicas(&self) -> RwLockReadGuard<'_, HashSet<AvailableServerReplica>> {
+        self.replicas.read().expect("Expected a read replica lock")
+    }
+
+    fn read_primary_replica(&self) -> Option<AvailableServerReplica> {
+        Self::filter_primary_replica(self.read_replicas().iter())
+    }
+
+    fn filter_primary_replica<'a, R: Replica + 'a>(replicas: impl IntoIterator<Item = &'a R>) -> Option<R> {
+        replicas.into_iter().filter(|replica| replica.is_primary()).max_by_key(|replica| replica.term().unwrap_or_default()).cloned()
+    }
+
     fn server_connection_failed_err(&self, errors: HashMap<Address, Error>) -> Error {
         let accessed_addresses =
             Addresses::from_addresses(self.read_replicas().iter().map(|replica| replica.address().clone()));