diff --git a/Cargo.lock b/Cargo.lock index 810141a..0e38d14 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -82,6 +82,26 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9555578bc9e57714c812a1f84e4fc5b4d21fcb063490c624de019f7464c91268" +[[package]] +name = "chacha20" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" +dependencies = [ + "cfg-if", + "cpufeatures", + "rand_core", +] + +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + [[package]] name = "equivalent" version = "1.0.2" @@ -198,6 +218,7 @@ dependencies = [ "cfg-if", "libc", "r-efi", + "rand_core", "wasip2", "wasip3", ] @@ -444,6 +465,23 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "rand" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" +dependencies = [ + "chacha20", + "getrandom", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + [[package]] name = "regex" version = "1.12.3" @@ -656,6 +694,7 @@ dependencies = [ "byteorder", "futures", "pin-project", + "rand", "snafu", "tokio", "tracing", diff --git a/Cargo.toml b/Cargo.toml index b5446ad..4a577f2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,6 +29,7 @@ snafu = "0.8.6" tokio = { version = "1.47.1", features = ["net", "rt", "time"] } tracing = "0.1.41" uuid = { version = "1.23.1", features = ["v4"], optional = true } +rand = "0.10.2" [dev-dependencies] tokio = { version = "1.47.1", features = ["macros", "rt-multi-thread"] } diff --git a/src/lib.rs b/src/lib.rs index fe60c0d..23f63d0 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -194,12 +194,14 @@ #![deny(missing_copy_implementations)] use error::Error; -use futures::{channel::oneshot, Stream}; -use snafu::{whatever as bail, ResultExt}; +use futures::{Stream, channel::oneshot}; +use rand::rng; +use rand::seq::SliceRandom; +use snafu::{whatever as bail, whatever}; use std::borrow::Cow; use std::net::SocketAddr; use std::time; -use tracing::{debug, instrument, trace}; +use tracing::{debug, instrument, trace, warn}; /// Per-operation ZooKeeper error types. pub mod error; @@ -265,7 +267,7 @@ impl Default for ZooKeeperBuilder { } impl ZooKeeperBuilder { - /// Connect to a ZooKeeper server instance at the given address. + /// Connect to a ZooKeeper server instance from the given quorum of addresses. /// /// A `ZooKeeper` instance is returned, along with a "watcher" that will provide notifications /// of any changes in state. @@ -275,13 +277,19 @@ impl ZooKeeperBuilder { /// during a disconnect may fail and have to be retried. pub async fn connect( self, - addr: &SocketAddr, + mut addrs: Vec, ) -> Result<(ZooKeeper, impl Stream), Error> { let (tx, rx) = futures::channel::mpsc::unbounded(); - let stream = tokio::net::TcpStream::connect(addr) - .await - .whatever_context("connect failed")?; - Ok((self.handshake(*addr, stream, tx).await?, rx)) + addrs.shuffle(&mut rng()); + for addr in addrs { + match tokio::net::TcpStream::connect(addr).await { + Ok(stream) => return Ok((self.handshake(addr, stream, tx).await?, rx)), + Err(err) => { + warn!("connection failed on address {}: {}", addr, err); + } + } + } + whatever!("Could not connect to any node in quorum") } /// Set the ZooKeeper [session expiry @@ -324,9 +332,9 @@ impl ZooKeeper { /// /// See [`ZooKeeperBuilder::connect`]. pub async fn connect( - addr: &SocketAddr, + addrs: Vec, ) -> Result<(Self, impl Stream), Error> { - ZooKeeperBuilder::default().connect(addr).await + ZooKeeperBuilder::default().connect(addrs).await } /// Create a node with the given `path` with `data` as its contents. @@ -769,8 +777,8 @@ mod tests { init_tracing_subscriber(); let builder = ZooKeeperBuilder::default(); - let connect_addr = "127.0.0.1:2181".parse().unwrap(); - let (zk, w) = builder.connect(&connect_addr).await.unwrap(); + let connect_addr: Vec = vec!["127.0.0.1:2181".parse().unwrap()]; + let (zk, w) = builder.connect(connect_addr).await.unwrap(); let (exists_w, stat) = zk.with_watcher().exists("/foo").await.unwrap(); assert_eq!(stat, None); let stat = zk.watch().exists("/foo").await.unwrap(); @@ -875,8 +883,8 @@ mod tests { #[tokio::test] async fn example() { - let connect_addr = "127.0.0.1:2181".parse().unwrap(); - let (zk, default_watcher) = ZooKeeper::connect(&connect_addr).await.unwrap(); + let connect_addr: Vec = vec!["127.0.0.1:2181".parse().unwrap()]; + let (zk, default_watcher) = ZooKeeper::connect(connect_addr).await.unwrap(); // let's first check if /example exists. the .watch() causes us to be notified // the next time the "exists" status of /example changes after the call. @@ -960,10 +968,9 @@ mod tests { async fn acl_test() { init_tracing_subscriber(); let builder = ZooKeeperBuilder::default(); + let connect_addr: Vec = vec!["127.0.0.1:2181".parse().unwrap()]; - let (zk, _) = (builder.connect(&"127.0.0.1:2181".parse().unwrap())) - .await - .unwrap(); + let (zk, _) = (builder.connect(connect_addr)).await.unwrap(); let _ = zk .create( "/acl_test", @@ -1024,10 +1031,8 @@ mod tests { Result::<_, Error>::Ok(res) } - let (zk, _) = builder - .connect(&"127.0.0.1:2181".parse().unwrap()) - .await - .unwrap(); + let connect_addr: Vec = vec!["127.0.0.1:2181".parse().unwrap()]; + let (zk, _) = builder.connect(connect_addr).await.unwrap(); let res = zk .multi() diff --git a/src/recipes/leader/mod.rs b/src/recipes/leader/mod.rs index 9656aa3..8a08985 100644 --- a/src/recipes/leader/mod.rs +++ b/src/recipes/leader/mod.rs @@ -449,6 +449,7 @@ async fn get_children( #[cfg(test)] mod tests { + use std::net::SocketAddr; use std::time::Duration; use super::*; @@ -481,11 +482,11 @@ mod tests { #[tokio::test] async fn election_works() { let builder = ZooKeeperBuilder::default(); - let connect_addr = "127.0.0.1:2181".parse().unwrap(); + let connect_addr: Vec = vec!["127.0.0.1:2181".parse().unwrap()]; init_tracing_subscriber(); - let (zk1, _w) = builder.connect(&connect_addr).await.unwrap(); + let (zk1, _w) = builder.connect(connect_addr.clone()).await.unwrap(); create_election_node(&zk1).await; let leader_election1 = LeaderElection::new(zk1, "/election", Acl::open_unsafe().to_vec()); let (mut rx1, jh1) = leader_election1.volunteer().await.unwrap(); @@ -493,7 +494,7 @@ mod tests { .await .expect("the first participant should be the leader"); - let (zk2, _w) = builder.connect(&connect_addr).await.unwrap(); + let (zk2, _w) = builder.connect(connect_addr).await.unwrap(); let leader_election2 = LeaderElection::new(zk2, "/election", Acl::open_unsafe().to_vec()); let (mut rx2, _jh2) = leader_election2.volunteer().await.unwrap();