From 624b4e0789df42bd762d8c07800fea1c398845c8 Mon Sep 17 00:00:00 2001
From: Cam Sweeney <cpsweene@gmail.com>
Date: Thu, 10 Mar 2022 11:02:19 -0800
Subject: [PATCH] support adding/removing clients from LBClient

---
 lbclient.go | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/lbclient.go b/lbclient.go
index 46d14b75cb..7dd4b9ab13 100644
--- a/lbclient.go
+++ b/lbclient.go
@@ -50,6 +50,7 @@ type LBClient struct {
 	cs []*lbClient
 
 	once sync.Once
+	mu   sync.Mutex
 }
 
 // DefaultLBClientTimeout is the default request timeout used by LBClient
@@ -91,6 +92,37 @@ func (cc *LBClient) init() {
 	}
 }
 
+// AddClient adds a new HostClient to the balanced clients
+// returns the new total number of clients
+func (cc *LBClient) AddClient(c *HostClient) int {
+	cc.mu.Lock()
+	cc.cs = append(cc.cs, &lbClient{
+		c:           c,
+		healthCheck: cc.HealthCheck,
+	})
+	cc.mu.Unlock()
+	return len(cc.cs)
+}
+
+// RemoveClients removes clients using the provided callback
+// if rc returns true, the passed client will be removed
+// returns the new total number of clients
+func (cc *LBClient) RemoveClients(rc func(*HostClient) bool) int {
+	cc.mu.Lock()
+	for idx, cs := range cc.cs {
+		// ignore non HostClient BalancingClients
+		hc, ok := cs.c.(*HostClient)
+		if !ok {
+			continue
+		}
+		if rc(hc) {
+			cc.cs = append(cc.cs[:idx], cc.cs[idx+1])
+		}
+	}
+	cc.mu.Unlock()
+	return len(cc.cs)
+}
+
 func (cc *LBClient) get() *lbClient {
 	cc.once.Do(cc.init)