rl544
133 lines
3.3 KiB
Go
133 lines
3.3 KiB
Go
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
|
|
}
|