1use async_trait::async_trait;
2use parking_lot::{Mutex as SyncMutex, RwLock};
3use serde::{Deserialize, Serialize};
4use std::collections::{HashMap, HashSet};
5use std::sync::Arc;
6use std::sync::atomic::{AtomicBool, Ordering};
7use tokio::sync::{Mutex as AsyncMutex, Notify, RwLock as AsyncRwLock, watch};
8
9use crate::channel::message::MessageChannel;
10use crate::client::{RpcClient, RpcClientHandle};
11use crate::codec::{BincodeCodec, Codec};
12use crate::discovery::Endpoint;
13use crate::error::{Result, RpcError, TransportError};
14use crate::loadbalancer::{LoadBalanceStrategy, LoadBalancer};
15use crate::streaming::{StreamId, StreamReceiver};
16
17#[async_trait]
18pub trait ClientFactory<T, C>: Send + Sync
19where
20 T: MessageChannel<C>,
21 C: Codec,
22{
23 async fn create(&self, endpoint: &Endpoint) -> Result<RpcClient<T, C>>;
24}
25
26struct ManagedClient<T, C>
27where
28 T: MessageChannel<C>,
29 C: Codec,
30{
31 client: Arc<RpcClient<T, C>>,
32 endpoint: Endpoint,
33 handle: AsyncMutex<Option<RpcClientHandle>>,
34}
35
36async fn wait_for_close(close_rx: &mut watch::Receiver<bool>) {
37 if *close_rx.borrow() {
38 return;
39 }
40 let _ = close_rx.changed().await;
41}
42
43type OwnedStreamIds = Arc<RwLock<HashSet<StreamId>>>;
44
45struct StreamAffinityGuard<S>
46where
47 S: LoadBalanceStrategy + 'static,
48{
49 load_balancer: Arc<LoadBalancer<S>>,
50 owned_stream_ids: OwnedStreamIds,
51 stream_id: StreamId,
52 armed: bool,
53}
54
55impl<S> StreamAffinityGuard<S>
56where
57 S: LoadBalanceStrategy + 'static,
58{
59 fn new(
60 load_balancer: Arc<LoadBalancer<S>>,
61 owned_stream_ids: OwnedStreamIds,
62 stream_id: StreamId,
63 ) -> Self {
64 owned_stream_ids.write().insert(stream_id);
65 Self {
66 load_balancer,
67 owned_stream_ids,
68 stream_id,
69 armed: true,
70 }
71 }
72
73 fn disarm(&mut self) {
74 self.armed = false;
75 }
76}
77
78impl<S> Drop for StreamAffinityGuard<S>
79where
80 S: LoadBalanceStrategy + 'static,
81{
82 fn drop(&mut self) {
83 if self.armed && self.owned_stream_ids.write().remove(&self.stream_id) {
84 self.load_balancer.release_stream(self.stream_id);
85 }
86 }
87}
88
89struct ActiveRequestGuard<S>
90where
91 S: LoadBalanceStrategy + 'static,
92{
93 load_balancer: Arc<LoadBalancer<S>>,
94 server_idx: usize,
95}
96
97impl<S> ActiveRequestGuard<S>
98where
99 S: LoadBalanceStrategy + 'static,
100{
101 fn new(load_balancer: Arc<LoadBalancer<S>>, server_idx: usize) -> Self {
102 load_balancer.acquire(server_idx);
103 Self {
104 load_balancer,
105 server_idx,
106 }
107 }
108}
109
110impl<S> Drop for ActiveRequestGuard<S>
111where
112 S: LoadBalanceStrategy + 'static,
113{
114 fn drop(&mut self) {
115 self.load_balancer.release(self.server_idx);
116 }
117}
118
119#[derive(Clone, Default)]
120enum LoadBalancedCloseState {
121 #[default]
122 Open,
123 Closing,
124 Closed(Result<()>),
125}
126
127struct LoadBalancedCloseFinalizer {
128 state: Arc<SyncMutex<LoadBalancedCloseState>>,
129 notify: Arc<Notify>,
130 completed: bool,
131}
132
133impl LoadBalancedCloseFinalizer {
134 fn complete(&mut self, result: Result<()>) {
135 *self.state.lock() = LoadBalancedCloseState::Closed(result);
136 self.completed = true;
137 self.notify.notify_waiters();
138 }
139}
140
141impl Drop for LoadBalancedCloseFinalizer {
142 fn drop(&mut self) {
143 if self.completed {
144 return;
145 }
146 *self.state.lock() = LoadBalancedCloseState::Closed(Err(RpcError::ClientError(
147 "Load-balanced close task was cancelled or panicked".to_string(),
148 )));
149 self.notify.notify_waiters();
150 }
151}
152
153pub struct LoadBalancedClient<T, C, S>
154where
155 T: MessageChannel<C>,
156 C: Codec,
157 S: LoadBalanceStrategy,
158{
159 load_balancer: Arc<LoadBalancer<S>>,
160 clients: RwLock<HashMap<usize, Arc<ManagedClient<T, C>>>>,
161 creation_locks: RwLock<HashMap<usize, Arc<AsyncMutex<()>>>>,
162 owned_stream_ids: OwnedStreamIds,
163 factory: Arc<dyn ClientFactory<T, C>>,
164 closed: AtomicBool,
165 close_cancel_tx: watch::Sender<bool>,
166 admission: AsyncRwLock<()>,
167 close_lock: AsyncMutex<()>,
168 close_state: Arc<SyncMutex<LoadBalancedCloseState>>,
169 close_notify: Arc<Notify>,
170}
171
172impl<T, C, S> LoadBalancedClient<T, C, S>
173where
174 T: MessageChannel<C> + 'static,
175 C: Codec + Clone + Default + 'static,
176 S: LoadBalanceStrategy + 'static,
177{
178 pub fn new(load_balancer: Arc<LoadBalancer<S>>, factory: Arc<dyn ClientFactory<T, C>>) -> Self {
179 let (close_cancel_tx, _) = watch::channel(false);
180 Self {
181 load_balancer,
182 clients: RwLock::new(HashMap::new()),
183 creation_locks: RwLock::new(HashMap::new()),
184 owned_stream_ids: Arc::new(RwLock::new(HashSet::new())),
185 factory,
186 closed: AtomicBool::new(false),
187 close_cancel_tx,
188 admission: AsyncRwLock::new(()),
189 close_lock: AsyncMutex::new(()),
190 close_state: Arc::new(SyncMutex::new(LoadBalancedCloseState::Open)),
191 close_notify: Arc::new(Notify::new()),
192 }
193 }
194
195 pub async fn init(&self) -> Result<()> {
196 self.ensure_open()?;
197 self.load_balancer.init().await?;
198 self.ensure_open()?;
199
200 for i in 0..self.load_balancer.server_count() {
201 if let Some(endpoint) = self.load_balancer.get_endpoint(i) {
202 let _ = self.get_or_create_client(i, &endpoint).await;
203 }
204 }
205
206 self.ensure_open()
207 }
208
209 fn ensure_open(&self) -> Result<()> {
210 if self.closed.load(Ordering::Acquire) {
211 return Err(RpcError::ConnectionClosed);
212 }
213 Ok(())
214 }
215
216 async fn close_entry(entry: &ManagedClient<T, C>) -> Result<()> {
217 let close_result = entry.client.close().await;
218 let handle = entry.handle.lock().await.take();
219 let join_result = match handle {
220 Some(handle) => handle.join().await,
221 None => Ok(()),
222 };
223
224 close_result?;
225
226 join_result.map_err(|error| {
227 RpcError::ClientError(format!("RPC client task failed while joining: {error}"))
228 })
229 }
230
231 fn spawn_entry_close(entry: Arc<ManagedClient<T, C>>) -> tokio::task::JoinHandle<Result<()>> {
232 tokio::spawn(async move { Self::close_entry(&entry).await })
233 }
234
235 fn close_unpublished_client(client: Arc<RpcClient<T, C>>) {
236 tokio::spawn(async move {
237 let _ = client.close().await;
238 });
239 }
240
241 fn cached_client(
242 &self,
243 server_idx: usize,
244 endpoint: &Endpoint,
245 ) -> Option<Arc<ManagedClient<T, C>>> {
246 if self.load_balancer.get_endpoint(server_idx).as_ref() != Some(endpoint) {
247 return None;
248 }
249 match self.clients.read().get(&server_idx).cloned() {
250 Some(entry) if entry.endpoint == *endpoint && entry.client.is_connected() => {
251 Some(entry)
252 }
253 _ => None,
254 }
255 }
256
257 async fn get_or_create_client(
258 &self,
259 server_idx: usize,
260 endpoint: &Endpoint,
261 ) -> Result<Arc<ManagedClient<T, C>>> {
262 self.ensure_open()?;
263 if let Some(entry) = self.cached_client(server_idx, endpoint) {
264 self.ensure_open()?;
265 return Ok(entry);
266 }
267
268 let mut close_rx = self.close_cancel_tx.subscribe();
269 let creation_lock = {
270 let mut creation_locks = self.creation_locks.write();
271 creation_locks
272 .entry(server_idx)
273 .or_insert_with(|| Arc::new(AsyncMutex::new(())))
274 .clone()
275 };
276 let _creation = tokio::select! {
277 biased;
278 _ = wait_for_close(&mut close_rx) => return Err(RpcError::ConnectionClosed),
279 creation = creation_lock.lock() => creation,
280 };
281 self.ensure_open()?;
282
283 if let Some(entry) = self.cached_client(server_idx, endpoint) {
284 return Ok(entry);
285 }
286
287 let stale = {
288 let current = self.clients.read().get(&server_idx).cloned();
289 match current {
290 Some(entry) => {
291 let mut clients = self.clients.write();
292 match clients.get(&server_idx) {
293 Some(current) if Arc::ptr_eq(current, &entry) => {
294 clients.remove(&server_idx)
295 }
296 _ => None,
297 }
298 }
299 None => None,
300 }
301 };
302 if let Some(stale) = stale {
303 let mut close_task = Self::spawn_entry_close(stale);
304 tokio::select! {
305 biased;
306 _ = wait_for_close(&mut close_rx) => return Err(RpcError::ConnectionClosed),
307 _ = &mut close_task => {}
308 }
309 }
310
311 self.ensure_open()?;
312 let client = tokio::select! {
313 biased;
314 result = self.factory.create(endpoint) => Arc::new(result?),
315 _ = wait_for_close(&mut close_rx) => return Err(RpcError::ConnectionClosed),
316 };
317
318 if self.closed.load(Ordering::Acquire) {
319 Self::close_unpublished_client(client);
320 return Err(RpcError::ConnectionClosed);
321 }
322
323 let admission = tokio::select! {
324 biased;
325 _ = wait_for_close(&mut close_rx) => {
326 Self::close_unpublished_client(client);
327 return Err(RpcError::ConnectionClosed);
328 }
329 admission = self.admission.read() => admission,
330 };
331 if let Err(error) = self.ensure_open() {
332 drop(admission);
333 Self::close_unpublished_client(client);
334 return Err(error);
335 }
336 if self.load_balancer.get_endpoint(server_idx).as_ref() != Some(endpoint) {
337 drop(admission);
338 Self::close_unpublished_client(client);
339 return Err(RpcError::ClientError(
340 "Server endpoint changed during client creation".to_string(),
341 ));
342 }
343
344 let handle = match client.try_start() {
345 Ok(handle) => handle,
346 Err(error) => {
347 drop(admission);
348 Self::close_unpublished_client(client);
349 return Err(error);
350 }
351 };
352 let entry = Arc::new(ManagedClient {
353 client,
354 endpoint: endpoint.clone(),
355 handle: AsyncMutex::new(Some(handle)),
356 });
357 let replaced = self.clients.write().insert(server_idx, Arc::clone(&entry));
358 drop(admission);
359
360 if let Some(replaced) = replaced {
361 drop(Self::spawn_entry_close(replaced));
362 }
363
364 Ok(entry)
365 }
366
367 async fn remove_and_close_client(
368 &self,
369 server_idx: usize,
370 expected: &Arc<ManagedClient<T, C>>,
371 ) {
372 let removed = {
373 let mut clients = self.clients.write();
374 match clients.get(&server_idx) {
375 Some(current) if Arc::ptr_eq(current, expected) => clients.remove(&server_idx),
376 _ => None,
377 }
378 };
379
380 if let Some(entry) = removed {
381 let _ = Self::spawn_entry_close(entry).await;
382 }
383 }
384
385 pub async fn close(&self) -> Result<()> {
386 self.closed.store(true, Ordering::Release);
387 self.close_cancel_tx.send_replace(true);
388
389 {
390 let _close = self.close_lock.lock().await;
391 let state = self.close_state.lock().clone();
392 match state {
393 LoadBalancedCloseState::Closed(result) => return result,
394 LoadBalancedCloseState::Closing => {}
395 LoadBalancedCloseState::Open => {
396 let _admission = self.admission.write().await;
397 let stream_ids = self.owned_stream_ids.write().drain().collect::<Vec<_>>();
398 for stream_id in stream_ids {
399 self.load_balancer.release_stream(stream_id);
400 }
401 let mut entries = self.clients.write().drain().collect::<Vec<_>>();
402 entries.sort_by_key(|(server_idx, _)| *server_idx);
403 self.creation_locks.write().clear();
404
405 *self.close_state.lock() = LoadBalancedCloseState::Closing;
406 let close_state = Arc::clone(&self.close_state);
407 let close_notify = Arc::clone(&self.close_notify);
408 tokio::spawn(async move {
409 let mut finalizer = LoadBalancedCloseFinalizer {
410 state: close_state,
411 notify: close_notify,
412 completed: false,
413 };
414 let close_tasks = entries
415 .into_iter()
416 .map(|(_, entry)| Self::spawn_entry_close(entry))
417 .collect::<Vec<_>>();
418 let mut first_error = None;
419 for close_task in close_tasks {
420 let result = match close_task.await {
421 Ok(result) => result,
422 Err(error) => Err(RpcError::ClientError(format!(
423 "RPC client close task failed while joining: {error}"
424 ))),
425 };
426 if first_error.is_none() {
427 first_error = result.err();
428 }
429 }
430 let result = match first_error {
431 Some(error) => Err(error),
432 None => Ok(()),
433 };
434 finalizer.complete(result);
435 });
436 }
437 }
438 }
439
440 loop {
441 let notified = self.close_notify.notified();
442 if let LoadBalancedCloseState::Closed(result) = &*self.close_state.lock() {
443 return result.clone();
444 }
445 notified.await;
446 }
447 }
448
449 fn child_error_is_terminal(entry: &ManagedClient<T, C>, error: &RpcError) -> bool {
450 !entry.client.is_connected()
451 || matches!(
452 error,
453 RpcError::ConnectionClosed
454 | RpcError::Transport(TransportError::ConnectionClosed)
455 | RpcError::Transport(TransportError::NotConnected)
456 )
457 }
458
459 pub async fn call<Req, Resp>(&self, method: &str, request: &Req) -> Result<Resp>
460 where
461 Req: Serialize,
462 Resp: for<'de> Deserialize<'de>,
463 {
464 self.ensure_open()?;
465 let config = self.load_balancer.config();
466 let max_attempts = if config.failover_enabled {
467 config.max_failover_attempts as usize + 1
468 } else {
469 1
470 };
471 let mut last_error = None;
472 let mut tried_servers = Vec::new();
473
474 for _ in 0..max_attempts {
475 self.ensure_open()?;
476 let server_idx = match self.load_balancer.select_excluding(&tried_servers) {
477 Some(idx) => idx,
478 None => break,
479 };
480 tried_servers.push(server_idx);
481
482 let endpoint = match self.load_balancer.get_endpoint(server_idx) {
483 Some(endpoint) => endpoint,
484 None => {
485 self.load_balancer.record_failure(server_idx);
486 last_error = Some(RpcError::ClientError("Server not found".to_string()));
487 continue;
488 }
489 };
490
491 let entry = match self.get_or_create_client(server_idx, &endpoint).await {
492 Ok(entry) => entry,
493 Err(error) => {
494 if self.closed.load(Ordering::Acquire) {
495 return Err(RpcError::ConnectionClosed);
496 }
497 self.load_balancer.record_failure(server_idx);
498 last_error = Some(error);
499 continue;
500 }
501 };
502
503 let active_request =
504 ActiveRequestGuard::new(Arc::clone(&self.load_balancer), server_idx);
505 let result = entry.client.call(method, request).await;
506 drop(active_request);
507
508 match result {
509 Ok(response) => {
510 self.load_balancer.record_success(server_idx);
511 return Ok(response);
512 }
513 Err(error) => {
514 self.load_balancer.record_failure(server_idx);
515 if Self::child_error_is_terminal(&entry, &error) {
516 self.remove_and_close_client(server_idx, &entry).await;
517 }
518 return Err(error);
519 }
520 }
521 }
522
523 Err(last_error.unwrap_or_else(|| RpcError::ClientError("No servers available".to_string())))
524 }
525
526 pub async fn call_server_stream<Req, Resp>(
527 &self,
528 method: &str,
529 request: &Req,
530 ) -> Result<StreamReceiver<Resp, C>>
531 where
532 Req: Serialize,
533 Resp: for<'de> Deserialize<'de>,
534 {
535 self.ensure_open()?;
536 let stream_id = crate::streaming::next_stream_id();
537
538 let server_idx = self
539 .load_balancer
540 .select_for_stream(stream_id)
541 .ok_or_else(|| RpcError::ClientError("No servers available".to_string()))?;
542 let mut affinity = StreamAffinityGuard::new(
543 Arc::clone(&self.load_balancer),
544 Arc::clone(&self.owned_stream_ids),
545 stream_id,
546 );
547
548 let endpoint = self
549 .load_balancer
550 .get_endpoint(server_idx)
551 .ok_or_else(|| RpcError::ClientError("Server not found".to_string()))?;
552
553 let entry = self.get_or_create_client(server_idx, &endpoint).await?;
554
555 match entry
556 .client
557 .call_server_stream_with_id(method, request, stream_id)
558 .await
559 {
560 Ok(mut receiver) => {
561 let owned_stream_ids = Arc::clone(&self.owned_stream_ids);
562 let load_balancer = Arc::clone(&self.load_balancer);
563 receiver.set_termination_callback(move |stream_id| {
564 if owned_stream_ids.write().remove(&stream_id) {
565 load_balancer.release_stream(stream_id);
566 }
567 });
568 affinity.disarm();
569 Ok(receiver)
570 }
571 Err(error) => Err(error),
572 }
573 }
574
575 pub async fn notify<Req: Serialize>(&self, method: &str, request: &Req) -> Result<()> {
576 self.ensure_open()?;
577 let server_idx = self
578 .load_balancer
579 .select()
580 .ok_or_else(|| RpcError::ClientError("No servers available".to_string()))?;
581
582 let endpoint = self
583 .load_balancer
584 .get_endpoint(server_idx)
585 .ok_or_else(|| RpcError::ClientError("Server not found".to_string()))?;
586
587 let entry = self.get_or_create_client(server_idx, &endpoint).await?;
588 entry.client.notify(method, request).await
589 }
590
591 pub fn load_balancer(&self) -> &LoadBalancer<S> {
592 &self.load_balancer
593 }
594
595 pub fn connection_count(&self) -> usize {
596 self.clients.read().len()
597 }
598
599 pub fn release_stream(&self, stream_id: crate::streaming::StreamId) {
600 if self.owned_stream_ids.write().remove(&stream_id) {
601 self.load_balancer.release_stream(stream_id);
602 }
603 }
604}
605
606impl<T, C, S> std::fmt::Debug for LoadBalancedClient<T, C, S>
607where
608 T: MessageChannel<C>,
609 C: Codec,
610 S: LoadBalanceStrategy + 'static,
611{
612 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
613 f.debug_struct("LoadBalancedClient")
614 .field("server_count", &self.load_balancer.server_count())
615 .field("available_count", &self.load_balancer.available_count())
616 .field("connection_count", &self.clients.read().len())
617 .field("strategy", &self.load_balancer.strategy_name())
618 .finish()
619 }
620}
621
622pub type DefaultLoadBalancedClient<T, S> = LoadBalancedClient<T, BincodeCodec, S>;
623
624#[cfg(test)]
625mod tests {
626 use super::*;
627 use crate::channel::message::MessageChannelAdapter;
628 use crate::discovery::StaticDiscovery;
629 use crate::error::{TransportError, TransportResult};
630 use crate::loadbalancer::{LoadBalancerConfig, RoundRobin};
631 use crate::message::Message;
632 use crate::streaming::StreamId;
633 use crate::transport::channel::{ChannelConfig, ChannelFrameTransport};
634 use parking_lot::Mutex;
635 use std::sync::atomic::AtomicUsize;
636 use std::time::Duration;
637
638 type TestChannel = MessageChannelAdapter<ChannelFrameTransport>;
639
640 struct MockClientFactory {
641 create_count: AtomicUsize,
642 created_endpoints: Mutex<Vec<Endpoint>>,
643 peers: Mutex<Vec<ChannelFrameTransport>>,
644 create_delay: Duration,
645 }
646
647 impl MockClientFactory {
648 fn new() -> Self {
649 Self::with_delay(Duration::ZERO)
650 }
651
652 fn with_delay(create_delay: Duration) -> Self {
653 Self {
654 create_count: AtomicUsize::new(0),
655 created_endpoints: Mutex::new(Vec::new()),
656 peers: Mutex::new(Vec::new()),
657 create_delay,
658 }
659 }
660 }
661
662 #[async_trait]
663 impl ClientFactory<TestChannel, BincodeCodec> for MockClientFactory {
664 async fn create(
665 &self,
666 endpoint: &Endpoint,
667 ) -> Result<RpcClient<TestChannel, BincodeCodec>> {
668 self.create_count.fetch_add(1, Ordering::Relaxed);
669 self.created_endpoints.lock().push(endpoint.clone());
670 if !self.create_delay.is_zero() {
671 tokio::time::sleep(self.create_delay).await;
672 }
673
674 let (client_transport, peer_transport) =
675 ChannelFrameTransport::create_pair("mock", ChannelConfig::default()).unwrap();
676 self.peers.lock().push(peer_transport);
677 Ok(RpcClient::new(MessageChannelAdapter::new(client_transport)))
678 }
679 }
680
681 #[derive(Debug)]
682 struct FailingChannel {
683 closed: AtomicBool,
684 send_count: Arc<AtomicUsize>,
685 sent_stream_ids: Arc<Mutex<Vec<StreamId>>>,
686 }
687
688 impl FailingChannel {
689 fn new(send_count: Arc<AtomicUsize>, sent_stream_ids: Arc<Mutex<Vec<StreamId>>>) -> Self {
690 Self {
691 closed: AtomicBool::new(false),
692 send_count,
693 sent_stream_ids,
694 }
695 }
696 }
697
698 #[async_trait]
699 impl MessageChannel<BincodeCodec> for FailingChannel {
700 async fn send(&self, message: &Message<BincodeCodec>) -> TransportResult<()> {
701 self.send_count.fetch_add(1, Ordering::SeqCst);
702 if let Some(stream_id) = message.metadata.stream_id {
703 self.sent_stream_ids.lock().push(stream_id);
704 }
705 Err(TransportError::ConnectionClosed)
706 }
707
708 async fn recv(&self) -> TransportResult<Message<BincodeCodec>> {
709 std::future::pending().await
710 }
711
712 fn is_connected(&self) -> bool {
713 !self.closed.load(Ordering::Acquire)
714 }
715
716 async fn close(&self) -> TransportResult<()> {
717 self.closed.store(true, Ordering::Release);
718 Err(TransportError::ConnectionClosed)
719 }
720 }
721
722 struct FailingClientFactory {
723 send_count: Arc<AtomicUsize>,
724 sent_stream_ids: Arc<Mutex<Vec<StreamId>>>,
725 }
726
727 impl FailingClientFactory {
728 fn new() -> Self {
729 Self {
730 send_count: Arc::new(AtomicUsize::new(0)),
731 sent_stream_ids: Arc::new(Mutex::new(Vec::new())),
732 }
733 }
734 }
735
736 #[async_trait]
737 impl ClientFactory<FailingChannel, BincodeCodec> for FailingClientFactory {
738 async fn create(
739 &self,
740 _endpoint: &Endpoint,
741 ) -> Result<RpcClient<FailingChannel, BincodeCodec>> {
742 Ok(RpcClient::new(FailingChannel::new(
743 Arc::clone(&self.send_count),
744 Arc::clone(&self.sent_stream_ids),
745 )))
746 }
747 }
748
749 #[derive(Debug)]
750 struct BlockingCloseControl {
751 close_count: AtomicUsize,
752 close_release: tokio::sync::Semaphore,
753 close_started: tokio::sync::Semaphore,
754 }
755
756 impl BlockingCloseControl {
757 fn new() -> Self {
758 Self {
759 close_count: AtomicUsize::new(0),
760 close_release: tokio::sync::Semaphore::new(0),
761 close_started: tokio::sync::Semaphore::new(0),
762 }
763 }
764 }
765
766 #[derive(Debug)]
767 struct BlockingCloseChannel {
768 control: Arc<BlockingCloseControl>,
769 }
770
771 #[async_trait]
772 impl MessageChannel<BincodeCodec> for BlockingCloseChannel {
773 async fn send(&self, _message: &Message<BincodeCodec>) -> TransportResult<()> {
774 Ok(())
775 }
776
777 async fn recv(&self) -> TransportResult<Message<BincodeCodec>> {
778 std::future::pending().await
779 }
780
781 fn is_connected(&self) -> bool {
782 true
783 }
784
785 async fn close(&self) -> TransportResult<()> {
786 self.control.close_count.fetch_add(1, Ordering::SeqCst);
787 self.control.close_started.add_permits(1);
788 self.control
789 .close_release
790 .acquire()
791 .await
792 .expect("close release semaphore must remain open")
793 .forget();
794 Err(TransportError::NotConnected)
795 }
796 }
797
798 struct BlockingCloseFactory {
799 control: Arc<BlockingCloseControl>,
800 }
801
802 #[async_trait]
803 impl ClientFactory<BlockingCloseChannel, BincodeCodec> for BlockingCloseFactory {
804 async fn create(
805 &self,
806 _endpoint: &Endpoint,
807 ) -> Result<RpcClient<BlockingCloseChannel, BincodeCodec>> {
808 Ok(RpcClient::new(BlockingCloseChannel {
809 control: Arc::clone(&self.control),
810 }))
811 }
812 }
813
814 struct BlockingCreateFactory {
815 started: Notify,
816 }
817
818 #[async_trait]
819 impl ClientFactory<TestChannel, BincodeCodec> for BlockingCreateFactory {
820 async fn create(
821 &self,
822 _endpoint: &Endpoint,
823 ) -> Result<RpcClient<TestChannel, BincodeCodec>> {
824 self.started.notify_one();
825 std::future::pending().await
826 }
827 }
828
829 struct FailingCreateFactory {
830 create_count: AtomicUsize,
831 }
832
833 #[async_trait]
834 impl ClientFactory<TestChannel, BincodeCodec> for FailingCreateFactory {
835 async fn create(
836 &self,
837 _endpoint: &Endpoint,
838 ) -> Result<RpcClient<TestChannel, BincodeCodec>> {
839 self.create_count.fetch_add(1, Ordering::SeqCst);
840 Err(RpcError::ClientError("client creation failed".to_string()))
841 }
842 }
843
844 #[derive(Debug)]
845 struct BlockingSendChannel {
846 send_started: Arc<Notify>,
847 }
848
849 #[async_trait]
850 impl MessageChannel<BincodeCodec> for BlockingSendChannel {
851 async fn send(&self, _message: &Message<BincodeCodec>) -> TransportResult<()> {
852 self.send_started.notify_one();
853 std::future::pending().await
854 }
855
856 async fn recv(&self) -> TransportResult<Message<BincodeCodec>> {
857 std::future::pending().await
858 }
859
860 fn is_connected(&self) -> bool {
861 true
862 }
863
864 async fn close(&self) -> TransportResult<()> {
865 Ok(())
866 }
867 }
868
869 struct BlockingSendFactory {
870 send_started: Arc<Notify>,
871 }
872
873 #[async_trait]
874 impl ClientFactory<BlockingSendChannel, BincodeCodec> for BlockingSendFactory {
875 async fn create(
876 &self,
877 _endpoint: &Endpoint,
878 ) -> Result<RpcClient<BlockingSendChannel, BincodeCodec>> {
879 Ok(RpcClient::new(BlockingSendChannel {
880 send_started: Arc::clone(&self.send_started),
881 }))
882 }
883 }
884
885 fn load_balancer(server_count: usize) -> Arc<LoadBalancer<RoundRobin>> {
886 load_balancer_with_config(server_count, LoadBalancerConfig::default())
887 }
888
889 fn load_balancer_with_config(
890 server_count: usize,
891 config: LoadBalancerConfig,
892 ) -> Arc<LoadBalancer<RoundRobin>> {
893 let endpoints = (0..server_count)
894 .map(|index| Endpoint::tcp_from_str(&format!("127.0.0.1:{}", 8001 + index)).unwrap())
895 .collect();
896 let discovery = Arc::new(StaticDiscovery::new(endpoints));
897 Arc::new(LoadBalancer::with_config(
898 discovery,
899 RoundRobin::new(),
900 config,
901 ))
902 }
903
904 #[tokio::test]
905 async fn test_lb_client_creation() {
906 let lb = load_balancer(2);
907 let factory = Arc::new(MockClientFactory::new());
908 let client = LoadBalancedClient::new(lb.clone(), factory.clone());
909
910 client.init().await.unwrap();
911
912 assert_eq!(client.load_balancer().server_count(), 2);
913 assert_eq!(client.connection_count(), 2);
914 assert_eq!(factory.create_count.load(Ordering::Relaxed), 2);
915 client.close().await.unwrap();
916 }
917
918 #[tokio::test]
919 async fn test_child_handle_is_retained_until_close() {
920 let factory = Arc::new(MockClientFactory::new());
921 let client = LoadBalancedClient::new(load_balancer(1), factory);
922 client.init().await.unwrap();
923
924 let entry = { client.clients.read().get(&0).cloned().unwrap() };
925 assert!(entry.handle.lock().await.is_some());
926
927 client.close().await.unwrap();
928
929 assert!(entry.handle.lock().await.is_none());
930 assert!(entry.client.is_closed());
931 }
932
933 #[tokio::test]
934 async fn test_close_drains_clients_and_is_idempotent() {
935 let factory = Arc::new(MockClientFactory::new());
936 let client = LoadBalancedClient::new(load_balancer(2), factory);
937 client.init().await.unwrap();
938
939 let entries = { client.clients.read().values().cloned().collect::<Vec<_>>() };
940 client.close().await.unwrap();
941
942 assert_eq!(client.connection_count(), 0);
943 for entry in entries {
944 assert!(entry.client.is_closed());
945 assert!(entry.handle.lock().await.is_none());
946 }
947 client.close().await.unwrap();
948 }
949
950 #[tokio::test]
951 async fn test_close_joins_handles_after_client_close_error() {
952 let client =
953 LoadBalancedClient::new(load_balancer(1), Arc::new(FailingClientFactory::new()));
954 client.init().await.unwrap();
955 let entry = { client.clients.read().get(&0).cloned().unwrap() };
956
957 let result = client.close().await;
958
959 assert!(matches!(result, Err(RpcError::ConnectionClosed)));
960 assert_eq!(client.connection_count(), 0);
961 assert!(entry.client.is_closed());
962 assert!(entry.handle.lock().await.is_none());
963 assert!(matches!(
964 client.close().await,
965 Err(RpcError::ConnectionClosed)
966 ));
967 }
968
969 #[tokio::test]
970 async fn test_cancelled_close_retries_finish_all_drained_children() {
971 let control = Arc::new(BlockingCloseControl::new());
972 let factory = Arc::new(BlockingCloseFactory {
973 control: Arc::clone(&control),
974 });
975 let client = Arc::new(LoadBalancedClient::new(load_balancer(2), factory));
976 client.init().await.unwrap();
977 let entries = client.clients.read().values().cloned().collect::<Vec<_>>();
978
979 let close_client = Arc::clone(&client);
980 let first = tokio::spawn(async move { close_client.close().await });
981 for _ in 0..2 {
982 tokio::time::timeout(Duration::from_secs(1), control.close_started.acquire())
983 .await
984 .unwrap()
985 .unwrap()
986 .forget();
987 }
988 assert_eq!(control.close_count.load(Ordering::SeqCst), 2);
989 assert_eq!(client.connection_count(), 0);
990 first.abort();
991 assert!(first.await.unwrap_err().is_cancelled());
992
993 control.close_release.add_permits(2);
994 let result = tokio::time::timeout(Duration::from_secs(1), client.close())
995 .await
996 .unwrap();
997
998 assert!(matches!(
999 result,
1000 Err(RpcError::Transport(TransportError::NotConnected))
1001 ));
1002 assert_eq!(control.close_count.load(Ordering::SeqCst), 2);
1003 for entry in entries {
1004 assert!(entry.client.is_closed());
1005 assert!(entry.handle.lock().await.is_none());
1006 }
1007 assert!(matches!(
1008 client.close().await,
1009 Err(RpcError::Transport(TransportError::NotConnected))
1010 ));
1011 }
1012
1013 #[tokio::test]
1014 async fn test_close_cancels_hung_client_creation() {
1015 let factory = Arc::new(BlockingCreateFactory {
1016 started: Notify::new(),
1017 });
1018 let client = Arc::new(LoadBalancedClient::new(load_balancer(1), factory.clone()));
1019 client.load_balancer().init().await.unwrap();
1020
1021 let call_client = Arc::clone(&client);
1022 let call = tokio::spawn(async move { call_client.call::<_, ()>("blocked", &()).await });
1023 tokio::time::timeout(Duration::from_secs(1), factory.started.notified())
1024 .await
1025 .unwrap();
1026
1027 tokio::time::timeout(Duration::from_secs(1), client.close())
1028 .await
1029 .unwrap()
1030 .unwrap();
1031 assert!(matches!(
1032 call.await.unwrap(),
1033 Err(RpcError::ConnectionClosed)
1034 ));
1035 }
1036
1037 #[tokio::test]
1038 async fn test_failover_removal_closes_and_joins_client() {
1039 let client =
1040 LoadBalancedClient::new(load_balancer(1), Arc::new(FailingClientFactory::new()));
1041 client.init().await.unwrap();
1042 let entry = { client.clients.read().get(&0).cloned().unwrap() };
1043
1044 let result: Result<()> = client.call("fail", &()).await;
1045
1046 assert!(matches!(result, Err(RpcError::ConnectionClosed)));
1047 assert_eq!(client.connection_count(), 0);
1048 assert!(entry.client.is_closed());
1049 assert!(entry.handle.lock().await.is_none());
1050 client.close().await.unwrap();
1051 }
1052
1053 #[tokio::test]
1054 async fn test_failover_disabled_attempts_one_client_creation() {
1055 let config = LoadBalancerConfig {
1056 failover_enabled: false,
1057 max_failover_attempts: 10,
1058 ..LoadBalancerConfig::default()
1059 };
1060 let factory = Arc::new(FailingCreateFactory {
1061 create_count: AtomicUsize::new(0),
1062 });
1063 let client = LoadBalancedClient::new(load_balancer_with_config(3, config), factory.clone());
1064 client.load_balancer().init().await.unwrap();
1065
1066 let result: Result<()> = client.call("create", &()).await;
1067
1068 assert!(matches!(result, Err(RpcError::ClientError(_))));
1069 assert_eq!(factory.create_count.load(Ordering::SeqCst), 1);
1070 client.close().await.unwrap();
1071 }
1072
1073 #[tokio::test]
1074 async fn test_call_error_is_not_retried_after_invocation() {
1075 let factory = Arc::new(FailingClientFactory::new());
1076 let client = LoadBalancedClient::new(load_balancer(2), factory.clone());
1077 client.init().await.unwrap();
1078
1079 let result: Result<()> = client.call("unknown_commit", &()).await;
1080
1081 assert!(matches!(result, Err(RpcError::ConnectionClosed)));
1082 assert_eq!(factory.send_count.load(Ordering::SeqCst), 1);
1083 assert_eq!(client.connection_count(), 1);
1084 let _ = client.close().await;
1085 }
1086
1087 #[tokio::test]
1088 async fn test_dropped_call_releases_active_request_count() {
1089 let send_started = Arc::new(Notify::new());
1090 let factory = Arc::new(BlockingSendFactory {
1091 send_started: Arc::clone(&send_started),
1092 });
1093 let client = Arc::new(LoadBalancedClient::new(load_balancer(1), factory));
1094 client.init().await.unwrap();
1095
1096 let call_client = Arc::clone(&client);
1097 let call = tokio::spawn(async move { call_client.call::<_, ()>("blocked", &()).await });
1098 tokio::time::timeout(Duration::from_secs(1), send_started.notified())
1099 .await
1100 .unwrap();
1101 assert_eq!(client.load_balancer().active_requests(0), 1);
1102
1103 call.abort();
1104 assert!(call.await.unwrap_err().is_cancelled());
1105 assert_eq!(client.load_balancer().active_requests(0), 0);
1106 client.close().await.unwrap();
1107 }
1108
1109 #[tokio::test]
1110 async fn test_dropped_stream_setup_releases_affinity_during_creation() {
1111 let lb = load_balancer(2);
1112 let factory = Arc::new(BlockingCreateFactory {
1113 started: Notify::new(),
1114 });
1115 let client = Arc::new(LoadBalancedClient::new(lb.clone(), factory.clone()));
1116 client.load_balancer().init().await.unwrap();
1117
1118 let stream_client = Arc::clone(&client);
1119 let setup = tokio::spawn(async move {
1120 stream_client
1121 .call_server_stream::<_, i32>("numbers", &())
1122 .await
1123 });
1124 tokio::time::timeout(Duration::from_secs(1), factory.started.notified())
1125 .await
1126 .unwrap();
1127 let stream_id = *client
1128 .owned_stream_ids
1129 .read()
1130 .iter()
1131 .next()
1132 .expect("stream affinity must be owned during creation");
1133
1134 setup.abort();
1135 assert!(matches!(setup.await, Err(error) if error.is_cancelled()));
1136 assert!(client.owned_stream_ids.read().is_empty());
1137 assert_eq!(lb.select_for_stream(stream_id), Some(1));
1138 lb.release_stream(stream_id);
1139 client.close().await.unwrap();
1140 }
1141
1142 #[tokio::test]
1143 async fn test_dropped_stream_setup_releases_affinity_during_send() {
1144 let lb = load_balancer(2);
1145 let send_started = Arc::new(Notify::new());
1146 let factory = Arc::new(BlockingSendFactory {
1147 send_started: Arc::clone(&send_started),
1148 });
1149 let client = Arc::new(LoadBalancedClient::new(lb.clone(), factory));
1150 client.init().await.unwrap();
1151 let entry = client.clients.read().get(&0).cloned().unwrap();
1152
1153 let stream_client = Arc::clone(&client);
1154 let setup = tokio::spawn(async move {
1155 stream_client
1156 .call_server_stream::<_, i32>("numbers", &())
1157 .await
1158 });
1159 tokio::time::timeout(Duration::from_secs(1), send_started.notified())
1160 .await
1161 .unwrap();
1162 let stream_id = *client
1163 .owned_stream_ids
1164 .read()
1165 .iter()
1166 .next()
1167 .expect("stream affinity must be owned during send");
1168
1169 setup.abort();
1170 assert!(matches!(setup.await, Err(error) if error.is_cancelled()));
1171 assert!(client.owned_stream_ids.read().is_empty());
1172 assert_eq!(entry.client.active_streams(), 0);
1173 assert_eq!(lb.select_for_stream(stream_id), Some(1));
1174 lb.release_stream(stream_id);
1175 client.close().await.unwrap();
1176 }
1177
1178 #[tokio::test]
1179 async fn test_server_stream_uses_affinity_stream_id() {
1180 let lb = load_balancer(2);
1181 let factory = Arc::new(MockClientFactory::new());
1182 let client = LoadBalancedClient::new(lb.clone(), factory.clone());
1183 client.init().await.unwrap();
1184 let peer_transport = { factory.peers.lock().remove(0) };
1185 let peer_channel = MessageChannelAdapter::new(peer_transport);
1186
1187 let stream: StreamReceiver<i32> = client.call_server_stream("numbers", &()).await.unwrap();
1188 let message = peer_channel.recv().await.unwrap();
1189 let stream_id = stream.stream_id();
1190
1191 assert_eq!(message.metadata.stream_id, Some(stream_id));
1192 assert_eq!(lb.select_for_stream(stream_id), Some(0));
1193 client.release_stream(stream_id);
1194 assert_eq!(lb.select_for_stream(stream_id), Some(1));
1195 lb.release_stream(stream_id);
1196 client.close().await.unwrap();
1197 }
1198
1199 #[tokio::test]
1200 async fn test_server_stream_error_releases_affinity() {
1201 let lb = load_balancer(2);
1202 let factory = Arc::new(FailingClientFactory::new());
1203 let client = LoadBalancedClient::new(lb.clone(), factory.clone());
1204 client.init().await.unwrap();
1205
1206 let result: Result<StreamReceiver<i32>> = client.call_server_stream("numbers", &()).await;
1207 let sent_stream_ids = { factory.sent_stream_ids.lock().clone() };
1208
1209 assert!(matches!(result, Err(RpcError::ConnectionClosed)));
1210 assert_eq!(sent_stream_ids.len(), 1);
1211 assert_eq!(lb.select_for_stream(sent_stream_ids[0]), Some(1));
1212 lb.release_stream(sent_stream_ids[0]);
1213 let _ = client.close().await;
1214 }
1215
1216 #[tokio::test]
1217 async fn test_server_stream_end_releases_affinity_automatically() {
1218 let lb = load_balancer(2);
1219 let factory = Arc::new(MockClientFactory::new());
1220 let client = LoadBalancedClient::new(lb.clone(), factory.clone());
1221 client.init().await.unwrap();
1222 let peer = MessageChannelAdapter::new(factory.peers.lock().remove(0));
1223
1224 let mut stream: StreamReceiver<i32> =
1225 client.call_server_stream("numbers", &()).await.unwrap();
1226 let request = peer.recv().await.unwrap();
1227 let stream_id = request.metadata.stream_id.unwrap();
1228 assert_eq!(lb.select_for_stream(stream_id), Some(0));
1229
1230 let end: Message = Message::stream_end(stream_id);
1231 peer.send(&end).await.unwrap();
1232 assert!(stream.recv().await.is_none());
1233 assert_eq!(lb.select_for_stream(stream_id), Some(1));
1234 lb.release_stream(stream_id);
1235 client.close().await.unwrap();
1236 }
1237
1238 #[tokio::test]
1239 async fn test_dropped_server_stream_releases_affinity_automatically() {
1240 let lb = load_balancer(2);
1241 let factory = Arc::new(MockClientFactory::new());
1242 let client = LoadBalancedClient::new(lb.clone(), factory);
1243 client.init().await.unwrap();
1244
1245 let stream: StreamReceiver<i32> = client.call_server_stream("numbers", &()).await.unwrap();
1246 let stream_id = stream.stream_id();
1247 assert_eq!(lb.select_for_stream(stream_id), Some(0));
1248 drop(stream);
1249 assert_eq!(lb.select_for_stream(stream_id), Some(1));
1250 lb.release_stream(stream_id);
1251 client.close().await.unwrap();
1252 }
1253
1254 #[tokio::test]
1255 async fn test_close_releases_only_owned_stream_affinities() {
1256 let lb = load_balancer(3);
1257 let first = LoadBalancedClient::new(lb.clone(), Arc::new(MockClientFactory::new()));
1258 let second = LoadBalancedClient::new(lb.clone(), Arc::new(MockClientFactory::new()));
1259 first.init().await.unwrap();
1260 second.init().await.unwrap();
1261
1262 let first_stream: StreamReceiver<i32> =
1263 first.call_server_stream("first", &()).await.unwrap();
1264 let second_stream: StreamReceiver<i32> =
1265 second.call_server_stream("second", &()).await.unwrap();
1266 let first_stream_id = first_stream.stream_id();
1267 let second_stream_id = second_stream.stream_id();
1268
1269 assert_eq!(lb.select_for_stream(first_stream_id), Some(0));
1270 assert_eq!(lb.select_for_stream(second_stream_id), Some(1));
1271 first.close().await.unwrap();
1272 assert_eq!(lb.select_for_stream(second_stream_id), Some(1));
1273 assert_eq!(lb.select_for_stream(first_stream_id), Some(2));
1274 lb.release_stream(first_stream_id);
1275
1276 second.close().await.unwrap();
1277 assert_eq!(lb.select_for_stream(second_stream_id), Some(0));
1278 lb.release_stream(second_stream_id);
1279 }
1280
1281 #[tokio::test]
1282 async fn test_endpoint_churn_replaces_connected_cached_client() {
1283 let load_balancer = load_balancer(1);
1284 let factory = Arc::new(MockClientFactory::new());
1285 let client = LoadBalancedClient::new(load_balancer.clone(), factory.clone());
1286 client.init().await.unwrap();
1287 let first = client.clients.read().get(&0).cloned().unwrap();
1288 let replacement = Endpoint::tcp_from_str("127.0.0.1:9001").unwrap();
1289
1290 load_balancer.update_endpoints(vec![replacement.clone()]);
1291 let second = client.get_or_create_client(0, &replacement).await.unwrap();
1292
1293 assert!(!Arc::ptr_eq(&first, &second));
1294 assert_eq!(second.endpoint, replacement);
1295 assert!(first.client.is_closed());
1296 assert!(first.handle.lock().await.is_none());
1297 assert_eq!(factory.create_count.load(Ordering::SeqCst), 2);
1298 assert_eq!(
1299 factory.created_endpoints.lock().as_slice(),
1300 &[
1301 Endpoint::tcp_from_str("127.0.0.1:8001").unwrap(),
1302 replacement
1303 ]
1304 );
1305 client.close().await.unwrap();
1306 }
1307
1308 #[tokio::test]
1309 async fn test_healthy_cached_client_bypasses_creation_mutex() {
1310 let factory = Arc::new(MockClientFactory::new());
1311 let client = LoadBalancedClient::new(load_balancer(1), factory);
1312 client.init().await.unwrap();
1313 let endpoint = client.load_balancer().get_endpoint(0).unwrap();
1314 let creation_lock = client.creation_locks.read().get(&0).cloned().unwrap();
1315 let _creation = creation_lock.lock().await;
1316
1317 let cached = tokio::time::timeout(
1318 Duration::from_millis(50),
1319 client.get_or_create_client(0, &endpoint),
1320 )
1321 .await
1322 .expect("healthy cached lookup waited for the creation mutex")
1323 .unwrap();
1324
1325 assert!(Arc::ptr_eq(&cached, client.clients.read().get(&0).unwrap()));
1326 drop(_creation);
1327 client.close().await.unwrap();
1328 }
1329
1330 #[tokio::test]
1331 async fn test_concurrent_creation_is_single_flight_per_server() {
1332 let factory = Arc::new(MockClientFactory::with_delay(Duration::from_millis(50)));
1333 let client = LoadBalancedClient::new(load_balancer(1), factory.clone());
1334 client.load_balancer().init().await.unwrap();
1335 let endpoint = client.load_balancer().get_endpoint(0).unwrap();
1336
1337 let (first, second, third) = tokio::join!(
1338 client.get_or_create_client(0, &endpoint),
1339 client.get_or_create_client(0, &endpoint),
1340 client.get_or_create_client(0, &endpoint),
1341 );
1342 let first = first.unwrap();
1343 let second = second.unwrap();
1344 let third = third.unwrap();
1345
1346 assert!(Arc::ptr_eq(&first, &second));
1347 assert!(Arc::ptr_eq(&first, &third));
1348 assert_eq!(factory.create_count.load(Ordering::Relaxed), 1);
1349 client.close().await.unwrap();
1350 }
1351
1352 #[tokio::test]
1353 async fn test_work_after_close_returns_connection_closed() {
1354 let factory = Arc::new(MockClientFactory::new());
1355 let client = LoadBalancedClient::new(load_balancer(1), factory);
1356 client.init().await.unwrap();
1357 client.close().await.unwrap();
1358
1359 let call_result: Result<()> = client.call("closed", &()).await;
1360 let stream_result: Result<StreamReceiver<(), BincodeCodec>> =
1361 client.call_server_stream("closed", &()).await;
1362 let notify_result = client.notify("closed", &()).await;
1363
1364 assert!(matches!(call_result, Err(RpcError::ConnectionClosed)));
1365 assert!(matches!(stream_result, Err(RpcError::ConnectionClosed)));
1366 assert!(matches!(notify_result, Err(RpcError::ConnectionClosed)));
1367 }
1368}