1#[cfg(feature = "discovery")]
2use std::sync::Arc;
3use std::time::Duration;
4
5use async_trait::async_trait;
6use rskit_discovery::Discovery;
7#[cfg(feature = "discovery")]
8use rskit_discovery::Watcher;
9use rskit_errors::AppResult;
10use tokio::sync::RwLock;
11use tokio::task::JoinHandle;
12use tonic::transport::Channel;
13use tracing::{debug, warn};
14
15use crate::channel::GrpcChannel;
16use crate::config::GrpcClientConfig;
17
18#[derive(Clone, Debug)]
20pub struct DiscoveryChannelConfig {
21 pub grpc: GrpcClientConfig,
23 pub resolve_interval: Duration,
26}
27
28impl DiscoveryChannelConfig {
29 pub fn new(grpc: GrpcClientConfig) -> Self {
31 Self {
32 grpc,
33 resolve_interval: Duration::from_secs(10),
34 }
35 }
36
37 #[must_use]
39 pub fn with_resolve_interval(mut self, interval: Duration) -> Self {
40 self.resolve_interval = interval;
41 self
42 }
43}
44
45impl From<GrpcClientConfig> for DiscoveryChannelConfig {
46 fn from(grpc: GrpcClientConfig) -> Self {
47 Self::new(grpc)
48 }
49}
50
51#[async_trait]
52trait ChannelConnector: Send + Sync {
53 async fn connect(&self, config: GrpcClientConfig) -> AppResult<GrpcChannel>;
54}
55
56#[derive(Default)]
57struct DefaultChannelConnector;
58
59#[async_trait]
60impl ChannelConnector for DefaultChannelConnector {
61 async fn connect(&self, config: GrpcClientConfig) -> AppResult<GrpcChannel> {
62 let channel = GrpcChannel::new(config);
63 channel.connect().await?;
64 Ok(channel)
65 }
66}
67
68pub struct DiscoveryChannel {
76 discovery: Arc<dyn Discovery>,
77 watcher: Option<Arc<dyn Watcher>>,
78 service_name: String,
79 config: DiscoveryChannelConfig,
80 connector: Arc<dyn ChannelConnector>,
81 current_target: Arc<RwLock<Option<String>>>,
83 channel: Arc<RwLock<Option<GrpcChannel>>>,
85 bg_handle: Arc<RwLock<Option<JoinHandle<()>>>>,
87}
88
89impl DiscoveryChannel {
90 pub fn new(
95 discovery: Arc<dyn Discovery>,
96 service_name: impl Into<String>,
97 config: GrpcClientConfig,
98 ) -> Self {
99 Self::build(
100 discovery,
101 None,
102 service_name,
103 DiscoveryChannelConfig::new(config),
104 Arc::new(DefaultChannelConnector),
105 )
106 }
107
108 pub fn with_watcher(
111 discovery: Arc<dyn Discovery>,
112 watcher: Option<Arc<dyn Watcher>>,
113 service_name: impl Into<String>,
114 config: DiscoveryChannelConfig,
115 ) -> Self {
116 Self::build(
117 discovery,
118 watcher,
119 service_name,
120 config,
121 Arc::new(DefaultChannelConnector),
122 )
123 }
124
125 fn build(
126 discovery: Arc<dyn Discovery>,
127 watcher: Option<Arc<dyn Watcher>>,
128 service_name: impl Into<String>,
129 config: DiscoveryChannelConfig,
130 connector: Arc<dyn ChannelConnector>,
131 ) -> Self {
132 let service_name = service_name.into();
133 debug!(
134 service = %service_name,
135 "Creating DiscoveryChannel target: {}",
136 config.grpc.target
137 );
138
139 Self {
140 discovery,
141 watcher,
142 service_name,
143 config,
144 connector,
145 current_target: Arc::new(RwLock::new(None)),
146 channel: Arc::new(RwLock::new(None)),
147 bg_handle: Arc::new(RwLock::new(None)),
148 }
149 }
150
151 #[cfg(test)]
152 fn with_connector(
153 discovery: Arc<dyn Discovery>,
154 watcher: Option<Arc<dyn Watcher>>,
155 service_name: impl Into<String>,
156 config: DiscoveryChannelConfig,
157 connector: Arc<dyn ChannelConnector>,
158 ) -> Self {
159 Self::build(discovery, watcher, service_name, config, connector)
160 }
161
162 pub async fn start_background(&self) -> AppResult<()> {
170 {
172 let guard = self.bg_handle.read().await;
173 if guard.is_some() {
174 return Ok(());
175 }
176 }
177
178 let discovery = Arc::clone(&self.discovery);
179 let service_name = self.service_name.clone();
180 let grpc_config = self.config.grpc.clone();
181 let connector = Arc::clone(&self.connector);
182 let current_target = Arc::clone(&self.current_target);
183 let channel = Arc::clone(&self.channel);
184 let resolve_interval = self.config.resolve_interval;
185
186 let handle = if let Some(watcher) = &self.watcher {
187 let mut rx = watcher.watch(&self.service_name).await?;
189
190 tokio::spawn(async move {
191 while let Some(instances) = rx.recv().await {
192 if instances.is_empty() {
193 debug!(
194 service = %service_name,
195 "watcher: received empty instance list, skipping"
196 );
197 continue;
198 }
199 let new_target = instances[0].endpoint();
200 Self::maybe_reconnect(
201 &service_name,
202 &grpc_config,
203 &new_target,
204 &connector,
205 ¤t_target,
206 &channel,
207 )
208 .await;
209 }
210 debug!(
211 service = %service_name,
212 "watcher channel closed, background task exiting"
213 );
214 })
215 } else {
216 tokio::spawn(async move {
218 loop {
219 tokio::time::sleep(resolve_interval).await;
220
221 match discovery.resolve(&service_name).await {
222 Ok(instances) if !instances.is_empty() => {
223 let new_target = instances[0].endpoint();
224 Self::maybe_reconnect(
225 &service_name,
226 &grpc_config,
227 &new_target,
228 &connector,
229 ¤t_target,
230 &channel,
231 )
232 .await;
233 }
234 Ok(_) => {
235 debug!(
236 service = %service_name,
237 "poll: no instances found"
238 );
239 }
240 Err(e) => {
241 warn!(
242 service = %service_name,
243 error = %e,
244 "poll: resolve failed"
245 );
246 }
247 }
248 }
249 })
250 };
251
252 *self.bg_handle.write().await = Some(handle);
253 Ok(())
254 }
255
256 async fn maybe_reconnect(
259 service_name: &str,
260 base_config: &GrpcClientConfig,
261 new_target: &str,
262 connector: &Arc<dyn ChannelConnector>,
263 current_target: &Arc<RwLock<Option<String>>>,
264 channel: &Arc<RwLock<Option<GrpcChannel>>>,
265 ) {
266 let old_target = current_target.read().await.clone();
267 if old_target.as_deref() == Some(new_target) {
268 return;
269 }
270
271 debug!(
272 service = %service_name,
273 old_target = ?old_target,
274 new_target = %new_target,
275 "target changed, reconnecting"
276 );
277
278 let gc = match connect_grpc_channel(connector, base_config, new_target).await {
279 Ok(gc) => gc,
280 Err(e) => {
281 warn!(
282 service = %service_name,
283 target = %new_target,
284 error = %e,
285 "background reconnect failed"
286 );
287 return;
288 }
289 };
290
291 *current_target.write().await = Some(new_target.to_owned());
292 *channel.write().await = Some(gc);
293 }
294
295 pub async fn resolve(&self) -> AppResult<String> {
299 debug!(service = %self.service_name, "Resolving service");
300
301 let instances = self
302 .discovery
303 .resolve(&self.service_name)
304 .await
305 .map_err(|e| {
306 warn!(
307 service = %self.service_name,
308 error = %e,
309 "Service discovery failed"
310 );
311 e
312 })?;
313
314 if instances.is_empty() {
315 return Err(rskit_errors::AppError::new(
316 rskit_errors::ErrorCode::ServiceUnavailable,
317 format!("no instances found for service: {}", self.service_name),
318 ));
319 }
320
321 let instance = &instances[0];
323 let target = instance.endpoint();
324
325 debug!(
326 service = %self.service_name,
327 target = %target,
328 "Resolved service to target"
329 );
330
331 Ok(target)
332 }
333
334 pub async fn channel(&self) -> AppResult<Channel> {
340 {
342 let ch = self.channel.read().await;
343 if let Some(gc) = ch.as_ref()
344 && let Ok(connected_ch) = gc.connected_channel().await
345 {
346 return Ok(connected_ch);
347 }
348 }
349
350 let target = {
352 let current = self.current_target.read().await;
353 current.clone()
354 };
355
356 let target = match target {
357 Some(t) => t,
358 None => {
359 self.resolve().await?
361 }
362 };
363
364 let config = config_for_target(&self.config.grpc, &target);
366 let gc = match self.connector.connect(config).await {
367 Ok(gc) => gc,
368 Err(e) => {
369 warn!(
372 service = %self.service_name,
373 target = %target,
374 error = %e,
375 "connection failed, triggering re-resolve"
376 );
377 if let Ok(new_target) = self.resolve().await
378 && new_target != target
379 && let Ok(gc2) =
380 connect_grpc_channel(&self.connector, &self.config.grpc, &new_target).await
381 {
382 let connected_ch = gc2.connected_channel().await?;
383 *self.current_target.write().await = Some(new_target);
384 *self.channel.write().await = Some(gc2);
385 return Ok(connected_ch);
386 }
387 return Err(e);
388 }
389 };
390
391 let connected_ch = gc.connected_channel().await?;
392
393 {
395 let mut ch = self.channel.write().await;
396 *ch = Some(gc);
397 }
398 *self.current_target.write().await = Some(target);
399
400 Ok(connected_ch)
401 }
402
403 pub async fn refresh(&self) -> AppResult<bool> {
408 let old_target = self.current_target.read().await.clone();
409 let new_target = self.resolve().await?;
410 let changed = old_target.as_ref() != Some(&new_target);
411
412 if changed {
413 debug!(
414 service = %self.service_name,
415 old_target = ?old_target,
416 new_target = %new_target,
417 "Service target changed, creating new channel"
418 );
419
420 let gc = connect_grpc_channel(&self.connector, &self.config.grpc, &new_target).await?;
422
423 let mut ch = self.channel.write().await;
424 *ch = Some(gc);
425 *self.current_target.write().await = Some(new_target);
426 }
427
428 Ok(changed)
429 }
430
431 pub async fn close(&mut self) -> AppResult<()> {
433 debug!(service = %self.service_name, "Closing DiscoveryChannel");
434
435 {
437 let mut bg = self.bg_handle.write().await;
438 if let Some(handle) = bg.take() {
439 handle.abort();
440 let _ = handle.await;
442 }
443 }
444
445 let mut ch = self.channel.write().await;
446 if let Some(mut gc) = ch.take() {
447 gc.close().await?;
448 }
449
450 let mut target = self.current_target.write().await;
451 *target = None;
452
453 Ok(())
454 }
455
456 pub fn service_name(&self) -> &str {
458 &self.service_name
459 }
460
461 pub async fn current_target(&self) -> Option<String> {
463 self.current_target.read().await.clone()
464 }
465
466 pub fn grpc_config(&self) -> &GrpcClientConfig {
468 &self.config.grpc
469 }
470
471 pub fn config(&self) -> &DiscoveryChannelConfig {
473 &self.config
474 }
475}
476
477fn config_for_target(base_config: &GrpcClientConfig, target: &str) -> GrpcClientConfig {
478 let mut config = base_config.clone();
479 config.target = target.to_string();
480 config
481}
482
483async fn connect_grpc_channel(
484 connector: &Arc<dyn ChannelConnector>,
485 base_config: &GrpcClientConfig,
486 target: &str,
487) -> AppResult<GrpcChannel> {
488 connector
489 .connect(config_for_target(base_config, target))
490 .await
491}
492
493#[cfg(test)]
494mod tests {
495 use std::collections::VecDeque;
496
497 use super::*;
498 use async_trait::async_trait;
499 use parking_lot::Mutex;
500 use rskit_discovery::{Watcher, instance::ServiceInstance};
501 use tokio::sync::mpsc;
502
503 struct MockDiscovery;
504
505 #[async_trait]
506 impl Discovery for MockDiscovery {
507 async fn resolve(&self, service: &str) -> AppResult<Vec<ServiceInstance>> {
508 if service == "valid-service" {
509 Ok(vec![ServiceInstance {
510 id: "test-instance".to_string(),
511 name: service.to_string(),
512 address: "localhost".to_string(),
513 port: 9090,
514 healthy: true,
515 weight: 1,
516 tags: vec![],
517 metadata: Default::default(),
518 }])
519 } else {
520 Err(rskit_errors::AppError::new(
521 rskit_errors::ErrorCode::ServiceUnavailable,
522 format!("service not found: {}", service),
523 ))
524 }
525 }
526 }
527
528 #[tokio::test]
529 async fn resolve_does_not_mutate_connected_target_cache() {
530 let discovery = Arc::new(MockDiscovery);
531 let ch = DiscoveryChannel::new(
532 discovery,
533 "valid-service",
534 GrpcClientConfig::new("localhost:50051"),
535 );
536
537 let target = ch.resolve().await.unwrap();
538 assert_eq!(target, "localhost:9090");
539 assert!(ch.current_target().await.is_none());
540 }
541
542 #[tokio::test]
543 async fn test_resolve_success() {
544 let discovery = Arc::new(MockDiscovery);
545 let ch = DiscoveryChannel::new(
546 discovery,
547 "valid-service",
548 GrpcClientConfig::new("localhost:50051"),
549 );
550
551 let target = ch.resolve().await;
552 assert!(target.is_ok());
553 assert_eq!(target.unwrap(), "localhost:9090");
554 }
555
556 #[tokio::test]
557 async fn test_resolve_failure() {
558 let discovery = Arc::new(MockDiscovery);
559 let ch = DiscoveryChannel::new(
560 discovery,
561 "unknown-service",
562 GrpcClientConfig::new("localhost:50051"),
563 );
564
565 let target = ch.resolve().await;
566 assert!(target.is_err());
567 }
568
569 #[tokio::test]
570 async fn test_with_watcher_constructor() {
571 let discovery = Arc::new(MockDiscovery);
572 let config = DiscoveryChannelConfig::new(GrpcClientConfig::new("localhost:50051"))
573 .with_resolve_interval(Duration::from_secs(5));
574
575 let ch = DiscoveryChannel::with_watcher(discovery, None, "valid-service", config);
577
578 let target = ch.resolve().await;
579 assert!(target.is_ok());
580 assert_eq!(target.unwrap(), "localhost:9090");
581 }
582
583 #[tokio::test]
584 async fn test_discovery_channel_config_defaults() {
585 let cfg = DiscoveryChannelConfig::new(GrpcClientConfig::new("localhost:50051"));
586 assert_eq!(cfg.resolve_interval, Duration::from_secs(10));
587 assert_eq!(cfg.grpc.target, "localhost:50051");
588 }
589
590 #[test]
591 fn discovery_channel_config_can_be_built_from_grpc_config() {
592 let cfg = DiscoveryChannelConfig::from(GrpcClientConfig::new("seed:50051"))
593 .with_resolve_interval(Duration::from_millis(25));
594
595 assert_eq!(cfg.grpc.target, "seed:50051");
596 assert_eq!(cfg.resolve_interval, Duration::from_millis(25));
597 }
598
599 #[tokio::test]
600 async fn test_close_without_background() {
601 let discovery = Arc::new(MockDiscovery);
602 let mut ch = DiscoveryChannel::new(
603 discovery,
604 "valid-service",
605 GrpcClientConfig::new("localhost:50051"),
606 );
607
608 let _ = ch.resolve().await;
610 let result = ch.close().await;
611 assert!(result.is_ok());
612 assert!(ch.current_target().await.is_none());
613 }
614
615 #[tokio::test]
616 async fn test_start_background_polling() {
617 tokio::time::pause();
618
619 let discovery = Arc::new(StaticDiscovery::new("127.0.0.1:9090"));
620 let connector = Arc::new(MockConnector::new([false, true]));
621 let config = DiscoveryChannelConfig::new(GrpcClientConfig::new("localhost:50051"))
622 .with_resolve_interval(Duration::from_millis(50));
623
624 let mut ch =
625 DiscoveryChannel::with_connector(discovery, None, "valid-service", config, connector);
626
627 ch.start_background().await.unwrap();
629
630 tokio::time::advance(Duration::from_millis(75)).await;
631 tokio::task::yield_now().await;
632 assert!(ch.current_target().await.is_none());
633
634 let mut target = None;
635 for _ in 0..4 {
636 tokio::time::advance(Duration::from_millis(50)).await;
637 tokio::task::yield_now().await;
638 target = ch.current_target().await;
639 if target.is_some() {
640 break;
641 }
642 }
643
644 assert_eq!(target, Some("127.0.0.1:9090".to_string()));
646
647 ch.close().await.unwrap();
649 }
650
651 #[tokio::test]
652 async fn refresh_connects_initial_target_without_false_no_change() {
653 let discovery = Arc::new(SequenceDiscovery::new(["127.0.0.1:9090"]));
654 let connector = Arc::new(MockConnector::new([true]));
655 let ch = DiscoveryChannel::with_connector(
656 discovery,
657 None,
658 "valid-service",
659 DiscoveryChannelConfig::new(GrpcClientConfig::new("localhost:50051")),
660 connector,
661 );
662
663 let changed = ch.refresh().await.unwrap();
664 assert!(changed);
665 assert_eq!(
666 ch.current_target().await,
667 Some("127.0.0.1:9090".to_string())
668 );
669 }
670
671 #[tokio::test]
672 async fn refresh_reports_no_change_after_successful_connection() {
673 let discovery = Arc::new(SequenceDiscovery::new(["127.0.0.1:9090", "127.0.0.1:9090"]));
674 let connector = Arc::new(MockConnector::new([true]));
675 let ch = DiscoveryChannel::with_connector(
676 discovery,
677 None,
678 "valid-service",
679 DiscoveryChannelConfig::new(GrpcClientConfig::new("localhost:50051")),
680 connector,
681 );
682
683 assert!(ch.refresh().await.unwrap());
684 assert!(!ch.refresh().await.unwrap());
685 }
686
687 #[tokio::test]
688 async fn resolve_errors_when_discovery_returns_no_instances() {
689 let discovery = Arc::new(EmptyDiscovery);
690 let ch = DiscoveryChannel::new(
691 discovery,
692 "empty-service",
693 GrpcClientConfig::new("localhost:50051"),
694 );
695
696 let err = ch.resolve().await.unwrap_err();
697
698 assert_eq!(err.code(), rskit_errors::ErrorCode::ServiceUnavailable);
699 assert!(err.to_string().contains("no instances"));
700 }
701
702 #[tokio::test]
703 async fn refresh_propagates_connector_failure_without_caching_target() {
704 let discovery = Arc::new(SequenceDiscovery::new(["127.0.0.1:9090"]));
705 let connector = Arc::new(MockConnector::new([false]));
706 let ch = DiscoveryChannel::with_connector(
707 discovery,
708 None,
709 "valid-service",
710 DiscoveryChannelConfig::new(GrpcClientConfig::new("localhost:50051")),
711 connector,
712 );
713
714 let err = ch.refresh().await.unwrap_err();
715
716 assert_eq!(err.code(), rskit_errors::ErrorCode::ServiceUnavailable);
717 assert!(ch.current_target().await.is_none());
718 }
719
720 #[tokio::test]
721 async fn watcher_background_reconnects_on_non_empty_updates_and_skips_empty_updates() {
722 let discovery = Arc::new(StaticDiscovery::new("127.0.0.1:9090"));
723 let (watcher, tx) = ChannelWatcher::new();
724 let connector = Arc::new(MockConnector::new([true]));
725 let mut ch = DiscoveryChannel::with_connector(
726 discovery,
727 Some(Arc::new(watcher)),
728 "valid-service",
729 DiscoveryChannelConfig::new(GrpcClientConfig::new("seed:50051")),
730 connector,
731 );
732
733 ch.start_background().await.unwrap();
734 tx.send(Vec::new()).await.expect("send empty update");
735 tokio::task::yield_now().await;
736 assert!(ch.current_target().await.is_none());
737
738 tx.send(vec![service_instance("valid-service", "127.0.0.1:9443")])
739 .await
740 .expect("send service update");
741 for _ in 0..4 {
742 tokio::task::yield_now().await;
743 if ch.current_target().await.is_some() {
744 break;
745 }
746 }
747
748 assert_eq!(
749 ch.current_target().await,
750 Some("127.0.0.1:9443".to_string())
751 );
752 assert_eq!(ch.service_name(), "valid-service");
753 assert_eq!(ch.grpc_config().target, "seed:50051");
754 assert_eq!(ch.config().grpc.target, "seed:50051");
755 ch.start_background().await.unwrap();
756 ch.close().await.unwrap();
757 assert!(ch.current_target().await.is_none());
758 }
759
760 struct StaticDiscovery {
761 target: String,
762 }
763
764 impl StaticDiscovery {
765 fn new(target: &str) -> Self {
766 Self {
767 target: target.to_string(),
768 }
769 }
770 }
771
772 #[async_trait]
773 impl Discovery for StaticDiscovery {
774 async fn resolve(&self, service: &str) -> AppResult<Vec<ServiceInstance>> {
775 Ok(vec![service_instance(service, &self.target)])
776 }
777 }
778
779 struct EmptyDiscovery;
780
781 #[async_trait]
782 impl Discovery for EmptyDiscovery {
783 async fn resolve(&self, _service: &str) -> AppResult<Vec<ServiceInstance>> {
784 Ok(Vec::new())
785 }
786 }
787
788 struct SequenceDiscovery {
789 targets: Mutex<VecDeque<String>>,
790 }
791
792 impl SequenceDiscovery {
793 fn new<const N: usize>(targets: [&str; N]) -> Self {
794 Self {
795 targets: Mutex::new(targets.into_iter().map(str::to_owned).collect()),
796 }
797 }
798 }
799
800 #[async_trait]
801 impl Discovery for SequenceDiscovery {
802 async fn resolve(&self, service: &str) -> AppResult<Vec<ServiceInstance>> {
803 let target = {
804 let mut targets = self.targets.lock();
805 match targets.len() {
806 0 => "127.0.0.1:9090".to_string(),
807 1 => targets.front().cloned().unwrap(),
808 _ => targets.pop_front().unwrap(),
809 }
810 };
811 Ok(vec![service_instance(service, &target)])
812 }
813 }
814
815 struct MockConnector {
816 outcomes: Mutex<VecDeque<bool>>,
817 }
818
819 impl MockConnector {
820 fn new<const N: usize>(outcomes: [bool; N]) -> Self {
821 Self {
822 outcomes: Mutex::new(outcomes.into_iter().collect()),
823 }
824 }
825 }
826
827 #[async_trait]
828 impl ChannelConnector for MockConnector {
829 async fn connect(&self, config: GrpcClientConfig) -> AppResult<GrpcChannel> {
830 let should_succeed = self.outcomes.lock().pop_front().unwrap_or(true);
831 if should_succeed {
832 Ok(GrpcChannel::new(config))
833 } else {
834 Err(rskit_errors::AppError::new(
835 rskit_errors::ErrorCode::ServiceUnavailable,
836 format!("mock connect failed for {}", config.target),
837 ))
838 }
839 }
840 }
841
842 fn service_instance(service: &str, endpoint: &str) -> ServiceInstance {
843 let (address, port) = endpoint.split_once(':').unwrap();
844 ServiceInstance {
845 id: format!("{service}-{endpoint}"),
846 name: service.to_string(),
847 address: address.to_string(),
848 port: port.parse().unwrap(),
849 healthy: true,
850 weight: 1,
851 tags: vec![],
852 metadata: Default::default(),
853 }
854 }
855
856 struct ChannelWatcher {
857 receiver: Mutex<Option<mpsc::Receiver<Vec<ServiceInstance>>>>,
858 }
859
860 impl ChannelWatcher {
861 fn new() -> (Self, mpsc::Sender<Vec<ServiceInstance>>) {
862 let (tx, rx) = mpsc::channel(4);
863 (
864 Self {
865 receiver: Mutex::new(Some(rx)),
866 },
867 tx,
868 )
869 }
870 }
871
872 #[async_trait]
873 impl Watcher for ChannelWatcher {
874 async fn watch(&self, _service: &str) -> AppResult<mpsc::Receiver<Vec<ServiceInstance>>> {
875 self.receiver.lock().take().ok_or_else(|| {
876 rskit_errors::AppError::new(
877 rskit_errors::ErrorCode::Internal,
878 "watcher already consumed",
879 )
880 })
881 }
882 }
883}