Add utility functions for HTTP header management and client IP extraction; implement metrics tracking in LoadBalancer

Author: itsfuad

.gitignore

@@ -1 +1,2 @@
-.env
+.env
+*.exe

balancer/balancer.go

@@ -1,35 +1,102 @@
 package balancer
 
-import "fmt"
-
 import (
+	"errors"
+	"fmt"
+	"io"
 	"log"
 	"net/http"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	sv "loadbalancer/server"
 	"loadbalancer/utils"
 )
 
 type LoadBalancer struct {
-	Servers  []*sv.Server
-	mu       sync.Mutex
-	Logger   *log.Logger
-	wg       sync.WaitGroup
-	shutdown bool
+	Servers     []*sv.Server
+	mu          sync.RWMutex
+	Logger      *log.Logger
+	wg          sync.WaitGroup
+	shutdown    bool
+	metrics     *Metrics
+	maxRetries  int
+}
+
+type Metrics struct {
+	TotalRequests uint64
+	FailedRequests uint64
+	ActiveConnections int64
+}
+
+func NewLoadBalancer(logger *log.Logger) *LoadBalancer {
+	if logger == nil {
+		logger = log.New(io.Discard, "", log.LstdFlags)
+	}
+	return &LoadBalancer{
+		Logger: logger,
+		metrics: &Metrics{},
+		maxRetries: 3,
+	}
 }
 
