feat: seek quorum for readonly session

Closes #30.

Author: kezhuw

Cargo.toml

@@ -46,3 +46,4 @@ assertor = "0.0.2"
 assert_matches = "1.5.0"
 tempfile = "3.6.0"
 rcgen = { version = "0.12.1", features = ["default", "x509-parser"] }
+serial_test = "3.0.0"

src/client/mod.rs

@@ -1577,7 +1577,7 @@ impl Connector {
         self
     }
 
-    /// Specifies whether readonly server is allowed.
+    /// Specifies whether readonly session is allowed.
     pub fn readonly(&mut self, readonly: bool) -> &mut Self {
         self.readonly = readonly;
         self
@@ -1717,7 +1717,7 @@ impl ClientBuilder {
         self
     }
 
-    /// Specifies whether readonly server is allowed.
+    /// Specifies whether readonly session is allowed.
     pub fn with_readonly(&mut self, readonly: bool) -> &mut ClientBuilder {
         self.connector.readonly = readonly;
         self

src/deadline.rs

@@ -0,0 +1,35 @@
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use tokio::time::{self, Instant, Sleep};
+
+pub struct Deadline {
+    sleep: Option<Sleep>,
+}
+
+impl Deadline {
+    pub fn never() -> Self {
+        Self { sleep: None }
+    }
+
+    pub fn until(deadline: Instant) -> Self {
+        Self { sleep: Some(time::sleep_until(deadline)) }
+    }
+
+    pub fn elapsed(&self) -> bool {
+        self.sleep.as_ref().map(|f| f.is_elapsed()).unwrap_or(false)
+    }
+}
+
+impl Future for Deadline {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        if self.sleep.is_none() {
+            return Poll::Pending;
+        }
+        let sleep = unsafe { self.map_unchecked_mut(|deadline| deadline.sleep.as_mut().unwrap_unchecked()) };
+        sleep.poll(cx)
+    }
+}

src/endpoint.rs

@@ -165,6 +165,10 @@ impl IterableEndpoints {
         Self { cycle: false, next: 0, endpoints: endpoints.into() }
     }
 
+    pub fn len(&self) -> usize {
+        self.endpoints.len()
+    }
+
     pub fn endpoints(&self) -> &[Endpoint] {
         &self.endpoints
     }
@@ -192,6 +196,14 @@ impl IterableEndpoints {
             self.next = 0;
         }
     }
+
+    pub fn peek(&self) -> Option<EndpointRef<'_>> {
+        let next = self.next;
+        if next >= self.endpoints.len() {
+            return None;
+        }
+        Some(self.endpoints[next].to_ref())
+    }
 }
 
 impl From<&[EndpointRef<'_>]> for IterableEndpoints {

src/lib.rs

@@ -1,6 +1,7 @@
 mod acl;
 mod chroot;
 mod client;
+mod deadline;
 mod endpoint;
 mod error;
 mod proto;

src/proto/connect.rs

@@ -1,4 +1,5 @@
 use bytes::BufMut;
+use derive_where::derive_where;
 
 use crate::record::{
     DeserializableRecord,
@@ -10,11 +11,13 @@ use crate::record::{
     UnsafeBuf,
 };
 
+#[derive_where(Debug)]
 pub struct ConnectRequest<'a> {
     pub protocol_version: i32,
     pub last_zxid_seen: i64,
     pub timeout: i32,
     pub session_id: i64,
+    #[derive_where(skip(Debug))]
     pub password: &'a [u8],
     pub readonly: bool,
 }
@@ -36,11 +39,13 @@ impl DynamicRecord for ConnectRequest<'_> {
     }
 }
 
+#[derive_where(Debug)]
 pub struct ConnectResponse<'a> {
     #[allow(dead_code)]
     pub protocol_version: i32,
     pub session_timeout: i32,
     pub session_id: i64,
+    #[derive_where(skip(Debug))]
     pub password: &'a [u8],
     pub readonly: bool,
 }
@@ -49,7 +54,7 @@ impl<'a> DeserializableRecord<'a> for ConnectResponse<'a> {
     type Error = DeserializeError;
 
     fn deserialize(buf: &mut ReadingBuf<'a>) -> Result<Self, Self::Error> {
-        let min_record_len = 4 + 4 + 8 + 4 + 1;
+        let min_record_len = 4 + 4 + 8 + 4;
         if buf.len() < min_record_len {
             return Err(DeserializeError::InsufficientBuf);
         }
@@ -64,12 +69,12 @@ impl<'a> DeserializableRecord<'a> for ConnectResponse<'a> {
             )));
         }
         let len = unsafe { buf.get_unchecked_i32() };
