diff --git a/client/client.go b/client/client.go index 7123124d..2be75b39 100644 --- a/client/client.go +++ b/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 { diff --git a/client/round_robin.go b/client/round_robin.go index 5a59afcd..5222a5ce 100644 --- a/client/round_robin.go +++ b/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) +} diff --git a/client/round_robin_test.go b/client/round_robin_test.go new file mode 100644 index 00000000..4099ca76 --- /dev/null +++ b/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) + } + } + } +} \ No newline at end of file