-func (lb *LoadBalancer) AddServer(url string) {
+func (lb *LoadBalancer) AddServer(url string) error {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+	
+	// Validate server URL
+	if url == "" {
+		return errors.New("server URL cannot be empty")
+	}
+	
+	// Check for duplicate servers
+	for _, server := range lb.Servers {
+		if server.URL == url {
+			return fmt.Errorf("server %s already exists", url)
+		}
+	}
+	
 	server := sv.NewServer(url, lb.Logger)
 	lb.Servers = append(lb.Servers, server)
-	//lb.Logger.Printf("Added server %s to the load balancer", url)
 	lb.Logger.Println(utils.Colorize("Added server "+url+" to the load balancer", utils.GREEN))
+	return nil
 }
 
-func (lb *LoadBalancer) GetLeastLoadedServer() *sv.Server {
+func (lb *LoadBalancer) RemoveServer(url string) error {
 	lb.mu.Lock()
 	defer lb.mu.Unlock()
+	
+	for i, server := range lb.Servers {
+		if server.URL == url {
+			// Wait for active connections to finish
+			for server.Load > 0 {
+				lb.mu.Unlock()
+				time.Sleep(100 * time.Millisecond)
+				lb.mu.Lock()
+			}
+			
+			// Remove server
+			lb.Servers = append(lb.Servers[:i], lb.Servers[i+1:]...)
+			lb.Logger.Println(utils.Colorize("Removed server "+url, utils.YELLOW))
+			return nil
+		}
+	}
+	
+	return fmt.Errorf("server %s not found", url)
+}
+
+func (lb *LoadBalancer) GetMetrics() *Metrics {
+	return &Metrics{
+		TotalRequests: atomic.LoadUint64(&lb.metrics.TotalRequests),
+		FailedRequests: atomic.LoadUint64(&lb.metrics.FailedRequests),
+		ActiveConnections: atomic.LoadInt64(&lb.metrics.ActiveConnections),
+	}
+}
+
+func (lb *LoadBalancer) GetLeastLoadedServer() *sv.Server {
+	lb.mu.RLock()
+	defer lb.mu.RUnlock()
 
 	var leastLoadedServer *sv.Server
 	for _, server := range lb.Servers {
@@ -50,13 +117,36 @@ func (lb *LoadBalancer) GetLeastLoadedServer() *sv.Server {
 }
 
 func (lb *LoadBalancer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	server := lb.GetLeastLoadedServer()
-	if server != nil {
-		server.HandleRequest(w, r)
-	} else {
-		http.Error(w, "No servers available", http.StatusServiceUnavailable)
-		lb.Logger.Println(utils.Colorize("No servers available to handle the request", utils.RED))
+	// Increment metrics
+	atomic.AddUint64(&lb.metrics.TotalRequests, 1)
+	atomic.AddInt64(&lb.metrics.ActiveConnections, 1)
+	defer atomic.AddInt64(&lb.metrics.ActiveConnections, -1)
+
+	// Check if we're shutting down
+	if lb.shutdown {
+		http.Error(w, "Service is shutting down", http.StatusServiceUnavailable)
+		return
 	}
+
+	// Try multiple servers if needed
+	var err error
+	for retry := 0; retry < lb.maxRetries; retry++ {
+		server := lb.GetLeastLoadedServer()
+		if server == nil {
+			continue
+		}
+
+		err = server.HandleRequest(w, r)
+		if err == nil {
+			return
+		}
+
+		lb.Logger.Printf("Request failed on server %s, attempt %d: %v", server.URL, retry+1, err)
+	}
+
+	// All retries failed
+	atomic.AddUint64(&lb.metrics.FailedRequests, 1)
+	http.Error(w, "All servers failed to process the request", http.StatusServiceUnavailable)
 }
 
 func (lb *LoadBalancer) StartHealthChecks(interval time.Duration) {

main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"encoding/json"
 	"fmt"
 	"log"
 	"net/http"
@@ -25,9 +26,7 @@ func main() {
 	}
 
 	// Create load balancer
-	lb := &balancer.LoadBalancer{
-		Logger: logger,
-	}
+	lb := balancer.NewLoadBalancer(logger)
 
 	// Add servers from configuration
 	for _, url := range config.Servers.URLs {
@@ -38,6 +37,12 @@ func main() {
 	healthCheckInterval := time.Duration(config.LoadBalancer.HealthCheckIntervalSeconds) * time.Second
 	lb.StartHealthChecks(healthCheckInterval)
 
+	// Add metrics endpoint
+	http.HandleFunc("/metrics", func(w http.ResponseWriter, r *http.Request) {
+		metrics := lb.GetMetrics()
+		json.NewEncoder(w).Encode(metrics)
+	})
+
 	// Setup graceful shutdown
 	stop := make(chan os.Signal, 1)
 	signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM)

main_test.go

@@ -0,0 +1,200 @@
+package main
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"loadbalancer/balancer"
+	"loadbalancer/config"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"testing"
+	"time"
+)
+
+// MockServer represents a test HTTP server
+type MockServer struct {
+	server *httptest.Server
+	URL    string
+}
+
+// setupMockServers creates test backend servers
+func setupMockServers(count int) []MockServer {
+	var servers []MockServer
+	for i := 0; i < count; i++ {
+		handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusOK)
+			fmt.Fprintf(w, "Mock Server %d", i+1)
+		})
+		server := httptest.NewServer(handler)
+		servers = append(servers, MockServer{
+			server: server,
+			URL:    server.URL,
+		})
+	}
+	return servers
+}
+
+// createTestConfig creates a temporary config file
+func createTestConfig(t *testing.T, servers []MockServer) string {
+	urls := make([]string, len(servers))
+	for i, s := range servers {
+		urls[i] = s.URL
+	}
+
+	configData := config.Config{
+		LoadBalancer: struct {
+			Port                       int "json:\"port\""
+			HealthCheckIntervalSeconds int "json:\"health_check_interval_seconds\""
+		}{
+			Port:                       8080,
+			HealthCheckIntervalSeconds: 5,
+		},
+		Servers: struct {
+			URLs []string "json:\"urls\""
+		}{
+			URLs: urls,
+		},
+	}
+
+	file, err := os.CreateTemp("", "test_config_*.json")
+	if err != nil {
+		t.Fatalf("Failed to create temp config file: %v", err)
+	}
+
+	encoder := json.NewEncoder(file)
+	if err := encoder.Encode(configData); err != nil {
+		t.Fatalf("Failed to write config data: %v", err)
+	}
+
+	return file.Name()
+}
+
+func TestLoadBalancerInitialization(t *testing.T) {
+	// Setup mock servers
+	mockServers := setupMockServers(2)
+	defer func() {
+		for _, s := range mockServers {
+			s.server.Close()
+		}
+	}()
+
+	// Create test config
+	configFile := createTestConfig(t, mockServers)
+	defer os.Remove(configFile)
+
+	// Test configuration loading
+	cfg, err := config.LoadConfig()
+	if err != nil {
+		t.Fatalf("Failed to load config: %v", err)
+	}
+
+	if len(cfg.Servers.URLs) != 2 {
+		t.Errorf("Expected 2 servers in config, got %d", len(cfg.Servers.URLs))
+	}
+
+	// Test load balancer creation
+	lb := balancer.NewLoadBalancer(nil)
+	if lb == nil {
+		t.Fatal("Failed to create load balancer")
+	}
+}
+
+func TestMetricsEndpoint(t *testing.T) {
+	// Create a test server with the metrics endpoint
+	lb := balancer.NewLoadBalancer(nil)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/metrics" {
+			metrics := lb.GetMetrics()
+			json.NewEncoder(w).Encode(metrics)
+		}
+	}))
+	defer server.Close()
+
+	// Test metrics endpoint
+	resp, err := http.Get(server.URL + "/metrics")
+	if err != nil {
+		t.Fatalf("Failed to get metrics: %v", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		t.Errorf("Expected status OK, got %v", resp.Status)
+	}
+
+	var metrics balancer.Metrics
+	if err := json.NewDecoder(resp.Body).Decode(&metrics); err != nil {
+		t.Fatalf("Failed to decode metrics: %v", err)
+	}
+}
+
+func TestGracefulShutdown(t *testing.T) {
+	lb := balancer.NewLoadBalancer(nil)
+	server := httptest.NewServer(lb)
+	defer server.Close()
+
+	// Create a channel to signal when shutdown is complete
+	done := make(chan bool)
+
+	// Simulate SIGTERM
+	p, err := os.FindProcess(os.Getpid())
+	if err != nil {
+		t.Fatalf("Failed to find process: %v", err)
+	}
+
+	go func() {
+		p.Signal(os.Interrupt)
+		done <- true
+	}()
+
+	// Wait for shutdown with timeout
+	select {
+	case <-done:
+		// Shutdown completed successfully
+	case <-time.After(5 * time.Second):
+		t.Fatal("Shutdown timeout")
+	}
+}
+
+func TestLoadBalancerRequestHandling(t *testing.T) {
+	// Setup mock backend servers
+	mockServers := setupMockServers(2)
+	defer func() {
+		for _, s := range mockServers {
+			s.server.Close()
+		}
+	}()
+
+	// Create load balancer
+	lb := balancer.NewLoadBalancer(nil)
+	for _, s := range mockServers {
+		if err := lb.AddServer(s.URL); err != nil {
+			t.Fatalf("Failed to add server: %v", err)
+		}
+	}
+
+	// Create test server with the load balancer
+	server := httptest.NewServer(lb)
+	defer server.Close()
+
+	// Test request handling
+	resp, err := http.Get(server.URL)
+	if err != nil {
+		t.Fatalf("Failed to send request: %v", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		t.Errorf("Expected status OK, got %v", resp.Status)
+	}
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		t.Fatalf("Failed to read response: %v", err)
+	}
+
+	if len(body) == 0 {
+		t.Error("Expected non-empty response body")
+	}
+}