feat: add Add/Delete methods to round-robin impl

client/client.go

@@ -96,6 +96,14 @@ func New(c *Config) (*Client, error) {
 	}, nil
 }
 
+func (c *Client) AddNode(addr string) {
+	c.roundRobin.Add(addr)
+}
+
+func (c *Client) DeleteNode(addr string) error {
+	return c.roundRobin.Delete(addr)
+}
+
 // Ping sends a dummy protocol messsage to the given host. This is useful to
 // measure RTT between hosts. It also can be used as aliveness check.
 func (c *Client) Ping(addr string) error {

client/round_robin.go

@@ -14,9 +14,12 @@
 
 package client
 
-import "sync"
+import (
+	"errors"
+	"sync"
+)
 
-// roundRobin implements quite simple round-robin algorithm to distribute load fairly between servers.
+// roundRobin implements quite simple round-robin scheduling algorithm to distribute load fairly between servers.
 type roundRobin struct {
 	sync.Mutex
 
@@ -45,3 +48,37 @@ func (r *roundRobin) Get() string {
 	r.current++
 	return addr
 }
+
+// Add adds a new address to the Round-Robin scheduler.
+func (r *roundRobin) Add(addr string) {
+	r.Lock()
+	defer r.Unlock()
+
+	r.addrs = append(r.addrs, addr)
+}
+
+// Delete deletes an address from the Round-Robin scheduler.
+func (r *roundRobin) Delete(addr string) error {
+	r.Lock()
+	defer r.Unlock()
+
+	for i, item := range r.addrs {
+		if item == addr {
+			if len(r.addrs) == 1 {
+				return errors.New("address cannot be removed")
+			}
+			r.addrs = append(r.addrs[:i], r.addrs[i+1:]...)
+			return nil
+		}
+	}
+	// not found
+	return nil
+}
+
+// Length returns the count of addresses
+func (r *roundRobin) Length() int {
+	r.Lock()
+	defer r.Unlock()
+
+	return len(r.addrs)
+}

client/round_robin_test.go

@@ -0,0 +1,91 @@
+// Copyright 2018-2020 Burak Sezer
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package client
+
+import (
+	"testing"
+)
+
+func TestRoundRobin(t *testing.T) {
+	addrs := []string{"127.0.0.1:2323", "127.0.0.1:4556", "127.0.0.1:7889"}
+	r := newRoundRobin(addrs)
+
+	t.Run("Get", func(t *testing.T) {
+		items := make(map[string]int)
+		for i := 0; i < r.Length(); i++ {
+			item := r.Get()
+			items[item]++
+		}
+		if len(items) != r.Length() {
+			t.Fatalf("Expected item count: %d. Got: %d", r.Length(), len(items))
+		}
+	})
+
+	t.Run("Add", func(t *testing.T) {
+		addr := "127.0.0.1:3320"
+		r.Add(addr)
+		items := make(map[string]int)
+		for i := 0; i < r.Length(); i++ {
+			item := r.Get()
+			items[item]++
+		}
+		if _, ok := items[addr]; !ok {
+			t.Fatalf("Addr not processed: %s", addr)
+		}
+		if len(items) != r.Length() {
+			t.Fatalf("Expected item count: %d. Got: %d", r.Length(), len(items))
+		}
+	})
+
+	t.Run("Delete", func(t *testing.T) {
+		addr := "127.0.0.1:7889"
+		if err := r.Delete(addr); err != nil {
+			t.Fatalf("Expected nil. Got: %v", err)
+		}
+
+		items := make(map[string]int)
+		for i := 0; i < r.Length(); i++ {
+			item := r.Get()
+			items[item]++
+		}
+		if _, ok := items[addr]; ok {
+			t.Fatalf("Address stil exists: %s", addr)
+		}
+		if len(items) != r.Length() {
+			t.Fatalf("Expected item count: %d. Got: %d", r.Length(), len(items))
+		}
+	})
+}
+
+
+
+func TestRoundRobin_Delete_NonExistent(t *testing.T) {
+	addrs := []string{"127.0.0.1:2323", "127.0.0.1:4556", "127.0.0.1:7889"}
+	r := newRoundRobin(addrs)
+
+	var fresh []string
+	fresh = append(fresh, addrs...)
+	for i, addr := range fresh {
+		if i+1 == len(addrs) {
+			if err := r.Delete(addr); err == nil {
+				t.Fatal("Expected an error. Got: nil")
+			}
+		} else {
+			if err := r.Delete(addr); err != nil {
+				t.Fatalf("Expected nil. Got: %v", err)
+			}
+		}
+	}
+}