Load Balancing

package main

import (
	"context"
	"fmt"
	"hash/fnv"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"os/signal"
	"sync"
	"sync/atomic"
	"syscall"
	"time"
)

// --- Types ---
type Backend struct {
	URL               *url.URL
	Healthy           atomic.Bool
	ActiveConnections atomic.Int64
	TotalRequests     atomic.Int64
	Weight            int
	ReverseProxy      *httputil.ReverseProxy
}

type Config struct {
	Backends           []BackendConfig
	HealthCheckPath    string
	HealthCheckInterval time.Duration
	HealthCheckTimeout  time.Duration
	Algorithm          string // "round-robin", "least-conn", "ip-hash"
}

type BackendConfig struct {
	URL    string
	Weight int
}

// --- Load Balancer ---
type LoadBalancer struct {
	backends []*Backend
	current  atomic.Uint64
	config   Config
}

func NewLoadBalancer(config Config) *LoadBalancer {
	lb := &LoadBalancer{config: config}

	for _, bc := range config.Backends {
		u, err := url.Parse(bc.URL)
		if err != nil {
			log.Fatalf("Invalid backend URL %s: %v", bc.URL, err)
		}

		proxy := httputil.NewSingleHostReverseProxy(u)
		proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
			log.Printf("Proxy error for %s: %v", u.String(), err)
			w.WriteHeader(http.StatusBadGateway)
			fmt.Fprintf(w, `{"error":"bad gateway"}`)
		}

		weight := bc.Weight
		if weight <= 0 {
			weight = 1
		}

		b := &Backend{
			URL:          u,
			Weight:       weight,
			ReverseProxy: proxy,
		}
		b.Healthy.Store(true)
		lb.backends = append(lb.backends, b)
	}

	return lb
}

func (lb *LoadBalancer) healthyBackends() []*Backend {
	var healthy []*Backend
	for _, b := range lb.backends {
		if b.Healthy.Load() {
			healthy = append(healthy, b)
		}
	}
	return healthy
}

func (lb *LoadBalancer) selectBackend(clientIP string) *Backend {
	healthy := lb.healthyBackends()
	if len(healthy) == 0 {
		return nil
	}

	switch lb.config.Algorithm {
	case "least-conn":
		var min *Backend
		for _, b := range healthy {
			if min == nil || b.ActiveConnections.Load() < min.ActiveConnections.Load() {
				min = b
			}
		}
		return min

	case "ip-hash":
		h := fnv.New32a()
		h.Write([]byte(clientIP))
		idx := h.Sum32() % uint32(len(healthy))
		return healthy[idx]

	default: // round-robin
		idx := lb.current.Add(1) - 1
		return healthy[idx%uint64(len(healthy))]
	}
}

func (lb *LoadBalancer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Status endpoint
	if r.URL.Path == "/lb/status" {
		lb.statusHandler(w)
		return
	}

	clientIP := r.RemoteAddr
	backend := lb.selectBackend(clientIP)
	if backend == nil {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusServiceUnavailable)
		fmt.Fprint(w, `{"error":"no healthy backends"}`)
		return
	}

	backend.ActiveConnections.Add(1)
	backend.TotalRequests.Add(1)
	defer backend.ActiveConnections.Add(-1)

	// Add proxy headers
	r.Header.Set("X-Forwarded-For", clientIP)
	r.Header.Set("X-Real-IP", clientIP)
	r.Header.Set("X-Forwarded-Proto", "http")

	start := time.Now()
	backend.ReverseProxy.ServeHTTP(w, r)
	log.Printf("%s %s -> %s [%v]", r.Method, r.URL.Path, backend.URL, time.Since(start))
}

func (lb *LoadBalancer) statusHandler(w http.ResponseWriter) {
	w.Header().Set("Content-Type", "application/json")
	fmt.Fprint(w, `{"backends":[`)
	for i, b := range lb.backends {
		if i > 0 {
			fmt.Fprint(w, ",")
		}
		fmt.Fprintf(w,
			`{"url":"%s","healthy":%t,"activeConnections":%d,"totalRequests":%d}`,
			b.URL, b.Healthy.Load(), b.ActiveConnections.Load(), b.TotalRequests.Load(),
		)
	}
	fmt.Fprint(w, "]}")
}

// --- Health Checking ---
func (lb *LoadBalancer) StartHealthChecks(ctx context.Context) {
	ticker := time.NewTicker(lb.config.HealthCheckInterval)
	defer ticker.Stop()

	client := &http.Client{Timeout: lb.config.HealthCheckTimeout}

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			var wg sync.WaitGroup
			for _, b := range lb.backends {
				wg.Add(1)
				go func(backend *Backend) {
					defer wg.Done()
					checkURL := backend.URL.String() + lb.config.HealthCheckPath

					resp, err := client.Get(checkURL)
					wasHealthy := backend.Healthy.Load()

					if err != nil || resp.StatusCode >= 500 {
						backend.Healthy.Store(false)
						if wasHealthy {
							log.Printf("Backend %s is now UNHEALTHY", backend.URL)
						}
					} else {
						backend.Healthy.Store(true)
						if !wasHealthy {
							log.Printf("Backend %s is now HEALTHY", backend.URL)
						}
					}
					if resp != nil {
						io.Copy(io.Discard, resp.Body)
						resp.Body.Close()
					}
				}(b)
			}
			wg.Wait()
		}
	}
}

func main() {
	config := Config{
		Backends: []BackendConfig{
			{URL: "http://localhost:3001", Weight: 2},
			{URL: "http://localhost:3002", Weight: 1},
			{URL: "http://localhost:3003", Weight: 1},
		},
		HealthCheckPath:     "/health",
		HealthCheckInterval: 5 * time.Second,
		HealthCheckTimeout:  3 * time.Second,
		Algorithm:           "round-robin",
	}

	lb := NewLoadBalancer(config)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	go lb.StartHealthChecks(ctx)

	srv := &http.Server{
		Addr:         ":8080",
		Handler:      lb,
		ReadTimeout:  10 * time.Second,
		WriteTimeout: 30 * time.Second,
	}

	go func() {
		log.Println("Load balancer listening on :8080")
		if err := srv.ListenAndServe(); err != http.ErrServerClosed {
			log.Fatal(err)
		}
	}()

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit

	cancel()
	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer shutdownCancel()
	srv.Shutdown(shutdownCtx)
	log.Println("Load balancer stopped")
}