serverpool.go

package main

import (
	"hash/fnv"
	"log"
	"math/rand"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"sync"
	"sync/atomic"
	"time"
)

// v1: passive health checks
func isBackendAlive(u *url.URL) bool {
	timeout := 2 * time.Second
	conn, err := net.DialTimeout("tcp", u.Host, timeout)
	if err != nil {
		return false
	}
	_ = conn.Close()
	return true
}

// ServerPool holds information about reachable backends
type ServerPool struct {
	backends []*Backend
	current  uint64
	// v17: marks if the last selection was canary
	lastPickWasCanary atomic.Bool
	// v19: protect backend slice for dynamic add/remove
	mu sync.RWMutex
}

// v19: safe snapshot of backends for iteration without holding lock
func (s *ServerPool) snapshot() []*Backend {
	s.mu.RLock()
	defer s.mu.RUnlock()
	out := make([]*Backend, len(s.backends))
	copy(out, s.backends)
	return out
}

// v12: find backend by exact URL string
func (s *ServerPool) FindBackend(urlStr string) *Backend {
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, b := range s.backends {
		if b.URL.String() == urlStr {
			return b
		}
	}
	return nil
}

// AddBackend adds a backend to the server pool
func (s *ServerPool) AddBackend(backend *Backend) {
	// Initialize reverse proxy for the backend
	backend.ReverseProxy = httputil.NewSingleHostReverseProxy(backend.URL)
    
	// v19: default health path if not set
	if backend.HealthPath == "" {
		backend.HealthPath = "/healthz"
	}
    
	// v15: propagate headers to backend
	rp := backend.ReverseProxy
	rp.Director = func(req *http.Request) {
		// preserve default director behavior:
		req.URL.Scheme = backend.URL.Scheme
		req.URL.Host = backend.URL.Host
		// preserve path/query automatically (we’re not rewriting)
		// X-Forwarded-For
		if ip, _, err := net.SplitHostPort(req.RemoteAddr); err == nil && ip != "" {
			prior := req.Header.Get("X-Forwarded-For")
			if prior == "" { req.Header.Set("X-Forwarded-For", ip) } else { req.Header.Set("X-Forwarded-For", prior+", "+ip) }
		}
		// X-Request-ID from context
		if v := req.Context().Value(ctxReqID); v != nil {
			req.Header.Set("X-Request-ID", v.(string))
		}
		// identify lb instance
		host, _ := os.Hostname()
		req.Header.Set("X-LB-Instance", host)
	}
    
	// v11: install ErrorHandler to do limited, budget-aware reroute
	rp.ErrorHandler = func(w http.ResponseWriter, req *http.Request, err error) {
		// Mark failure on the current backend (breaker inputs)
		backend.FailureCount.Add(1)
		backend.LastFailure.Store(time.Now().UnixNano())
		if backend.FailureCount.Load() >= failureThreshold {
			backend.BreakerState.Store(int32(Open))
		}

		// Respect context deadline
		ctx := req.Context()
		if ctx.Err() != nil {
			// Deadline/cancelled: finish with 504
			http.Error(w, "Gateway Timeout", http.StatusGatewayTimeout)
			return
		}

		// Retry only idempotent methods
		m := req.Method
		if !(m == http.MethodGet || m == http.MethodHead || m == http.MethodOptions) {
			http.Error(w, "Bad Gateway", http.StatusBadGateway)
			return
		}

		// Attempts from context
		attempts := getAttempts(ctx)
		if attempts >= maxRetryAttempts {
			http.Error(w, "Bad Gateway", http.StatusBadGateway)
			return
		}

		// Small jitter (bounded by remaining context budget)
		// If almost out of time, skip sleep.
		if dl, ok := ctx.Deadline(); ok {
			rem := time.Until(dl)
			if rem > retryJitterMin {
				jitter := retryJitterMin + time.Duration(rand.Int63n(int64(retryJitterMax-retryJitterMin)))
				if jitter < rem {
					select {
					case <-time.After(jitter):
					case <-ctx.Done():
						http.Error(w, "Gateway Timeout", http.StatusGatewayTimeout)
						return
					}
				}
			}
		}

		// Choose an alternative backend via current ServerPool & client key
		sp := serverPool.Load()
		if sp == nil {
			http.Error(w, "Service Unavailable", http.StatusServiceUnavailable)
			return
		}
		ck := getClientKey(ctx)
		alt := sp.GetNextPeer(ck)
		if alt == nil || alt == backend {
			// No alternative or same target — give up
			http.Error(w, "Bad Gateway", http.StatusBadGateway)
			return
		}

		// Bump attempts and recurse to the alternative ReverseProxy
		ctx2 := withAttempts(ctx, attempts+1)
		alt.ReverseProxy.ServeHTTP(w, req.WithContext(ctx2))
	}
    
	s.mu.Lock()
	s.backends = append(s.backends, backend)
	s.mu.Unlock()
}

