add new file

lib/lb/algo.go

@@ -0,0 +1,96 @@
+package lb
+
+import (
+	"errors"
+	"sync"
+)
+
+func GetLbAlgo(algo string) Algo {
+	// switch
+	return NewRoundRobin()
+}
+
+type Algo interface {
+	Next() (interface{}, error)
+	Append(i interface{}) error
+	Remove(i interface{}) error
+	Empty() bool
+}
+
+// rotation
+type roundRobin struct {
+	head *server
+	now  *server
+	sync.RWMutex
+}
+
+type server struct {
+	self interface{}
+	next *server
+}
+
+func NewRoundRobin() *roundRobin {
+	return &roundRobin{}
+}
+
+func (r *roundRobin) Append(i interface{}) error {
+	r.Lock()
+	defer r.Unlock()
+	if r.head == nil {
+		r.head = &server{self: i}
+		return nil
+	}
+	r.now = r.head
+	for {
+		if r.now.next == nil {
+			r.now.next = &server{self: i}
+			break
+		}
+		r.now = r.now.next
+	}
+	return nil
+}
+
+func (r *roundRobin) Remove(i interface{}) error {
+	r.Lock()
+	defer r.Unlock()
+	o := r.head
+	var last *server
+	for {
+		if o == nil {
+			return errors.New("not round")
+		}
+		if o.self == i {
+			if last == nil {
+				r.head = o.next
+			} else {
+				last.next = o.next
+			}
+			r.now = r.head
+			return nil
+		}
+		last = o
+		o = o.next
+	}
+}
+
+func (r *roundRobin) Next() (interface{}, error) {
+	r.Lock()
+	defer r.Unlock()
+	if r.head == nil {
+		return nil, errors.New("not found component")
+	}
+	if r.now == nil {
+		r.now = r.head
+	}
+	i := r.now
+	r.now = r.now.next
+	return i.self, nil
+}
+
+func (r *roundRobin) Empty() bool {
+	if r.head != nil {
+		return false
+	}
+	return true
+}

lib/lb/lb.go

@@ -0,0 +1,59 @@
+package lb
+
+import (
+	"errors"
+	"sync"
+)
+
+func NewLoadBalancer() *LoadBalancer {
+	return &LoadBalancer{
+		instances: make(map[string]Algo, 0),
+	}
+}
+
+type LoadBalancer struct {
+	instances map[string]Algo
+	Algo      string
+	sync.RWMutex
+}
+
+func (lb *LoadBalancer) SetClient(id string, instance interface{}) error {
+	lb.Lock()
+	defer lb.Unlock()
+	var l Algo
+	var ok bool
+	if l, ok = lb.instances[id]; !ok {
+		l = GetLbAlgo(lb.Algo)
+		lb.instances[id] = l
+	}
+	return l.Append(instance)
+}
+
+func (lb *LoadBalancer) RemoveClient(id string, instance interface{}) error {
+	lb.Lock()
+	defer lb.Unlock()
+	var l Algo
+	var ok bool
+	if l, ok = lb.instances[id]; !ok {
+		return errors.New("not found Client")
+	}
+	err := l.Remove(instance)
+	if l.Empty() {
+		delete(lb.instances, id)
+	}
+	return err
+}
+
+func (lb *LoadBalancer) GetClient(id string) (interface{}, error) {
+	lb.Lock()
+	l, ok := lb.instances[id]
+	lb.Unlock()
+	if !ok {
+		return nil, errors.New("client can not found")
+	}
+	i, err := l.Next()
+	if err != nil {
+		return nil, err
+	}
+	return i, nil
+}

lib/lb/lb_test.go

@@ -0,0 +1,49 @@
+package lb
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+type test struct {
+	id int
+}
+
+func TestLb(t *testing.T) {
+	lb := NewLoadBalancer()
+
+	for i := 1; i <= 100; i++ {
+		err := lb.SetClient("test", &test{id: i})
+		assert.NoError(t, err)
+	}
+
+	m := make(map[int]bool)
+
+	for i := 1; i <= 100; i++ {
+		tt, err := lb.GetClient("test")
+		assert.NoError(t, err)
+		if _, ok := m[tt.(*test).id]; ok {
+			t.Fail()
+		}
+		m[tt.(*test).id] = true
+	}
+
+	for i := 1; i <= 50; i++ {
+		tt, err := lb.GetClient("test")
+		assert.NoError(t, err)
+		err = lb.RemoveClient("test", tt)
+		assert.NoError(t, err)
+	}
+
+	m = make(map[int]bool)
+	for i := 1; i <= 50; i++ {
+		tt, err := lb.GetClient("test")
+		assert.NoError(t, err)
+		if _, ok := m[tt.(*test).id]; ok {
+			t.Fail()
+		}
+		m[tt.(*test).id] = true
+	}
+
+
+}