pub

internal/proxy/http.go

@@ -0,0 +1,132 @@
+package proxy
+
+import (
+	"fmt"
+	"log"
+	"net"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+	"sync"
+	"sync/atomic"
+	"time"
+)
+
+// Backend represents a target server with its health state
+type Backend struct {
+	URL   *url.URL
+	Alive bool
+	mu    sync.RWMutex
+}
+
+func (b *Backend) SetAlive(alive bool) {
+	b.mu.Lock()
+	b.Alive = alive
+	b.mu.Unlock()
+}
+
+func (b *Backend) IsAlive() bool {
+	b.mu.RLock()
+	alive := b.Alive
+	b.mu.RUnlock()
+	return alive
+}
+
+// LoadBalancer holds a slice of backends and a counter for Round-Robin
+type LoadBalancer struct {
+	backends []*Backend
+	current  uint64
+}
+
+// Next returns the next healthy backend URL using Round-Robin
+func (lb *LoadBalancer) Next() *url.URL {
+	// Loop at most len(backends) times to find a healthy one
+	for i := 0; i < len(lb.backends); i++ {
+		next := atomic.AddUint64(&lb.current, 1)
+		idx := next % uint64(len(lb.backends))
+		b := lb.backends[idx]
+
+		if b.IsAlive() {
+			return b.URL
+		}
+	}
+	return nil // No healthy backends available
+}
+
+// StartHealthCheck begins a background routine to ping backends
+func (lb *LoadBalancer) StartHealthCheck() {
+	go func() {
+		for {
+			for _, b := range lb.backends {
+				go func(backend *Backend) {
+					// TCP Dial is a safe, universal health check for HTTP servers
+					host := backend.URL.Host
+					if backend.URL.Port() == "" {
+						if backend.URL.Scheme == "https" {
+							host += ":443"
+						} else {
+							host += ":80"
+						}
+					}
+					
+					conn, err := net.DialTimeout("tcp", host, 2*time.Second)
+					if err != nil {
+						if backend.IsAlive() {
+							log.Printf("[HealthCheck] Backend %s went DOWN: %v", backend.URL.String(), err)
+							backend.SetAlive(false)
+						}
+					} else {
+						conn.Close()
+						if !backend.IsAlive() {
+							log.Printf("[HealthCheck] Backend %s came UP", backend.URL.String())
+							backend.SetAlive(true)
+						}
+					}
+				}(b)
+			}
+			time.Sleep(5 * time.Second) // Check every 5 seconds
+		}
+	}()
+}
+
+// NewLoadBalancedHTTPProxy creates a ReverseProxy that load balances across multiple backends
+func NewLoadBalancedHTTPProxy(backends []string) (*httputil.ReverseProxy, error) {
+	if len(backends) == 0 {
+		return nil, fmt.Errorf("no backends provided")
+	}
+
+	lb := &LoadBalancer{}
+	for _, b := range backends {
+		u, err := url.Parse(b)
+		if err != nil {
+			return nil, fmt.Errorf("invalid backend url %s: %w", b, err)
+		}
+		// Initially mark all as alive
+		lb.backends = append(lb.backends, &Backend{URL: u, Alive: true})
+	}
+
+	// Start health checking in the background
+	lb.StartHealthCheck()
+
+	proxy := &httputil.ReverseProxy{
+		Director: func(req *http.Request) {
+			target := lb.Next()
+			if target == nil {
+				// We can't return an error directly from Director, but setting Scheme to empty
+				// will cause the RoundTripper to fail, returning 502 Bad Gateway by default.
+				return
+			}
+			req.URL.Scheme = target.Scheme
+			req.URL.Host = target.Host
+			// Note: We don't overwrite req.URL.Path here, as it's modified by the router
+			req.Header.Set("X-Forwarded-Host", req.Header.Get("Host"))
+		},
+		// Optional custom error handler when backends are unreachable
+		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
+			log.Printf("Proxy error: %v", err)
+			http.Error(w, "502 Bad Gateway - No healthy backends available", http.StatusBadGateway)
+		},
+	}
+
+	return proxy, nil
+}