// v5: breaker check
func (b *Backend) CanServe() bool {
	// v16: skip serving if in quarantine window
	if until := time.Unix(0, b.QuarantineUntil.Load()); until.After(time.Now()) {
		return false
	}

	// v18: per-backend concurrency cap
	if cap := b.MaxInFlight.Load(); cap > 0 && b.InFlight.Load() >= cap {
		return false
	}

	// v18: warm-up gating (probabilistic admit while warming)
	if until := b.WarmUntil.Load(); until > 0 {
		now := time.Now()
		end := time.Unix(0, until)
		if now.Before(end) {
			// linear ramp from warmMinProb -> 1.0 over warmupDuration
			elapsed := warmupDuration - time.Until(end)
			p := float64(warmMinProb) + (1.0-float64(warmMinProb))*(elapsed.Seconds()/warmupDuration.Seconds())
			if p > 1.0 { p = 1.0 }
			if rand.Float64() > p {
				return false
			}
		} else {
			// warm-up finished
			b.WarmUntil.Store(0)
		}
	}
	state := breakerState(b.BreakerState.Load())
	switch state {
	case Closed:
		return b.IsAlive()
	case Open:
		// check if enough time has passed for half-open
		lastFail := time.Unix(0, b.LastFailure.Load())
		if time.Since(lastFail) > openTimeout {
			// transition to HalfOpen
			if b.BreakerState.CompareAndSwap(int32(Open), int32(HalfOpen)) {
				return true
			}
		}
		return false
	case HalfOpen:
		// allow one trial (concurrent protection: only first will flip to Closed if success)
		return true
	}
	return false
}

// NextIndex returns the next index in round robin fashion
func (s *ServerPool) NextIndex() int {
	backs := s.snapshot()
	if len(backs) == 0 {
		return 0
	}
	return int(atomic.AddUint64(&s.current, uint64(1)) % uint64(len(backs)))
}

// v4: Least-Connections selection among alive backends
func (s *ServerPool) pickLeastConn() *Backend {
	backs := s.snapshot()
	if len(backs) == 0 {
		return nil
	}
	// find a starting offset for tie-breaking (RR-ish)
	start := s.NextIndex()
	var best *Backend
	var bestVal int64 = 1 << 62 // large

	l := len(backs) + start
	for i := start; i < l; i++ {
		b := backs[i%len(backs)]
		if !b.CanServe() {
			continue
		}
		v := b.InFlight.Load()
		if v < bestVal {
			bestVal, best = v, b
		}
	}
	return best
}

// v6: Weighted Round Robin selection
func (s *ServerPool) pickWeightedRR() *Backend {
	backs := s.snapshot()
	if len(backs) == 0 {
		return nil
	}
	var best *Backend
	total := 0
	for _, b := range backs {
		if !b.CanServe() || b.Weight <= 0 {
			continue
		}
		b.currentWeight += b.Weight
		total += b.Weight
		if best == nil || b.currentWeight > best.currentWeight {
			best = b
		}
	}
	if best == nil {
		return nil
	}
	best.currentWeight -= total
	return best
}

