feat:支持负载均衡

Author: chuntaojun

src/core/config/consumer.rs

@@ -45,6 +45,7 @@ pub struct ServiceRouterPluginConfig {
 #[derive(Deserialize, Debug)]
 #[serde(rename_all = "camelCase", deny_unknown_fields)]
 pub struct LoadBalancerConfig {
+    pub default_policy: String,
     pub plugins: Option<Vec<String>>,
 }
 

src/core/context.rs

@@ -21,6 +21,7 @@ use crate::core::engine::Engine;
 use crate::core::model::error::{ErrorCode, PolarisError};
 
 pub struct SDKContext {
+    pub conf: Arc<Configuration>,
     engine: Arc<Engine>,
 }
 
@@ -53,12 +54,14 @@ impl SDKContext {
     // create_by_configuration
     pub fn create_by_configuration(cfg: Configuration) -> Result<SDKContext, PolarisError> {
         let start_time = std::time::Instant::now();
-        let ret = Engine::new(cfg);
+        let cfg = Arc::new(cfg);
+        let ret = Engine::new(cfg.clone());
         tracing::info!("create engine cost: {:?}", start_time.elapsed());
         if ret.is_err() {
             return Err(ret.err().unwrap());
         }
         return Ok(Self {
+            conf: cfg,
             engine: Arc::new(ret.ok().unwrap()),
         });
     }

src/core/engine.rs

@@ -17,6 +17,7 @@ use std::collections::HashMap;
 use std::sync::Arc;
 
 use tokio::runtime::{Builder, Runtime};
+use tokio::sync::RwLock;
 
 use crate::config::req::{
     CreateConfigFileRequest, GetConfigFileRequest, PublishConfigFileRequest,
@@ -35,8 +36,10 @@ use crate::discovery::req::{
 };
 
 use super::model::config::ConfigFile;
+use super::model::naming::ServiceInstances;
 use super::plugin::cache::{Filter, ResourceCache, ResourceListener};
 use super::plugin::connector::Connector;
+use super::plugin::loadbalance::LoadBalancer;
 use super::plugin::location::{LocationProvider, LocationSupplier};
 
 pub struct Engine
@@ -47,10 +50,11 @@ where
     local_cache: Arc<Box<dyn ResourceCache>>,
     server_connector: Arc<Box<dyn Connector>>,
     location_provider: Arc<LocationProvider>,
+    load_balancer: Arc<RwLock<HashMap<String, Arc<Box<dyn LoadBalancer>>>>>,
 }
 
 impl Engine {
-    pub fn new(conf: Configuration) -> Result<Self, PolarisError> {
+    pub fn new(arc_conf: Arc<Configuration>) -> Result<Self, PolarisError> {
         let runtime = Arc::new(
             Builder::new_multi_thread()
                 .enable_all()
@@ -60,7 +64,6 @@ impl Engine {
                 .unwrap(),
         );
 
-        let arc_conf = Arc::new(conf);
         let client_id = crate::core::plugin::plugins::acquire_client_id(arc_conf.clone());
 
         // 初始化 extensions
@@ -101,6 +104,7 @@ impl Engine {
             local_cache: Arc::new(local_cache),
             server_connector: server_connector,
             location_provider: Arc::new(location_provider),
+            load_balancer: Arc::new(RwLock::new(extension.load_loadbalancers())),
         })
     }
 
@@ -220,17 +224,12 @@ impl Engine {
 
         let svc_ins = ret.unwrap();
 
-        if only_available {
-            Ok(InstancesResponse {
-                service_info: svc_ins.service,
-                instances: svc_ins.available_instances,
-            })
-        } else {
-            Ok(InstancesResponse {
-                service_info: svc_ins.service,
-                instances: svc_ins.instances,
-            })
-        }
+        Ok(InstancesResponse {
+            instances: ServiceInstances::new(
+                svc_ins.get_service_info(),
+                svc_ins.list_instances(only_available).await,
+            ),
+        })
     }
 
     /// get_service_rule 获取服务规则
@@ -340,6 +339,11 @@ impl Engine {
         };
     }
 
+    pub async fn lookup_loadbalancer(&self, name: &str) -> Option<Arc<Box<dyn LoadBalancer>>> {
+        let lb = self.load_balancer.read().await;
+        lb.get(name).map(|lb| lb.clone())
+    }
+
     /// register_resource_listener 注册资源监听器
     pub async fn register_resource_listener(&self, listener: Arc<dyn ResourceListener>) {
         self.local_cache.register_resource_listener(listener).await;

src/core/model/cache.rs

@@ -326,6 +326,8 @@ pub struct ServiceInstancesCacheItem {
     initialized: Arc<AtomicBool>,
     pub svc_info: Service,
     pub value: Arc<RwLock<Vec<Instance>>>,
+    pub available_instances: Arc<RwLock<Vec<Instance>>>,
+    pub total_weight: u64,
     pub revision: String,
 }
 
@@ -344,11 +346,16 @@ impl ServiceInstancesCacheItem {
             initialized: Arc::new(AtomicBool::new(false)),
             svc_info: Service::default(),
             value: Arc::new(RwLock::new(Vec::new())),
+            available_instances: Arc::new(RwLock::new(Vec::new())),
+            total_weight: 0,
             revision: String::new(),
         }
     }
 
-    pub async fn list_instances(&self) -> Vec<Instance> {
+    pub async fn list_instances(&self, only_available: bool) -> Vec<Instance> {
+        if only_available {
+            return self.available_instances.read().await.clone();
+        }
         self.value.read().await.clone()
     }
 
@@ -375,6 +382,8 @@ impl Clone for ServiceInstancesCacheItem {
             initialized: self.initialized.clone(),
             svc_info: self.svc_info.clone(),
             value: self.value.clone(),
+            available_instances: self.available_instances.clone(),
+            total_weight: self.total_weight,
             revision: self.revision.clone(),
         }
     }

src/core/model/error.rs

@@ -78,7 +78,7 @@ where
 
 impl Display for PolarisError {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "{:?}", self)
+        write!(f, "err_code={:?} err_msg={}", self.err_code, self.err_msg)
     }
 }
 

src/core/model/naming.rs

@@ -45,11 +45,37 @@ pub struct ServiceInfo {
     pub revision: String,
 }
 
-#[derive(Default)]
+#[derive(Default, Clone, Debug)]
 pub struct ServiceInstances {
     pub service: ServiceInfo,
     pub instances: Vec<Instance>,
-    pub available_instances: Vec<Instance>,
+    pub total_weight: u64,
+}
+
+impl ServiceInstances {
+    pub fn get_cache_key(&self) -> String {
+        format!(
+            "Instance-namespace={}-service={}",
+            self.service.namespace, self.service.name
+        )
+    }
+
+    pub fn new(svc_info: ServiceInfo, all_ins: Vec<Instance>) -> Self {
+        let mut total_weight: u64 = 0;
+        for (_, val) in all_ins.iter().enumerate() {
+            total_weight += val.weight as u64;
+        }
+
+        Self {
+            service: svc_info,
+            instances: all_ins,
+            total_weight,
+        }
+    }
+
+    pub fn get_total_weight(&self) -> u64 {
+        self.total_weight
+    }
 }
 
 #[derive(Default, Debug, Clone)]

src/core/plugin/cache.rs

@@ -14,7 +14,9 @@
 // specific language governing permissions and limitations under the License.
 
 use crate::core::config::config::Configuration;
-use crate::core::model::cache::{EventType, ResourceEventKey, ServerEvent};
+use crate::core::model::cache::{
+    EventType, ResourceEventKey, ServerEvent, ServiceInstancesCacheItem,
+};
 use crate::core::model::config::{ConfigFile, ConfigGroup};
 use crate::core::model::error::PolarisError;
 use crate::core::model::naming::{ServiceInstances, ServiceRule, Services};
@@ -69,7 +71,7 @@ pub trait ResourceCache: Plugin {
     async fn load_service_instances(
         &self,
         filter: Filter,
-    ) -> Result<ServiceInstances, PolarisError>;
+    ) -> Result<ServiceInstancesCacheItem, PolarisError>;
     // 加载配置文件
     async fn load_config_file(&self, filter: Filter) -> Result<ConfigFile, PolarisError>;
     // 加载配置文件组
@@ -113,7 +115,7 @@ impl ResourceCache for NoopResourceCache {
     async fn load_service_instances(
         &self,
         filter: Filter,
-    ) -> Result<ServiceInstances, PolarisError> {
+    ) -> Result<ServiceInstancesCacheItem, PolarisError> {
         todo!()
     }
 

src/core/plugin/loadbalance.rs

@@ -14,6 +14,7 @@
 // specific language governing permissions and limitations under the License.
 
 use crate::core::model::{
+    error::PolarisError,
     loadbalance::Criteria,
     naming::{Instance, ServiceInstances},
 };
@@ -26,5 +27,9 @@ where
     Self: Plugin,
 {
     /// choose_instance 选择一个实例
-    fn choose_instance(&self, criteria: Criteria, instances: ServiceInstances) -> Option<Instance>;
+    fn choose_instance(
+        &self,
+        criteria: Criteria,
+        instances: ServiceInstances,
+    ) -> Result<Instance, PolarisError>;
 }

src/core/plugin/plugins.rs

@@ -23,6 +23,9 @@ use crate::core::plugin::router::ServiceRouter;
 use crate::plugins::cache::memory::memory::MemoryCache;
 use crate::plugins::connector::grpc::connector::GrpcConnector;
 use crate::plugins::filter::configcrypto::crypto::ConfigFileCryptoFilter;
+use crate::plugins::loadbalance::random::random::WeightRandomLoadbalancer;
+use crate::plugins::loadbalance::ringhash::ringhash::ConsistentHashLoadBalancer;
+use crate::plugins::loadbalance::roundrobin::roundrobin::WeightedRoundRobinBalancer;
 use std::collections::HashMap;
 use std::fmt::Display;
 use std::hash::Hash;
@@ -35,6 +38,7 @@ use tokio::runtime::Runtime;
 use super::cache::InitResourceCacheOption;
 use super::connector::InitConnectorOption;
 use super::filter::DiscoverFilter;
+use super::loadbalance::LoadBalancer;
 
 static SEQ: AtomicU64 = AtomicU64::new(1);
 
@@ -173,6 +177,15 @@ impl Extensions {
         self.config_filters = Some(Arc::new(filters));
         Ok(())
     }
+
+    pub fn load_loadbalancers(&mut self) -> HashMap<String, Arc<Box<dyn LoadBalancer>>> {
+        let mut loadbalancers = HashMap::<String, Arc<Box<dyn LoadBalancer>>>::new();
+        for (name, supplier) in self.plugin_container.load_balancers.iter() {
+            let lb = supplier();
+            loadbalancers.insert(name.clone(), Arc::new(lb));
+        }
+        loadbalancers
+    }
 }
 
 pub struct PluginContainer {
@@ -181,6 +194,7 @@ pub struct PluginContainer {
     caches: HashMap<String, fn(InitResourceCacheOption) -> Box<dyn ResourceCache>>,
     discover_filters:
         HashMap<String, fn(serde_yaml::Value) -> Result<Box<dyn DiscoverFilter>, PolarisError>>,
+    load_balancers: HashMap<String, fn() -> Box<dyn LoadBalancer>>,
 }
 
 impl Default for PluginContainer {
@@ -190,6 +204,7 @@ impl Default for PluginContainer {
             routers: Default::default(),
             caches: Default::default(),
             discover_filters: Default::default(),
+            load_balancers: Default::default(),
         };
 
         return c;
@@ -201,6 +216,7 @@ impl PluginContainer {
         self.register_resource_cache();
         self.register_connector();
         self.register_discover_filter();
+        self.register_load_balancer();
     }
 
     fn register_connector(&mut self) {
@@ -227,6 +243,18 @@ impl PluginContainer {
         }
     }
 
+    fn register_load_balancer(&mut self) {
+        let vec = vec![
+            WeightRandomLoadbalancer::builder,
+            ConsistentHashLoadBalancer::builder,
+            WeightedRoundRobinBalancer::builder,
+        ];
+        for c in vec {
+            let (supplier, name) = c();
+            self.load_balancers.insert(name, supplier);
+        }
+    }
+
     fn get_connector_supplier(
         &self,
         name: &str,

src/discovery/api.rs

@@ -95,7 +95,7 @@ where
     async fn get_one_instance(
         &self,
         req: GetOneInstanceRequest,
-    ) -> Result<InstancesResponse, PolarisError>;
+    ) -> Result<InstanceResponse, PolarisError>;
 
     /// get_health_instance 拉取健康实例
     async fn get_health_instance(