From 5381c1236579459470fb340f53d1fc56fdcf8f8d Mon Sep 17 00:00:00 2001
From: holmes1412 <liyingxin1412@gmail.com>
Date: Mon, 25 Jan 2021 16:53:55 +0800
Subject: [PATCH] add UPSWeightedRandomPolicy __add_server() and
 __remove_server()

---
 src/manager/UPSPolicy.cc | 34 +++++++++++++++++++++++++++++-----
 src/manager/UPSPolicy.h  |  5 ++++-
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/src/manager/UPSPolicy.cc b/src/manager/UPSPolicy.cc
index dc408178..bc8981e4 100644
--- a/src/manager/UPSPolicy.cc
+++ b/src/manager/UPSPolicy.cc
@@ -511,8 +511,8 @@ void UPSGroupPolicy::__add_server(EndpointAddress *addr)
 	}
 
 	pthread_mutex_lock(&group->mutex);
-	this->recover_one_server(addr);
 	addr->group = group;
+	this->recover_one_server(addr);
 	if (addr->params.server_type == 0)
 	{
 		group->mains.push_back(addr);
@@ -614,14 +614,39 @@ const EndpointAddress *UPSGroupPolicy::consistent_hash_with_group(unsigned int h
 	return this->check_and_get(addr);
 }
 
+void UPSWeightedRandomPolicy::__add_server(EndpointAddress *addr)
+{
+	UPSGroupPolicy::__add_server(addr);
+	if (addr->params.server_type == 0)
+	{
+		this->total_weight += addr->params.weight;
+		this->servers.push_back(addr);
+	}
+	return;
+}
+
+int UPSWeightedRandomPolicy::__remove_server(const std::string& address)
+{
+	const auto map_it = this->server_map.find(address);
+
+	if (map_it != this->server_map.cend())
+	{
+		for (EndpointAddress *addr : map_it->second)
+		{
+			if (addr->params.server_type == 0)
+				this->total_weight -= addr->params.weight;
+		}
+	}
+
+	return UPSGroupPolicy::__remove_server(address);
+}
+
 const EndpointAddress *UPSWeightedRandomPolicy::first_stradegy(const ParsedURI& uri)
 {
 	int x = 0;
 	int s = 0;
 	size_t idx;
-	int temp_weight;
-
-	temp_weight = this->total_weight;
+	int temp_weight = this->total_weight;
 
 	if (temp_weight > 0)
 		x = rand() % temp_weight;
@@ -661,7 +686,6 @@ const EndpointAddress *UPSWeightedRandomPolicy::another_stradegy(const ParsedURI
 	return this->check_and_get(addr);
 }
 
-
 void UPSWeightedRandomPolicy::recover_one_server(const EndpointAddress *addr)
 {
 	this->nalives++;
diff --git a/src/manager/UPSPolicy.h b/src/manager/UPSPolicy.h
index 71e8cfb5..58534990 100644
--- a/src/manager/UPSPolicy.h
+++ b/src/manager/UPSPolicy.h
@@ -168,10 +168,11 @@ private:
 	}
 	// override: select() add_server() remove_server()
 	virtual bool select(const ParsedURI& uri, EndpointAddress **addr);
+
+protected:
 	virtual void __add_server(EndpointAddress *addr);
 	virtual int __remove_server(const std::string& address);
 
-protected:
 	const EndpointAddress *consistent_hash_with_group(unsigned int hash) const;
 
 	inline const EndpointAddress *check_and_get(const EndpointAddress *addr) const // check_get_weak
@@ -212,6 +213,8 @@ protected:
 private:
 	virtual void recover_one_server(const EndpointAddress *addr);
 	virtual void fuse_one_server(const EndpointAddress *addr);
+	virtual void __add_server(EndpointAddress *addr);
+	virtual int __remove_server(const std::string& address);
 };
 
 using select_t = std::function<unsigned int (const char *, const char *, const char *)>;