@@ -38,6 +38,7 @@ use crate::{
38
38
use base64:: write:: EncoderWriter as Base64Encoder ;
39
39
use bytes:: BytesMut ;
40
40
use serde:: Serialize ;
41
+ use serde_json:: Value ;
41
42
use std:: {
42
43
error, fmt,
43
44
fmt:: Debug ,
@@ -288,7 +289,7 @@ impl Default for TransportBuilder {
288
289
/// A connection to an Elasticsearch node, used to send an API request
289
290
#[ derive( Debug , Clone ) ]
290
291
pub struct Connection {
291
- url : Url ,
292
+ url : Arc < Url > ,
292
293
}
293
294
294
295
impl Connection {
@@ -303,8 +304,14 @@ impl Connection {
303
304
url
304
305
} ;
305
306
307
+ let url = Arc :: new ( url) ;
308
+
306
309
Self { url }
307
310
}
311
+
312
+ pub fn url ( & self ) -> Arc < Url > {
313
+ self . url . clone ( )
314
+ }
308
315
}
309
316
310
317
/// A HTTP transport responsible for making the API requests to Elasticsearch,
@@ -365,27 +372,22 @@ impl Transport {
365
372
Ok ( transport)
366
373
}
367
374
368
- /// Creates an asynchronous request that can be awaited
369
- pub async fn send < B , Q > (
375
+ pub fn request_builder < B , Q > (
370
376
& self ,
377
+ connection : & Connection ,
371
378
method : Method ,
372
379
path : & str ,
373
380
headers : HeaderMap ,
374
381
query_string : Option < & Q > ,
375
382
body : Option < B > ,
376
383
timeout : Option < Duration > ,
377
- ) -> Result < Response , Error >
384
+ ) -> Result < reqwest :: RequestBuilder , Error >
378
385
where
379
386
B : Body ,
380
387
Q : Serialize + ?Sized ,
381
388
{
382
- if self . conn_pool . reseedable ( ) {
383
- // Reseed nodes
384
- println ! ( "Reseeding!" ) ;
385
- }
386
- let connection = self . conn_pool . next ( ) ;
387
- let url = connection. url . join ( path. trim_start_matches ( '/' ) ) ?;
388
389
let reqwest_method = self . method ( method) ;
390
+ let url = connection. url . join ( path. trim_start_matches ( '/' ) ) ?;
389
391
let mut request_builder = self . client . request ( reqwest_method, url) ;
390
392
391
393
if let Some ( t) = timeout {
@@ -442,6 +444,71 @@ impl Transport {
442
444
if let Some ( q) = query_string {
443
445
request_builder = request_builder. query ( q) ;
444
446
}
447
+ Ok ( request_builder)
448
+ }
449
+
450
+ /// Creates an asynchronous request that can be awaited
451
+ pub async fn send < B , Q > (
452
+ & self ,
453
+ method : Method ,
454
+ path : & str ,
455
+ headers : HeaderMap ,
456
+ query_string : Option < & Q > ,
457
+ body : Option < B > ,
458
+ timeout : Option < Duration > ,
459
+ ) -> Result < Response , Error >
460
+ where
461
+ B : Body ,
462
+ Q : Serialize + ?Sized ,
463
+ {
464
+ let connection = self . conn_pool . next ( ) ;
465
+
466
+ // Threads will execute against old connection pool during reseed
467
+ if self . conn_pool . reseedable ( ) {
468
+ // Set as reseeding prevents another thread from attempting
469
+ // to reseed during es request and reseed
470
+ self . conn_pool . reseeding ( ) ;
471
+
472
+ let scheme = & connection. url . scheme ( ) ;
473
+ // Build node info request
474
+ let node_request = self . request_builder (
475
+ & connection,
476
+ Method :: Get ,
477
+ "_nodes/_all/http" ,
478
+ headers. clone ( ) ,
479
+ None :: < & Q > ,
480
+ None :: < B > ,
481
+ timeout,
482
+ ) ?;
483
+ let resp = node_request. send ( ) . await ?;
484
+ let json: Value = resp. json ( ) . await ?;
485
+ let connections: Vec < Connection > = json[ "nodes" ]
486
+ . as_object ( )
487
+ . unwrap ( )
488
+ . iter ( )
489
+ . map ( |h| {
490
+ let url = format ! (
491
+ "{}://{}" ,
492
+ scheme,
493
+ h. 1 [ "http" ] [ "publish_address" ] . as_str( ) . unwrap( )
494
+ ) ;
495
+ let url = Url :: parse ( & url) . unwrap ( ) ;
496
+ Connection :: new ( url)
497
+ } )
498
+ . collect ( ) ;
499
+ self . conn_pool . reseed ( connections) ;
500
+ }
501
+
502
+ let connection = self . conn_pool . next ( ) ;
503
+ let request_builder = self . request_builder (
504
+ & connection,
505
+ method,
506
+ path,
507
+ headers,
508
+ query_string,
509
+ body,
510
+ timeout,
511
+ ) ?;
445
512
446
513
let response = request_builder. send ( ) . await ;
447
514
match response {
@@ -471,6 +538,9 @@ pub trait ConnectionPool: Debug + dyn_clone::DynClone + Sync + Send {
471
538
false
472
539
}
473
540
541
+ // NOOP
542
+ fn reseeding ( & self ) { }
543
+
474
544
// NOOP by default
475
545
fn reseed ( & self , _connection : Vec < Connection > ) { }
476
546
}
@@ -629,70 +699,80 @@ impl ConnectionPool for CloudConnectionPool {
629
699
630
700
/// A Connection Pool that manages a static connection of nodes
631
701
#[ derive( Debug , Clone ) ]
632
- pub struct MultiNodeConnectionPool < TStrategy = RoundRobin > {
702
+ pub struct MultiNodeConnectionPool < LoadBalancingStrategy = RoundRobin > {
633
703
inner : Arc < RwLock < MultiNodeConnectionPoolInner > > ,
634
- wait : Option < Duration > ,
635
- strategy : TStrategy ,
704
+ reseed_frequency : Option < Duration > ,
705
+ load_balancing_strategy : LoadBalancingStrategy ,
636
706
}
637
707
638
708
#[ derive( Debug , Clone ) ]
639
709
pub struct MultiNodeConnectionPoolInner {
710
+ reseeding : bool ,
640
711
last_update : Option < Instant > ,
641
712
connections : Vec < Connection > ,
642
713
}
643
714
644
715
impl < TStrategy > ConnectionPool for MultiNodeConnectionPool < TStrategy >
645
716
where
646
- TStrategy : Strategy + Clone ,
717
+ TStrategy : LoadBalancingStrategy + Clone ,
647
718
{
648
719
fn next ( & self ) -> Connection {
649
720
let inner = self . inner . read ( ) . expect ( "lock poisoned" ) ;
650
- self . strategy . try_next ( & inner. connections ) . unwrap ( )
721
+ self . load_balancing_strategy
722
+ . try_next ( & inner. connections )
723
+ . unwrap ( )
651
724
}
652
725
653
726
fn reseedable ( & self ) -> bool {
654
727
let inner = self . inner . read ( ) . expect ( "lock poisoned" ) ;
655
- let wait = match self . wait {
728
+ let reseed_frequency = match self . reseed_frequency {
656
729
Some ( wait) => wait,
657
730
None => return false ,
658
731
} ;
659
732
let last_update_is_stale = inner
660
733
. last_update
661
734
. as_ref ( )
662
- . map ( |last_update| last_update. elapsed ( ) > wait) ;
663
- last_update_is_stale. unwrap_or ( true )
735
+ . map ( |last_update| last_update. elapsed ( ) > reseed_frequency) ;
736
+ last_update_is_stale. unwrap_or ( true ) && !inner. reseeding
737
+ }
738
+
739
+ fn reseeding ( & self ) {
740
+ let mut inner = self . inner . write ( ) . expect ( "Lock Poisoned" ) ;
741
+ inner. reseeding = true
664
742
}
665
743
666
744
fn reseed ( & self , mut connection : Vec < Connection > ) {
667
745
let mut inner = self . inner . write ( ) . expect ( "lock poisoned" ) ;
668
746
inner. last_update = Some ( Instant :: now ( ) ) ;
669
747
inner. connections . clear ( ) ;
670
748
inner. connections . append ( & mut connection) ;
749
+ inner. reseeding = false ;
671
750
}
672
751
}
673
752
674
753
impl MultiNodeConnectionPool < RoundRobin > {
675
754
/** Use a round-robin strategy for balancing traffic over the given set of nodes. */
676
- pub fn round_robin ( urls : Vec < Url > , wait : Option < Duration > ) -> Self {
755
+ pub fn round_robin ( urls : Vec < Url > , reseed_frequency : Option < Duration > ) -> Self {
677
756
let connections = urls. into_iter ( ) . map ( Connection :: new) . collect ( ) ;
678
757
679
758
let inner: Arc < RwLock < MultiNodeConnectionPoolInner > > =
680
759
Arc :: new ( RwLock :: new ( MultiNodeConnectionPoolInner {
760
+ reseeding : false ,
681
761
last_update : None ,
682
762
connections,
683
763
} ) ) ;
684
764
685
- let strategy = RoundRobin :: default ( ) ;
765
+ let load_balancing_strategy = RoundRobin :: default ( ) ;
686
766
Self {
687
767
inner,
688
- strategy ,
689
- wait ,
768
+ load_balancing_strategy ,
769
+ reseed_frequency ,
690
770
}
691
771
}
692
772
}
693
773
694
774
/** The strategy selects an address from a given collection. */
695
- pub trait Strategy : Send + Sync + Debug {
775
+ pub trait LoadBalancingStrategy : Send + Sync + Debug {
696
776
/** Try get the next connection. */
697
777
fn try_next < ' a > ( & self , connections : & ' a [ Connection ] ) -> Result < Connection , Error > ;
698
778
}
@@ -711,7 +791,7 @@ impl Default for RoundRobin {
711
791
}
712
792
}
713
793
714
- impl Strategy for RoundRobin {
794
+ impl LoadBalancingStrategy for RoundRobin {
715
795
fn try_next < ' a > ( & self , connections : & ' a [ Connection ] ) -> Result < Connection , Error > {
716
796
if connections. is_empty ( ) {
717
797
Err ( crate :: error:: lib ( "Connection list empty" ) )
0 commit comments