From 3b804d0af21b412b8be2ed15eedac116a82faef9 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Wed, 17 Jul 2024 14:07:20 +0700 Subject: [PATCH 01/42] init structure of balancer --- ydb/src/client_builder.rs | 10 +++++++--- ydb/src/load_balancer.rs | 39 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/ydb/src/client_builder.rs b/ydb/src/client_builder.rs index d5f91cb4..52962483 100644 --- a/ydb/src/client_builder.rs +++ b/ydb/src/client_builder.rs @@ -158,7 +158,7 @@ fn token_static_password(uri: &str, mut client_builder: ClientBuilder) -> YdbRes password, endpoint, client_builder.database.clone(), - ) + ), }; client_builder.credentials = credencials_ref(creds); @@ -238,8 +238,12 @@ impl ClientBuilder { interceptor.with_interceptor(DiscoveryPessimizationInterceptor::new(discovery.clone())); let load_balancer = SharedLoadBalancer::new(discovery.as_ref().as_ref()); - let connection_manager = - GrpcConnectionManager::new(load_balancer, db_cred.database.clone(), interceptor, self.cert_path); + let connection_manager = GrpcConnectionManager::new( + load_balancer, + db_cred.database.clone(), + interceptor, + self.cert_path, + ); Client::new(db_cred, discovery, connection_manager) } diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index 4564e5c2..ca3f1e4a 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -182,6 +182,45 @@ pub(crate) async fn update_load_balancer( } } +pub(crate) struct NearestDCBalancer { + discovery_state: Arc, + waiter: Arc, + random_balancer: Option, +} + +impl NearestDCBalancer { + pub(crate) fn new(backup_with_random_balancer: bool) -> Self { + let mut random_balancer = None; + if backup_with_random_balancer { + random_balancer = Some(RandomLoadBalancer::new()) + } + Self { + discovery_state: Arc::new(DiscoveryState::default()), + waiter: Arc::new(WaiterImpl::new()), + random_balancer, + } + } +} + +#[async_trait::async_trait] +impl Waiter for NearestDCBalancer { + async fn wait(&self) -> YdbResult<()> { + todo!() + } +} + +impl LoadBalancer for NearestDCBalancer { + fn endpoint(&self, service: Service) -> YdbResult { + todo!() + } + fn set_discovery_state(&mut self, discovery_state: &Arc) -> YdbResult<()> { + todo!() + } + fn waiter(&self) -> Box { + todo!() + } +} + #[cfg(test)] mod test { use super::*; From 44e86a49c14cea672969f236224535e76aafaae1 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 28 Jul 2024 22:50:46 +0700 Subject: [PATCH 02/42] Support location inside NodeInfo --- ydb/src/discovery.rs | 19 +++++++++++++++---- ydb/src/grpc_wrapper/raw_discovery_client.rs | 2 ++ 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/ydb/src/discovery.rs b/ydb/src/discovery.rs index 017b836b..c5d32d0b 100644 --- a/ydb/src/discovery.rs +++ b/ydb/src/discovery.rs @@ -97,11 +97,19 @@ impl Default for DiscoveryState { #[derive(Clone, Debug, PartialEq)] pub(crate) struct NodeInfo { pub(crate) uri: Uri, + pub(crate) location: String, } impl NodeInfo { pub(crate) fn new(uri: Uri) -> Self { - Self { uri } + Self { + uri, + location: String::new(), + } + } + + pub(crate) fn new_with_location(uri: Uri, location: String) -> Self { + Self { uri, location } } } @@ -324,14 +332,17 @@ impl DiscoverySharedState { fn list_endpoints_to_node_infos(list: Vec) -> YdbResult> { list.into_iter() - .map(|item| match Self::endpoint_info_to_uri(item) { - Ok(uri) => YdbResult::::Ok(NodeInfo::new(uri)), + .map(|item| match Self::endpoint_info_to_uri(&item) { + Ok(uri) => YdbResult::::Ok(NodeInfo::new_with_location( + uri, + item.location.clone(), + )), Err(err) => YdbResult::::Err(err), }) .try_collect() } - fn endpoint_info_to_uri(endpoint_info: EndpointInfo) -> YdbResult { + fn endpoint_info_to_uri(endpoint_info: &EndpointInfo) -> YdbResult { let authority: Authority = Authority::from_str(format!("{}:{}", endpoint_info.fqdn, endpoint_info.port).as_str())?; diff --git a/ydb/src/grpc_wrapper/raw_discovery_client.rs b/ydb/src/grpc_wrapper/raw_discovery_client.rs index 5a0996aa..e4d14f9e 100644 --- a/ydb/src/grpc_wrapper/raw_discovery_client.rs +++ b/ydb/src/grpc_wrapper/raw_discovery_client.rs @@ -35,6 +35,7 @@ impl GrpcDiscoveryClient { fqdn: item.address, port: item.port, ssl: item.ssl, + location: item.location }) .collect_vec(); Ok(res) @@ -45,6 +46,7 @@ pub(crate) struct EndpointInfo { pub(crate) fqdn: String, pub(crate) port: u32, pub(crate) ssl: bool, + pub(crate) location: String, } impl GrpcServiceForDiscovery for GrpcDiscoveryClient { From b65b0d083d58972667bf9f4ed70edbb57f1010a8 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 28 Jul 2024 22:52:25 +0700 Subject: [PATCH 03/42] Balancer logic (need to optimize) --- ydb/src/load_balancer.rs | 161 ++++++++++++++++++++++++++++++++++----- 1 file changed, 143 insertions(+), 18 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index ca3f1e4a..2eeaf9c4 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -1,9 +1,14 @@ -use crate::discovery::{Discovery, DiscoveryState}; +use crate::discovery::{Discovery, DiscoveryState, NodeInfo}; use crate::errors::*; use http::Uri; +use rand::thread_rng; +use tokio::runtime::Handle; +use tokio_util::sync::CancellationToken; use crate::grpc_wrapper::raw_services::Service; use crate::waiter::{Waiter, WaiterImpl}; +use std::collections::HashMap; +use std::net::ToSocketAddrs; use std::sync::{Arc, RwLock}; use tokio::sync::watch::Receiver; @@ -183,41 +188,161 @@ pub(crate) async fn update_load_balancer( } pub(crate) struct NearestDCBalancer { - discovery_state: Arc, - waiter: Arc, - random_balancer: Option, + random_balancer: RandomLoadBalancer, } impl NearestDCBalancer { - pub(crate) fn new(backup_with_random_balancer: bool) -> Self { - let mut random_balancer = None; - if backup_with_random_balancer { - random_balancer = Some(RandomLoadBalancer::new()) - } - Self { - discovery_state: Arc::new(DiscoveryState::default()), - waiter: Arc::new(WaiterImpl::new()), - random_balancer, - } + pub(crate) fn new() -> Self { + let random_balancer = RandomLoadBalancer::new(); + Self { random_balancer } } } #[async_trait::async_trait] impl Waiter for NearestDCBalancer { async fn wait(&self) -> YdbResult<()> { - todo!() + self.random_balancer.wait().await } } +const NODES_PER_DC: usize = 5; + impl LoadBalancer for NearestDCBalancer { fn endpoint(&self, service: Service) -> YdbResult { - todo!() + let nodes = self.random_balancer.discovery_state.get_nodes(&service); + match nodes { + None => Err(YdbError::Custom(format!( + "no endpoints for service: '{}'", + service + ))), + Some(nodes) => self.get_local_endpoint(nodes), + } } + fn set_discovery_state(&mut self, discovery_state: &Arc) -> YdbResult<()> { - todo!() + self.random_balancer.set_discovery_state(discovery_state) } fn waiter(&self) -> Box { - todo!() + self.random_balancer.waiter() + } +} + +impl NearestDCBalancer { + fn get_local_endpoint(&self, nodes: &Vec) -> YdbResult { + let mut dc_to_nodes = self.split_endpoints_by_location(nodes); + let mut to_check = Vec::with_capacity(NODES_PER_DC * dc_to_nodes.keys().len()); + + dc_to_nodes + .iter_mut() + .for_each(|(_, eps)| to_check.append(self.get_random_endpoints(eps))); + + if to_check.len() == 0 { + return Err(YdbError::Custom(format!("no endpoints"))); + } + + self.find_fastest_endpoint(&to_check) + } + + fn split_endpoints_by_location(&self, nodes: &Vec) -> HashMap> { + // no clones =( ? + let mut dc_to_eps = HashMap::>::new(); + nodes.into_iter().for_each(|info| { + if let Some(nodes) = dc_to_eps.get_mut(&info.location) { + nodes.push(info.clone()); + } else { + dc_to_eps.insert(info.location.clone(), vec![info.clone()]); + } + }); + dc_to_eps + } + + fn get_random_endpoints<'a>( + &'a self, + dc_endpoints: &'a mut Vec, + ) -> &mut Vec { + use rand::seq::SliceRandom; + dc_endpoints.shuffle(&mut thread_rng()); + dc_endpoints.truncate(NODES_PER_DC); + dc_endpoints + } + + fn find_fastest_endpoint(&self, to_check: &Vec) -> YdbResult { + let addr_to_node = self.addr_to_node(to_check); + if addr_to_node.len() == 0 { + return Err(YdbError::Custom(format!("no ips for endpoints"))); + } + + let addrs = addr_to_node.keys(); + + let fastest_address = self.find_fastest_address(addrs.collect()); + + Ok(addr_to_node[&fastest_address].uri.clone()) + } + + fn addr_to_node(&self, nodes: &Vec) -> HashMap { + let mut addr_to_node = HashMap::::with_capacity(2 * nodes.len()); // ipv4 + ipv6 + + nodes.into_iter().for_each(|info| { + let host: &str; + let port: u16; + match info.uri.host() { + Some(uri_host) => host = uri_host, + None => return, + } + + match info.uri.port() { + Some(uri_port) => port = uri_port.as_u16(), + None => return, + } + + let _ = (host, port).to_socket_addrs().and_then(|addrs| { + for addr in addrs { + addr_to_node.insert(addr.to_string(), info.clone()); + } + Ok(()) + }); + }); + + addr_to_node + } + fn find_fastest_address(&self, addrs: Vec<&String>) -> String { + let token = CancellationToken::new(); + let (fire, _) = tokio::sync::broadcast::channel::<()>(1); + let (addr_sender, mut addr_reciever) = tokio::sync::mpsc::channel::(1); + + for addr in addrs { + let mut wait_for_start = fire.subscribe(); + let stop_measure = token.clone(); + let addr = addr.clone(); + let addr_sender = addr_sender.clone(); + + tokio::spawn(async move { + let _ = wait_for_start.recv().await; + tokio::select! { + connection_result = tokio::net::TcpStream::connect(addr.clone()) =>{ + if connection_result.is_ok(){ + addr_sender.send(addr); //????? + } + } + + _ = stop_measure.cancelled() => { + return ; + } + } + }); + } + + tokio::task::block_in_place(move || { + Handle::current().block_on(async { + match addr_reciever.recv().await { + Some(address) => { + token.cancel(); + address + } + None => unreachable!(), + } + }) + }) } } From 76b9dbc5582a132399635400bb04243d711177b8 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Mon, 29 Jul 2024 14:47:59 +0700 Subject: [PATCH 04/42] Fire to start computation, await on sending addr ready, Pretty vars --- ydb/src/load_balancer.rs | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index 2eeaf9c4..32a99699 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -306,22 +306,24 @@ impl NearestDCBalancer { addr_to_node } fn find_fastest_address(&self, addrs: Vec<&String>) -> String { - let token = CancellationToken::new(); - let (fire, _) = tokio::sync::broadcast::channel::<()>(1); + let stop_measure = CancellationToken::new(); + let (start_measure, _) = tokio::sync::broadcast::channel::<()>(1); let (addr_sender, mut addr_reciever) = tokio::sync::mpsc::channel::(1); for addr in addrs { - let mut wait_for_start = fire.subscribe(); - let stop_measure = token.clone(); - let addr = addr.clone(); - let addr_sender = addr_sender.clone(); - + let (mut wait_for_start, stop_measure, addr, addr_sender) = ( + start_measure.subscribe(), + stop_measure.clone(), + addr.clone(), + addr_sender.clone(), + ); + tokio::spawn(async move { let _ = wait_for_start.recv().await; tokio::select! { connection_result = tokio::net::TcpStream::connect(addr.clone()) =>{ if connection_result.is_ok(){ - addr_sender.send(addr); //????? + addr_sender.send(addr).await.unwrap(); } } @@ -333,10 +335,11 @@ impl NearestDCBalancer { } tokio::task::block_in_place(move || { + let _ = start_measure.send(()); Handle::current().block_on(async { match addr_reciever.recv().await { Some(address) => { - token.cancel(); + stop_measure.cancel(); address } None => unreachable!(), From 6a999e31052a475759199f81813d13575b992977 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Wed, 31 Jul 2024 13:15:42 +0700 Subject: [PATCH 05/42] Structured concurrency --- ydb/src/load_balancer.rs | 39 +++++++++++++++++++++++++++------------ 1 file changed, 27 insertions(+), 12 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index 32a99699..7fca04bc 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -2,7 +2,11 @@ use crate::discovery::{Discovery, DiscoveryState, NodeInfo}; use crate::errors::*; use http::Uri; use rand::thread_rng; +use tokio::io::AsyncWriteExt; +use tokio::net::TcpStream; use tokio::runtime::Handle; +use tokio::sync::mpsc; +use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; use crate::grpc_wrapper::raw_services::Service; @@ -308,7 +312,9 @@ impl NearestDCBalancer { fn find_fastest_address(&self, addrs: Vec<&String>) -> String { let stop_measure = CancellationToken::new(); let (start_measure, _) = tokio::sync::broadcast::channel::<()>(1); - let (addr_sender, mut addr_reciever) = tokio::sync::mpsc::channel::(1); + let (addr_sender, addr_reciever) = tokio::sync::mpsc::channel::(1); + + let mut nursery = JoinSet::new(); for addr in addrs { let (mut wait_for_start, stop_measure, addr, addr_sender) = ( @@ -317,13 +323,14 @@ impl NearestDCBalancer { addr.clone(), addr_sender.clone(), ); - - tokio::spawn(async move { + + nursery.spawn(async move { let _ = wait_for_start.recv().await; tokio::select! { - connection_result = tokio::net::TcpStream::connect(addr.clone()) =>{ + connection_result = TcpStream::connect(addr.clone()) =>{ if connection_result.is_ok(){ - addr_sender.send(addr).await.unwrap(); + let _ = connection_result.unwrap().shutdown().await; + addr_sender.send(addr).await.unwrap(); } } @@ -337,16 +344,24 @@ impl NearestDCBalancer { tokio::task::block_in_place(move || { let _ = start_measure.send(()); Handle::current().block_on(async { - match addr_reciever.recv().await { - Some(address) => { - stop_measure.cancel(); - address - } - None => unreachable!(), - } + let addr = self.wait_for_address(addr_reciever).await; + stop_measure.cancel(); + self.join_all(&mut nursery).await; + addr }) }) } + + async fn wait_for_address(&self, mut addr_reciever: mpsc::Receiver) -> String { + match addr_reciever.recv().await { + Some(address) => address, + None => unreachable!(), + } + } + + async fn join_all(&self, awaitable: &mut JoinSet<()>) { + while let Some(_) = awaitable.join_next().await {} + } } #[cfg(test)] From aa441be42500f6a29e479cd51119d1d97413c406 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Wed, 31 Jul 2024 13:21:10 +0700 Subject: [PATCH 06/42] Construct NodeInfo with location --- ydb/src/discovery.rs | 13 +++---------- ydb/src/load_balancer.rs | 6 +++--- 2 files changed, 6 insertions(+), 13 deletions(-) diff --git a/ydb/src/discovery.rs b/ydb/src/discovery.rs index c5d32d0b..b5471e31 100644 --- a/ydb/src/discovery.rs +++ b/ydb/src/discovery.rs @@ -101,14 +101,7 @@ pub(crate) struct NodeInfo { } impl NodeInfo { - pub(crate) fn new(uri: Uri) -> Self { - Self { - uri, - location: String::new(), - } - } - - pub(crate) fn new_with_location(uri: Uri, location: String) -> Self { + pub(crate) fn new(uri: Uri, location: String) -> Self { Self { uri, location } } } @@ -149,7 +142,7 @@ pub struct StaticDiscovery { impl StaticDiscovery { pub fn new_from_str<'a, T: Into<&'a str>>(endpoint: T) -> YdbResult { let endpoint = Uri::from_str(endpoint.into())?; - let nodes = vec![NodeInfo::new(endpoint)]; + let nodes = vec![NodeInfo::new(endpoint, String::new())]; let state = DiscoveryState::new(std::time::Instant::now(), nodes); let state = Arc::new(state); @@ -333,7 +326,7 @@ impl DiscoverySharedState { fn list_endpoints_to_node_infos(list: Vec) -> YdbResult> { list.into_iter() .map(|item| match Self::endpoint_info_to_uri(&item) { - Ok(uri) => YdbResult::::Ok(NodeInfo::new_with_location( + Ok(uri) => YdbResult::::Ok(NodeInfo::new( uri, item.location.clone(), )), diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index 7fca04bc..cf32f674 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -409,7 +409,7 @@ mod test { let new_discovery_state = Arc::new(DiscoveryState::default().with_node_info( Table, - NodeInfo::new(Uri::from_str("http://test.com").unwrap()), + NodeInfo::new(Uri::from_str("http://test.com").unwrap(), String::new()), )); let (first_update_sender, first_update_receiver) = tokio::sync::oneshot::channel(); @@ -473,8 +473,8 @@ mod test { let load_balancer = RandomLoadBalancer { discovery_state: Arc::new( DiscoveryState::default() - .with_node_info(Table, NodeInfo::new(one.clone())) - .with_node_info(Table, NodeInfo::new(two.clone())), + .with_node_info(Table, NodeInfo::new(one.clone(), String::new())) + .with_node_info(Table, NodeInfo::new(two.clone(), String::new())), ), waiter: Arc::new(WaiterImpl::new()), }; From 15316176605ff684c4d109f95e84eb8136a17379 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Wed, 31 Jul 2024 13:38:48 +0700 Subject: [PATCH 07/42] Balancr backup strategy options --- ydb/src/load_balancer.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index cf32f674..ffe45881 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -191,6 +191,13 @@ pub(crate) async fn update_load_balancer( } } +// What will balancer do if there is no available endpoints at local dc +pub(crate) enum NearestDCStrategy { + Error, // Just throw error + Random, // Random endpoint from other dcs + Next, // Find next fastest dc +} + pub(crate) struct NearestDCBalancer { random_balancer: RandomLoadBalancer, } From 8072d92dd02a7cb847a2e028000a1315896ee3bd Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 4 Aug 2024 17:59:21 +0700 Subject: [PATCH 08/42] first draft --- ydb/src/discovery.rs | 16 +- ydb/src/grpc_wrapper/raw_discovery_client.rs | 14 +- ydb/src/load_balancer.rs | 175 ++++++++++++------- 3 files changed, 141 insertions(+), 64 deletions(-) diff --git a/ydb/src/discovery.rs b/ydb/src/discovery.rs index b5471e31..f506bf6f 100644 --- a/ydb/src/discovery.rs +++ b/ydb/src/discovery.rs @@ -63,6 +63,10 @@ impl DiscoveryState { Some(&self.nodes) } + pub(crate) fn get_all_nodes(&self) -> Option<&Vec> { + Some(&self.nodes) + } + pub(crate) fn is_empty(&self) -> bool { self.nodes.len() == 0 } @@ -98,11 +102,16 @@ impl Default for DiscoveryState { pub(crate) struct NodeInfo { pub(crate) uri: Uri, pub(crate) location: String, + pub(crate) service: Vec, } impl NodeInfo { - pub(crate) fn new(uri: Uri, location: String) -> Self { - Self { uri, location } + pub(crate) fn new(uri: Uri, location: String, service: Vec) -> Self { + Self { + uri, + location, + service, + } } } @@ -142,7 +151,7 @@ pub struct StaticDiscovery { impl StaticDiscovery { pub fn new_from_str<'a, T: Into<&'a str>>(endpoint: T) -> YdbResult { let endpoint = Uri::from_str(endpoint.into())?; - let nodes = vec![NodeInfo::new(endpoint, String::new())]; + let nodes = vec![NodeInfo::new(endpoint, String::new(), Vec::new())]; let state = DiscoveryState::new(std::time::Instant::now(), nodes); let state = Arc::new(state); @@ -329,6 +338,7 @@ impl DiscoverySharedState { Ok(uri) => YdbResult::::Ok(NodeInfo::new( uri, item.location.clone(), + item.service.clone(), )), Err(err) => YdbResult::::Err(err), }) diff --git a/ydb/src/grpc_wrapper/raw_discovery_client.rs b/ydb/src/grpc_wrapper/raw_discovery_client.rs index e4d14f9e..96fc28f8 100644 --- a/ydb/src/grpc_wrapper/raw_discovery_client.rs +++ b/ydb/src/grpc_wrapper/raw_discovery_client.rs @@ -3,9 +3,12 @@ use crate::grpc_wrapper::raw_services::{GrpcServiceForDiscovery, Service}; use crate::grpc_wrapper::runtime_interceptors::InterceptedChannel; use crate::YdbResult; use itertools::Itertools; +use serde_json::from_str; use ydb_grpc::ydb_proto::discovery::v1::discovery_service_client::DiscoveryServiceClient; use ydb_grpc::ydb_proto::discovery::{ListEndpointsRequest, ListEndpointsResult}; +use super::raw_services; + pub struct GrpcDiscoveryClient { service: DiscoveryServiceClient, } @@ -28,6 +31,8 @@ impl GrpcDiscoveryClient { }; let resp = self.service.list_endpoints(req).await?; let result: ListEndpointsResult = grpc_read_operation_result(resp)?; + + use std::str::FromStr; let res = result .endpoints .into_iter() @@ -35,7 +40,13 @@ impl GrpcDiscoveryClient { fqdn: item.address, port: item.port, ssl: item.ssl, - location: item.location + location: item.location, + service: item + .service + .into_iter() + .map(|s| Service::from_str(s.as_str())) + .map(|res| res.unwrap()) + .collect(), }) .collect_vec(); Ok(res) @@ -47,6 +58,7 @@ pub(crate) struct EndpointInfo { pub(crate) port: u32, pub(crate) ssl: bool, pub(crate) location: String, + pub(crate) service: Vec, } impl GrpcServiceForDiscovery for GrpcDiscoveryClient { diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index ffe45881..c9be3435 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -1,6 +1,7 @@ use crate::discovery::{Discovery, DiscoveryState, NodeInfo}; use crate::errors::*; use http::Uri; +use itertools::Itertools; use rand::thread_rng; use tokio::io::AsyncWriteExt; use tokio::net::TcpStream; @@ -12,7 +13,9 @@ use tokio_util::sync::CancellationToken; use crate::grpc_wrapper::raw_services::Service; use crate::waiter::{Waiter, WaiterImpl}; use std::collections::HashMap; +use std::fmt::format; use std::net::ToSocketAddrs; +use std::str::FromStr; use std::sync::{Arc, RwLock}; use tokio::sync::watch::Receiver; @@ -192,7 +195,7 @@ pub(crate) async fn update_load_balancer( } // What will balancer do if there is no available endpoints at local dc -pub(crate) enum NearestDCStrategy { +pub(crate) enum FallbackStrategy { Error, // Just throw error Random, // Random endpoint from other dcs Next, // Find next fastest dc @@ -200,12 +203,20 @@ pub(crate) enum NearestDCStrategy { pub(crate) struct NearestDCBalancer { random_balancer: RandomLoadBalancer, + fallback_strategy: FallbackStrategy, //Default - Error + allowed_endpoints: Vec, + location: String, // vec of locations sorted by speed. Next strategy??? } impl NearestDCBalancer { pub(crate) fn new() -> Self { let random_balancer = RandomLoadBalancer::new(); - Self { random_balancer } + Self { + random_balancer, + fallback_strategy: FallbackStrategy::Random, + allowed_endpoints: Vec::new(), + location: String::new(), + } } } @@ -216,52 +227,84 @@ impl Waiter for NearestDCBalancer { } } -const NODES_PER_DC: usize = 5; - impl LoadBalancer for NearestDCBalancer { fn endpoint(&self, service: Service) -> YdbResult { - let nodes = self.random_balancer.discovery_state.get_nodes(&service); - match nodes { - None => Err(YdbError::Custom(format!( - "no endpoints for service: '{}'", - service - ))), - Some(nodes) => self.get_local_endpoint(nodes), - } + self.get_endpoint(service) } fn set_discovery_state(&mut self, discovery_state: &Arc) -> YdbResult<()> { - self.random_balancer.set_discovery_state(discovery_state) + self.random_balancer.set_discovery_state(discovery_state)?; + self.adjust_local_dc()?; + self.adjust_allowed_endpoints() } + fn waiter(&self) -> Box { self.random_balancer.waiter() } } +const NODES_PER_DC: usize = 5; + impl NearestDCBalancer { - fn get_local_endpoint(&self, nodes: &Vec) -> YdbResult { + fn get_endpoint(&self, service: Service) -> YdbResult { + for ep in self.allowed_endpoints.iter() { + if ep.service.contains(&service) { + return YdbResult::Ok(ep.uri.clone()); + } + } + + match self.fallback_strategy { + FallbackStrategy::Error => Err(YdbError::custom(format!( + "no available endpoints for service {} in local dc", + service + ))), + FallbackStrategy::Random => self.random_balancer.endpoint(service), + FallbackStrategy::Next => todo!(), + } + } + + fn adjust_local_dc(&mut self) -> YdbResult<()> { + let nodes = self.get_nodes()?; let mut dc_to_nodes = self.split_endpoints_by_location(nodes); let mut to_check = Vec::with_capacity(NODES_PER_DC * dc_to_nodes.keys().len()); - dc_to_nodes .iter_mut() - .for_each(|(_, eps)| to_check.append(self.get_random_endpoints(eps))); + .for_each(|(_, endpoints)| to_check.append(self.get_random_endpoints(endpoints))); + let local_dc = self.find_local_dc(&to_check)?; + self.location = local_dc; + Ok(()) + } - if to_check.len() == 0 { - return Err(YdbError::Custom(format!("no endpoints"))); - } + fn adjust_allowed_endpoints(&mut self) -> YdbResult<()> { + self.allowed_endpoints = self + .get_nodes()? + .into_iter() + .filter(|ep| ep.location == self.location) + .map(|ep| ep.clone()) + .collect_vec(); + Ok(()) + } - self.find_fastest_endpoint(&to_check) + fn get_nodes(&self) -> YdbResult<&Vec> { + let nodes = self.random_balancer.discovery_state.get_all_nodes(); + match nodes { + None => Err(YdbError::Custom(format!( + "no endpoints on discovery update" + ))), + Some(nodes) => Ok(nodes), + } } - fn split_endpoints_by_location(&self, nodes: &Vec) -> HashMap> { - // no clones =( ? - let mut dc_to_eps = HashMap::>::new(); + fn split_endpoints_by_location<'a>( + &'a self, + nodes: &'a Vec, + ) -> HashMap> { + let mut dc_to_eps = HashMap::>::new(); nodes.into_iter().for_each(|info| { if let Some(nodes) = dc_to_eps.get_mut(&info.location) { - nodes.push(info.clone()); + nodes.push(info); } else { - dc_to_eps.insert(info.location.clone(), vec![info.clone()]); + dc_to_eps.insert(info.location.clone(), vec![info]); } }); dc_to_eps @@ -269,30 +312,28 @@ impl NearestDCBalancer { fn get_random_endpoints<'a>( &'a self, - dc_endpoints: &'a mut Vec, - ) -> &mut Vec { + dc_endpoints: &'a mut Vec<&'a NodeInfo>, + ) -> &mut Vec<&NodeInfo> { use rand::seq::SliceRandom; dc_endpoints.shuffle(&mut thread_rng()); dc_endpoints.truncate(NODES_PER_DC); dc_endpoints } - fn find_fastest_endpoint(&self, to_check: &Vec) -> YdbResult { + fn find_local_dc(&self, to_check: &[&NodeInfo]) -> YdbResult { let addr_to_node = self.addr_to_node(to_check); - if addr_to_node.len() == 0 { - return Err(YdbError::Custom(format!("no ips for endpoints"))); + if addr_to_node.is_empty() { + return Err(YdbError::Custom(format!("no ip addresses for endpoints"))); } - let addrs = addr_to_node.keys(); - - let fastest_address = self.find_fastest_address(addrs.collect()); - - Ok(addr_to_node[&fastest_address].uri.clone()) + match self.find_fastest_address(addrs.collect()) { + Some(fastest_address) => Ok(addr_to_node[&fastest_address].location.clone()), + None => Err(YdbError::custom("could not find fastest address")), + } } - fn addr_to_node(&self, nodes: &Vec) -> HashMap { - let mut addr_to_node = HashMap::::with_capacity(2 * nodes.len()); // ipv4 + ipv6 - + fn addr_to_node<'a>(&'a self, nodes: &[&'a NodeInfo]) -> HashMap { + let mut addr_to_node = HashMap::::with_capacity(2 * nodes.len()); // ipv4 + ipv6 nodes.into_iter().for_each(|info| { let host: &str; let port: u16; @@ -300,27 +341,25 @@ impl NearestDCBalancer { Some(uri_host) => host = uri_host, None => return, } - match info.uri.port() { Some(uri_port) => port = uri_port.as_u16(), None => return, } - let _ = (host, port).to_socket_addrs().and_then(|addrs| { for addr in addrs { - addr_to_node.insert(addr.to_string(), info.clone()); + addr_to_node.insert(addr.to_string(), info); } Ok(()) }); }); - addr_to_node } - fn find_fastest_address(&self, addrs: Vec<&String>) -> String { + + fn find_fastest_address(&self, addrs: Vec<&String>) -> Option { let stop_measure = CancellationToken::new(); let (start_measure, _) = tokio::sync::broadcast::channel::<()>(1); - let (addr_sender, addr_reciever) = tokio::sync::mpsc::channel::(1); - + let (addr_sender, mut addr_reciever) = tokio::sync::mpsc::channel::>(1); + let addr_count = addrs.len(); let mut nursery = JoinSet::new(); for addr in addrs { @@ -335,14 +374,16 @@ impl NearestDCBalancer { let _ = wait_for_start.recv().await; tokio::select! { connection_result = TcpStream::connect(addr.clone()) =>{ - if connection_result.is_ok(){ - let _ = connection_result.unwrap().shutdown().await; - addr_sender.send(addr).await.unwrap(); + match connection_result{ + Ok(mut connection) => { + let _ = connection.shutdown().await; + addr_sender.send(Some(addr)).await; + }, + Err(_) => {addr_sender.send(None).await;}, } } - _ = stop_measure.cancelled() => { - return ; + (); } } }); @@ -351,18 +392,28 @@ impl NearestDCBalancer { tokio::task::block_in_place(move || { let _ = start_measure.send(()); Handle::current().block_on(async { - let addr = self.wait_for_address(addr_reciever).await; - stop_measure.cancel(); - self.join_all(&mut nursery).await; - addr + for _ in 0..addr_count { + match self.wait_for_single_address(&mut addr_reciever).await { + Some(address) => { + stop_measure.cancel(); + self.join_all(&mut nursery).await; + return Some(address); + } + None => continue, + } + } + None }) }) } - async fn wait_for_address(&self, mut addr_reciever: mpsc::Receiver) -> String { + async fn wait_for_single_address( + &self, + addr_reciever: &mut mpsc::Receiver>, + ) -> Option { match addr_reciever.recv().await { - Some(address) => address, - None => unreachable!(), + Some(maybe_address) => maybe_address, + None => unreachable!(), // no channel closing while awaiting address } } @@ -416,7 +467,11 @@ mod test { let new_discovery_state = Arc::new(DiscoveryState::default().with_node_info( Table, - NodeInfo::new(Uri::from_str("http://test.com").unwrap(), String::new()), + NodeInfo::new( + Uri::from_str("http://test.com").unwrap(), + String::new(), + Vec::new(), + ), )); let (first_update_sender, first_update_receiver) = tokio::sync::oneshot::channel(); @@ -480,8 +535,8 @@ mod test { let load_balancer = RandomLoadBalancer { discovery_state: Arc::new( DiscoveryState::default() - .with_node_info(Table, NodeInfo::new(one.clone(), String::new())) - .with_node_info(Table, NodeInfo::new(two.clone(), String::new())), + .with_node_info(Table, NodeInfo::new(one.clone(), String::new(), Vec::new())) + .with_node_info(Table, NodeInfo::new(two.clone(), String::new(), Vec::new())), ), waiter: Arc::new(WaiterImpl::new()), }; From 1156ef7530db08686c06cbe3d05d582712962d84 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 4 Aug 2024 18:05:57 +0700 Subject: [PATCH 09/42] less diff --- ydb/src/client_builder.rs | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/ydb/src/client_builder.rs b/ydb/src/client_builder.rs index 52962483..234ed6ad 100644 --- a/ydb/src/client_builder.rs +++ b/ydb/src/client_builder.rs @@ -158,7 +158,7 @@ fn token_static_password(uri: &str, mut client_builder: ClientBuilder) -> YdbRes password, endpoint, client_builder.database.clone(), - ), + ) }; client_builder.credentials = credencials_ref(creds); @@ -238,12 +238,7 @@ impl ClientBuilder { interceptor.with_interceptor(DiscoveryPessimizationInterceptor::new(discovery.clone())); let load_balancer = SharedLoadBalancer::new(discovery.as_ref().as_ref()); - let connection_manager = GrpcConnectionManager::new( - load_balancer, - db_cred.database.clone(), - interceptor, - self.cert_path, - ); + let connection_manager = GrpcConnectionManager::new(load_balancer, db_cred.database.clone(), interceptor, self.cert_path); Client::new(db_cred, discovery, connection_manager) } From 4a9f1d95b33d278eed37888e579c1a6af0258568 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 4 Aug 2024 18:06:32 +0700 Subject: [PATCH 10/42] less diff --- ydb/src/client_builder.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ydb/src/client_builder.rs b/ydb/src/client_builder.rs index 234ed6ad..34ae5dbd 100644 --- a/ydb/src/client_builder.rs +++ b/ydb/src/client_builder.rs @@ -238,7 +238,8 @@ impl ClientBuilder { interceptor.with_interceptor(DiscoveryPessimizationInterceptor::new(discovery.clone())); let load_balancer = SharedLoadBalancer::new(discovery.as_ref().as_ref()); - let connection_manager = GrpcConnectionManager::new(load_balancer, db_cred.database.clone(), interceptor, self.cert_path); + let connection_manager = + GrpcConnectionManager::new(load_balancer, db_cred.database.clone(), interceptor, self.cert_path); Client::new(db_cred, discovery, connection_manager) } From fcfa8c965df262bbf7004ef1f30844e66089d7d5 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 4 Aug 2024 18:06:56 +0700 Subject: [PATCH 11/42] less diff --- ydb/src/client_builder.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ydb/src/client_builder.rs b/ydb/src/client_builder.rs index 34ae5dbd..d5f91cb4 100644 --- a/ydb/src/client_builder.rs +++ b/ydb/src/client_builder.rs @@ -238,7 +238,7 @@ impl ClientBuilder { interceptor.with_interceptor(DiscoveryPessimizationInterceptor::new(discovery.clone())); let load_balancer = SharedLoadBalancer::new(discovery.as_ref().as_ref()); - let connection_manager = + let connection_manager = GrpcConnectionManager::new(load_balancer, db_cred.database.clone(), interceptor, self.cert_path); Client::new(db_cred, discovery, connection_manager) From 03d82cee51bd7a66fd6a2c3a41e17ab9d4f972d2 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 4 Aug 2024 18:18:21 +0700 Subject: [PATCH 12/42] info about location inside error trace --- ydb/src/load_balancer.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index c9be3435..121a281f 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -255,8 +255,9 @@ impl NearestDCBalancer { match self.fallback_strategy { FallbackStrategy::Error => Err(YdbError::custom(format!( - "no available endpoints for service {} in local dc", - service + "no available endpoints for service:{} in local dc:{}", + service, + self.location ))), FallbackStrategy::Random => self.random_balancer.endpoint(service), FallbackStrategy::Next => todo!(), From 55eb7f0a2f73ea7d6d5e9a36bf425599462c062c Mon Sep 17 00:00:00 2001 From: kshprenger Date: Thu, 8 Aug 2024 12:16:33 +0700 Subject: [PATCH 13/42] config --- ydb/src/load_balancer.rs | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index 121a281f..bfd730a4 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -194,26 +194,29 @@ pub(crate) async fn update_load_balancer( } } +pub(crate) struct BalancerConfig { + fallback_strategy: FallbackStrategy, +} + // What will balancer do if there is no available endpoints at local dc pub(crate) enum FallbackStrategy { Error, // Just throw error Random, // Random endpoint from other dcs - Next, // Find next fastest dc } pub(crate) struct NearestDCBalancer { random_balancer: RandomLoadBalancer, - fallback_strategy: FallbackStrategy, //Default - Error + config: BalancerConfig, allowed_endpoints: Vec, - location: String, // vec of locations sorted by speed. Next strategy??? + location: String, } impl NearestDCBalancer { - pub(crate) fn new() -> Self { + pub(crate) fn new(config: BalancerConfig) -> Self { let random_balancer = RandomLoadBalancer::new(); Self { random_balancer, - fallback_strategy: FallbackStrategy::Random, + config, allowed_endpoints: Vec::new(), location: String::new(), } @@ -253,14 +256,12 @@ impl NearestDCBalancer { } } - match self.fallback_strategy { + match self.config.fallback_strategy { FallbackStrategy::Error => Err(YdbError::custom(format!( "no available endpoints for service:{} in local dc:{}", - service, - self.location + service, self.location ))), FallbackStrategy::Random => self.random_balancer.endpoint(service), - FallbackStrategy::Next => todo!(), } } From 25d8e4b8e17813c96c497c9f09f066364e5e625f Mon Sep 17 00:00:00 2001 From: kshprenger Date: Thu, 8 Aug 2024 12:23:20 +0700 Subject: [PATCH 14/42] no need in services yet --- ydb/src/discovery.rs | 17 ++++------------- ydb/src/grpc_wrapper/raw_discovery_client.rs | 11 ----------- ydb/src/load_balancer.rs | 16 ++++------------ 3 files changed, 8 insertions(+), 36 deletions(-) diff --git a/ydb/src/discovery.rs b/ydb/src/discovery.rs index f506bf6f..77418044 100644 --- a/ydb/src/discovery.rs +++ b/ydb/src/discovery.rs @@ -102,16 +102,11 @@ impl Default for DiscoveryState { pub(crate) struct NodeInfo { pub(crate) uri: Uri, pub(crate) location: String, - pub(crate) service: Vec, } impl NodeInfo { - pub(crate) fn new(uri: Uri, location: String, service: Vec) -> Self { - Self { - uri, - location, - service, - } + pub(crate) fn new(uri: Uri, location: String) -> Self { + Self { uri, location } } } @@ -151,7 +146,7 @@ pub struct StaticDiscovery { impl StaticDiscovery { pub fn new_from_str<'a, T: Into<&'a str>>(endpoint: T) -> YdbResult { let endpoint = Uri::from_str(endpoint.into())?; - let nodes = vec![NodeInfo::new(endpoint, String::new(), Vec::new())]; + let nodes = vec![NodeInfo::new(endpoint, String::new())]; let state = DiscoveryState::new(std::time::Instant::now(), nodes); let state = Arc::new(state); @@ -335,11 +330,7 @@ impl DiscoverySharedState { fn list_endpoints_to_node_infos(list: Vec) -> YdbResult> { list.into_iter() .map(|item| match Self::endpoint_info_to_uri(&item) { - Ok(uri) => YdbResult::::Ok(NodeInfo::new( - uri, - item.location.clone(), - item.service.clone(), - )), + Ok(uri) => YdbResult::::Ok(NodeInfo::new(uri, item.location.clone())), Err(err) => YdbResult::::Err(err), }) .try_collect() diff --git a/ydb/src/grpc_wrapper/raw_discovery_client.rs b/ydb/src/grpc_wrapper/raw_discovery_client.rs index 96fc28f8..acbe0489 100644 --- a/ydb/src/grpc_wrapper/raw_discovery_client.rs +++ b/ydb/src/grpc_wrapper/raw_discovery_client.rs @@ -3,12 +3,9 @@ use crate::grpc_wrapper::raw_services::{GrpcServiceForDiscovery, Service}; use crate::grpc_wrapper::runtime_interceptors::InterceptedChannel; use crate::YdbResult; use itertools::Itertools; -use serde_json::from_str; use ydb_grpc::ydb_proto::discovery::v1::discovery_service_client::DiscoveryServiceClient; use ydb_grpc::ydb_proto::discovery::{ListEndpointsRequest, ListEndpointsResult}; -use super::raw_services; - pub struct GrpcDiscoveryClient { service: DiscoveryServiceClient, } @@ -32,7 +29,6 @@ impl GrpcDiscoveryClient { let resp = self.service.list_endpoints(req).await?; let result: ListEndpointsResult = grpc_read_operation_result(resp)?; - use std::str::FromStr; let res = result .endpoints .into_iter() @@ -41,12 +37,6 @@ impl GrpcDiscoveryClient { port: item.port, ssl: item.ssl, location: item.location, - service: item - .service - .into_iter() - .map(|s| Service::from_str(s.as_str())) - .map(|res| res.unwrap()) - .collect(), }) .collect_vec(); Ok(res) @@ -58,7 +48,6 @@ pub(crate) struct EndpointInfo { pub(crate) port: u32, pub(crate) ssl: bool, pub(crate) location: String, - pub(crate) service: Vec, } impl GrpcServiceForDiscovery for GrpcDiscoveryClient { diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index bfd730a4..25706739 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -13,9 +13,7 @@ use tokio_util::sync::CancellationToken; use crate::grpc_wrapper::raw_services::Service; use crate::waiter::{Waiter, WaiterImpl}; use std::collections::HashMap; -use std::fmt::format; use std::net::ToSocketAddrs; -use std::str::FromStr; use std::sync::{Arc, RwLock}; use tokio::sync::watch::Receiver; @@ -251,9 +249,7 @@ const NODES_PER_DC: usize = 5; impl NearestDCBalancer { fn get_endpoint(&self, service: Service) -> YdbResult { for ep in self.allowed_endpoints.iter() { - if ep.service.contains(&service) { - return YdbResult::Ok(ep.uri.clone()); - } + return YdbResult::Ok(ep.uri.clone()); } match self.config.fallback_strategy { @@ -469,11 +465,7 @@ mod test { let new_discovery_state = Arc::new(DiscoveryState::default().with_node_info( Table, - NodeInfo::new( - Uri::from_str("http://test.com").unwrap(), - String::new(), - Vec::new(), - ), + NodeInfo::new(Uri::from_str("http://test.com").unwrap(), String::new()), )); let (first_update_sender, first_update_receiver) = tokio::sync::oneshot::channel(); @@ -537,8 +529,8 @@ mod test { let load_balancer = RandomLoadBalancer { discovery_state: Arc::new( DiscoveryState::default() - .with_node_info(Table, NodeInfo::new(one.clone(), String::new(), Vec::new())) - .with_node_info(Table, NodeInfo::new(two.clone(), String::new(), Vec::new())), + .with_node_info(Table, NodeInfo::new(one.clone(), String::new())) + .with_node_info(Table, NodeInfo::new(two.clone(), String::new())), ), waiter: Arc::new(WaiterImpl::new()), }; From e97bb2ac6c37f70dfded88b81c1a415dea68efe5 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Thu, 8 Aug 2024 12:30:55 +0700 Subject: [PATCH 15/42] lint checks --- ydb/src/load_balancer.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index 25706739..89f0d428 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -375,9 +375,9 @@ impl NearestDCBalancer { match connection_result{ Ok(mut connection) => { let _ = connection.shutdown().await; - addr_sender.send(Some(addr)).await; + let _ = addr_sender.send(Some(addr)).await; }, - Err(_) => {addr_sender.send(None).await;}, + Err(_) => {let _ = addr_sender.send(None).await;}, } } _ = stop_measure.cancelled() => { From aa9f207e5f711b92db376fa17921d344625ea9ab Mon Sep 17 00:00:00 2001 From: kshprenger Date: Thu, 8 Aug 2024 12:55:47 +0700 Subject: [PATCH 16/42] balancer config --- ydb/src/load_balancer.rs | 35 +++++++++++++++++++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index 89f0d428..af6a21ff 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -21,6 +21,7 @@ use tokio::sync::watch::Receiver; pub(crate) trait LoadBalancer: Send + Sync + Waiter { fn endpoint(&self, service: Service) -> YdbResult; fn set_discovery_state(&mut self, discovery_state: &Arc) -> YdbResult<()>; + fn set_config(&mut self, config: BalancerConfig) -> YdbResult<()>; fn waiter(&self) -> Box; // need for wait ready in without read lock } @@ -38,7 +39,11 @@ pub(crate) struct SharedLoadBalancer { impl SharedLoadBalancer { pub(crate) fn new(discovery: &dyn Discovery) -> Self { - Self::new_with_balancer_and_updater(Box::new(RandomLoadBalancer::new()), discovery) + Self::new_with_balancer_and_updater( + Box::new(RandomLoadBalancer::new()), + BalancerConfig::default(), + discovery, + ) } pub(crate) fn new_with_balancer(load_balancer: Box) -> Self { @@ -49,9 +54,11 @@ impl SharedLoadBalancer { pub(crate) fn new_with_balancer_and_updater( load_balancer: Box, + config: BalancerConfig, discovery: &dyn Discovery, ) -> Self { let mut shared_lb = Self::new_with_balancer(load_balancer); + let _ = shared_lb.set_config(config); let shared_lb_updater = shared_lb.clone(); let discovery_receiver = discovery.subscribe(); let _ = shared_lb.set_discovery_state(&discovery.state()); @@ -71,6 +78,10 @@ impl LoadBalancer for SharedLoadBalancer { self.inner.write()?.set_discovery_state(discovery_state) } + fn set_config(&mut self, config: BalancerConfig) -> YdbResult<()> { + self.inner.write()?.set_config(config) + } + fn waiter(&self) -> Box { return self.inner.read().unwrap().waiter(); } @@ -106,6 +117,12 @@ impl LoadBalancer for StaticLoadBalancer { )) } + fn set_config(&mut self, _: BalancerConfig) -> YdbResult<()> { + Err(YdbError::Custom( + "static balancer does not have config".into(), + )) + } + fn waiter(&self) -> Box { let waiter = WaiterImpl::new(); waiter.set_received(Ok(())); @@ -166,6 +183,12 @@ impl LoadBalancer for RandomLoadBalancer { Ok(()) } + fn set_config(&mut self, _: BalancerConfig) -> YdbResult<()> { + Err(YdbError::Custom( + "random balancer does not have config".into(), + )) + } + fn waiter(&self) -> Box { Box::new(self.waiter.clone()) } @@ -192,13 +215,16 @@ pub(crate) async fn update_load_balancer( } } +#[derive(Default)] pub(crate) struct BalancerConfig { fallback_strategy: FallbackStrategy, } +#[derive(Default)] // What will balancer do if there is no available endpoints at local dc pub(crate) enum FallbackStrategy { - Error, // Just throw error + #[default] + Error, // Just throw error Random, // Random endpoint from other dcs } @@ -239,6 +265,11 @@ impl LoadBalancer for NearestDCBalancer { self.adjust_allowed_endpoints() } + fn set_config(&mut self, config: BalancerConfig) -> YdbResult<()> { + self.config = config; + Ok(()) + } + fn waiter(&self) -> Box { self.random_balancer.waiter() } From 85c86c9ac3d542ff21140face87a2a2219373454 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Thu, 8 Aug 2024 13:00:10 +0700 Subject: [PATCH 17/42] non experimental default --- ydb/src/load_balancer.rs | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index af6a21ff..b18d822b 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -215,19 +215,24 @@ pub(crate) async fn update_load_balancer( } } -#[derive(Default)] pub(crate) struct BalancerConfig { fallback_strategy: FallbackStrategy, } -#[derive(Default)] // What will balancer do if there is no available endpoints at local dc pub(crate) enum FallbackStrategy { - #[default] - Error, // Just throw error + Error, // Just throw error Random, // Random endpoint from other dcs } +impl Default for BalancerConfig { + fn default() -> Self { + BalancerConfig { + fallback_strategy: FallbackStrategy::Error, + } + } +} + pub(crate) struct NearestDCBalancer { random_balancer: RandomLoadBalancer, config: BalancerConfig, From 7f2fc72622460d775865643d95a3f5d95e29a3ee Mon Sep 17 00:00:00 2001 From: kshprenger Date: Fri, 9 Aug 2024 13:19:17 +0700 Subject: [PATCH 18/42] naming --- ydb/src/load_balancer.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index b18d822b..145c23a1 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -236,7 +236,7 @@ impl Default for BalancerConfig { pub(crate) struct NearestDCBalancer { random_balancer: RandomLoadBalancer, config: BalancerConfig, - allowed_endpoints: Vec, + preferred_endpoints: Vec, location: String, } @@ -246,7 +246,7 @@ impl NearestDCBalancer { Self { random_balancer, config, - allowed_endpoints: Vec::new(), + preferred_endpoints: Vec::new(), location: String::new(), } } @@ -267,7 +267,7 @@ impl LoadBalancer for NearestDCBalancer { fn set_discovery_state(&mut self, discovery_state: &Arc) -> YdbResult<()> { self.random_balancer.set_discovery_state(discovery_state)?; self.adjust_local_dc()?; - self.adjust_allowed_endpoints() + self.adjust_preferred_endpoints() } fn set_config(&mut self, config: BalancerConfig) -> YdbResult<()> { @@ -284,7 +284,7 @@ const NODES_PER_DC: usize = 5; impl NearestDCBalancer { fn get_endpoint(&self, service: Service) -> YdbResult { - for ep in self.allowed_endpoints.iter() { + for ep in self.preferred_endpoints.iter() { return YdbResult::Ok(ep.uri.clone()); } @@ -309,8 +309,8 @@ impl NearestDCBalancer { Ok(()) } - fn adjust_allowed_endpoints(&mut self) -> YdbResult<()> { - self.allowed_endpoints = self + fn adjust_preferred_endpoints(&mut self) -> YdbResult<()> { + self.preferred_endpoints = self .get_nodes()? .into_iter() .filter(|ep| ep.location == self.location) From 9b233ecafef1a609e8b68c7f03f2a5bce3f97073 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Fri, 9 Aug 2024 13:23:11 +0700 Subject: [PATCH 19/42] random - default fallback --- ydb/src/load_balancer.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index 145c23a1..17784af5 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -228,7 +228,7 @@ pub(crate) enum FallbackStrategy { impl Default for BalancerConfig { fn default() -> Self { BalancerConfig { - fallback_strategy: FallbackStrategy::Error, + fallback_strategy: FallbackStrategy::Random, } } } From 2c8f650694a15d6c0a304165883894777e6e32ba Mon Sep 17 00:00:00 2001 From: kshprenger Date: Fri, 9 Aug 2024 13:27:20 +0700 Subject: [PATCH 20/42] static config --- ydb/src/load_balancer.rs | 25 ------------------------- 1 file changed, 25 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index 17784af5..85e116c5 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -21,7 +21,6 @@ use tokio::sync::watch::Receiver; pub(crate) trait LoadBalancer: Send + Sync + Waiter { fn endpoint(&self, service: Service) -> YdbResult; fn set_discovery_state(&mut self, discovery_state: &Arc) -> YdbResult<()>; - fn set_config(&mut self, config: BalancerConfig) -> YdbResult<()>; fn waiter(&self) -> Box; // need for wait ready in without read lock } @@ -41,7 +40,6 @@ impl SharedLoadBalancer { pub(crate) fn new(discovery: &dyn Discovery) -> Self { Self::new_with_balancer_and_updater( Box::new(RandomLoadBalancer::new()), - BalancerConfig::default(), discovery, ) } @@ -54,11 +52,9 @@ impl SharedLoadBalancer { pub(crate) fn new_with_balancer_and_updater( load_balancer: Box, - config: BalancerConfig, discovery: &dyn Discovery, ) -> Self { let mut shared_lb = Self::new_with_balancer(load_balancer); - let _ = shared_lb.set_config(config); let shared_lb_updater = shared_lb.clone(); let discovery_receiver = discovery.subscribe(); let _ = shared_lb.set_discovery_state(&discovery.state()); @@ -78,10 +74,6 @@ impl LoadBalancer for SharedLoadBalancer { self.inner.write()?.set_discovery_state(discovery_state) } - fn set_config(&mut self, config: BalancerConfig) -> YdbResult<()> { - self.inner.write()?.set_config(config) - } - fn waiter(&self) -> Box { return self.inner.read().unwrap().waiter(); } @@ -117,12 +109,6 @@ impl LoadBalancer for StaticLoadBalancer { )) } - fn set_config(&mut self, _: BalancerConfig) -> YdbResult<()> { - Err(YdbError::Custom( - "static balancer does not have config".into(), - )) - } - fn waiter(&self) -> Box { let waiter = WaiterImpl::new(); waiter.set_received(Ok(())); @@ -183,12 +169,6 @@ impl LoadBalancer for RandomLoadBalancer { Ok(()) } - fn set_config(&mut self, _: BalancerConfig) -> YdbResult<()> { - Err(YdbError::Custom( - "random balancer does not have config".into(), - )) - } - fn waiter(&self) -> Box { Box::new(self.waiter.clone()) } @@ -270,11 +250,6 @@ impl LoadBalancer for NearestDCBalancer { self.adjust_preferred_endpoints() } - fn set_config(&mut self, config: BalancerConfig) -> YdbResult<()> { - self.config = config; - Ok(()) - } - fn waiter(&self) -> Box { self.random_balancer.waiter() } From e4abe1d1ca985067579cf233981bbc170bff18b1 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 11 Aug 2024 18:36:03 +0700 Subject: [PATCH 21/42] all waiter --- ydb/src/waiter.rs | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/ydb/src/waiter.rs b/ydb/src/waiter.rs index 3ad38285..c7fc8e41 100644 --- a/ydb/src/waiter.rs +++ b/ydb/src/waiter.rs @@ -65,3 +65,24 @@ impl Waiter for Arc { self.as_ref().wait().await } } + +pub(crate) struct AllWaiter { + waiters: Vec>, +} + +impl AllWaiter { + pub fn new(waiters: Vec>) -> Self { + Self { waiters } + } +} + +// AllWaiter should point to group of Arc Waiters +#[async_trait::async_trait] +impl Waiter for AllWaiter { + async fn wait(&self) -> YdbResult<()> { + for waiter in self.waiters.iter() { + waiter.wait().await? + } + Ok(()) + } +} From 97eb2cce059415809103768776658e14ff134556 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 11 Aug 2024 18:37:12 +0700 Subject: [PATCH 22/42] share state between self and child Balancer and wait alll if needed + generalize child balancer --- ydb/src/load_balancer.rs | 72 +++++++++++++++++++++++++++++----------- 1 file changed, 53 insertions(+), 19 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index 85e116c5..937845b1 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -11,7 +11,7 @@ use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; use crate::grpc_wrapper::raw_services::Service; -use crate::waiter::{Waiter, WaiterImpl}; +use crate::waiter::{AllWaiter, Waiter, WaiterImpl}; use std::collections::HashMap; use std::net::ToSocketAddrs; use std::sync::{Arc, RwLock}; @@ -38,10 +38,7 @@ pub(crate) struct SharedLoadBalancer { impl SharedLoadBalancer { pub(crate) fn new(discovery: &dyn Discovery) -> Self { - Self::new_with_balancer_and_updater( - Box::new(RandomLoadBalancer::new()), - discovery, - ) + Self::new_with_balancer_and_updater(Box::new(RandomLoadBalancer::new()), discovery) } pub(crate) fn new_with_balancer(load_balancer: Box) -> Self { @@ -197,45 +194,65 @@ pub(crate) async fn update_load_balancer( pub(crate) struct BalancerConfig { fallback_strategy: FallbackStrategy, + fallback_balancer: Option>, } // What will balancer do if there is no available endpoints at local dc +#[derive(PartialEq, Eq)] pub(crate) enum FallbackStrategy { - Error, // Just throw error - Random, // Random endpoint from other dcs + Error, // Just throw error + BalanceWithOther, // Use another balancer } impl Default for BalancerConfig { fn default() -> Self { BalancerConfig { - fallback_strategy: FallbackStrategy::Random, + fallback_strategy: FallbackStrategy::BalanceWithOther, + fallback_balancer: Some(Box::new(RandomLoadBalancer::new())), } } } pub(crate) struct NearestDCBalancer { - random_balancer: RandomLoadBalancer, + discovery_state: Arc, + waiter: Arc, config: BalancerConfig, preferred_endpoints: Vec, location: String, } impl NearestDCBalancer { - pub(crate) fn new(config: BalancerConfig) -> Self { - let random_balancer = RandomLoadBalancer::new(); - Self { - random_balancer, + pub(crate) fn new(config: BalancerConfig) -> YdbResult { + match config.fallback_balancer.as_ref() { + Some(_) => { + if config.fallback_strategy == FallbackStrategy::Error { + return Err(YdbError::Custom( + "fallback strategy is Error but balancer was provided".to_string(), + )); + } + } + None => { + if config.fallback_strategy == FallbackStrategy::BalanceWithOther { + return Err(YdbError::Custom( + "no fallback balancer was provided".to_string(), + )); + } + } + } + Ok(Self { + discovery_state: Arc::new(DiscoveryState::default()), + waiter: Arc::new(WaiterImpl::new()), config, preferred_endpoints: Vec::new(), location: String::new(), - } + }) } } #[async_trait::async_trait] impl Waiter for NearestDCBalancer { async fn wait(&self) -> YdbResult<()> { - self.random_balancer.wait().await + self.waiter().wait().await } } @@ -245,13 +262,24 @@ impl LoadBalancer for NearestDCBalancer { } fn set_discovery_state(&mut self, discovery_state: &Arc) -> YdbResult<()> { - self.random_balancer.set_discovery_state(discovery_state)?; + match self.config.fallback_balancer.as_mut() { + Some(balancer) => balancer.set_discovery_state(discovery_state)?, + None => (), + } + self.discovery_state = discovery_state.clone(); + if !self.discovery_state.is_empty() { + self.waiter.set_received(Ok(())) + } self.adjust_local_dc()?; self.adjust_preferred_endpoints() } fn waiter(&self) -> Box { - self.random_balancer.waiter() + let self_waiter = Box::new(self.waiter.clone()); + match self.config.fallback_balancer.as_ref() { + Some(balancer) => Box::new(AllWaiter::new(vec![self_waiter, balancer.waiter()])), + None => self_waiter, + } } } @@ -268,7 +296,13 @@ impl NearestDCBalancer { "no available endpoints for service:{} in local dc:{}", service, self.location ))), - FallbackStrategy::Random => self.random_balancer.endpoint(service), + FallbackStrategy::BalanceWithOther => { + self.config + .fallback_balancer + .as_ref() + .unwrap() // unwrap is safe [checks inside ::new()] + .endpoint(service) + } } } @@ -295,7 +329,7 @@ impl NearestDCBalancer { } fn get_nodes(&self) -> YdbResult<&Vec> { - let nodes = self.random_balancer.discovery_state.get_all_nodes(); + let nodes = self.discovery_state.get_all_nodes(); match nodes { None => Err(YdbError::Custom(format!( "no endpoints on discovery update" From 3b9a9445ae9b8626d19154ce19c0532ed8933731 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 11 Aug 2024 19:05:51 +0700 Subject: [PATCH 23/42] better concurrent wait on waiters awaits --- ydb/src/waiter.rs | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/ydb/src/waiter.rs b/ydb/src/waiter.rs index c7fc8e41..1b3eb3d9 100644 --- a/ydb/src/waiter.rs +++ b/ydb/src/waiter.rs @@ -1,4 +1,7 @@ -use crate::YdbResult; +use crate::{YdbError, YdbResult}; +use itertools::Itertools; +use std::future::Future; +use std::process::Output; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use tokio::sync::watch; @@ -80,9 +83,15 @@ impl AllWaiter { #[async_trait::async_trait] impl Waiter for AllWaiter { async fn wait(&self) -> YdbResult<()> { - for waiter in self.waiters.iter() { - waiter.wait().await? - } + let awaitables = self + .waiters + .iter() + .map(|waiter| waiter.wait()) + .collect::>(); + futures_util::future::join_all(awaitables) + .await + .into_iter() + .collect::, YdbError>>()?; // If any waiter produced error - return it, otherwise - Ok Ok(()) } } From e139cef0e029eebf32a95944fd1366b24fc16d86 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 11 Aug 2024 19:08:22 +0700 Subject: [PATCH 24/42] reduce imports --- ydb/src/waiter.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/ydb/src/waiter.rs b/ydb/src/waiter.rs index 1b3eb3d9..8ae1a7d9 100644 --- a/ydb/src/waiter.rs +++ b/ydb/src/waiter.rs @@ -1,7 +1,4 @@ use crate::{YdbError, YdbResult}; -use itertools::Itertools; -use std::future::Future; -use std::process::Output; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use tokio::sync::watch; From 30caee313a35e1a759570917e904317168d5450c Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 18 Aug 2024 18:15:20 +0700 Subject: [PATCH 25/42] non-blocking balancer with full cancellation control and tracing --- ydb/src/load_balancer.rs | 233 ++++++++++++++++++++++++++------------- 1 file changed, 159 insertions(+), 74 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index 937845b1..7f45896b 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -1,21 +1,23 @@ use crate::discovery::{Discovery, DiscoveryState, NodeInfo}; use crate::errors::*; +use crate::grpc_wrapper::raw_services::Service; +use crate::waiter::{AllWaiter, Waiter, WaiterImpl}; use http::Uri; use itertools::Itertools; use rand::thread_rng; +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; +use std::time::Duration; use tokio::io::AsyncWriteExt; use tokio::net::TcpStream; -use tokio::runtime::Handle; -use tokio::sync::mpsc; -use tokio::task::JoinSet; +use tokio::sync::watch::{Receiver, Sender}; +use tokio::sync::Mutex; +use tokio::sync::{mpsc, watch}; +use tokio::task::{JoinHandle, JoinSet}; +use tokio::time::Timeout; use tokio_util::sync::CancellationToken; - -use crate::grpc_wrapper::raw_services::Service; -use crate::waiter::{AllWaiter, Waiter, WaiterImpl}; -use std::collections::HashMap; -use std::net::ToSocketAddrs; -use std::sync::{Arc, RwLock}; -use tokio::sync::watch::Receiver; +use tracing::info; +use tracing::{error, warn}; #[mockall::automock] pub(crate) trait LoadBalancer: Send + Sync + Waiter { @@ -215,10 +217,12 @@ impl Default for BalancerConfig { pub(crate) struct NearestDCBalancer { discovery_state: Arc, + state_sender: Sender>, + adjust_local_dc_process_control: CancellationToken, waiter: Arc, config: BalancerConfig, preferred_endpoints: Vec, - location: String, + location: Arc>, } impl NearestDCBalancer { @@ -239,16 +243,41 @@ impl NearestDCBalancer { } } } + + let discovery_state = Arc::new(DiscoveryState::default()); + let self_location = Arc::new(Mutex::new(String::new())); + let location_updater = self_location.clone(); + let (state_sender, state_reciever) = watch::channel(discovery_state.clone()); + let adjust_local_dc_process_control = CancellationToken::new(); + let adjust_local_dc_process_control_clone = adjust_local_dc_process_control.clone(); + + tokio::spawn(async move { + Self::adjust_local_dc( + location_updater, + state_reciever, + adjust_local_dc_process_control_clone, + ) + .await + }); + Ok(Self { - discovery_state: Arc::new(DiscoveryState::default()), + discovery_state, + state_sender, + adjust_local_dc_process_control, waiter: Arc::new(WaiterImpl::new()), config, preferred_endpoints: Vec::new(), - location: String::new(), + location: Arc::new(Mutex::new(String::new())), }) } } +impl Drop for NearestDCBalancer { + fn drop(&mut self) { + self.adjust_local_dc_process_control.cancel(); + } +} + #[async_trait::async_trait] impl Waiter for NearestDCBalancer { async fn wait(&self) -> YdbResult<()> { @@ -267,11 +296,12 @@ impl LoadBalancer for NearestDCBalancer { None => (), } self.discovery_state = discovery_state.clone(); + let _ = self.state_sender.send(discovery_state.clone()); + self.adjust_preferred_endpoints()?; if !self.discovery_state.is_empty() { self.waiter.set_received(Ok(())) } - self.adjust_local_dc()?; - self.adjust_preferred_endpoints() + Ok(()) } fn waiter(&self) -> Box { @@ -284,6 +314,7 @@ impl LoadBalancer for NearestDCBalancer { } const NODES_PER_DC: usize = 5; +const PING_TIMEOUT_SECS: u64 = 60; impl NearestDCBalancer { fn get_endpoint(&self, service: Service) -> YdbResult { @@ -293,8 +324,8 @@ impl NearestDCBalancer { match self.config.fallback_strategy { FallbackStrategy::Error => Err(YdbError::custom(format!( - "no available endpoints for service:{} in local dc:{}", - service, self.location + "no available endpoints for service:{}", + service ))), FallbackStrategy::BalanceWithOther => { self.config @@ -306,30 +337,60 @@ impl NearestDCBalancer { } } - fn adjust_local_dc(&mut self) -> YdbResult<()> { - let nodes = self.get_nodes()?; - let mut dc_to_nodes = self.split_endpoints_by_location(nodes); - let mut to_check = Vec::with_capacity(NODES_PER_DC * dc_to_nodes.keys().len()); - dc_to_nodes - .iter_mut() - .for_each(|(_, endpoints)| to_check.append(self.get_random_endpoints(endpoints))); - let local_dc = self.find_local_dc(&to_check)?; - self.location = local_dc; - Ok(()) - } - fn adjust_preferred_endpoints(&mut self) -> YdbResult<()> { - self.preferred_endpoints = self - .get_nodes()? + let location = match self.location.try_lock() { + Ok(location_guard) => (*location_guard).clone(), + Err(_) => { + info!("could not acquire lock on location"); + "".into() + } + }; + self.preferred_endpoints = Self::extract_nodes(&self.discovery_state)? .into_iter() - .filter(|ep| ep.location == self.location) + .filter(|ep| ep.location == *location) .map(|ep| ep.clone()) .collect_vec(); Ok(()) } - fn get_nodes(&self) -> YdbResult<&Vec> { - let nodes = self.discovery_state.get_all_nodes(); + async fn adjust_local_dc( + self_location: Arc>, + mut state_reciever: watch::Receiver>, + stop_ping_process: CancellationToken, + ) { + loop { + let new_state = state_reciever.borrow_and_update().clone(); + match Self::extract_nodes(&new_state) { + Ok(some_nodes) => { + let mut dc_to_nodes = Self::split_endpoints_by_location(some_nodes); + let mut to_check = Vec::with_capacity(NODES_PER_DC * dc_to_nodes.keys().len()); + dc_to_nodes.iter_mut().for_each(|(_, endpoints)| { + to_check.append(Self::get_random_endpoints(endpoints)) + }); + match Self::find_local_dc(&to_check).await { + Ok(dc) => { + info!("found new local dc:{}", dc); + *self_location.lock().await = dc; // fast lock + } + Err(err) => { + error!("error on search local dc:{}", err); + continue; + } + } + } + Err(_) => continue, + } + if state_reciever.changed().await.is_err() { + return; + } + if stop_ping_process.is_cancelled() { + return; + } + } + } + + fn extract_nodes(from_state: &Arc) -> YdbResult<&Vec> { + let nodes = from_state.get_all_nodes(); match nodes { None => Err(YdbError::Custom(format!( "no endpoints on discovery update" @@ -339,7 +400,6 @@ impl NearestDCBalancer { } fn split_endpoints_by_location<'a>( - &'a self, nodes: &'a Vec, ) -> HashMap> { let mut dc_to_eps = HashMap::>::new(); @@ -353,41 +413,49 @@ impl NearestDCBalancer { dc_to_eps } - fn get_random_endpoints<'a>( - &'a self, - dc_endpoints: &'a mut Vec<&'a NodeInfo>, - ) -> &mut Vec<&NodeInfo> { + fn get_random_endpoints<'a>(dc_endpoints: &'a mut Vec<&'a NodeInfo>) -> &mut Vec<&NodeInfo> { use rand::seq::SliceRandom; dc_endpoints.shuffle(&mut thread_rng()); dc_endpoints.truncate(NODES_PER_DC); dc_endpoints } - fn find_local_dc(&self, to_check: &[&NodeInfo]) -> YdbResult { - let addr_to_node = self.addr_to_node(to_check); + async fn find_local_dc(to_check: &[&NodeInfo]) -> YdbResult { + let addr_to_node = Self::addr_to_node(to_check); if addr_to_node.is_empty() { return Err(YdbError::Custom(format!("no ip addresses for endpoints"))); } let addrs = addr_to_node.keys(); - match self.find_fastest_address(addrs.collect()) { - Some(fastest_address) => Ok(addr_to_node[&fastest_address].location.clone()), - None => Err(YdbError::custom("could not find fastest address")), + match Self::find_fastest_address(addrs.collect()).await { + Ok(fastest_address) => Ok(addr_to_node[&fastest_address].location.clone()), + Err(err) => { + error!("could not find fastest address:{}", err); + Err(err) + } } } - fn addr_to_node<'a>(&'a self, nodes: &[&'a NodeInfo]) -> HashMap { + fn addr_to_node<'a>(nodes: &[&'a NodeInfo]) -> HashMap { let mut addr_to_node = HashMap::::with_capacity(2 * nodes.len()); // ipv4 + ipv6 nodes.into_iter().for_each(|info| { let host: &str; let port: u16; match info.uri.host() { Some(uri_host) => host = uri_host, - None => return, + None => { + warn!("no host for uri:{}", info.uri); + return; + } } match info.uri.port() { Some(uri_port) => port = uri_port.as_u16(), - None => return, + None => { + warn!("no port for uri:{}", info.uri); + return; + } } + + use std::net::ToSocketAddrs; let _ = (host, port).to_socket_addrs().and_then(|addrs| { for addr in addrs { addr_to_node.insert(addr.to_string(), info); @@ -398,11 +466,14 @@ impl NearestDCBalancer { addr_to_node } - fn find_fastest_address(&self, addrs: Vec<&String>) -> Option { - let stop_measure = CancellationToken::new(); + async fn find_fastest_address(addrs: Vec<&String>) -> YdbResult { + // Cancellation flow: timeout -> address collector -> address producers + let interrupt_via_timeout = CancellationToken::new(); + let interrupt_collector_future = interrupt_via_timeout.child_token(); + let stop_measure = interrupt_collector_future.child_token(); // (*) + let (start_measure, _) = tokio::sync::broadcast::channel::<()>(1); let (addr_sender, mut addr_reciever) = tokio::sync::mpsc::channel::>(1); - let addr_count = addrs.len(); let mut nursery = JoinSet::new(); for addr in addrs { @@ -432,39 +503,53 @@ impl NearestDCBalancer { }); } - tokio::task::block_in_place(move || { - let _ = start_measure.send(()); - Handle::current().block_on(async { - for _ in 0..addr_count { - match self.wait_for_single_address(&mut addr_reciever).await { - Some(address) => { - stop_measure.cancel(); - self.join_all(&mut nursery).await; - return Some(address); + let wait_first_some_or_cancel = async { + loop { + tokio::select! { + biased; // check timeout first + _ = interrupt_collector_future.cancelled() =>{ + Self::join_all(&mut nursery).await; // children will be cancelled due to tokens chaining, see (*) + return YdbResult::Err("cancelled".into()) + } + address_reciever_option = addr_reciever.recv() =>{ + match address_reciever_option { + Some(address_option) => { + match address_option { + Some(address) =>{ + interrupt_collector_future.cancel(); // Cancel other producing children + Self::join_all(&mut nursery).await; + return YdbResult::Ok(address); + }, + None => continue, // Some producer sent blank address -> wait others + } + }, + None => return YdbResult::Err("no fastest address".into()), // Channel closed, all producers have done measures } - None => continue, } } - None - }) - }) - } + } + }; - async fn wait_for_single_address( - &self, - addr_reciever: &mut mpsc::Receiver>, - ) -> Option { - match addr_reciever.recv().await { - Some(maybe_address) => maybe_address, - None => unreachable!(), // no channel closing while awaiting address + let _ = start_measure.send(()); + + match tokio::time::timeout( + Duration::from_secs(PING_TIMEOUT_SECS), + wait_first_some_or_cancel, + ) + .await + { + Ok(address_option) => address_option, + Err(_) => { + interrupt_via_timeout.cancel(); + YdbResult::Err("timeout while detecting fastest address".into()) + } } } - async fn join_all(&self, awaitable: &mut JoinSet<()>) { - while let Some(_) = awaitable.join_next().await {} + async fn join_all(awaitable: &mut JoinSet<()>) { + while awaitable.join_next().await.is_some() {} } } - #[cfg(test)] mod test { use super::*; From dd8bf2f8ddf6ef487e52627e6945ff6a03df5f7c Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 18 Aug 2024 18:20:38 +0700 Subject: [PATCH 26/42] optimize imports --- ydb/src/load_balancer.rs | 32 ++++++++++++++++++-------------- 1 file changed, 18 insertions(+), 14 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index 7f45896b..d61616f6 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -5,19 +5,23 @@ use crate::waiter::{AllWaiter, Waiter, WaiterImpl}; use http::Uri; use itertools::Itertools; use rand::thread_rng; -use std::collections::HashMap; -use std::sync::{Arc, RwLock}; -use std::time::Duration; -use tokio::io::AsyncWriteExt; -use tokio::net::TcpStream; -use tokio::sync::watch::{Receiver, Sender}; -use tokio::sync::Mutex; -use tokio::sync::{mpsc, watch}; -use tokio::task::{JoinHandle, JoinSet}; -use tokio::time::Timeout; +use std::{ + collections::HashMap, + sync::{Arc, RwLock}, + time::Duration, +}; +use tokio::{ + io::AsyncWriteExt, + net::TcpStream, + sync::{ + broadcast, mpsc, watch, + watch::{Receiver, Sender}, + Mutex, + }, + task::JoinSet, +}; use tokio_util::sync::CancellationToken; -use tracing::info; -use tracing::{error, warn}; +use tracing::{error, info, warn}; #[mockall::automock] pub(crate) trait LoadBalancer: Send + Sync + Waiter { @@ -472,8 +476,8 @@ impl NearestDCBalancer { let interrupt_collector_future = interrupt_via_timeout.child_token(); let stop_measure = interrupt_collector_future.child_token(); // (*) - let (start_measure, _) = tokio::sync::broadcast::channel::<()>(1); - let (addr_sender, mut addr_reciever) = tokio::sync::mpsc::channel::>(1); + let (start_measure, _) = broadcast::channel::<()>(1); + let (addr_sender, mut addr_reciever) = mpsc::channel::>(1); let mut nursery = JoinSet::new(); for addr in addrs { From 4be16242b440e1b01ef2c1250775ccca9843bb00 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Tue, 20 Aug 2024 13:26:53 +0700 Subject: [PATCH 27/42] client waits for balancer set up its state both in waiter and on endpoint() --- ydb/src/load_balancer.rs | 131 ++++++++++++++++++++++----------------- 1 file changed, 74 insertions(+), 57 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index d61616f6..60f951f4 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -5,6 +5,7 @@ use crate::waiter::{AllWaiter, Waiter, WaiterImpl}; use http::Uri; use itertools::Itertools; use rand::thread_rng; +use std::borrow::{Borrow, BorrowMut}; use std::{ collections::HashMap, sync::{Arc, RwLock}, @@ -203,6 +204,11 @@ pub(crate) struct BalancerConfig { fallback_balancer: Option>, } +#[derive(Default)] +struct BalancerState { + preferred_endpoints: Vec, +} + // What will balancer do if there is no available endpoints at local dc #[derive(PartialEq, Eq)] pub(crate) enum FallbackStrategy { @@ -222,11 +228,10 @@ impl Default for BalancerConfig { pub(crate) struct NearestDCBalancer { discovery_state: Arc, state_sender: Sender>, - adjust_local_dc_process_control: CancellationToken, + ping_token: CancellationToken, waiter: Arc, config: BalancerConfig, - preferred_endpoints: Vec, - location: Arc>, + balancer_state: Arc>, } impl NearestDCBalancer { @@ -249,17 +254,22 @@ impl NearestDCBalancer { } let discovery_state = Arc::new(DiscoveryState::default()); - let self_location = Arc::new(Mutex::new(String::new())); - let location_updater = self_location.clone(); + let balancer_state = Arc::new(Mutex::new(BalancerState::default())); + let balancer_state_updater = balancer_state.clone(); let (state_sender, state_reciever) = watch::channel(discovery_state.clone()); - let adjust_local_dc_process_control = CancellationToken::new(); - let adjust_local_dc_process_control_clone = adjust_local_dc_process_control.clone(); + + let ping_token = CancellationToken::new(); + let ping_token_clone = ping_token.clone(); + + let waiter = Arc::new(WaiterImpl::new()); + let waiter_clone = waiter.clone(); tokio::spawn(async move { Self::adjust_local_dc( - location_updater, + balancer_state_updater, state_reciever, - adjust_local_dc_process_control_clone, + ping_token_clone, + waiter_clone, ) .await }); @@ -267,18 +277,17 @@ impl NearestDCBalancer { Ok(Self { discovery_state, state_sender, - adjust_local_dc_process_control, - waiter: Arc::new(WaiterImpl::new()), + ping_token, + waiter, config, - preferred_endpoints: Vec::new(), - location: Arc::new(Mutex::new(String::new())), + balancer_state, }) } } impl Drop for NearestDCBalancer { fn drop(&mut self) { - self.adjust_local_dc_process_control.cancel(); + self.ping_token.cancel(); } } @@ -299,12 +308,9 @@ impl LoadBalancer for NearestDCBalancer { Some(balancer) => balancer.set_discovery_state(discovery_state)?, None => (), } + println!("sending signal to start ping"); self.discovery_state = discovery_state.clone(); let _ = self.state_sender.send(discovery_state.clone()); - self.adjust_preferred_endpoints()?; - if !self.discovery_state.is_empty() { - self.waiter.set_received(Ok(())) - } Ok(()) } @@ -322,49 +328,39 @@ const PING_TIMEOUT_SECS: u64 = 60; impl NearestDCBalancer { fn get_endpoint(&self, service: Service) -> YdbResult { - for ep in self.preferred_endpoints.iter() { - return YdbResult::Ok(ep.uri.clone()); - } - - match self.config.fallback_strategy { - FallbackStrategy::Error => Err(YdbError::custom(format!( - "no available endpoints for service:{}", - service - ))), - FallbackStrategy::BalanceWithOther => { - self.config - .fallback_balancer - .as_ref() - .unwrap() // unwrap is safe [checks inside ::new()] - .endpoint(service) + match self.balancer_state.try_lock() { + Ok(state_guard) => { + for ep in state_guard.borrow().preferred_endpoints.iter() { + return YdbResult::Ok(ep.uri.clone()); + } + match self.config.fallback_strategy { + FallbackStrategy::Error => Err(YdbError::custom(format!( + "no available endpoints for service:{}", + service + ))), + FallbackStrategy::BalanceWithOther => { + info!("trying another balancer..."); + self.config + .fallback_balancer + .as_ref() + .unwrap() // unwrap is safe [checks inside ::new()] + .endpoint(service) + } + } } + Err(_) => Err(YdbError::Custom("balancer is updating its state".into())), } } - fn adjust_preferred_endpoints(&mut self) -> YdbResult<()> { - let location = match self.location.try_lock() { - Ok(location_guard) => (*location_guard).clone(), - Err(_) => { - info!("could not acquire lock on location"); - "".into() - } - }; - self.preferred_endpoints = Self::extract_nodes(&self.discovery_state)? - .into_iter() - .filter(|ep| ep.location == *location) - .map(|ep| ep.clone()) - .collect_vec(); - Ok(()) - } - async fn adjust_local_dc( - self_location: Arc>, + balancer_state: Arc>, mut state_reciever: watch::Receiver>, stop_ping_process: CancellationToken, + waiter: Arc, ) { loop { - let new_state = state_reciever.borrow_and_update().clone(); - match Self::extract_nodes(&new_state) { + let new_discovery_state = state_reciever.borrow_and_update().clone(); + match Self::extract_nodes(&new_discovery_state) { Ok(some_nodes) => { let mut dc_to_nodes = Self::split_endpoints_by_location(some_nodes); let mut to_check = Vec::with_capacity(NODES_PER_DC * dc_to_nodes.keys().len()); @@ -373,10 +369,13 @@ impl NearestDCBalancer { }); match Self::find_local_dc(&to_check).await { Ok(dc) => { + println!("found new local dc:{}", dc); info!("found new local dc:{}", dc); - *self_location.lock().await = dc; // fast lock + Self::adjust_preferred_endpoints(&balancer_state, some_nodes, dc).await; + waiter.set_received(Ok(())); } Err(err) => { + println!("error on search local dc:{}", err); error!("error on search local dc:{}", err); continue; } @@ -384,15 +383,29 @@ impl NearestDCBalancer { } Err(_) => continue, } - if state_reciever.changed().await.is_err() { - return; - } - if stop_ping_process.is_cancelled() { + if stop_ping_process.is_cancelled() || state_reciever.changed().await.is_err() { return; } } } + async fn adjust_preferred_endpoints( + balancer_state: &Arc>, + new_nodes: &Vec, + local_dc: String, + ) { + info!("adjusting endpoints"); + let new_preferred_endpoints = new_nodes + .into_iter() + .filter(|ep| ep.location == local_dc) + .map(|ep| ep.clone()) + .collect_vec(); + println!("new preferred endpoints:{:?}", new_preferred_endpoints); + (balancer_state.lock().await) // fast lock + .borrow_mut() + .preferred_endpoints = new_preferred_endpoints; + } + fn extract_nodes(from_state: &Arc) -> YdbResult<&Vec> { let nodes = from_state.get_all_nodes(); match nodes { @@ -433,6 +446,7 @@ impl NearestDCBalancer { match Self::find_fastest_address(addrs.collect()).await { Ok(fastest_address) => Ok(addr_to_node[&fastest_address].location.clone()), Err(err) => { + println!("could not find fastest address:{}", err); error!("could not find fastest address:{}", err); Err(err) } @@ -447,6 +461,7 @@ impl NearestDCBalancer { match info.uri.host() { Some(uri_host) => host = uri_host, None => { + println!("no host for uri:{}", info.uri); warn!("no host for uri:{}", info.uri); return; } @@ -454,6 +469,7 @@ impl NearestDCBalancer { match info.uri.port() { Some(uri_port) => port = uri_port.as_u16(), None => { + println!("no port for uri:{}", info.uri); warn!("no port for uri:{}", info.uri); return; } @@ -558,6 +574,7 @@ impl NearestDCBalancer { mod test { use super::*; use crate::discovery::NodeInfo; + use crate::grpc_wrapper::raw_services; use crate::grpc_wrapper::raw_services::Service::Table; use mockall::predicate; use std::collections::HashMap; From 2212b32b7abc3ff81cca3c2fcaf732436579ebd9 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sat, 31 Aug 2024 15:27:23 +0300 Subject: [PATCH 28/42] pretty + tests on internal non-async functions --- ydb/src/load_balancer.rs | 77 ++++++++++++++++++++++++++++++++++------ ydb/src/waiter.rs | 3 +- 2 files changed, 67 insertions(+), 13 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index 60f951f4..ecc8167c 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -240,7 +240,8 @@ impl NearestDCBalancer { Some(_) => { if config.fallback_strategy == FallbackStrategy::Error { return Err(YdbError::Custom( - "fallback strategy is Error but balancer was provided".to_string(), + "fallback strategy is \"Error\" but fallback balancer was provided" + .to_string(), )); } } @@ -308,7 +309,6 @@ impl LoadBalancer for NearestDCBalancer { Some(balancer) => balancer.set_discovery_state(discovery_state)?, None => (), } - println!("sending signal to start ping"); self.discovery_state = discovery_state.clone(); let _ = self.state_sender.send(discovery_state.clone()); Ok(()) @@ -339,7 +339,7 @@ impl NearestDCBalancer { service ))), FallbackStrategy::BalanceWithOther => { - info!("trying another balancer..."); + info!("trying fallback balancer..."); self.config .fallback_balancer .as_ref() @@ -348,7 +348,9 @@ impl NearestDCBalancer { } } } - Err(_) => Err(YdbError::Custom("balancer is updating its state".into())), + Err(_) => Err(YdbError::Custom( + "balancer is updating its state".to_string(), + )), } } @@ -369,13 +371,11 @@ impl NearestDCBalancer { }); match Self::find_local_dc(&to_check).await { Ok(dc) => { - println!("found new local dc:{}", dc); info!("found new local dc:{}", dc); Self::adjust_preferred_endpoints(&balancer_state, some_nodes, dc).await; waiter.set_received(Ok(())); } Err(err) => { - println!("error on search local dc:{}", err); error!("error on search local dc:{}", err); continue; } @@ -400,7 +400,6 @@ impl NearestDCBalancer { .filter(|ep| ep.location == local_dc) .map(|ep| ep.clone()) .collect_vec(); - println!("new preferred endpoints:{:?}", new_preferred_endpoints); (balancer_state.lock().await) // fast lock .borrow_mut() .preferred_endpoints = new_preferred_endpoints; @@ -446,7 +445,6 @@ impl NearestDCBalancer { match Self::find_fastest_address(addrs.collect()).await { Ok(fastest_address) => Ok(addr_to_node[&fastest_address].location.clone()), Err(err) => { - println!("could not find fastest address:{}", err); error!("could not find fastest address:{}", err); Err(err) } @@ -461,7 +459,6 @@ impl NearestDCBalancer { match info.uri.host() { Some(uri_host) => host = uri_host, None => { - println!("no host for uri:{}", info.uri); warn!("no host for uri:{}", info.uri); return; } @@ -469,7 +466,6 @@ impl NearestDCBalancer { match info.uri.port() { Some(uri_port) => port = uri_port.as_u16(), None => { - println!("no port for uri:{}", info.uri); warn!("no port for uri:{}", info.uri); return; } @@ -574,9 +570,9 @@ impl NearestDCBalancer { mod test { use super::*; use crate::discovery::NodeInfo; - use crate::grpc_wrapper::raw_services; use crate::grpc_wrapper::raw_services::Service::Table; use mockall::predicate; + use ntest::assert_true; use std::collections::HashMap; use std::str::FromStr; use std::sync::atomic::AtomicUsize; @@ -701,4 +697,63 @@ mod test { assert!(*map.get(two.to_string().as_str()).unwrap() > 30); Ok(()) } + + #[test] + fn split_by_location() -> YdbResult<()> { + let nodes = vec![ + NodeInfo::new(Uri::from_str("http://one:213")?, "A".to_string()), + NodeInfo::new(Uri::from_str("http://two:213")?, "A".to_string()), + NodeInfo::new(Uri::from_str("http://three:213")?, "B".to_string()), + NodeInfo::new(Uri::from_str("http://four:213")?, "B".to_string()), + NodeInfo::new(Uri::from_str("http://five:213")?, "C".to_string()), + ]; + let splitted = NearestDCBalancer::split_endpoints_by_location(&nodes); + assert_eq!(splitted.keys().len(), 3); + assert_eq!(splitted["A"].len(), 2); + assert_eq!(splitted["B"].len(), 2); + assert_eq!(splitted["C"].len(), 1); + Ok(()) + } + + #[test] + fn choose_random_endpoints() -> YdbResult<()> { + let nodes = vec![ + NodeInfo::new(Uri::from_str("http://one:213")?, "C".to_string()), + NodeInfo::new(Uri::from_str("http://two:213")?, "C".to_string()), + NodeInfo::new(Uri::from_str("http://three:213")?, "C".to_string()), + NodeInfo::new(Uri::from_str("http://four:213")?, "C".to_string()), + NodeInfo::new(Uri::from_str("http://five:213")?, "C".to_string()), + NodeInfo::new(Uri::from_str("http://seven:213")?, "C".to_string()), + NodeInfo::new(Uri::from_str("http://eight:213")?, "C".to_string()), + NodeInfo::new(Uri::from_str("http://nine:213")?, "C".to_string()), + ]; + + let mut refs = nodes.iter().collect_vec(); + let nodes_clone = refs.clone(); + let random_subset = NearestDCBalancer::get_random_endpoints(&mut refs); + + assert_eq!(random_subset.len(), NODES_PER_DC); + for node in random_subset { + assert_true!(nodes_clone.contains(node)) + } + + Ok(()) + } + + #[test] + fn extract_addrs_and_map_them() -> YdbResult<()> { + + let one = NodeInfo::new(Uri::from_str("http://localhost:123")?, "C".to_string()); + let two = NodeInfo::new(Uri::from_str("http://localhost:321")?, "C".to_string()); + let nodes = vec![&one, &two]; + let map = NearestDCBalancer::addr_to_node(&nodes); + + assert_eq!(map.keys().len(), 4); // ipv4 + ipv6 on each + assert_true!(map.keys().contains(&"127.0.0.1:123".to_string())); + assert_true!(map.keys().contains(&"[::1]:123".to_string())); + assert!(map["127.0.0.1:123"].eq(&one)); + assert!(map["127.0.0.1:123"].eq(map["[::1]:123"])); + + Ok(()) + } } diff --git a/ydb/src/waiter.rs b/ydb/src/waiter.rs index 8ae1a7d9..204037ec 100644 --- a/ydb/src/waiter.rs +++ b/ydb/src/waiter.rs @@ -76,7 +76,6 @@ impl AllWaiter { } } -// AllWaiter should point to group of Arc Waiters #[async_trait::async_trait] impl Waiter for AllWaiter { async fn wait(&self) -> YdbResult<()> { @@ -88,7 +87,7 @@ impl Waiter for AllWaiter { futures_util::future::join_all(awaitables) .await .into_iter() - .collect::, YdbError>>()?; // If any waiter produced error - return it, otherwise - Ok + .collect::, YdbError>>()?; // If some waiters produced error - return first, otherwise - Ok Ok(()) } } From a9e787cc531cc52a02ac034b46bfdafba84166c5 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sat, 31 Aug 2024 16:00:53 +0300 Subject: [PATCH 29/42] ci restart From dde7b9bf389465b2ad3acc08f836269f01c93f4b Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 1 Sep 2024 19:14:08 +0300 Subject: [PATCH 30/42] Fix deadlock on buffer send by producers, fixed livelock on infinite loop if continue from error. Added tests --- ydb/src/load_balancer.rs | 456 +++++++++++++++++++++++++++++++++++++-- 1 file changed, 440 insertions(+), 16 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index ecc8167c..2b6ee58d 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -361,6 +361,16 @@ impl NearestDCBalancer { waiter: Arc, ) { loop { + tokio::select! { + _ = stop_ping_process.cancelled() => { + return + } + result = state_reciever.changed() =>{ + if result.is_err(){ // sender have been dropped + return + } + } + } let new_discovery_state = state_reciever.borrow_and_update().clone(); match Self::extract_nodes(&new_discovery_state) { Ok(some_nodes) => { @@ -383,9 +393,6 @@ impl NearestDCBalancer { } Err(_) => continue, } - if stop_ping_process.is_cancelled() || state_reciever.changed().await.is_err() { - return; - } } } @@ -442,7 +449,9 @@ impl NearestDCBalancer { return Err(YdbError::Custom(format!("no ip addresses for endpoints"))); } let addrs = addr_to_node.keys(); - match Self::find_fastest_address(addrs.collect()).await { + match Self::find_fastest_address(addrs.collect(), Duration::from_secs(PING_TIMEOUT_SECS)) + .await + { Ok(fastest_address) => Ok(addr_to_node[&fastest_address].location.clone()), Err(err) => { error!("could not find fastest address:{}", err); @@ -470,7 +479,7 @@ impl NearestDCBalancer { return; } } - + println!("host:{}, port:{}", host, port); use std::net::ToSocketAddrs; let _ = (host, port).to_socket_addrs().and_then(|addrs| { for addr in addrs { @@ -482,14 +491,17 @@ impl NearestDCBalancer { addr_to_node } - async fn find_fastest_address(addrs: Vec<&String>) -> YdbResult { + async fn find_fastest_address(addrs: Vec<&String>, timeout: Duration) -> YdbResult { + println!("addrs:{:?}", addrs); + // Cancellation flow: timeout -> address collector -> address producers let interrupt_via_timeout = CancellationToken::new(); let interrupt_collector_future = interrupt_via_timeout.child_token(); let stop_measure = interrupt_collector_future.child_token(); // (*) let (start_measure, _) = broadcast::channel::<()>(1); - let (addr_sender, mut addr_reciever) = mpsc::channel::>(1); + let buffer_cap = if addrs.len() > 0 { addrs.len() } else { 1 }; + let (addr_sender, mut addr_reciever) = mpsc::channel::>(buffer_cap); let mut nursery = JoinSet::new(); for addr in addrs { @@ -547,14 +559,14 @@ impl NearestDCBalancer { }; let _ = start_measure.send(()); + println!("start measure"); - match tokio::time::timeout( - Duration::from_secs(PING_TIMEOUT_SECS), - wait_first_some_or_cancel, - ) - .await - { - Ok(address_option) => address_option, + match tokio::time::timeout(timeout, wait_first_some_or_cancel).await { + Ok(address_option) => { + println!("got fastest address:{:?}", address_option); + + address_option + } Err(_) => { interrupt_via_timeout.cancel(); YdbResult::Err("timeout while detecting fastest address".into()) @@ -578,6 +590,8 @@ mod test { use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering::Relaxed; use std::time::Duration; + use tokio::net::TcpListener; + use tokio::time::timeout; use tracing::trace; #[test] @@ -742,9 +756,8 @@ mod test { #[test] fn extract_addrs_and_map_them() -> YdbResult<()> { - let one = NodeInfo::new(Uri::from_str("http://localhost:123")?, "C".to_string()); - let two = NodeInfo::new(Uri::from_str("http://localhost:321")?, "C".to_string()); + let two = NodeInfo::new(Uri::from_str("http://localhost:321")?, "C".to_string()); let nodes = vec![&one, &two]; let map = NearestDCBalancer::addr_to_node(&nodes); @@ -756,4 +769,415 @@ mod test { Ok(()) } + + #[tokio::test] + async fn detect_fastest_addr_just_some() -> YdbResult<()> { + let l1 = TcpListener::bind("127.0.0.1:0").await?; + let l2 = TcpListener::bind("127.0.0.1:0").await?; + let l3 = TcpListener::bind("127.0.0.1:0").await?; + + let l1_addr = l1.local_addr()?; + let l2_addr = l2.local_addr()?; + let l3_addr = l3.local_addr()?; + + println!("Listener №1 on: {}", l1_addr); + println!("Listener №2 on: {}", l2_addr); + println!("Listener №3 on: {}", l3_addr); + + let nodes = vec![ + l1_addr.to_string(), + l2_addr.to_string(), + l3_addr.to_string(), + ]; + + for _ in 0..100 { + let addr = NearestDCBalancer::find_fastest_address( + nodes.iter().collect_vec(), + Duration::from_secs(PING_TIMEOUT_SECS), + ) + .await?; + assert!(nodes.contains(&addr)) + } + + Ok(()) + } + + #[tokio::test] + async fn detect_fastest_addr_with_fault() -> YdbResult<()> { + let l1 = TcpListener::bind("127.0.0.1:0").await?; + let l2 = TcpListener::bind("127.0.0.1:0").await?; + let l3 = TcpListener::bind("127.0.0.1:0").await?; + + let l1_addr = l1.local_addr()?; + let l2_addr = l2.local_addr()?; + let l3_addr = l3.local_addr()?; + + println!("Listener №1 on: {}", l1_addr); + println!("Listener №2 on: {}", l2_addr); + println!("Listener №3 on: {}", l3_addr); + + let nodes = vec![ + l1_addr.to_string(), + l2_addr.to_string(), + l3_addr.to_string(), + ]; + + drop(l1); + + for _ in 0..100 { + let addr = NearestDCBalancer::find_fastest_address( + nodes.iter().collect_vec(), + Duration::from_secs(PING_TIMEOUT_SECS), + ) + .await?; + assert!(nodes.contains(&addr) && addr != l1_addr.to_string()) + } + + Ok(()) + } + + #[tokio::test] + async fn detect_fastest_addr_one_alive() -> YdbResult<()> { + let l1 = TcpListener::bind("127.0.0.1:0").await?; + let l2 = TcpListener::bind("127.0.0.1:0").await?; + let l3 = TcpListener::bind("127.0.0.1:0").await?; + + let l1_addr = l1.local_addr()?; + let l2_addr = l2.local_addr()?; + let l3_addr = l3.local_addr()?; + + println!("Listener №1 on: {}", l1_addr); + println!("Listener №2 on: {}", l2_addr); + println!("Listener №3 on: {}", l3_addr); + + let nodes = vec![ + l1_addr.to_string(), + l2_addr.to_string(), + l3_addr.to_string(), + ]; + + drop(l1); + drop(l2); + + for _ in 0..100 { + let addr = NearestDCBalancer::find_fastest_address( + nodes.iter().collect_vec(), + Duration::from_secs(PING_TIMEOUT_SECS), + ) + .await?; + assert!(addr == l3_addr.to_string()) + } + + Ok(()) + } + + #[tokio::test] + async fn detect_fastest_addr_timeout() -> YdbResult<()> { + let l1 = TcpListener::bind("127.0.0.1:0").await?; + let l2 = TcpListener::bind("127.0.0.1:0").await?; + let l3 = TcpListener::bind("127.0.0.1:0").await?; + + let l1_addr = l1.local_addr()?; + let l2_addr = l2.local_addr()?; + let l3_addr = l3.local_addr()?; + + println!("Listener №1 on: {}", l1_addr); + println!("Listener №2 on: {}", l2_addr); + println!("Listener №3 on: {}", l3_addr); + + let nodes = vec![ + l1_addr.to_string(), + l2_addr.to_string(), + l3_addr.to_string(), + ]; + + drop(l1); + drop(l2); + drop(l3); + + let result = NearestDCBalancer::find_fastest_address( + nodes.iter().collect_vec(), + Duration::from_secs(3), + ) + .await; + match result { + Ok(_) => unreachable!(), + Err(err) => { + assert_eq!( + err.to_string(), + "Custom(\"timeout while detecting fastest address\")" + ); + } + } + Ok(()) + } + + #[tokio::test] + async fn no_addr_timeout() -> YdbResult<()> { + let result = + NearestDCBalancer::find_fastest_address(Vec::new(), Duration::from_secs(3)).await; + match result { + Ok(_) => unreachable!(), + Err(err) => { + assert_eq!( + err.to_string(), + "Custom(\"timeout while detecting fastest address\")" + ); + } + } + Ok(()) + } + + #[tokio::test] + async fn detect_fastest_addr() -> YdbResult<()> { + let l1 = TcpListener::bind("127.0.0.1:0").await?; + let l2 = TcpListener::bind("127.0.0.1:0").await?; + let l3 = TcpListener::bind("127.0.0.1:0").await?; + + let l1_addr = l1.local_addr()?; + let l2_addr = l2.local_addr()?; + let l3_addr = l3.local_addr()?; + + println!("Listener №1 on: {}", l1_addr); + println!("Listener №2 on: {}", l2_addr); + println!("Listener №3 on: {}", l3_addr); + + let nodes = vec![ + l1_addr.to_string(), + l2_addr.to_string(), + l3_addr.to_string(), + ]; + + drop(l1); + drop(l2); + drop(l3); + + let result = NearestDCBalancer::find_fastest_address( + nodes.iter().collect_vec(), + Duration::from_secs(3), + ) + .await; + match result { + Ok(_) => unreachable!(), + Err(err) => { + assert_eq!( + err.to_string(), + "Custom(\"timeout while detecting fastest address\")" + ); + } + } + Ok(()) + } + + #[tokio::test] + async fn adjusting_dc() -> YdbResult<()> { + let l1 = TcpListener::bind("127.0.0.1:0").await?; + let l2 = TcpListener::bind("127.0.0.1:0").await?; + let l3 = TcpListener::bind("127.0.0.1:0").await?; + + let l1_addr = l1.local_addr()?; + let l2_addr = l2.local_addr()?; + let l3_addr = l3.local_addr()?; + + println!("Listener №1 on: {}", l1_addr); + println!("Listener №2 on: {}", l2_addr); + println!("Listener №3 on: {}", l3_addr); + + let discovery_state = Arc::new(DiscoveryState::default()); + let balancer_state = Arc::new(Mutex::new(BalancerState::default())); + let balancer_state_updater = balancer_state.clone(); + let (state_sender, state_reciever) = watch::channel(discovery_state.clone()); + + let ping_token = CancellationToken::new(); + let ping_token_clone = ping_token.clone(); + + let waiter = Arc::new(WaiterImpl::new()); + let waiter_clone = waiter.clone(); + + let updater = tokio::spawn(async move { + NearestDCBalancer::adjust_local_dc( + balancer_state_updater, + state_reciever, + ping_token_clone, + waiter_clone, + ) + .await + }); + + let updated_state = Arc::new( + DiscoveryState::default() + .with_node_info( + Table, + NodeInfo::new( + Uri::from_str(l1_addr.to_string().as_str()).unwrap(), + "A".to_string(), + ), + ) + .with_node_info( + Table, + NodeInfo::new( + Uri::from_str(l2_addr.to_string().as_str()).unwrap(), + "B".to_string(), + ), + ) + .with_node_info( + Table, + NodeInfo::new( + Uri::from_str(l2_addr.to_string().as_str()).unwrap(), + "C".to_string(), + ), + ), + ); + assert!( + (balancer_state.lock().await) + .borrow() + .preferred_endpoints + .len() + == 0 // no endpoints + ); + let _ = state_sender.send(updated_state); + tokio::time::sleep(Duration::from_secs(2)).await; + assert_true!(timeout(Duration::from_secs(3), waiter.wait()).await.is_ok()); // should not wait + assert!( + (balancer_state.lock().await) + .borrow() + .preferred_endpoints + .len() + == 1 // only one endpoint in each dc + ); + let updated_state_next = Arc::new( + DiscoveryState::default() + .with_node_info( + Table, + NodeInfo::new( + Uri::from_str(l1_addr.to_string().as_str()).unwrap(), + "A".to_string(), + ), + ) + .with_node_info( + Table, + NodeInfo::new( + Uri::from_str(l2_addr.to_string().as_str()).unwrap(), + "A".to_string(), + ), + ), + ); + let _ = state_sender.send(updated_state_next); + tokio::time::sleep(Duration::from_secs(2)).await; + assert_true!(timeout(Duration::from_secs(3), waiter.wait()).await.is_ok()); // should not wait + assert!( + (balancer_state.lock().await) + .borrow() + .preferred_endpoints + .len() + == 2 // both endpoints in same dc + ); + ping_token.cancel(); // reciever stops wait for state change + let _ = tokio::join!(updater); // should join + Ok(()) + } + + #[tokio::test] + async fn nearest_dc_balancer_integration_with_error_fallback() -> YdbResult<()> { + let balancer = NearestDCBalancer::new(BalancerConfig { + fallback_strategy: FallbackStrategy::Error, + fallback_balancer: None, + }) + .unwrap(); + + let sh = SharedLoadBalancer::new_with_balancer(Box::new(balancer)); + + match sh.endpoint(Table) { + Ok(_) => unreachable!(), + Err(err) => assert_eq!( + err.to_string(), + "Custom(\"no available endpoints for service:table_service\")".to_string() + ), + } + Ok(()) + } + + #[tokio::test] + async fn nearest_dc_balancer_integration_with_other_fallback_error() -> YdbResult<()> { + let balancer = NearestDCBalancer::new(BalancerConfig { + fallback_strategy: FallbackStrategy::BalanceWithOther, + fallback_balancer: Some(Box::new(RandomLoadBalancer::new())), + }) + .unwrap(); + + let sh = SharedLoadBalancer::new_with_balancer(Box::new(balancer)); + + match sh.endpoint(Table) { + Ok(_) => unreachable!(), + Err(err) => assert_eq!( + err.to_string(), + "Custom(\"empty endpoint list for service: table_service\")".to_string() + ), + } + Ok(()) + } + + #[tokio::test] + async fn nearest_dc_balancer_integration() -> YdbResult<()> { + let l1 = TcpListener::bind("127.0.0.1:0").await?; + let l2 = TcpListener::bind("127.0.0.1:0").await?; + + let l1_addr = l1.local_addr()?; + let l2_addr = l2.local_addr()?; + + println!("Listener №1 on: {}", l1_addr); + println!("Listener №2 on: {}", l2_addr); + + let balancer = NearestDCBalancer::new(BalancerConfig { + fallback_strategy: FallbackStrategy::Error, + fallback_balancer: None, + }) + .unwrap(); + + let sh = SharedLoadBalancer::new_with_balancer(Box::new(balancer)); + let self_updater = sh.clone(); + let (state_sender, state_reciever) = + watch::channel::>(Arc::new(DiscoveryState::default())); + + tokio::spawn(async move { update_load_balancer(self_updater, state_reciever).await }); + + match sh.endpoint(Table) { + Ok(_) => unreachable!(), + Err(err) => assert_eq!( + err.to_string(), + "Custom(\"no available endpoints for service:table_service\")".to_string() + ), + } + + let updated_state = Arc::new( + DiscoveryState::default() + .with_node_info( + Table, + NodeInfo::new( + Uri::from_str(l1_addr.to_string().as_str()).unwrap(), + "A".to_string(), + ), + ) + .with_node_info( + Table, + NodeInfo::new( + Uri::from_str(l2_addr.to_string().as_str()).unwrap(), + "A".to_string(), + ), + ), + ); + + let _ = state_sender.send(updated_state); + + sh.wait().await?; + + match sh.endpoint(Table) { + Ok(uri) => { + let addr = uri.host().unwrap(); + assert!(addr == "127.0.0.1" || addr == "[::1]") + } + Err(err) => unreachable!("{}", err.to_string()), + } + Ok(()) + } } From 103a410d925c6020eeb367c49398eb4a1eb56e3e Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 1 Sep 2024 19:15:33 +0300 Subject: [PATCH 31/42] clear --- ydb/src/load_balancer.rs | 5 ----- 1 file changed, 5 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index 2b6ee58d..f68ba517 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -479,7 +479,6 @@ impl NearestDCBalancer { return; } } - println!("host:{}, port:{}", host, port); use std::net::ToSocketAddrs; let _ = (host, port).to_socket_addrs().and_then(|addrs| { for addr in addrs { @@ -492,7 +491,6 @@ impl NearestDCBalancer { } async fn find_fastest_address(addrs: Vec<&String>, timeout: Duration) -> YdbResult { - println!("addrs:{:?}", addrs); // Cancellation flow: timeout -> address collector -> address producers let interrupt_via_timeout = CancellationToken::new(); @@ -559,12 +557,9 @@ impl NearestDCBalancer { }; let _ = start_measure.send(()); - println!("start measure"); match tokio::time::timeout(timeout, wait_first_some_or_cancel).await { Ok(address_option) => { - println!("got fastest address:{:?}", address_option); - address_option } Err(_) => { From 94f00cfa113b28b36c275f4b72f4338c35951779 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sat, 7 Sep 2024 15:57:59 +0300 Subject: [PATCH 32/42] Move fallback balancer inside fallback strategy --- ydb/src/load_balancer.rs | 74 +++++++++++++--------------------------- 1 file changed, 23 insertions(+), 51 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index f68ba517..3d61551a 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -201,7 +201,6 @@ pub(crate) async fn update_load_balancer( pub(crate) struct BalancerConfig { fallback_strategy: FallbackStrategy, - fallback_balancer: Option>, } #[derive(Default)] @@ -210,17 +209,17 @@ struct BalancerState { } // What will balancer do if there is no available endpoints at local dc -#[derive(PartialEq, Eq)] pub(crate) enum FallbackStrategy { - Error, // Just throw error - BalanceWithOther, // Use another balancer + Error, // Just throw error + BalanceWithOther(Box), // Use another balancer } impl Default for BalancerConfig { fn default() -> Self { BalancerConfig { - fallback_strategy: FallbackStrategy::BalanceWithOther, - fallback_balancer: Some(Box::new(RandomLoadBalancer::new())), + fallback_strategy: FallbackStrategy::BalanceWithOther(Box::new( + RandomLoadBalancer::new(), + )), } } } @@ -236,24 +235,6 @@ pub(crate) struct NearestDCBalancer { impl NearestDCBalancer { pub(crate) fn new(config: BalancerConfig) -> YdbResult { - match config.fallback_balancer.as_ref() { - Some(_) => { - if config.fallback_strategy == FallbackStrategy::Error { - return Err(YdbError::Custom( - "fallback strategy is \"Error\" but fallback balancer was provided" - .to_string(), - )); - } - } - None => { - if config.fallback_strategy == FallbackStrategy::BalanceWithOther { - return Err(YdbError::Custom( - "no fallback balancer was provided".to_string(), - )); - } - } - } - let discovery_state = Arc::new(DiscoveryState::default()); let balancer_state = Arc::new(Mutex::new(BalancerState::default())); let balancer_state_updater = balancer_state.clone(); @@ -305,9 +286,11 @@ impl LoadBalancer for NearestDCBalancer { } fn set_discovery_state(&mut self, discovery_state: &Arc) -> YdbResult<()> { - match self.config.fallback_balancer.as_mut() { - Some(balancer) => balancer.set_discovery_state(discovery_state)?, - None => (), + match self.config.fallback_strategy.borrow_mut() { + FallbackStrategy::BalanceWithOther(balancer) => { + balancer.set_discovery_state(discovery_state)? + } + FallbackStrategy::Error => (), } self.discovery_state = discovery_state.clone(); let _ = self.state_sender.send(discovery_state.clone()); @@ -316,9 +299,11 @@ impl LoadBalancer for NearestDCBalancer { fn waiter(&self) -> Box { let self_waiter = Box::new(self.waiter.clone()); - match self.config.fallback_balancer.as_ref() { - Some(balancer) => Box::new(AllWaiter::new(vec![self_waiter, balancer.waiter()])), - None => self_waiter, + match self.config.fallback_strategy.borrow() { + FallbackStrategy::BalanceWithOther(balancer) => { + Box::new(AllWaiter::new(vec![self_waiter, balancer.waiter()])) + } + FallbackStrategy::Error => self_waiter, } } } @@ -333,18 +318,14 @@ impl NearestDCBalancer { for ep in state_guard.borrow().preferred_endpoints.iter() { return YdbResult::Ok(ep.uri.clone()); } - match self.config.fallback_strategy { + match self.config.fallback_strategy.borrow() { FallbackStrategy::Error => Err(YdbError::custom(format!( "no available endpoints for service:{}", service ))), - FallbackStrategy::BalanceWithOther => { + FallbackStrategy::BalanceWithOther(balancer) => { info!("trying fallback balancer..."); - self.config - .fallback_balancer - .as_ref() - .unwrap() // unwrap is safe [checks inside ::new()] - .endpoint(service) + balancer.endpoint(service) } } } @@ -491,7 +472,6 @@ impl NearestDCBalancer { } async fn find_fastest_address(addrs: Vec<&String>, timeout: Duration) -> YdbResult { - // Cancellation flow: timeout -> address collector -> address producers let interrupt_via_timeout = CancellationToken::new(); let interrupt_collector_future = interrupt_via_timeout.child_token(); @@ -559,9 +539,7 @@ impl NearestDCBalancer { let _ = start_measure.send(()); match tokio::time::timeout(timeout, wait_first_some_or_cancel).await { - Ok(address_option) => { - address_option - } + Ok(address_option) => address_option, Err(_) => { interrupt_via_timeout.cancel(); YdbResult::Err("timeout while detecting fastest address".into()) @@ -1076,7 +1054,6 @@ mod test { async fn nearest_dc_balancer_integration_with_error_fallback() -> YdbResult<()> { let balancer = NearestDCBalancer::new(BalancerConfig { fallback_strategy: FallbackStrategy::Error, - fallback_balancer: None, }) .unwrap(); @@ -1094,11 +1071,7 @@ mod test { #[tokio::test] async fn nearest_dc_balancer_integration_with_other_fallback_error() -> YdbResult<()> { - let balancer = NearestDCBalancer::new(BalancerConfig { - fallback_strategy: FallbackStrategy::BalanceWithOther, - fallback_balancer: Some(Box::new(RandomLoadBalancer::new())), - }) - .unwrap(); + let balancer = NearestDCBalancer::new(BalancerConfig::default()).unwrap(); let sh = SharedLoadBalancer::new_with_balancer(Box::new(balancer)); @@ -1125,7 +1098,6 @@ mod test { let balancer = NearestDCBalancer::new(BalancerConfig { fallback_strategy: FallbackStrategy::Error, - fallback_balancer: None, }) .unwrap(); @@ -1161,11 +1133,11 @@ mod test { ), ), ); - + let _ = state_sender.send(updated_state); - + sh.wait().await?; - + match sh.endpoint(Table) { Ok(uri) => { let addr = uri.host().unwrap(); From 9f92b50647af48caebf2e00bf4a6ceb40445404a Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sat, 7 Sep 2024 16:19:03 +0300 Subject: [PATCH 33/42] Choose random endpoint from prefered --- ydb/src/load_balancer.rs | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index 3d61551a..a62e0e88 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -4,7 +4,8 @@ use crate::grpc_wrapper::raw_services::Service; use crate::waiter::{AllWaiter, Waiter, WaiterImpl}; use http::Uri; use itertools::Itertools; -use rand::thread_rng; +use rand::seq::IteratorRandom; +use rand::{seq::SliceRandom, thread_rng}; use std::borrow::{Borrow, BorrowMut}; use std::{ collections::HashMap, @@ -315,8 +316,13 @@ impl NearestDCBalancer { fn get_endpoint(&self, service: Service) -> YdbResult { match self.balancer_state.try_lock() { Ok(state_guard) => { - for ep in state_guard.borrow().preferred_endpoints.iter() { - return YdbResult::Ok(ep.uri.clone()); + match state_guard + .borrow() + .preferred_endpoints + .choose(&mut thread_rng()) + { + Some(ep) => return YdbResult::Ok(ep.uri.clone()), + None => (), } match self.config.fallback_strategy.borrow() { FallbackStrategy::Error => Err(YdbError::custom(format!( @@ -418,7 +424,6 @@ impl NearestDCBalancer { } fn get_random_endpoints<'a>(dc_endpoints: &'a mut Vec<&'a NodeInfo>) -> &mut Vec<&NodeInfo> { - use rand::seq::SliceRandom; dc_endpoints.shuffle(&mut thread_rng()); dc_endpoints.truncate(NODES_PER_DC); dc_endpoints From 57d26273597ae083b0175fadbf0eed3de249aa2a Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sat, 7 Sep 2024 16:21:47 +0300 Subject: [PATCH 34/42] Iter over only values in dc_to_nodes --- ydb/src/load_balancer.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs index a62e0e88..f5ccda89 100644 --- a/ydb/src/load_balancer.rs +++ b/ydb/src/load_balancer.rs @@ -363,7 +363,7 @@ impl NearestDCBalancer { Ok(some_nodes) => { let mut dc_to_nodes = Self::split_endpoints_by_location(some_nodes); let mut to_check = Vec::with_capacity(NODES_PER_DC * dc_to_nodes.keys().len()); - dc_to_nodes.iter_mut().for_each(|(_, endpoints)| { + dc_to_nodes.values_mut().for_each(|endpoints| { to_check.append(Self::get_random_endpoints(endpoints)) }); match Self::find_local_dc(&to_check).await { From 5661d8501071de4438ce5f1c7dd11c13300824c3 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 15 Sep 2024 19:36:27 +0300 Subject: [PATCH 35/42] Divide into multiple files --- ydb/src/load_balancer.rs | 1155 ------------------ ydb/src/load_balancer/balancer_test.rs | 602 +++++++++ ydb/src/load_balancer/mod.rs | 46 + ydb/src/load_balancer/nearest_dc_balancer.rs | 393 ++++++ ydb/src/load_balancer/random_balancer.rs | 68 ++ ydb/src/load_balancer/shared_balancer.rs | 60 + ydb/src/load_balancer/static_balancer.rs | 43 + 7 files changed, 1212 insertions(+), 1155 deletions(-) delete mode 100644 ydb/src/load_balancer.rs create mode 100644 ydb/src/load_balancer/balancer_test.rs create mode 100644 ydb/src/load_balancer/mod.rs create mode 100644 ydb/src/load_balancer/nearest_dc_balancer.rs create mode 100644 ydb/src/load_balancer/random_balancer.rs create mode 100644 ydb/src/load_balancer/shared_balancer.rs create mode 100644 ydb/src/load_balancer/static_balancer.rs diff --git a/ydb/src/load_balancer.rs b/ydb/src/load_balancer.rs deleted file mode 100644 index f5ccda89..00000000 --- a/ydb/src/load_balancer.rs +++ /dev/null @@ -1,1155 +0,0 @@ -use crate::discovery::{Discovery, DiscoveryState, NodeInfo}; -use crate::errors::*; -use crate::grpc_wrapper::raw_services::Service; -use crate::waiter::{AllWaiter, Waiter, WaiterImpl}; -use http::Uri; -use itertools::Itertools; -use rand::seq::IteratorRandom; -use rand::{seq::SliceRandom, thread_rng}; -use std::borrow::{Borrow, BorrowMut}; -use std::{ - collections::HashMap, - sync::{Arc, RwLock}, - time::Duration, -}; -use tokio::{ - io::AsyncWriteExt, - net::TcpStream, - sync::{ - broadcast, mpsc, watch, - watch::{Receiver, Sender}, - Mutex, - }, - task::JoinSet, -}; -use tokio_util::sync::CancellationToken; -use tracing::{error, info, warn}; - -#[mockall::automock] -pub(crate) trait LoadBalancer: Send + Sync + Waiter { - fn endpoint(&self, service: Service) -> YdbResult; - fn set_discovery_state(&mut self, discovery_state: &Arc) -> YdbResult<()>; - fn waiter(&self) -> Box; // need for wait ready in without read lock -} - -#[async_trait::async_trait] -impl Waiter for MockLoadBalancer { - async fn wait(&self) -> YdbResult<()> { - Ok(()) - } -} - -#[derive(Clone)] -pub(crate) struct SharedLoadBalancer { - inner: Arc>>, -} - -impl SharedLoadBalancer { - pub(crate) fn new(discovery: &dyn Discovery) -> Self { - Self::new_with_balancer_and_updater(Box::new(RandomLoadBalancer::new()), discovery) - } - - pub(crate) fn new_with_balancer(load_balancer: Box) -> Self { - Self { - inner: Arc::new(RwLock::new(load_balancer)), - } - } - - pub(crate) fn new_with_balancer_and_updater( - load_balancer: Box, - discovery: &dyn Discovery, - ) -> Self { - let mut shared_lb = Self::new_with_balancer(load_balancer); - let shared_lb_updater = shared_lb.clone(); - let discovery_receiver = discovery.subscribe(); - let _ = shared_lb.set_discovery_state(&discovery.state()); - tokio::spawn( - async move { update_load_balancer(shared_lb_updater, discovery_receiver).await }, - ); - shared_lb - } -} - -impl LoadBalancer for SharedLoadBalancer { - fn endpoint(&self, service: Service) -> YdbResult { - self.inner.read()?.endpoint(service) - } - - fn set_discovery_state(&mut self, discovery_state: &Arc) -> YdbResult<()> { - self.inner.write()?.set_discovery_state(discovery_state) - } - - fn waiter(&self) -> Box { - return self.inner.read().unwrap().waiter(); - } -} - -#[async_trait::async_trait] -impl Waiter for SharedLoadBalancer { - async fn wait(&self) -> YdbResult<()> { - let waiter = self.inner.read()?.waiter(); - return waiter.wait().await; - } -} - -pub(crate) struct StaticLoadBalancer { - endpoint: Uri, -} - -impl StaticLoadBalancer { - #[allow(dead_code)] - pub(crate) fn new(endpoint: Uri) -> Self { - Self { endpoint } - } -} - -impl LoadBalancer for StaticLoadBalancer { - fn endpoint(&self, _: Service) -> YdbResult { - Ok(self.endpoint.clone()) - } - - fn set_discovery_state(&mut self, _: &Arc) -> YdbResult<()> { - Err(YdbError::Custom( - "static balancer no way to update state".into(), - )) - } - - fn waiter(&self) -> Box { - let waiter = WaiterImpl::new(); - waiter.set_received(Ok(())); - Box::new(waiter) - } -} - -#[async_trait::async_trait] -impl Waiter for StaticLoadBalancer { - async fn wait(&self) -> YdbResult<()> { - Ok(()) - } -} - -#[derive(Clone)] -pub(crate) struct RandomLoadBalancer { - discovery_state: Arc, - waiter: Arc, -} - -impl RandomLoadBalancer { - pub(crate) fn new() -> Self { - Self { - discovery_state: Arc::new(DiscoveryState::default()), - waiter: Arc::new(WaiterImpl::new()), - } - } -} - -impl LoadBalancer for RandomLoadBalancer { - fn endpoint(&self, service: Service) -> YdbResult { - let nodes = self.discovery_state.get_nodes(&service); - match nodes { - None => Err(YdbError::Custom(format!( - "no endpoints for service: '{}'", - service - ))), - Some(nodes) => { - if !nodes.is_empty() { - let index = rand::random::() % nodes.len(); - let node = &nodes[index % nodes.len()]; - Ok(node.uri.clone()) - } else { - Err(YdbError::Custom(format!( - "empty endpoint list for service: {}", - service - ))) - } - } - } - } - - fn set_discovery_state(&mut self, discovery_state: &Arc) -> YdbResult<()> { - self.discovery_state = discovery_state.clone(); - if !self.discovery_state.is_empty() { - self.waiter.set_received(Ok(())) - } - Ok(()) - } - - fn waiter(&self) -> Box { - Box::new(self.waiter.clone()) - } -} - -#[async_trait::async_trait] -impl Waiter for RandomLoadBalancer { - async fn wait(&self) -> YdbResult<()> { - self.waiter.wait().await - } -} - -pub(crate) async fn update_load_balancer( - mut lb: impl LoadBalancer, - mut receiver: Receiver>, -) { - loop { - // clone for prevent block send side while update current lb - let state = receiver.borrow_and_update().clone(); - let _ = lb.set_discovery_state(&state); - if receiver.changed().await.is_err() { - break; - } - } -} - -pub(crate) struct BalancerConfig { - fallback_strategy: FallbackStrategy, -} - -#[derive(Default)] -struct BalancerState { - preferred_endpoints: Vec, -} - -// What will balancer do if there is no available endpoints at local dc -pub(crate) enum FallbackStrategy { - Error, // Just throw error - BalanceWithOther(Box), // Use another balancer -} - -impl Default for BalancerConfig { - fn default() -> Self { - BalancerConfig { - fallback_strategy: FallbackStrategy::BalanceWithOther(Box::new( - RandomLoadBalancer::new(), - )), - } - } -} - -pub(crate) struct NearestDCBalancer { - discovery_state: Arc, - state_sender: Sender>, - ping_token: CancellationToken, - waiter: Arc, - config: BalancerConfig, - balancer_state: Arc>, -} - -impl NearestDCBalancer { - pub(crate) fn new(config: BalancerConfig) -> YdbResult { - let discovery_state = Arc::new(DiscoveryState::default()); - let balancer_state = Arc::new(Mutex::new(BalancerState::default())); - let balancer_state_updater = balancer_state.clone(); - let (state_sender, state_reciever) = watch::channel(discovery_state.clone()); - - let ping_token = CancellationToken::new(); - let ping_token_clone = ping_token.clone(); - - let waiter = Arc::new(WaiterImpl::new()); - let waiter_clone = waiter.clone(); - - tokio::spawn(async move { - Self::adjust_local_dc( - balancer_state_updater, - state_reciever, - ping_token_clone, - waiter_clone, - ) - .await - }); - - Ok(Self { - discovery_state, - state_sender, - ping_token, - waiter, - config, - balancer_state, - }) - } -} - -impl Drop for NearestDCBalancer { - fn drop(&mut self) { - self.ping_token.cancel(); - } -} - -#[async_trait::async_trait] -impl Waiter for NearestDCBalancer { - async fn wait(&self) -> YdbResult<()> { - self.waiter().wait().await - } -} - -impl LoadBalancer for NearestDCBalancer { - fn endpoint(&self, service: Service) -> YdbResult { - self.get_endpoint(service) - } - - fn set_discovery_state(&mut self, discovery_state: &Arc) -> YdbResult<()> { - match self.config.fallback_strategy.borrow_mut() { - FallbackStrategy::BalanceWithOther(balancer) => { - balancer.set_discovery_state(discovery_state)? - } - FallbackStrategy::Error => (), - } - self.discovery_state = discovery_state.clone(); - let _ = self.state_sender.send(discovery_state.clone()); - Ok(()) - } - - fn waiter(&self) -> Box { - let self_waiter = Box::new(self.waiter.clone()); - match self.config.fallback_strategy.borrow() { - FallbackStrategy::BalanceWithOther(balancer) => { - Box::new(AllWaiter::new(vec![self_waiter, balancer.waiter()])) - } - FallbackStrategy::Error => self_waiter, - } - } -} - -const NODES_PER_DC: usize = 5; -const PING_TIMEOUT_SECS: u64 = 60; - -impl NearestDCBalancer { - fn get_endpoint(&self, service: Service) -> YdbResult { - match self.balancer_state.try_lock() { - Ok(state_guard) => { - match state_guard - .borrow() - .preferred_endpoints - .choose(&mut thread_rng()) - { - Some(ep) => return YdbResult::Ok(ep.uri.clone()), - None => (), - } - match self.config.fallback_strategy.borrow() { - FallbackStrategy::Error => Err(YdbError::custom(format!( - "no available endpoints for service:{}", - service - ))), - FallbackStrategy::BalanceWithOther(balancer) => { - info!("trying fallback balancer..."); - balancer.endpoint(service) - } - } - } - Err(_) => Err(YdbError::Custom( - "balancer is updating its state".to_string(), - )), - } - } - - async fn adjust_local_dc( - balancer_state: Arc>, - mut state_reciever: watch::Receiver>, - stop_ping_process: CancellationToken, - waiter: Arc, - ) { - loop { - tokio::select! { - _ = stop_ping_process.cancelled() => { - return - } - result = state_reciever.changed() =>{ - if result.is_err(){ // sender have been dropped - return - } - } - } - let new_discovery_state = state_reciever.borrow_and_update().clone(); - match Self::extract_nodes(&new_discovery_state) { - Ok(some_nodes) => { - let mut dc_to_nodes = Self::split_endpoints_by_location(some_nodes); - let mut to_check = Vec::with_capacity(NODES_PER_DC * dc_to_nodes.keys().len()); - dc_to_nodes.values_mut().for_each(|endpoints| { - to_check.append(Self::get_random_endpoints(endpoints)) - }); - match Self::find_local_dc(&to_check).await { - Ok(dc) => { - info!("found new local dc:{}", dc); - Self::adjust_preferred_endpoints(&balancer_state, some_nodes, dc).await; - waiter.set_received(Ok(())); - } - Err(err) => { - error!("error on search local dc:{}", err); - continue; - } - } - } - Err(_) => continue, - } - } - } - - async fn adjust_preferred_endpoints( - balancer_state: &Arc>, - new_nodes: &Vec, - local_dc: String, - ) { - info!("adjusting endpoints"); - let new_preferred_endpoints = new_nodes - .into_iter() - .filter(|ep| ep.location == local_dc) - .map(|ep| ep.clone()) - .collect_vec(); - (balancer_state.lock().await) // fast lock - .borrow_mut() - .preferred_endpoints = new_preferred_endpoints; - } - - fn extract_nodes(from_state: &Arc) -> YdbResult<&Vec> { - let nodes = from_state.get_all_nodes(); - match nodes { - None => Err(YdbError::Custom(format!( - "no endpoints on discovery update" - ))), - Some(nodes) => Ok(nodes), - } - } - - fn split_endpoints_by_location<'a>( - nodes: &'a Vec, - ) -> HashMap> { - let mut dc_to_eps = HashMap::>::new(); - nodes.into_iter().for_each(|info| { - if let Some(nodes) = dc_to_eps.get_mut(&info.location) { - nodes.push(info); - } else { - dc_to_eps.insert(info.location.clone(), vec![info]); - } - }); - dc_to_eps - } - - fn get_random_endpoints<'a>(dc_endpoints: &'a mut Vec<&'a NodeInfo>) -> &mut Vec<&NodeInfo> { - dc_endpoints.shuffle(&mut thread_rng()); - dc_endpoints.truncate(NODES_PER_DC); - dc_endpoints - } - - async fn find_local_dc(to_check: &[&NodeInfo]) -> YdbResult { - let addr_to_node = Self::addr_to_node(to_check); - if addr_to_node.is_empty() { - return Err(YdbError::Custom(format!("no ip addresses for endpoints"))); - } - let addrs = addr_to_node.keys(); - match Self::find_fastest_address(addrs.collect(), Duration::from_secs(PING_TIMEOUT_SECS)) - .await - { - Ok(fastest_address) => Ok(addr_to_node[&fastest_address].location.clone()), - Err(err) => { - error!("could not find fastest address:{}", err); - Err(err) - } - } - } - - fn addr_to_node<'a>(nodes: &[&'a NodeInfo]) -> HashMap { - let mut addr_to_node = HashMap::::with_capacity(2 * nodes.len()); // ipv4 + ipv6 - nodes.into_iter().for_each(|info| { - let host: &str; - let port: u16; - match info.uri.host() { - Some(uri_host) => host = uri_host, - None => { - warn!("no host for uri:{}", info.uri); - return; - } - } - match info.uri.port() { - Some(uri_port) => port = uri_port.as_u16(), - None => { - warn!("no port for uri:{}", info.uri); - return; - } - } - use std::net::ToSocketAddrs; - let _ = (host, port).to_socket_addrs().and_then(|addrs| { - for addr in addrs { - addr_to_node.insert(addr.to_string(), info); - } - Ok(()) - }); - }); - addr_to_node - } - - async fn find_fastest_address(addrs: Vec<&String>, timeout: Duration) -> YdbResult { - // Cancellation flow: timeout -> address collector -> address producers - let interrupt_via_timeout = CancellationToken::new(); - let interrupt_collector_future = interrupt_via_timeout.child_token(); - let stop_measure = interrupt_collector_future.child_token(); // (*) - - let (start_measure, _) = broadcast::channel::<()>(1); - let buffer_cap = if addrs.len() > 0 { addrs.len() } else { 1 }; - let (addr_sender, mut addr_reciever) = mpsc::channel::>(buffer_cap); - let mut nursery = JoinSet::new(); - - for addr in addrs { - let (mut wait_for_start, stop_measure, addr, addr_sender) = ( - start_measure.subscribe(), - stop_measure.clone(), - addr.clone(), - addr_sender.clone(), - ); - - nursery.spawn(async move { - let _ = wait_for_start.recv().await; - tokio::select! { - connection_result = TcpStream::connect(addr.clone()) =>{ - match connection_result{ - Ok(mut connection) => { - let _ = connection.shutdown().await; - let _ = addr_sender.send(Some(addr)).await; - }, - Err(_) => {let _ = addr_sender.send(None).await;}, - } - } - _ = stop_measure.cancelled() => { - (); - } - } - }); - } - - let wait_first_some_or_cancel = async { - loop { - tokio::select! { - biased; // check timeout first - _ = interrupt_collector_future.cancelled() =>{ - Self::join_all(&mut nursery).await; // children will be cancelled due to tokens chaining, see (*) - return YdbResult::Err("cancelled".into()) - } - address_reciever_option = addr_reciever.recv() =>{ - match address_reciever_option { - Some(address_option) => { - match address_option { - Some(address) =>{ - interrupt_collector_future.cancel(); // Cancel other producing children - Self::join_all(&mut nursery).await; - return YdbResult::Ok(address); - }, - None => continue, // Some producer sent blank address -> wait others - } - }, - None => return YdbResult::Err("no fastest address".into()), // Channel closed, all producers have done measures - } - } - } - } - }; - - let _ = start_measure.send(()); - - match tokio::time::timeout(timeout, wait_first_some_or_cancel).await { - Ok(address_option) => address_option, - Err(_) => { - interrupt_via_timeout.cancel(); - YdbResult::Err("timeout while detecting fastest address".into()) - } - } - } - - async fn join_all(awaitable: &mut JoinSet<()>) { - while awaitable.join_next().await.is_some() {} - } -} -#[cfg(test)] -mod test { - use super::*; - use crate::discovery::NodeInfo; - use crate::grpc_wrapper::raw_services::Service::Table; - use mockall::predicate; - use ntest::assert_true; - use std::collections::HashMap; - use std::str::FromStr; - use std::sync::atomic::AtomicUsize; - use std::sync::atomic::Ordering::Relaxed; - use std::time::Duration; - use tokio::net::TcpListener; - use tokio::time::timeout; - use tracing::trace; - - #[test] - fn shared_load_balancer() -> YdbResult<()> { - let endpoint_counter = Arc::new(AtomicUsize::new(0)); - let test_uri = Uri::from_str("http://test.com")?; - - let mut lb_mock = MockLoadBalancer::new(); - let endpoint_counter_mock = endpoint_counter.clone(); - let test_uri_mock = test_uri.clone(); - - lb_mock.expect_endpoint().returning(move |_service| { - endpoint_counter_mock.fetch_add(1, Relaxed); - Ok(test_uri_mock.clone()) - }); - - let s1 = SharedLoadBalancer::new_with_balancer(Box::new(lb_mock)); - - #[allow(clippy::redundant_clone)] - let s2 = s1.clone(); - - assert_eq!(test_uri, s1.endpoint(Table)?); - assert_eq!(test_uri, s2.endpoint(Table)?); - assert_eq!(endpoint_counter.load(Relaxed), 2); - Ok(()) - } - - #[tokio::test] - async fn update_load_balancer_test() -> YdbResult<()> { - let original_discovery_state = Arc::new(DiscoveryState::default()); - let (sender, receiver) = tokio::sync::watch::channel(original_discovery_state.clone()); - - let new_discovery_state = Arc::new(DiscoveryState::default().with_node_info( - Table, - NodeInfo::new(Uri::from_str("http://test.com").unwrap(), String::new()), - )); - - let (first_update_sender, first_update_receiver) = tokio::sync::oneshot::channel(); - let (second_update_sender, second_update_receiver) = tokio::sync::oneshot::channel(); - let (updater_finished_sender, updater_finished_receiver) = - tokio::sync::oneshot::channel::<()>(); - - let mut first_update_sender = Some(first_update_sender); - let mut second_update_sender = Some(second_update_sender); - let mut lb_mock = MockLoadBalancer::new(); - lb_mock - .expect_set_discovery_state() - .with(predicate::eq(original_discovery_state.clone())) - .times(1) - .returning(move |_| { - trace!("first set"); - first_update_sender.take().unwrap().send(()).unwrap(); - Ok(()) - }); - - lb_mock - .expect_set_discovery_state() - .with(predicate::eq(new_discovery_state.clone())) - .times(1) - .returning(move |_| { - trace!("second set"); - second_update_sender.take().unwrap().send(()).unwrap(); - Ok(()) - }); - - let shared_lb = SharedLoadBalancer::new_with_balancer(Box::new(lb_mock)); - - tokio::spawn(async move { - trace!("updater start"); - update_load_balancer(shared_lb, receiver).await; - trace!("updater finished"); - updater_finished_sender.send(()).unwrap(); - }); - - tokio::spawn(async move { - first_update_receiver.await.unwrap(); - sender.send(new_discovery_state).unwrap(); - second_update_receiver.await.unwrap(); - drop(sender); - }); - - tokio::select! { - _ = updater_finished_receiver =>{} - _ = tokio::time::sleep(Duration::from_secs(10)) => { - panic!("test failed"); - } - } - // updater_finished_receiver.await.unwrap(); - Ok(()) - } - - #[test] - fn random_load_balancer() -> YdbResult<()> { - let one = Uri::from_str("http://one:213")?; - let two = Uri::from_str("http://two:213")?; - let load_balancer = RandomLoadBalancer { - discovery_state: Arc::new( - DiscoveryState::default() - .with_node_info(Table, NodeInfo::new(one.clone(), String::new())) - .with_node_info(Table, NodeInfo::new(two.clone(), String::new())), - ), - waiter: Arc::new(WaiterImpl::new()), - }; - - let mut map = HashMap::new(); - map.insert(one.to_string(), 0); - map.insert(two.to_string(), 0); - - for _ in 0..100 { - let u = load_balancer.endpoint(Table)?; - let val = *map.get_mut(u.to_string().as_str()).unwrap(); - map.insert(u.to_string(), val + 1); - } - - assert_eq!(map.len(), 2); - assert!(*map.get(one.to_string().as_str()).unwrap() > 30); - assert!(*map.get(two.to_string().as_str()).unwrap() > 30); - Ok(()) - } - - #[test] - fn split_by_location() -> YdbResult<()> { - let nodes = vec![ - NodeInfo::new(Uri::from_str("http://one:213")?, "A".to_string()), - NodeInfo::new(Uri::from_str("http://two:213")?, "A".to_string()), - NodeInfo::new(Uri::from_str("http://three:213")?, "B".to_string()), - NodeInfo::new(Uri::from_str("http://four:213")?, "B".to_string()), - NodeInfo::new(Uri::from_str("http://five:213")?, "C".to_string()), - ]; - let splitted = NearestDCBalancer::split_endpoints_by_location(&nodes); - assert_eq!(splitted.keys().len(), 3); - assert_eq!(splitted["A"].len(), 2); - assert_eq!(splitted["B"].len(), 2); - assert_eq!(splitted["C"].len(), 1); - Ok(()) - } - - #[test] - fn choose_random_endpoints() -> YdbResult<()> { - let nodes = vec![ - NodeInfo::new(Uri::from_str("http://one:213")?, "C".to_string()), - NodeInfo::new(Uri::from_str("http://two:213")?, "C".to_string()), - NodeInfo::new(Uri::from_str("http://three:213")?, "C".to_string()), - NodeInfo::new(Uri::from_str("http://four:213")?, "C".to_string()), - NodeInfo::new(Uri::from_str("http://five:213")?, "C".to_string()), - NodeInfo::new(Uri::from_str("http://seven:213")?, "C".to_string()), - NodeInfo::new(Uri::from_str("http://eight:213")?, "C".to_string()), - NodeInfo::new(Uri::from_str("http://nine:213")?, "C".to_string()), - ]; - - let mut refs = nodes.iter().collect_vec(); - let nodes_clone = refs.clone(); - let random_subset = NearestDCBalancer::get_random_endpoints(&mut refs); - - assert_eq!(random_subset.len(), NODES_PER_DC); - for node in random_subset { - assert_true!(nodes_clone.contains(node)) - } - - Ok(()) - } - - #[test] - fn extract_addrs_and_map_them() -> YdbResult<()> { - let one = NodeInfo::new(Uri::from_str("http://localhost:123")?, "C".to_string()); - let two = NodeInfo::new(Uri::from_str("http://localhost:321")?, "C".to_string()); - let nodes = vec![&one, &two]; - let map = NearestDCBalancer::addr_to_node(&nodes); - - assert_eq!(map.keys().len(), 4); // ipv4 + ipv6 on each - assert_true!(map.keys().contains(&"127.0.0.1:123".to_string())); - assert_true!(map.keys().contains(&"[::1]:123".to_string())); - assert!(map["127.0.0.1:123"].eq(&one)); - assert!(map["127.0.0.1:123"].eq(map["[::1]:123"])); - - Ok(()) - } - - #[tokio::test] - async fn detect_fastest_addr_just_some() -> YdbResult<()> { - let l1 = TcpListener::bind("127.0.0.1:0").await?; - let l2 = TcpListener::bind("127.0.0.1:0").await?; - let l3 = TcpListener::bind("127.0.0.1:0").await?; - - let l1_addr = l1.local_addr()?; - let l2_addr = l2.local_addr()?; - let l3_addr = l3.local_addr()?; - - println!("Listener №1 on: {}", l1_addr); - println!("Listener №2 on: {}", l2_addr); - println!("Listener №3 on: {}", l3_addr); - - let nodes = vec![ - l1_addr.to_string(), - l2_addr.to_string(), - l3_addr.to_string(), - ]; - - for _ in 0..100 { - let addr = NearestDCBalancer::find_fastest_address( - nodes.iter().collect_vec(), - Duration::from_secs(PING_TIMEOUT_SECS), - ) - .await?; - assert!(nodes.contains(&addr)) - } - - Ok(()) - } - - #[tokio::test] - async fn detect_fastest_addr_with_fault() -> YdbResult<()> { - let l1 = TcpListener::bind("127.0.0.1:0").await?; - let l2 = TcpListener::bind("127.0.0.1:0").await?; - let l3 = TcpListener::bind("127.0.0.1:0").await?; - - let l1_addr = l1.local_addr()?; - let l2_addr = l2.local_addr()?; - let l3_addr = l3.local_addr()?; - - println!("Listener №1 on: {}", l1_addr); - println!("Listener №2 on: {}", l2_addr); - println!("Listener №3 on: {}", l3_addr); - - let nodes = vec![ - l1_addr.to_string(), - l2_addr.to_string(), - l3_addr.to_string(), - ]; - - drop(l1); - - for _ in 0..100 { - let addr = NearestDCBalancer::find_fastest_address( - nodes.iter().collect_vec(), - Duration::from_secs(PING_TIMEOUT_SECS), - ) - .await?; - assert!(nodes.contains(&addr) && addr != l1_addr.to_string()) - } - - Ok(()) - } - - #[tokio::test] - async fn detect_fastest_addr_one_alive() -> YdbResult<()> { - let l1 = TcpListener::bind("127.0.0.1:0").await?; - let l2 = TcpListener::bind("127.0.0.1:0").await?; - let l3 = TcpListener::bind("127.0.0.1:0").await?; - - let l1_addr = l1.local_addr()?; - let l2_addr = l2.local_addr()?; - let l3_addr = l3.local_addr()?; - - println!("Listener №1 on: {}", l1_addr); - println!("Listener №2 on: {}", l2_addr); - println!("Listener №3 on: {}", l3_addr); - - let nodes = vec![ - l1_addr.to_string(), - l2_addr.to_string(), - l3_addr.to_string(), - ]; - - drop(l1); - drop(l2); - - for _ in 0..100 { - let addr = NearestDCBalancer::find_fastest_address( - nodes.iter().collect_vec(), - Duration::from_secs(PING_TIMEOUT_SECS), - ) - .await?; - assert!(addr == l3_addr.to_string()) - } - - Ok(()) - } - - #[tokio::test] - async fn detect_fastest_addr_timeout() -> YdbResult<()> { - let l1 = TcpListener::bind("127.0.0.1:0").await?; - let l2 = TcpListener::bind("127.0.0.1:0").await?; - let l3 = TcpListener::bind("127.0.0.1:0").await?; - - let l1_addr = l1.local_addr()?; - let l2_addr = l2.local_addr()?; - let l3_addr = l3.local_addr()?; - - println!("Listener №1 on: {}", l1_addr); - println!("Listener №2 on: {}", l2_addr); - println!("Listener №3 on: {}", l3_addr); - - let nodes = vec![ - l1_addr.to_string(), - l2_addr.to_string(), - l3_addr.to_string(), - ]; - - drop(l1); - drop(l2); - drop(l3); - - let result = NearestDCBalancer::find_fastest_address( - nodes.iter().collect_vec(), - Duration::from_secs(3), - ) - .await; - match result { - Ok(_) => unreachable!(), - Err(err) => { - assert_eq!( - err.to_string(), - "Custom(\"timeout while detecting fastest address\")" - ); - } - } - Ok(()) - } - - #[tokio::test] - async fn no_addr_timeout() -> YdbResult<()> { - let result = - NearestDCBalancer::find_fastest_address(Vec::new(), Duration::from_secs(3)).await; - match result { - Ok(_) => unreachable!(), - Err(err) => { - assert_eq!( - err.to_string(), - "Custom(\"timeout while detecting fastest address\")" - ); - } - } - Ok(()) - } - - #[tokio::test] - async fn detect_fastest_addr() -> YdbResult<()> { - let l1 = TcpListener::bind("127.0.0.1:0").await?; - let l2 = TcpListener::bind("127.0.0.1:0").await?; - let l3 = TcpListener::bind("127.0.0.1:0").await?; - - let l1_addr = l1.local_addr()?; - let l2_addr = l2.local_addr()?; - let l3_addr = l3.local_addr()?; - - println!("Listener №1 on: {}", l1_addr); - println!("Listener №2 on: {}", l2_addr); - println!("Listener №3 on: {}", l3_addr); - - let nodes = vec![ - l1_addr.to_string(), - l2_addr.to_string(), - l3_addr.to_string(), - ]; - - drop(l1); - drop(l2); - drop(l3); - - let result = NearestDCBalancer::find_fastest_address( - nodes.iter().collect_vec(), - Duration::from_secs(3), - ) - .await; - match result { - Ok(_) => unreachable!(), - Err(err) => { - assert_eq!( - err.to_string(), - "Custom(\"timeout while detecting fastest address\")" - ); - } - } - Ok(()) - } - - #[tokio::test] - async fn adjusting_dc() -> YdbResult<()> { - let l1 = TcpListener::bind("127.0.0.1:0").await?; - let l2 = TcpListener::bind("127.0.0.1:0").await?; - let l3 = TcpListener::bind("127.0.0.1:0").await?; - - let l1_addr = l1.local_addr()?; - let l2_addr = l2.local_addr()?; - let l3_addr = l3.local_addr()?; - - println!("Listener №1 on: {}", l1_addr); - println!("Listener №2 on: {}", l2_addr); - println!("Listener №3 on: {}", l3_addr); - - let discovery_state = Arc::new(DiscoveryState::default()); - let balancer_state = Arc::new(Mutex::new(BalancerState::default())); - let balancer_state_updater = balancer_state.clone(); - let (state_sender, state_reciever) = watch::channel(discovery_state.clone()); - - let ping_token = CancellationToken::new(); - let ping_token_clone = ping_token.clone(); - - let waiter = Arc::new(WaiterImpl::new()); - let waiter_clone = waiter.clone(); - - let updater = tokio::spawn(async move { - NearestDCBalancer::adjust_local_dc( - balancer_state_updater, - state_reciever, - ping_token_clone, - waiter_clone, - ) - .await - }); - - let updated_state = Arc::new( - DiscoveryState::default() - .with_node_info( - Table, - NodeInfo::new( - Uri::from_str(l1_addr.to_string().as_str()).unwrap(), - "A".to_string(), - ), - ) - .with_node_info( - Table, - NodeInfo::new( - Uri::from_str(l2_addr.to_string().as_str()).unwrap(), - "B".to_string(), - ), - ) - .with_node_info( - Table, - NodeInfo::new( - Uri::from_str(l2_addr.to_string().as_str()).unwrap(), - "C".to_string(), - ), - ), - ); - assert!( - (balancer_state.lock().await) - .borrow() - .preferred_endpoints - .len() - == 0 // no endpoints - ); - let _ = state_sender.send(updated_state); - tokio::time::sleep(Duration::from_secs(2)).await; - assert_true!(timeout(Duration::from_secs(3), waiter.wait()).await.is_ok()); // should not wait - assert!( - (balancer_state.lock().await) - .borrow() - .preferred_endpoints - .len() - == 1 // only one endpoint in each dc - ); - let updated_state_next = Arc::new( - DiscoveryState::default() - .with_node_info( - Table, - NodeInfo::new( - Uri::from_str(l1_addr.to_string().as_str()).unwrap(), - "A".to_string(), - ), - ) - .with_node_info( - Table, - NodeInfo::new( - Uri::from_str(l2_addr.to_string().as_str()).unwrap(), - "A".to_string(), - ), - ), - ); - let _ = state_sender.send(updated_state_next); - tokio::time::sleep(Duration::from_secs(2)).await; - assert_true!(timeout(Duration::from_secs(3), waiter.wait()).await.is_ok()); // should not wait - assert!( - (balancer_state.lock().await) - .borrow() - .preferred_endpoints - .len() - == 2 // both endpoints in same dc - ); - ping_token.cancel(); // reciever stops wait for state change - let _ = tokio::join!(updater); // should join - Ok(()) - } - - #[tokio::test] - async fn nearest_dc_balancer_integration_with_error_fallback() -> YdbResult<()> { - let balancer = NearestDCBalancer::new(BalancerConfig { - fallback_strategy: FallbackStrategy::Error, - }) - .unwrap(); - - let sh = SharedLoadBalancer::new_with_balancer(Box::new(balancer)); - - match sh.endpoint(Table) { - Ok(_) => unreachable!(), - Err(err) => assert_eq!( - err.to_string(), - "Custom(\"no available endpoints for service:table_service\")".to_string() - ), - } - Ok(()) - } - - #[tokio::test] - async fn nearest_dc_balancer_integration_with_other_fallback_error() -> YdbResult<()> { - let balancer = NearestDCBalancer::new(BalancerConfig::default()).unwrap(); - - let sh = SharedLoadBalancer::new_with_balancer(Box::new(balancer)); - - match sh.endpoint(Table) { - Ok(_) => unreachable!(), - Err(err) => assert_eq!( - err.to_string(), - "Custom(\"empty endpoint list for service: table_service\")".to_string() - ), - } - Ok(()) - } - - #[tokio::test] - async fn nearest_dc_balancer_integration() -> YdbResult<()> { - let l1 = TcpListener::bind("127.0.0.1:0").await?; - let l2 = TcpListener::bind("127.0.0.1:0").await?; - - let l1_addr = l1.local_addr()?; - let l2_addr = l2.local_addr()?; - - println!("Listener №1 on: {}", l1_addr); - println!("Listener №2 on: {}", l2_addr); - - let balancer = NearestDCBalancer::new(BalancerConfig { - fallback_strategy: FallbackStrategy::Error, - }) - .unwrap(); - - let sh = SharedLoadBalancer::new_with_balancer(Box::new(balancer)); - let self_updater = sh.clone(); - let (state_sender, state_reciever) = - watch::channel::>(Arc::new(DiscoveryState::default())); - - tokio::spawn(async move { update_load_balancer(self_updater, state_reciever).await }); - - match sh.endpoint(Table) { - Ok(_) => unreachable!(), - Err(err) => assert_eq!( - err.to_string(), - "Custom(\"no available endpoints for service:table_service\")".to_string() - ), - } - - let updated_state = Arc::new( - DiscoveryState::default() - .with_node_info( - Table, - NodeInfo::new( - Uri::from_str(l1_addr.to_string().as_str()).unwrap(), - "A".to_string(), - ), - ) - .with_node_info( - Table, - NodeInfo::new( - Uri::from_str(l2_addr.to_string().as_str()).unwrap(), - "A".to_string(), - ), - ), - ); - - let _ = state_sender.send(updated_state); - - sh.wait().await?; - - match sh.endpoint(Table) { - Ok(uri) => { - let addr = uri.host().unwrap(); - assert!(addr == "127.0.0.1" || addr == "[::1]") - } - Err(err) => unreachable!("{}", err.to_string()), - } - Ok(()) - } -} diff --git a/ydb/src/load_balancer/balancer_test.rs b/ydb/src/load_balancer/balancer_test.rs new file mode 100644 index 00000000..1ef3cd50 --- /dev/null +++ b/ydb/src/load_balancer/balancer_test.rs @@ -0,0 +1,602 @@ +use super::*; +use super::{ + update_load_balancer, BalancerConfig, LoadBalancer, MockLoadBalancer, NearestDCBalancer, + RandomLoadBalancer, SharedLoadBalancer, +}; +use crate::discovery::NodeInfo; +use crate::grpc_wrapper::raw_services::Service::Table; +use crate::waiter::WaiterImpl; +use crate::YdbResult; +use http::Uri; +use itertools::Itertools; +use mockall::predicate; +use nearest_dc_balancer::{BalancerState, NODES_PER_DC, PING_TIMEOUT_SECS}; +use ntest::assert_true; +use std::borrow::Borrow; +use std::collections::HashMap; +use std::str::FromStr; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering::Relaxed; +use std::sync::Arc; +use std::time::Duration; +use tokio::net::TcpListener; +use tokio::sync::{watch, Mutex}; +use tokio::time::timeout; +use tokio_util::sync::CancellationToken; +use tracing::trace; + +#[test] +fn shared_load_balancer() -> YdbResult<()> { + let endpoint_counter = Arc::new(AtomicUsize::new(0)); + let test_uri = Uri::from_str("http://test.com")?; + + let mut lb_mock = MockLoadBalancer::new(); + let endpoint_counter_mock = endpoint_counter.clone(); + let test_uri_mock = test_uri.clone(); + + lb_mock.expect_endpoint().returning(move |_service| { + endpoint_counter_mock.fetch_add(1, Relaxed); + Ok(test_uri_mock.clone()) + }); + + let s1 = SharedLoadBalancer::new_with_balancer(Box::new(lb_mock)); + + #[allow(clippy::redundant_clone)] + let s2 = s1.clone(); + + assert_eq!(test_uri, s1.endpoint(Table)?); + assert_eq!(test_uri, s2.endpoint(Table)?); + assert_eq!(endpoint_counter.load(Relaxed), 2); + Ok(()) +} + +#[tokio::test] +async fn update_load_balancer_test() -> YdbResult<()> { + let original_discovery_state = Arc::new(DiscoveryState::default()); + let (sender, receiver) = tokio::sync::watch::channel(original_discovery_state.clone()); + + let new_discovery_state = Arc::new(DiscoveryState::default().with_node_info( + Table, + NodeInfo::new(Uri::from_str("http://test.com").unwrap(), String::new()), + )); + + let (first_update_sender, first_update_receiver) = tokio::sync::oneshot::channel(); + let (second_update_sender, second_update_receiver) = tokio::sync::oneshot::channel(); + let (updater_finished_sender, updater_finished_receiver) = + tokio::sync::oneshot::channel::<()>(); + + let mut first_update_sender = Some(first_update_sender); + let mut second_update_sender = Some(second_update_sender); + let mut lb_mock = MockLoadBalancer::new(); + lb_mock + .expect_set_discovery_state() + .with(predicate::eq(original_discovery_state.clone())) + .times(1) + .returning(move |_| { + trace!("first set"); + first_update_sender.take().unwrap().send(()).unwrap(); + Ok(()) + }); + + lb_mock + .expect_set_discovery_state() + .with(predicate::eq(new_discovery_state.clone())) + .times(1) + .returning(move |_| { + trace!("second set"); + second_update_sender.take().unwrap().send(()).unwrap(); + Ok(()) + }); + + let shared_lb = SharedLoadBalancer::new_with_balancer(Box::new(lb_mock)); + + tokio::spawn(async move { + trace!("updater start"); + update_load_balancer(shared_lb, receiver).await; + trace!("updater finished"); + updater_finished_sender.send(()).unwrap(); + }); + + tokio::spawn(async move { + first_update_receiver.await.unwrap(); + sender.send(new_discovery_state).unwrap(); + second_update_receiver.await.unwrap(); + drop(sender); + }); + + tokio::select! { + _ = updater_finished_receiver =>{} + _ = tokio::time::sleep(Duration::from_secs(10)) => { + panic!("test failed"); + } + } + // updater_finished_receiver.await.unwrap(); + Ok(()) +} + +#[test] +fn random_load_balancer() -> YdbResult<()> { + let one = Uri::from_str("http://one:213")?; + let two = Uri::from_str("http://two:213")?; + let load_balancer = RandomLoadBalancer { + discovery_state: Arc::new( + DiscoveryState::default() + .with_node_info(Table, NodeInfo::new(one.clone(), String::new())) + .with_node_info(Table, NodeInfo::new(two.clone(), String::new())), + ), + waiter: Arc::new(WaiterImpl::new()), + }; + + let mut map = HashMap::new(); + map.insert(one.to_string(), 0); + map.insert(two.to_string(), 0); + + for _ in 0..100 { + let u = load_balancer.endpoint(Table)?; + let val = *map.get_mut(u.to_string().as_str()).unwrap(); + map.insert(u.to_string(), val + 1); + } + + assert_eq!(map.len(), 2); + assert!(*map.get(one.to_string().as_str()).unwrap() > 30); + assert!(*map.get(two.to_string().as_str()).unwrap() > 30); + Ok(()) +} + +#[test] +fn split_by_location() -> YdbResult<()> { + let nodes = vec![ + NodeInfo::new(Uri::from_str("http://one:213")?, "A".to_string()), + NodeInfo::new(Uri::from_str("http://two:213")?, "A".to_string()), + NodeInfo::new(Uri::from_str("http://three:213")?, "B".to_string()), + NodeInfo::new(Uri::from_str("http://four:213")?, "B".to_string()), + NodeInfo::new(Uri::from_str("http://five:213")?, "C".to_string()), + ]; + let splitted = NearestDCBalancer::split_endpoints_by_location(&nodes); + assert_eq!(splitted.keys().len(), 3); + assert_eq!(splitted["A"].len(), 2); + assert_eq!(splitted["B"].len(), 2); + assert_eq!(splitted["C"].len(), 1); + Ok(()) +} + +#[test] +fn choose_random_endpoints() -> YdbResult<()> { + let nodes = vec![ + NodeInfo::new(Uri::from_str("http://one:213")?, "C".to_string()), + NodeInfo::new(Uri::from_str("http://two:213")?, "C".to_string()), + NodeInfo::new(Uri::from_str("http://three:213")?, "C".to_string()), + NodeInfo::new(Uri::from_str("http://four:213")?, "C".to_string()), + NodeInfo::new(Uri::from_str("http://five:213")?, "C".to_string()), + NodeInfo::new(Uri::from_str("http://seven:213")?, "C".to_string()), + NodeInfo::new(Uri::from_str("http://eight:213")?, "C".to_string()), + NodeInfo::new(Uri::from_str("http://nine:213")?, "C".to_string()), + ]; + + let mut refs = nodes.iter().collect_vec(); + let nodes_clone = refs.clone(); + let random_subset = NearestDCBalancer::get_random_endpoints(&mut refs); + + assert_eq!(random_subset.len(), NODES_PER_DC); + for node in random_subset { + assert_true!(nodes_clone.contains(node)) + } + + Ok(()) +} + +#[test] +fn extract_addrs_and_map_them() -> YdbResult<()> { + let one = NodeInfo::new(Uri::from_str("http://localhost:123")?, "C".to_string()); + let two = NodeInfo::new(Uri::from_str("http://localhost:321")?, "C".to_string()); + let nodes = vec![&one, &two]; + let map = NearestDCBalancer::addr_to_node(&nodes); + + assert_eq!(map.keys().len(), 4); // ipv4 + ipv6 on each + assert_true!(map.keys().contains(&"127.0.0.1:123".to_string())); + assert_true!(map.keys().contains(&"[::1]:123".to_string())); + assert!(map["127.0.0.1:123"].eq(&one)); + assert!(map["127.0.0.1:123"].eq(map["[::1]:123"])); + + Ok(()) +} + +#[tokio::test] +async fn detect_fastest_addr_just_some() -> YdbResult<()> { + let l1 = TcpListener::bind("127.0.0.1:0").await?; + let l2 = TcpListener::bind("127.0.0.1:0").await?; + let l3 = TcpListener::bind("127.0.0.1:0").await?; + + let l1_addr = l1.local_addr()?; + let l2_addr = l2.local_addr()?; + let l3_addr = l3.local_addr()?; + + println!("Listener №1 on: {}", l1_addr); + println!("Listener №2 on: {}", l2_addr); + println!("Listener №3 on: {}", l3_addr); + + let nodes = vec![ + l1_addr.to_string(), + l2_addr.to_string(), + l3_addr.to_string(), + ]; + + for _ in 0..100 { + let addr = NearestDCBalancer::find_fastest_address( + nodes.iter().collect_vec(), + Duration::from_secs(PING_TIMEOUT_SECS), + ) + .await?; + assert!(nodes.contains(&addr)) + } + + Ok(()) +} + +#[tokio::test] +async fn detect_fastest_addr_with_fault() -> YdbResult<()> { + let l1 = TcpListener::bind("127.0.0.1:0").await?; + let l2 = TcpListener::bind("127.0.0.1:0").await?; + let l3 = TcpListener::bind("127.0.0.1:0").await?; + + let l1_addr = l1.local_addr()?; + let l2_addr = l2.local_addr()?; + let l3_addr = l3.local_addr()?; + + println!("Listener №1 on: {}", l1_addr); + println!("Listener №2 on: {}", l2_addr); + println!("Listener №3 on: {}", l3_addr); + + let nodes = vec![ + l1_addr.to_string(), + l2_addr.to_string(), + l3_addr.to_string(), + ]; + + drop(l1); + + for _ in 0..100 { + let addr = NearestDCBalancer::find_fastest_address( + nodes.iter().collect_vec(), + Duration::from_secs(PING_TIMEOUT_SECS), + ) + .await?; + assert!(nodes.contains(&addr) && addr != l1_addr.to_string()) + } + + Ok(()) +} + +#[tokio::test] +async fn detect_fastest_addr_one_alive() -> YdbResult<()> { + let l1 = TcpListener::bind("127.0.0.1:0").await?; + let l2 = TcpListener::bind("127.0.0.1:0").await?; + let l3 = TcpListener::bind("127.0.0.1:0").await?; + + let l1_addr = l1.local_addr()?; + let l2_addr = l2.local_addr()?; + let l3_addr = l3.local_addr()?; + + println!("Listener №1 on: {}", l1_addr); + println!("Listener №2 on: {}", l2_addr); + println!("Listener №3 on: {}", l3_addr); + + let nodes = vec![ + l1_addr.to_string(), + l2_addr.to_string(), + l3_addr.to_string(), + ]; + + drop(l1); + drop(l2); + + for _ in 0..100 { + let addr = NearestDCBalancer::find_fastest_address( + nodes.iter().collect_vec(), + Duration::from_secs(PING_TIMEOUT_SECS), + ) + .await?; + assert!(addr == l3_addr.to_string()) + } + + Ok(()) +} + +#[tokio::test] +async fn detect_fastest_addr_timeout() -> YdbResult<()> { + let l1 = TcpListener::bind("127.0.0.1:0").await?; + let l2 = TcpListener::bind("127.0.0.1:0").await?; + let l3 = TcpListener::bind("127.0.0.1:0").await?; + + let l1_addr = l1.local_addr()?; + let l2_addr = l2.local_addr()?; + let l3_addr = l3.local_addr()?; + + println!("Listener №1 on: {}", l1_addr); + println!("Listener №2 on: {}", l2_addr); + println!("Listener №3 on: {}", l3_addr); + + let nodes = vec![ + l1_addr.to_string(), + l2_addr.to_string(), + l3_addr.to_string(), + ]; + + drop(l1); + drop(l2); + drop(l3); + + let result = + NearestDCBalancer::find_fastest_address(nodes.iter().collect_vec(), Duration::from_secs(3)) + .await; + match result { + Ok(_) => unreachable!(), + Err(err) => { + assert_eq!( + err.to_string(), + "Custom(\"timeout while detecting fastest address\")" + ); + } + } + Ok(()) +} + +#[tokio::test] +async fn no_addr_timeout() -> YdbResult<()> { + let result = NearestDCBalancer::find_fastest_address(Vec::new(), Duration::from_secs(3)).await; + match result { + Ok(_) => unreachable!(), + Err(err) => { + assert_eq!( + err.to_string(), + "Custom(\"timeout while detecting fastest address\")" + ); + } + } + Ok(()) +} + +#[tokio::test] +async fn detect_fastest_addr() -> YdbResult<()> { + let l1 = TcpListener::bind("127.0.0.1:0").await?; + let l2 = TcpListener::bind("127.0.0.1:0").await?; + let l3 = TcpListener::bind("127.0.0.1:0").await?; + + let l1_addr = l1.local_addr()?; + let l2_addr = l2.local_addr()?; + let l3_addr = l3.local_addr()?; + + println!("Listener №1 on: {}", l1_addr); + println!("Listener №2 on: {}", l2_addr); + println!("Listener №3 on: {}", l3_addr); + + let nodes = vec![ + l1_addr.to_string(), + l2_addr.to_string(), + l3_addr.to_string(), + ]; + + drop(l1); + drop(l2); + drop(l3); + + let result = + NearestDCBalancer::find_fastest_address(nodes.iter().collect_vec(), Duration::from_secs(3)) + .await; + match result { + Ok(_) => unreachable!(), + Err(err) => { + assert_eq!( + err.to_string(), + "Custom(\"timeout while detecting fastest address\")" + ); + } + } + Ok(()) +} + +#[tokio::test] +async fn adjusting_dc() -> YdbResult<()> { + let l1 = TcpListener::bind("127.0.0.1:0").await?; + let l2 = TcpListener::bind("127.0.0.1:0").await?; + let l3 = TcpListener::bind("127.0.0.1:0").await?; + + let l1_addr = l1.local_addr()?; + let l2_addr = l2.local_addr()?; + let l3_addr = l3.local_addr()?; + + println!("Listener №1 on: {}", l1_addr); + println!("Listener №2 on: {}", l2_addr); + println!("Listener №3 on: {}", l3_addr); + + let discovery_state = Arc::new(DiscoveryState::default()); + let balancer_state = Arc::new(Mutex::new(BalancerState::default())); + let balancer_state_updater = balancer_state.clone(); + let (state_sender, state_reciever) = watch::channel(discovery_state.clone()); + + let ping_token = CancellationToken::new(); + let ping_token_clone = ping_token.clone(); + + let waiter = Arc::new(WaiterImpl::new()); + let waiter_clone = waiter.clone(); + + let updater = tokio::spawn(async move { + NearestDCBalancer::adjust_local_dc( + balancer_state_updater, + state_reciever, + ping_token_clone, + waiter_clone, + ) + .await + }); + + let updated_state = Arc::new( + DiscoveryState::default() + .with_node_info( + Table, + NodeInfo::new( + Uri::from_str(l1_addr.to_string().as_str()).unwrap(), + "A".to_string(), + ), + ) + .with_node_info( + Table, + NodeInfo::new( + Uri::from_str(l2_addr.to_string().as_str()).unwrap(), + "B".to_string(), + ), + ) + .with_node_info( + Table, + NodeInfo::new( + Uri::from_str(l2_addr.to_string().as_str()).unwrap(), + "C".to_string(), + ), + ), + ); + assert!( + (balancer_state.lock().await) + .borrow() + .preferred_endpoints + .len() + == 0 // no endpoints + ); + let _ = state_sender.send(updated_state); + tokio::time::sleep(Duration::from_secs(2)).await; + assert_true!(timeout(Duration::from_secs(3), waiter.wait()).await.is_ok()); // should not wait + assert!( + (balancer_state.lock().await) + .borrow() + .preferred_endpoints + .len() + == 1 // only one endpoint in each dc + ); + let updated_state_next = Arc::new( + DiscoveryState::default() + .with_node_info( + Table, + NodeInfo::new( + Uri::from_str(l1_addr.to_string().as_str()).unwrap(), + "A".to_string(), + ), + ) + .with_node_info( + Table, + NodeInfo::new( + Uri::from_str(l2_addr.to_string().as_str()).unwrap(), + "A".to_string(), + ), + ), + ); + let _ = state_sender.send(updated_state_next); + tokio::time::sleep(Duration::from_secs(2)).await; + assert_true!(timeout(Duration::from_secs(3), waiter.wait()).await.is_ok()); // should not wait + assert!( + (balancer_state.lock().await) + .borrow() + .preferred_endpoints + .len() + == 2 // both endpoints in same dc + ); + ping_token.cancel(); // reciever stops wait for state change + let _ = tokio::join!(updater); // should join + Ok(()) +} + +#[tokio::test] +async fn nearest_dc_balancer_integration_with_error_fallback() -> YdbResult<()> { + let balancer = NearestDCBalancer::new(BalancerConfig { + fallback_strategy: FallbackStrategy::Error, + }) + .unwrap(); + + let sh = SharedLoadBalancer::new_with_balancer(Box::new(balancer)); + + match sh.endpoint(Table) { + Ok(_) => unreachable!(), + Err(err) => assert_eq!( + err.to_string(), + "Custom(\"no available endpoints for service:table_service\")".to_string() + ), + } + Ok(()) +} + +#[tokio::test] +async fn nearest_dc_balancer_integration_with_other_fallback_error() -> YdbResult<()> { + let balancer = NearestDCBalancer::new(BalancerConfig::default()).unwrap(); + + let sh = SharedLoadBalancer::new_with_balancer(Box::new(balancer)); + + match sh.endpoint(Table) { + Ok(_) => unreachable!(), + Err(err) => assert_eq!( + err.to_string(), + "Custom(\"empty endpoint list for service: table_service\")".to_string() + ), + } + Ok(()) +} + +#[tokio::test] +async fn nearest_dc_balancer_integration() -> YdbResult<()> { + let l1 = TcpListener::bind("127.0.0.1:0").await?; + let l2 = TcpListener::bind("127.0.0.1:0").await?; + + let l1_addr = l1.local_addr()?; + let l2_addr = l2.local_addr()?; + + println!("Listener №1 on: {}", l1_addr); + println!("Listener №2 on: {}", l2_addr); + + let balancer = NearestDCBalancer::new(BalancerConfig { + fallback_strategy: FallbackStrategy::Error, + }) + .unwrap(); + + let sh = SharedLoadBalancer::new_with_balancer(Box::new(balancer)); + let self_updater = sh.clone(); + let (state_sender, state_reciever) = + watch::channel::>(Arc::new(DiscoveryState::default())); + + tokio::spawn(async move { update_load_balancer(self_updater, state_reciever).await }); + + match sh.endpoint(Table) { + Ok(_) => unreachable!(), + Err(err) => assert_eq!( + err.to_string(), + "Custom(\"no available endpoints for service:table_service\")".to_string() + ), + } + + let updated_state = Arc::new( + DiscoveryState::default() + .with_node_info( + Table, + NodeInfo::new( + Uri::from_str(l1_addr.to_string().as_str()).unwrap(), + "A".to_string(), + ), + ) + .with_node_info( + Table, + NodeInfo::new( + Uri::from_str(l2_addr.to_string().as_str()).unwrap(), + "A".to_string(), + ), + ), + ); + + let _ = state_sender.send(updated_state); + + sh.wait().await?; + + match sh.endpoint(Table) { + Ok(uri) => { + let addr = uri.host().unwrap(); + assert!(addr == "127.0.0.1" || addr == "[::1]") + } + Err(err) => unreachable!("{}", err.to_string()), + } + Ok(()) +} diff --git a/ydb/src/load_balancer/mod.rs b/ydb/src/load_balancer/mod.rs new file mode 100644 index 00000000..64d6049a --- /dev/null +++ b/ydb/src/load_balancer/mod.rs @@ -0,0 +1,46 @@ +use std::sync::Arc; + +use http::Uri; +use tokio::sync::watch::Receiver; + +use crate::{grpc_wrapper::raw_services::Service, DiscoveryState, Waiter, YdbResult}; + +#[cfg(test)] +pub mod balancer_test; +pub mod nearest_dc_balancer; +pub mod random_balancer; +pub mod shared_balancer; +pub mod static_balancer; + +pub(crate) use nearest_dc_balancer::{BalancerConfig, FallbackStrategy, NearestDCBalancer}; +pub(crate) use random_balancer::RandomLoadBalancer; +pub(crate) use shared_balancer::SharedLoadBalancer; +pub(crate) use static_balancer::StaticLoadBalancer; + +#[mockall::automock] +pub(crate) trait LoadBalancer: Send + Sync + Waiter { + fn endpoint(&self, service: Service) -> YdbResult; + fn set_discovery_state(&mut self, discovery_state: &Arc) -> YdbResult<()>; + fn waiter(&self) -> Box; // need for wait ready in without read lock +} + +#[async_trait::async_trait] +impl Waiter for MockLoadBalancer { + async fn wait(&self) -> YdbResult<()> { + Ok(()) + } +} + +pub(crate) async fn update_load_balancer( + mut lb: impl LoadBalancer, + mut receiver: Receiver>, +) { + loop { + // clone for prevent block send side while update current lb + let state = receiver.borrow_and_update().clone(); + let _ = lb.set_discovery_state(&state); + if receiver.changed().await.is_err() { + break; + } + } +} diff --git a/ydb/src/load_balancer/nearest_dc_balancer.rs b/ydb/src/load_balancer/nearest_dc_balancer.rs new file mode 100644 index 00000000..97350cb5 --- /dev/null +++ b/ydb/src/load_balancer/nearest_dc_balancer.rs @@ -0,0 +1,393 @@ +use std::{ + borrow::{Borrow, BorrowMut}, + collections::HashMap, + sync::Arc, +}; + +use http::Uri; +use itertools::Itertools; +use rand::{seq::SliceRandom, thread_rng}; +use std::time::Duration; +use tokio::{ + io::AsyncWriteExt, + net::TcpStream, + sync::{ + broadcast, mpsc, + watch::{self, Sender}, + Mutex, + }, + task::JoinSet, +}; +use tokio_util::sync::CancellationToken; +use tracing::{error, info, warn}; + +use crate::{ + discovery::NodeInfo, + grpc_wrapper::raw_services::Service, + waiter::{AllWaiter, WaiterImpl}, + DiscoveryState, Waiter, YdbError, YdbResult, +}; + +use super::{random_balancer::RandomLoadBalancer, LoadBalancer}; + +pub(crate) struct BalancerConfig { + pub(super) fallback_strategy: FallbackStrategy, +} + +#[derive(Default)] +pub(super) struct BalancerState { + pub(super) preferred_endpoints: Vec, +} + +// What will balancer do if there is no available endpoints at local dc +pub enum FallbackStrategy { + Error, // Just throw error + BalanceWithOther(Box), // Use another balancer +} + +impl Default for BalancerConfig { + fn default() -> Self { + BalancerConfig { + fallback_strategy: FallbackStrategy::BalanceWithOther(Box::new( + RandomLoadBalancer::new(), + )), + } + } +} + +pub(crate) struct NearestDCBalancer { + discovery_state: Arc, + state_sender: Sender>, + ping_token: CancellationToken, + waiter: Arc, + config: BalancerConfig, + balancer_state: Arc>, +} + +impl NearestDCBalancer { + pub(crate) fn new(config: BalancerConfig) -> YdbResult { + let discovery_state = Arc::new(DiscoveryState::default()); + let balancer_state = Arc::new(Mutex::new(BalancerState::default())); + let balancer_state_updater = balancer_state.clone(); + let (state_sender, state_reciever) = watch::channel(discovery_state.clone()); + + let ping_token = CancellationToken::new(); + let ping_token_clone = ping_token.clone(); + + let waiter = Arc::new(WaiterImpl::new()); + let waiter_clone = waiter.clone(); + + tokio::spawn(async move { + Self::adjust_local_dc( + balancer_state_updater, + state_reciever, + ping_token_clone, + waiter_clone, + ) + .await + }); + + Ok(Self { + discovery_state, + state_sender, + ping_token, + waiter, + config, + balancer_state, + }) + } +} + +impl Drop for NearestDCBalancer { + fn drop(&mut self) { + self.ping_token.cancel(); + } +} + +#[async_trait::async_trait] +impl Waiter for NearestDCBalancer { + async fn wait(&self) -> YdbResult<()> { + self.waiter().wait().await + } +} + +impl LoadBalancer for NearestDCBalancer { + fn endpoint(&self, service: Service) -> YdbResult { + self.get_endpoint(service) + } + + fn set_discovery_state(&mut self, discovery_state: &Arc) -> YdbResult<()> { + match self.config.fallback_strategy.borrow_mut() { + FallbackStrategy::BalanceWithOther(balancer) => { + balancer.set_discovery_state(discovery_state)? + } + FallbackStrategy::Error => (), + } + self.discovery_state = discovery_state.clone(); + let _ = self.state_sender.send(discovery_state.clone()); + Ok(()) + } + + fn waiter(&self) -> Box { + let self_waiter = Box::new(self.waiter.clone()); + match self.config.fallback_strategy.borrow() { + FallbackStrategy::BalanceWithOther(balancer) => { + Box::new(AllWaiter::new(vec![self_waiter, balancer.waiter()])) + } + FallbackStrategy::Error => self_waiter, + } + } +} + +pub(super) const NODES_PER_DC: usize = 5; +pub(super) const PING_TIMEOUT_SECS: u64 = 60; + +impl NearestDCBalancer { + fn get_endpoint(&self, service: Service) -> YdbResult { + match self.balancer_state.try_lock() { + Ok(state_guard) => { + match state_guard + .borrow() + .preferred_endpoints + .choose(&mut thread_rng()) + { + Some(ep) => return YdbResult::Ok(ep.uri.clone()), + None => (), + } + match self.config.fallback_strategy.borrow() { + FallbackStrategy::Error => Err(YdbError::custom(format!( + "no available endpoints for service:{}", + service + ))), + FallbackStrategy::BalanceWithOther(balancer) => { + info!("trying fallback balancer..."); + balancer.endpoint(service) + } + } + } + Err(_) => Err(YdbError::Custom( + "balancer is updating its state".to_string(), + )), + } + } + + pub(super) async fn adjust_local_dc( + balancer_state: Arc>, + mut state_reciever: watch::Receiver>, + stop_ping_process: CancellationToken, + waiter: Arc, + ) { + loop { + tokio::select! { + _ = stop_ping_process.cancelled() => { + return + } + result = state_reciever.changed() =>{ + if result.is_err(){ // sender have been dropped + return + } + } + } + let new_discovery_state = state_reciever.borrow_and_update().clone(); + match Self::extract_nodes(&new_discovery_state) { + Ok(some_nodes) => { + let mut dc_to_nodes = Self::split_endpoints_by_location(some_nodes); + let mut to_check = Vec::with_capacity(NODES_PER_DC * dc_to_nodes.keys().len()); + dc_to_nodes.values_mut().for_each(|endpoints| { + to_check.append(Self::get_random_endpoints(endpoints)) + }); + match Self::find_local_dc(&to_check).await { + Ok(dc) => { + info!("found new local dc:{}", dc); + Self::adjust_preferred_endpoints(&balancer_state, some_nodes, dc).await; + waiter.set_received(Ok(())); + } + Err(err) => { + error!("error on search local dc:{}", err); + continue; + } + } + } + Err(_) => continue, + } + } + } + + async fn adjust_preferred_endpoints( + balancer_state: &Arc>, + new_nodes: &Vec, + local_dc: String, + ) { + info!("adjusting endpoints"); + let new_preferred_endpoints = new_nodes + .into_iter() + .filter(|ep| ep.location == local_dc) + .map(|ep| ep.clone()) + .collect_vec(); + (balancer_state.lock().await) // fast lock + .borrow_mut() + .preferred_endpoints = new_preferred_endpoints; + } + + pub(super) fn extract_nodes(from_state: &Arc) -> YdbResult<&Vec> { + let nodes = from_state.get_all_nodes(); + match nodes { + None => Err(YdbError::Custom(format!( + "no endpoints on discovery update" + ))), + Some(nodes) => Ok(nodes), + } + } + + pub(super) fn split_endpoints_by_location<'a>( + nodes: &'a Vec, + ) -> HashMap> { + let mut dc_to_eps = HashMap::>::new(); + nodes.into_iter().for_each(|info| { + if let Some(nodes) = dc_to_eps.get_mut(&info.location) { + nodes.push(info); + } else { + dc_to_eps.insert(info.location.clone(), vec![info]); + } + }); + dc_to_eps + } + + pub(super) fn get_random_endpoints<'a>( + dc_endpoints: &'a mut Vec<&'a NodeInfo>, + ) -> &mut Vec<&NodeInfo> { + dc_endpoints.shuffle(&mut thread_rng()); + dc_endpoints.truncate(NODES_PER_DC); + dc_endpoints + } + + pub(super) async fn find_local_dc(to_check: &[&NodeInfo]) -> YdbResult { + let addr_to_node = Self::addr_to_node(to_check); + if addr_to_node.is_empty() { + return Err(YdbError::Custom(format!("no ip addresses for endpoints"))); + } + let addrs = addr_to_node.keys(); + match Self::find_fastest_address(addrs.collect(), Duration::from_secs(PING_TIMEOUT_SECS)) + .await + { + Ok(fastest_address) => Ok(addr_to_node[&fastest_address].location.clone()), + Err(err) => { + error!("could not find fastest address:{}", err); + Err(err) + } + } + } + + pub(super) fn addr_to_node<'a>(nodes: &[&'a NodeInfo]) -> HashMap { + let mut addr_to_node = HashMap::::with_capacity(2 * nodes.len()); // ipv4 + ipv6 + nodes.into_iter().for_each(|info| { + let host: &str; + let port: u16; + match info.uri.host() { + Some(uri_host) => host = uri_host, + None => { + warn!("no host for uri:{}", info.uri); + return; + } + } + match info.uri.port() { + Some(uri_port) => port = uri_port.as_u16(), + None => { + warn!("no port for uri:{}", info.uri); + return; + } + } + use std::net::ToSocketAddrs; + let _ = (host, port).to_socket_addrs().and_then(|addrs| { + for addr in addrs { + addr_to_node.insert(addr.to_string(), info); + } + Ok(()) + }); + }); + addr_to_node + } + + pub(super) async fn find_fastest_address( + addrs: Vec<&String>, + timeout: Duration, + ) -> YdbResult { + // Cancellation flow: timeout -> address collector -> address producers + let interrupt_via_timeout = CancellationToken::new(); + let interrupt_collector_future = interrupt_via_timeout.child_token(); + let stop_measure = interrupt_collector_future.child_token(); // (*) + + let (start_measure, _) = broadcast::channel::<()>(1); + let buffer_cap = if addrs.len() > 0 { addrs.len() } else { 1 }; + let (addr_sender, mut addr_reciever) = mpsc::channel::>(buffer_cap); + let mut nursery = JoinSet::new(); + + for addr in addrs { + let (mut wait_for_start, stop_measure, addr, addr_sender) = ( + start_measure.subscribe(), + stop_measure.clone(), + addr.clone(), + addr_sender.clone(), + ); + + nursery.spawn(async move { + let _ = wait_for_start.recv().await; + tokio::select! { + connection_result = TcpStream::connect(addr.clone()) =>{ + match connection_result{ + Ok(mut connection) => { + let _ = connection.shutdown().await; + let _ = addr_sender.send(Some(addr)).await; + }, + Err(_) => {let _ = addr_sender.send(None).await;}, + } + } + _ = stop_measure.cancelled() => { + (); + } + } + }); + } + + let wait_first_some_or_cancel = async { + loop { + tokio::select! { + biased; // check timeout first + _ = interrupt_collector_future.cancelled() =>{ + Self::join_all(&mut nursery).await; // children will be cancelled due to tokens chaining, see (*) + return YdbResult::Err("cancelled".into()) + } + address_reciever_option = addr_reciever.recv() =>{ + match address_reciever_option { + Some(address_option) => { + match address_option { + Some(address) =>{ + interrupt_collector_future.cancel(); // Cancel other producing children + Self::join_all(&mut nursery).await; + return YdbResult::Ok(address); + }, + None => continue, // Some producer sent blank address -> wait others + } + }, + None => return YdbResult::Err("no fastest address".into()), // Channel closed, all producers have done measures + } + } + } + } + }; + + let _ = start_measure.send(()); + + match tokio::time::timeout(timeout, wait_first_some_or_cancel).await { + Ok(address_option) => address_option, + Err(_) => { + interrupt_via_timeout.cancel(); + YdbResult::Err("timeout while detecting fastest address".into()) + } + } + } + + async fn join_all(awaitable: &mut JoinSet<()>) { + while awaitable.join_next().await.is_some() {} + } +} diff --git a/ydb/src/load_balancer/random_balancer.rs b/ydb/src/load_balancer/random_balancer.rs new file mode 100644 index 00000000..2fa7f92b --- /dev/null +++ b/ydb/src/load_balancer/random_balancer.rs @@ -0,0 +1,68 @@ +use std::sync::Arc; + +use http::Uri; + +use crate::{ + grpc_wrapper::raw_services::Service, waiter::WaiterImpl, DiscoveryState, Waiter, YdbError, + YdbResult, +}; + +use super::LoadBalancer; + +#[derive(Clone)] +pub(crate) struct RandomLoadBalancer { + pub(super) discovery_state: Arc, + pub(super) waiter: Arc, +} + +impl RandomLoadBalancer { + pub(crate) fn new() -> Self { + Self { + discovery_state: Arc::new(DiscoveryState::default()), + waiter: Arc::new(WaiterImpl::new()), + } + } +} + +impl LoadBalancer for RandomLoadBalancer { + fn endpoint(&self, service: Service) -> YdbResult { + let nodes = self.discovery_state.get_nodes(&service); + match nodes { + None => Err(YdbError::Custom(format!( + "no endpoints for service: '{}'", + service + ))), + Some(nodes) => { + if !nodes.is_empty() { + let index = rand::random::() % nodes.len(); + let node = &nodes[index % nodes.len()]; + Ok(node.uri.clone()) + } else { + Err(YdbError::Custom(format!( + "empty endpoint list for service: {}", + service + ))) + } + } + } + } + + fn set_discovery_state(&mut self, discovery_state: &Arc) -> YdbResult<()> { + self.discovery_state = discovery_state.clone(); + if !self.discovery_state.is_empty() { + self.waiter.set_received(Ok(())) + } + Ok(()) + } + + fn waiter(&self) -> Box { + Box::new(self.waiter.clone()) + } +} + +#[async_trait::async_trait] +impl Waiter for RandomLoadBalancer { + async fn wait(&self) -> YdbResult<()> { + self.waiter.wait().await + } +} diff --git a/ydb/src/load_balancer/shared_balancer.rs b/ydb/src/load_balancer/shared_balancer.rs new file mode 100644 index 00000000..c81b91f3 --- /dev/null +++ b/ydb/src/load_balancer/shared_balancer.rs @@ -0,0 +1,60 @@ +use std::sync::{Arc, RwLock}; + +use http::Uri; + +use crate::{grpc_wrapper::raw_services::Service, Discovery, DiscoveryState, Waiter, YdbResult}; + +use super::{random_balancer::RandomLoadBalancer, update_load_balancer, LoadBalancer}; + +#[derive(Clone)] +pub(crate) struct SharedLoadBalancer { + inner: Arc>>, +} + +impl SharedLoadBalancer { + pub(crate) fn new(discovery: &dyn Discovery) -> Self { + Self::new_with_balancer_and_updater(Box::new(RandomLoadBalancer::new()), discovery) + } + + pub(crate) fn new_with_balancer(load_balancer: Box) -> Self { + Self { + inner: Arc::new(RwLock::new(load_balancer)), + } + } + + pub(crate) fn new_with_balancer_and_updater( + load_balancer: Box, + discovery: &dyn Discovery, + ) -> Self { + let mut shared_lb = Self::new_with_balancer(load_balancer); + let shared_lb_updater = shared_lb.clone(); + let discovery_receiver = discovery.subscribe(); + let _ = shared_lb.set_discovery_state(&discovery.state()); + tokio::spawn( + async move { update_load_balancer(shared_lb_updater, discovery_receiver).await }, + ); + shared_lb + } +} + +impl LoadBalancer for SharedLoadBalancer { + fn endpoint(&self, service: Service) -> YdbResult { + self.inner.read()?.endpoint(service) + } + + fn set_discovery_state(&mut self, discovery_state: &Arc) -> YdbResult<()> { + self.inner.write()?.set_discovery_state(discovery_state) + } + + fn waiter(&self) -> Box { + return self.inner.read().unwrap().waiter(); + } +} + +#[async_trait::async_trait] +impl Waiter for SharedLoadBalancer { + async fn wait(&self) -> YdbResult<()> { + let waiter = self.inner.read()?.waiter(); + return waiter.wait().await; + } +} diff --git a/ydb/src/load_balancer/static_balancer.rs b/ydb/src/load_balancer/static_balancer.rs new file mode 100644 index 00000000..c7112de3 --- /dev/null +++ b/ydb/src/load_balancer/static_balancer.rs @@ -0,0 +1,43 @@ +use std::sync::Arc; + +use http::Uri; + +use crate::{grpc_wrapper::raw_services::Service, waiter::WaiterImpl, DiscoveryState, Waiter, YdbError, YdbResult}; + +use super::LoadBalancer; + +pub(crate) struct StaticLoadBalancer { + endpoint: Uri, +} + +impl StaticLoadBalancer { + #[allow(dead_code)] + pub(crate) fn new(endpoint: Uri) -> Self { + Self { endpoint } + } +} + +impl LoadBalancer for StaticLoadBalancer { + fn endpoint(&self, _: Service) -> YdbResult { + Ok(self.endpoint.clone()) + } + + fn set_discovery_state(&mut self, _: &Arc) -> YdbResult<()> { + Err(YdbError::Custom( + "static balancer no way to update state".into(), + )) + } + + fn waiter(&self) -> Box { + let waiter = WaiterImpl::new(); + waiter.set_received(Ok(())); + Box::new(waiter) + } +} + +#[async_trait::async_trait] +impl Waiter for StaticLoadBalancer { + async fn wait(&self) -> YdbResult<()> { + Ok(()) + } +} From fa2e570afa745b1107a353e67446e23705e8fdcf Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 15 Sep 2024 19:45:05 +0300 Subject: [PATCH 36/42] Address producers do nothing on ping error --- ydb/src/load_balancer/nearest_dc_balancer.rs | 27 ++++++++------------ 1 file changed, 11 insertions(+), 16 deletions(-) diff --git a/ydb/src/load_balancer/nearest_dc_balancer.rs b/ydb/src/load_balancer/nearest_dc_balancer.rs index 97350cb5..2998449e 100644 --- a/ydb/src/load_balancer/nearest_dc_balancer.rs +++ b/ydb/src/load_balancer/nearest_dc_balancer.rs @@ -319,7 +319,7 @@ impl NearestDCBalancer { let (start_measure, _) = broadcast::channel::<()>(1); let buffer_cap = if addrs.len() > 0 { addrs.len() } else { 1 }; - let (addr_sender, mut addr_reciever) = mpsc::channel::>(buffer_cap); + let (addr_sender, mut addr_reciever) = mpsc::channel::(buffer_cap); let mut nursery = JoinSet::new(); for addr in addrs { @@ -337,13 +337,13 @@ impl NearestDCBalancer { match connection_result{ Ok(mut connection) => { let _ = connection.shutdown().await; - let _ = addr_sender.send(Some(addr)).await; + let _ = addr_sender.send(addr).await; }, - Err(_) => {let _ = addr_sender.send(None).await;}, + Err(_) => (), // Just do nothing if there is error on connection attempt } } _ = stop_measure.cancelled() => { - (); + (); // Also do nothing if there is request to stop pings (balancer already got fastest address) } } }); @@ -354,20 +354,15 @@ impl NearestDCBalancer { tokio::select! { biased; // check timeout first _ = interrupt_collector_future.cancelled() =>{ - Self::join_all(&mut nursery).await; // children will be cancelled due to tokens chaining, see (*) + Self::join_all(&mut nursery).await; // Children will be cancelled due to tokens chaining, see (*) return YdbResult::Err("cancelled".into()) } - address_reciever_option = addr_reciever.recv() =>{ - match address_reciever_option { - Some(address_option) => { - match address_option { - Some(address) =>{ - interrupt_collector_future.cancel(); // Cancel other producing children - Self::join_all(&mut nursery).await; - return YdbResult::Ok(address); - }, - None => continue, // Some producer sent blank address -> wait others - } + address_option = addr_reciever.recv() =>{ + match address_option { + Some(address) => { + interrupt_collector_future.cancel(); // Cancel other producing children + Self::join_all(&mut nursery).await; + return YdbResult::Ok(address); }, None => return YdbResult::Err("no fastest address".into()), // Channel closed, all producers have done measures } From a25f949a8ab87af8e55ae9542de8e6b44153928d Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 15 Sep 2024 19:56:01 +0300 Subject: [PATCH 37/42] Check exact map --- ydb/src/load_balancer/balancer_test.rs | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/ydb/src/load_balancer/balancer_test.rs b/ydb/src/load_balancer/balancer_test.rs index 1ef3cd50..e862a3f9 100644 --- a/ydb/src/load_balancer/balancer_test.rs +++ b/ydb/src/load_balancer/balancer_test.rs @@ -145,18 +145,31 @@ fn random_load_balancer() -> YdbResult<()> { #[test] fn split_by_location() -> YdbResult<()> { - let nodes = vec![ + let (one, two, three, four, five) = ( NodeInfo::new(Uri::from_str("http://one:213")?, "A".to_string()), NodeInfo::new(Uri::from_str("http://two:213")?, "A".to_string()), NodeInfo::new(Uri::from_str("http://three:213")?, "B".to_string()), NodeInfo::new(Uri::from_str("http://four:213")?, "B".to_string()), NodeInfo::new(Uri::from_str("http://five:213")?, "C".to_string()), + ); + + let nodes = vec![ + one.clone(), + two.clone(), + three.clone(), + four.clone(), + five.clone(), ]; + let splitted = NearestDCBalancer::split_endpoints_by_location(&nodes); - assert_eq!(splitted.keys().len(), 3); - assert_eq!(splitted["A"].len(), 2); - assert_eq!(splitted["B"].len(), 2); - assert_eq!(splitted["C"].len(), 1); + assert_eq!( + splitted, + HashMap::from([ + ("A".to_string(), vec![&one, &two]), + ("B".to_string(), vec![&three, &four]), + ("C".to_string(), vec![&five]), + ]) + ); Ok(()) } From 0b83dad9bbb998fc071050e93d6dc551b4f92b53 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 15 Sep 2024 20:01:49 +0300 Subject: [PATCH 38/42] Better random shuffle test --- ydb/src/load_balancer/balancer_test.rs | 31 +++++++++++++++++++------- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/ydb/src/load_balancer/balancer_test.rs b/ydb/src/load_balancer/balancer_test.rs index e862a3f9..5e6bc493 100644 --- a/ydb/src/load_balancer/balancer_test.rs +++ b/ydb/src/load_balancer/balancer_test.rs @@ -175,25 +175,40 @@ fn split_by_location() -> YdbResult<()> { #[test] fn choose_random_endpoints() -> YdbResult<()> { - let nodes = vec![ + let (one, two, three, four, five, six, seven, eight, nine) = ( NodeInfo::new(Uri::from_str("http://one:213")?, "C".to_string()), NodeInfo::new(Uri::from_str("http://two:213")?, "C".to_string()), NodeInfo::new(Uri::from_str("http://three:213")?, "C".to_string()), NodeInfo::new(Uri::from_str("http://four:213")?, "C".to_string()), NodeInfo::new(Uri::from_str("http://five:213")?, "C".to_string()), + NodeInfo::new(Uri::from_str("http://six:213")?, "C".to_string()), NodeInfo::new(Uri::from_str("http://seven:213")?, "C".to_string()), NodeInfo::new(Uri::from_str("http://eight:213")?, "C".to_string()), NodeInfo::new(Uri::from_str("http://nine:213")?, "C".to_string()), + ); + + let nodes_big = vec![ + one.clone(), + two.clone(), + three.clone(), + four.clone(), + five.clone(), + six.clone(), + seven.clone(), + eight.clone(), + nine.clone(), ]; - let mut refs = nodes.iter().collect_vec(); - let nodes_clone = refs.clone(); - let random_subset = NearestDCBalancer::get_random_endpoints(&mut refs); + let nodes_small = vec![one.clone(), two.clone(), three.clone()]; - assert_eq!(random_subset.len(), NODES_PER_DC); - for node in random_subset { - assert_true!(nodes_clone.contains(node)) - } + let mut refs_big = nodes_big.iter().collect_vec(); + let mut refs_small = nodes_small.iter().collect_vec(); + + let random_subset_big = NearestDCBalancer::get_random_endpoints(&mut refs_big); + let random_subset_small = NearestDCBalancer::get_random_endpoints(&mut refs_small); + + assert_eq!(random_subset_big.len(), NODES_PER_DC); + assert_eq!(random_subset_small.len(), nodes_small.len()); Ok(()) } From a225958271e600c3e0a3ec0fef09ca142ce2ad05 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 15 Sep 2024 20:04:31 +0300 Subject: [PATCH 39/42] No clones --- ydb/src/load_balancer/balancer_test.rs | 23 ++++++----------------- 1 file changed, 6 insertions(+), 17 deletions(-) diff --git a/ydb/src/load_balancer/balancer_test.rs b/ydb/src/load_balancer/balancer_test.rs index 5e6bc493..8e4ab2e6 100644 --- a/ydb/src/load_balancer/balancer_test.rs +++ b/ydb/src/load_balancer/balancer_test.rs @@ -187,28 +187,17 @@ fn choose_random_endpoints() -> YdbResult<()> { NodeInfo::new(Uri::from_str("http://nine:213")?, "C".to_string()), ); - let nodes_big = vec![ - one.clone(), - two.clone(), - three.clone(), - four.clone(), - five.clone(), - six.clone(), - seven.clone(), - eight.clone(), - nine.clone(), + let mut nodes_big = vec![ + &one, &two, &three, &four, &five, &six, &seven, &eight, &nine, ]; - let nodes_small = vec![one.clone(), two.clone(), three.clone()]; - - let mut refs_big = nodes_big.iter().collect_vec(); - let mut refs_small = nodes_small.iter().collect_vec(); + let mut nodes_small = vec![&one, &two, &three]; - let random_subset_big = NearestDCBalancer::get_random_endpoints(&mut refs_big); - let random_subset_small = NearestDCBalancer::get_random_endpoints(&mut refs_small); + let random_subset_big = NearestDCBalancer::get_random_endpoints(&mut nodes_big); + let random_subset_small = NearestDCBalancer::get_random_endpoints(&mut nodes_small); assert_eq!(random_subset_big.len(), NODES_PER_DC); - assert_eq!(random_subset_small.len(), nodes_small.len()); + assert_eq!(random_subset_small.len(), 3); Ok(()) } From 6499c412423434a35017ef85c41ceced22583e06 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 15 Sep 2024 20:29:13 +0300 Subject: [PATCH 40/42] Remove useless test --- ydb/src/load_balancer/balancer_test.rs | 38 -------------------------- 1 file changed, 38 deletions(-) diff --git a/ydb/src/load_balancer/balancer_test.rs b/ydb/src/load_balancer/balancer_test.rs index 8e4ab2e6..2b922582 100644 --- a/ydb/src/load_balancer/balancer_test.rs +++ b/ydb/src/load_balancer/balancer_test.rs @@ -373,44 +373,6 @@ async fn no_addr_timeout() -> YdbResult<()> { Ok(()) } -#[tokio::test] -async fn detect_fastest_addr() -> YdbResult<()> { - let l1 = TcpListener::bind("127.0.0.1:0").await?; - let l2 = TcpListener::bind("127.0.0.1:0").await?; - let l3 = TcpListener::bind("127.0.0.1:0").await?; - - let l1_addr = l1.local_addr()?; - let l2_addr = l2.local_addr()?; - let l3_addr = l3.local_addr()?; - - println!("Listener №1 on: {}", l1_addr); - println!("Listener №2 on: {}", l2_addr); - println!("Listener №3 on: {}", l3_addr); - - let nodes = vec![ - l1_addr.to_string(), - l2_addr.to_string(), - l3_addr.to_string(), - ]; - - drop(l1); - drop(l2); - drop(l3); - - let result = - NearestDCBalancer::find_fastest_address(nodes.iter().collect_vec(), Duration::from_secs(3)) - .await; - match result { - Ok(_) => unreachable!(), - Err(err) => { - assert_eq!( - err.to_string(), - "Custom(\"timeout while detecting fastest address\")" - ); - } - } - Ok(()) -} #[tokio::test] async fn adjusting_dc() -> YdbResult<()> { From 4c2071d76ce3546bf8086bf025be63093d6fa0d0 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 15 Sep 2024 21:23:05 +0300 Subject: [PATCH 41/42] pub(crate) --- ydb/src/load_balancer/balancer_test.rs | 1 - ydb/src/load_balancer/nearest_dc_balancer.rs | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/ydb/src/load_balancer/balancer_test.rs b/ydb/src/load_balancer/balancer_test.rs index 2b922582..1953fd24 100644 --- a/ydb/src/load_balancer/balancer_test.rs +++ b/ydb/src/load_balancer/balancer_test.rs @@ -548,7 +548,6 @@ async fn nearest_dc_balancer_integration() -> YdbResult<()> { ), } - let updated_state = Arc::new( DiscoveryState::default() .with_node_info( Table, diff --git a/ydb/src/load_balancer/nearest_dc_balancer.rs b/ydb/src/load_balancer/nearest_dc_balancer.rs index 2998449e..1ec039ee 100644 --- a/ydb/src/load_balancer/nearest_dc_balancer.rs +++ b/ydb/src/load_balancer/nearest_dc_balancer.rs @@ -40,7 +40,7 @@ pub(super) struct BalancerState { } // What will balancer do if there is no available endpoints at local dc -pub enum FallbackStrategy { +pub(crate) enum FallbackStrategy { Error, // Just throw error BalanceWithOther(Box), // Use another balancer } From 079f46c8440c7f6581f4d6caec409aa03ce32796 Mon Sep 17 00:00:00 2001 From: kshprenger Date: Sun, 15 Sep 2024 22:01:34 +0300 Subject: [PATCH 42/42] fix --- ydb/src/load_balancer/balancer_test.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ydb/src/load_balancer/balancer_test.rs b/ydb/src/load_balancer/balancer_test.rs index 1953fd24..0d664396 100644 --- a/ydb/src/load_balancer/balancer_test.rs +++ b/ydb/src/load_balancer/balancer_test.rs @@ -373,7 +373,6 @@ async fn no_addr_timeout() -> YdbResult<()> { Ok(()) } - #[tokio::test] async fn adjusting_dc() -> YdbResult<()> { let l1 = TcpListener::bind("127.0.0.1:0").await?; @@ -547,7 +546,7 @@ async fn nearest_dc_balancer_integration() -> YdbResult<()> { "Custom(\"no available endpoints for service:table_service\")".to_string() ), } - + let updated_state = Arc::new( DiscoveryState::default() .with_node_info( Table,