Add load balancer for backend nodes

backend/loadbalancer.go

@@ -0,0 +1,71 @@
+package backend
+
+import (
+	"fmt"
+	"sync/atomic"
+
+	"github.com/ksysoev/wasabi"
+)
+
+var ErrNotEnoughBackends = fmt.Errorf("load balancer requires at least 2 backends")
+
+const minRequiredBackends = 2
+
+type LoadBalancerNode struct {
+	backend wasabi.RequestHandler
+	counter atomic.Int32
+}
+
+type LoadBalancer struct {
+	backends []*LoadBalancerNode
+}
+
+// NewLoadBalancer creates a new instance of LoadBalancer with the given backends.
+// It takes a slice of RequestHandler as a parameter and returns a new instance of LoadBalancer.
+func NewLoadBalancer(backends []wasabi.RequestHandler) (*LoadBalancer, error) {
+	if len(backends) < minRequiredBackends {
+		return nil, ErrNotEnoughBackends
+	}
+
+	nodes := make([]*LoadBalancerNode, len(backends))
+
+	for i, backend := range backends {
+		nodes[i] = &LoadBalancerNode{
+			backend: backend,
+			counter: atomic.Int32{},
+		}
+	}
+
+	return &LoadBalancer{
+		backends: nodes,
+	}, nil
+}
+
+// Handle handles the incoming request by sending it to the least busy backend and returning the response.
+// It takes a connection and a request as parameters and returns an error if any.
+func (lb *LoadBalancer) Handle(conn wasabi.Connection, r wasabi.Request) error {
+	backend := lb.getLeastBusyNode()
+
+	backend.counter.Add(1)
+	defer backend.counter.Add(-1)
+
+	return backend.backend.Handle(conn, r)
+}
+
+// getLeastBusyNode returns the least busy backend node.
+// It returns the least busy backend node.
+func (lb *LoadBalancer) getLeastBusyNode() *LoadBalancerNode {
+	minRequests := lb.backends[0].counter.Load()
+	minBackend := lb.backends[0]
+
+	for _, b := range lb.backends[1:] {
+		counter := b.counter.Load()
+
+		if counter < minRequests {
+			minRequests = counter
+			minBackend = b
+		}
+	}
+
+	return minBackend
+}

backend/loadbalancer_test.go

@@ -0,0 +1,90 @@
+package backend
+
+import (
+	"testing"
+
+	"github.com/ksysoev/wasabi"
+	"github.com/ksysoev/wasabi/mocks"
+)
+
+func TestNewLoadBalancer(t *testing.T) {
+	backends := []wasabi.RequestHandler{
+		mocks.NewMockBackend(t),
+	}
+
+	_, err := NewLoadBalancer(backends)
+	if err != ErrNotEnoughBackends {
+		t.Errorf("Expected error to be 'load balancer requires at least 2 backends', but got %v", err)
+	}
+
+	backends = append(backends, mocks.NewMockBackend(t))
+
+	lb, err := NewLoadBalancer(backends)
+	if err != nil {
+		t.Fatalf("Failed to create load balancer: %v", err)
+	}
+
+	if len(lb.backends) != len(backends) {
+		t.Errorf("Expected %d backends, but got %d", len(backends), len(lb.backends))
+	}
+
+	for i, backend := range lb.backends {
+		if backend.backend != backends[i] {
+			t.Errorf("Expected backend at index %d to be %v, but got %v", i, backends[i], backend.backend)
+		}
+
+		if backend.counter.Load() != 0 {
+			t.Errorf("Expected backend counter at index %d to be 0, but got %d", i, backend.counter.Load())
+		}
+	}
+}
+
+func TestLoadBalancer_getLeastBusyNode(t *testing.T) {
+	backends := []wasabi.RequestHandler{
+		// Mock backends for testing
+		mocks.NewMockBackend(t),
+		mocks.NewMockBackend(t),
+	}
+
+	lb, err := NewLoadBalancer(backends)
+	if err != nil {
+		t.Fatalf("Failed to create load balancer: %v", err)
+	}
+
+	// Increment the counter of the second backend
+	lb.backends[0].counter.Add(10)
+
+	// Get the least busy node
+	leastBusyNode := lb.getLeastBusyNode()
+
+	// Check if the least busy node is the second backend
+	if leastBusyNode != lb.backends[1] {
+		t.Errorf("Expected least busy node to be the second backend, but got %v", leastBusyNode)
+	}
+}
+func TestLoadBalancer_Handle(t *testing.T) {
+	firstBackend := mocks.NewMockBackend(t)
+	// Create mock backends for testing
+	backends := []wasabi.RequestHandler{
+		firstBackend,
+		mocks.NewMockBackend(t),
+		mocks.NewMockBackend(t),
+	}
+
+	lb, err := NewLoadBalancer(backends)
+	if err != nil {
+		t.Fatalf("Failed to create load balancer: %v", err)
+	}
+
+	// Create mock connection and request
+	mockConn := mocks.NewMockConnection(t)
+	mockRequest := mocks.NewMockRequest(t)
+
+	firstBackend.EXPECT().Handle(mockConn, mockRequest).Return(nil)
+
+	// Call the Handle method
+	err = lb.Handle(mockConn, mockRequest)
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+}