-        if len <= 0 || len != (buf.len() - 1) as i32 {
+        if len <= 0 || len >= buf.len() as i32 {
             return Err(DeserializeError::UnmarshalError(format!("invalid session password length {len}")));
         }
         let len = len as usize;
         let password = unsafe { buf.get_unchecked(..len) };
-        let readonly = unsafe { *buf.get_unchecked(len) };
+        let readonly = if buf.is_empty() { 0 } else { unsafe { *buf.get_unchecked(len) } };
         if readonly != 0 && readonly != 1 {
             return Err(DeserializeError::UnmarshalError(format!("invalid session readonly value {readonly}")));
         }

src/session/connection.rs

@@ -1,12 +1,22 @@
-use std::io::{ErrorKind, IoSlice, Result};
+use std::io::{Error, ErrorKind, IoSlice, Result};
 use std::pin::Pin;
 use std::ptr;
+use std::sync::Arc;
 use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
+use std::time::Duration;
 
 use bytes::buf::BufMut;
-use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
+use ignore_result::Ignore;
+use rustls::pki_types::ServerName;
+use rustls::ClientConfig;
+use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufStream, ReadBuf};
 use tokio::net::TcpStream;
+use tokio::{select, time};
 use tokio_rustls::client::TlsStream;
+use tokio_rustls::TlsConnector;
+
+use crate::deadline::Deadline;
+use crate::endpoint::{EndpointRef, IterableEndpoints};
 
 const NOOP_VTABLE: RawWakerVTable =
     RawWakerVTable::new(|_| RawWaker::new(ptr::null(), &NOOP_VTABLE), |_| {}, |_| {}, |_| {});
@@ -126,4 +136,90 @@ impl Connection {
             Poll::Ready(result) => result,
         }
     }
+
+    pub async fn command(self, cmd: &str) -> Result<String> {
+        let mut stream = BufStream::new(self);
+        stream.write_all(cmd.as_bytes()).await?;
+        stream.flush().await?;
+        let mut line = String::new();
+        stream.read_line(&mut line).await?;
+        stream.shutdown().await.ignore();
+        Ok(line)
+    }
+
+    pub async fn command_isro(self) -> Result<bool> {
+        let r = self.command("isro").await?;
+        if r == "rw" {
+            Ok(true)
+        } else {
+            Ok(false)
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct Connector {
+    tls: TlsConnector,
+    timeout: Duration,
+}
+
+impl Connector {
+    pub fn new(config: impl Into<Arc<ClientConfig>>) -> Self {
+        Self { tls: TlsConnector::from(config.into()), timeout: Duration::from_secs(10) }
+    }
+
+    pub fn timeout(&self) -> Duration {
+        self.timeout
+    }
+
+    pub fn set_timeout(&mut self, timeout: Duration) {
+        self.timeout = timeout;
+    }
+
+    pub async fn connect(&self, endpoint: EndpointRef<'_>, deadline: &mut Deadline) -> Result<Connection> {
+        select! {
+            _ = unsafe { Pin::new_unchecked(deadline) } => Err(Error::new(ErrorKind::TimedOut, "deadline exceed")),
+            _ = time::sleep(self.timeout) => Err(Error::new(ErrorKind::TimedOut, format!("connection timeout{:?} exceed", self.timeout))),
+            r = TcpStream::connect((endpoint.host, endpoint.port)) => {
+                match r {
+                    Err(err) => Err(err),
+                    Ok(sock) => {
+                        let connection = if endpoint.tls {
+                            let domain = ServerName::try_from(endpoint.host).unwrap().to_owned();
+                            let stream = self.tls.connect(domain, sock).await?;
+                            Connection::new_tls(stream)
+                        } else {
+                            Connection::new_raw(sock)
+                        };
+                        Ok(connection)
+                    },
+                }
+            },
+        }
+    }
+
+    pub async fn seek_for_writable(self, endpoints: &mut IterableEndpoints) -> Option<EndpointRef<'_>> {
+        let n = endpoints.len();
+        let max_timeout = Duration::from_secs(60);
+        let mut i = 0;
+        let mut timeout = Duration::from_millis(100);
+        let mut deadline = Deadline::never();
+        while let Some(endpoint) = endpoints.peek() {
+            i += 1;
+            if let Ok(conn) = self.connect(endpoint, &mut deadline).await {
+                if let Ok(true) = conn.command_isro().await {
+                    return Some(unsafe { std::mem::transmute(endpoint) });
+                }
+            }
+            endpoints.step();
+            if i % n == 0 {
+                log::debug!("ZooKeeper fails to contact writable server from endpoints {:?}", endpoints.endpoints());
+                time::sleep(timeout).await;
+                timeout = max_timeout.min(timeout * 2);
+            } else {
+                time::sleep(Duration::from_millis(5)).await;
+            }
+        }
+        None
+    }
 }