// v7: Power-of-Two-Choices by EWMA latency
func (s *ServerPool) pickP2CEWMA() *Backend {
	backs := s.snapshot()
	n := len(backs)
	if n == 0 {
		return nil
	}

	// sample two distinct indices, try a few times to find CanServe()
	pickOne := func() *Backend {
		for attempts := 0; attempts < 4; attempts++ {
			b := backs[rand.Intn(n)]
			if b.CanServe() {
				return b
			}
		}
		return nil
	}

	a := pickOne()
	b := pickOne()
	if a == nil && b == nil {
		return nil
	}
	if a == nil {
		return b
	}
	if b == nil {
		return a
	}

	// choose lower EWMA (lower = faster)
	aw := a.getEWMA()
	bw := b.getEWMA()

	// treat 0 as "unknown" -> bias slightly against unknown if the other has data
	if aw == 0 && bw > 0 {
		return b
	}
	if bw == 0 && aw > 0 {
		return a
	}
	if aw <= bw {
		return a
	}
	return b
}

// v9: pick by IP-hash (affinity), honoring CanServe()
func (s *ServerPool) pickSticky(key string) *Backend {
	backs := s.snapshot()
	n := len(backs)
	if n == 0 {
		return nil
	}
	h := fnv.New32a()
	_, _ = h.Write([]byte(key))
	start := int(h.Sum32()) % n

	// scan from hashed index for a CanServe backend
	l := n + start
	for i := start; i < l; i++ {
		b := backs[i%n]
		if b.CanServe() {
			return b
		}
	}
	return nil
}

// GetNextPeer returns the next active peer to take a connection
func (s *ServerPool) GetNextPeer(clientKey string) *Backend {	
	// v1: empty-pool safety
	backs := s.snapshot()
	if len(backs) == 0 {
		return nil
	}

	// v17: canary routing (takes precedence if eligible)
	if b := s.pickCanary(); b != nil {
		return b
	}

	// v9: sticky session selection
	// v14: use runtime toggles
	if policyEnableSticky.Load() {
		// v10: use clientKey parameter instead of tmpKey field
		if clientKey != "" {
			if b := s.pickSticky(clientKey); b != nil {
				return b
			}
			// fall through to normal policies if none CanServe
		}
	}
	
	// v7: integrate P2C by EWMA
	// v14: use runtime toggles
	if policyUseP2CEWMA.Load() {
		if b := s.pickP2CEWMA(); b != nil {
			return b
		}
		return nil
	}
	
	// v6: integrate Weighted Round Robin
	// v14: use runtime toggles
	if policyUseWeightedRR.Load() {
		if b := s.pickWeightedRR(); b != nil {
			return b
		}
		return nil
	}
	
	// v4: integrate optional Least-Connections
	// v14: use runtime toggles
	if policyUseLeastConn.Load() {
		if b := s.pickLeastConn(); b != nil {
			return b
		}
		return nil
	}
	
	// existing Round-Robin over alive
	next := s.NextIndex()
	backs = s.snapshot()
	l := len(backs) + next // start from next and move a full cycle
	
	for i := next; i < l; i++ {
		idx := i % len(backs) // take an index by modding with length
		
		// if we have a backend that can serve, use it and store if its not the original one
		if backs[idx].CanServe() {
			if i != next {
				atomic.StoreUint64(&s.current, uint64(idx))
			}
			return backs[idx]
		}
	}
	
	return nil
}

// v1: passive health checks
func (s *ServerPool) HealthCheck() {
	for _, b := range s.snapshot() {
		alive := isBackendAlive(b.URL)
		b.SetAlive(alive)
		status := "up"
		if !alive {
			status = "down"
		}
		// v3: metrics
		if alive {
			healthUp.WithLabelValues(b.URL.String()).Set(1)
		} else {
			healthUp.WithLabelValues(b.URL.String()).Set(0)
		}
		// v5: breaker state metrics
		breakerStateGauge.WithLabelValues(b.URL.String()).Set(float64(b.BreakerState.Load()))
		log.Printf("%s [%s]\n", b.URL, status)
	}
}

// v17: try canary selection based on runtime config
func (s *ServerPool) pickCanary() *Backend {
	if !canaryEnabled.Load() {
		return nil
	}
	pct := canaryPercent.Load()
	if pct <= 0 {
		return nil
	}
	tgtAny := canaryTarget.Load()
	if tgtAny == nil {
		return nil
	}
	tgt := tgtAny.(string)
	// probabilistic gate
	if rand.Intn(100) >= int(pct) {
		return nil
	}
	b := s.FindBackend(tgt)
	if b == nil || !b.CanServe() {
		return nil
	}
	// mark canary pick
	s.lastPickWasCanary.Store(true)
	return b
}