Load balancing: when, which kind, which algorithm

// LeastConnections picks the backend with the fewest requests in flight. The scan starts one
// place later on every pick, so ties spread over the backends.
type LeastConnections struct{ start int }

// Pick implements Picker.
func (p *LeastConnections) Pick(v View, _ uint64) (int, bool) {
  return p.min(v, func(i int) int64 { return int64(v.InFlight(i)) })
}

func (p *LeastConnections) min(v View, score func(i int) int64) (int, bool) {
  n := v.Len()
  best, bestScore := -1, int64(0)
  for k := range n {
    i := (p.start + k) % n
    if !v.Up(i) {
      continue
    }
    if s := score(i); best < 0 || s < bestScore {
      best, bestScore = i, s
    }
  }
  p.start = (p.start + 1) % n
  return best, best >= 0
}

// PowerOfTwo samples two different backends at random and sends the request to the one with fewer
// requests in flight.
type PowerOfTwo struct {
  rng *rand.Rand
  up  []int
}

// NewPowerOfTwo returns a power-of-two-choices picker with a seeded random source.
func NewPowerOfTwo(seed uint64) *PowerOfTwo {
  return &PowerOfTwo{rng: rand.New(rand.NewPCG(seed, 0x9e3779b97f4a7c15))}
}

// Pick implements Picker.
func (p *PowerOfTwo) Pick(v View, _ uint64) (int, bool) {
  p.up = p.up[:0]
  for i := range v.Len() {
    if v.Up(i) {
      p.up = append(p.up, i)
    }
  }
  switch len(p.up) {
  case 0:
    return 0, false
  case 1:
    return p.up[0], true
  }
  a := p.rng.IntN(len(p.up))
  b := p.rng.IntN(len(p.up) - 1)
  if b >= a {
    b++
  }
  if v.InFlight(p.up[b]) < v.InFlight(p.up[a]) {
    return p.up[b], true
  }
  return p.up[a], true
}

// WeightedRoundRobin is the smooth weighted round robin of nginx: over one cycle each backend gets
// requests in proportion to its weight, and the picks of a heavy backend are spread out.
type WeightedRoundRobin struct{ current []int }

// Pick implements Picker.
func (p *WeightedRoundRobin) Pick(v View, _ uint64) (int, bool) {
  if len(p.current) != v.Len() {
    p.current = make([]int, v.Len())
  }
  total, best := 0, -1
  for i := range v.Len() {
    if !v.Up(i) {
      continue
    }
    p.current[i] += v.Weight(i)
    total += v.Weight(i)
    if best < 0 || p.current[i] > p.current[best] {
      best = i
    }
  }
  if best < 0 {
    return 0, false
  }
  p.current[best] -= total
  return best, true
}

// Pick implements Picker. A backend that is down is skipped; its keys go to the next backend
// clockwise, and only those keys move.
func (c *ConsistentHash) Pick(v View, key uint64) (int, bool) {
  h := Mix(key)
  at, _ := slices.BinarySearchFunc(c.ring, h, func(p point, h uint64) int {
    switch {
    case p.hash < h:
      return -1
    case p.hash > h:
      return 1
    }
    return 0
  })
  for k := range len(c.ring) {
    p := c.ring[(at+k)%len(c.ring)]
    if p.backend < v.Len() && v.Up(p.backend) {
      return p.backend, true
    }
  }
  return 0, false
}

// HealthCheck probes every backend's /healthz each interval until ctx ends. A backend goes down
// after fall failed probes in a row, and up again after rise good probes in a row.
func (p *Proxy) HealthCheck(ctx context.Context, interval time.Duration, fall, rise int) {
  client := &http.Client{Transport: p.transport, Timeout: interval}
  bad := make([]int, len(p.backends))
  good := make([]int, len(p.backends))
  t := time.NewTicker(interval)
  defer t.Stop()
  for {
    select {
    case <-ctx.Done():
      return
    case <-t.C:
    }
    for i, b := range p.backends {
      if p.probe(ctx, client, b) {
        bad[i], good[i] = 0, good[i]+1
        if good[i] >= rise && b.down.Swap(false) {
          p.log.Info("backend up", "backend", b.URL.Host)
        }
        continue
      }
      good[i], bad[i] = 0, bad[i]+1
      if bad[i] >= fall && !b.down.Swap(true) {
        p.log.Warn("backend down", "backend", b.URL.Host)
      }
    }
  }
}

// RoundTrip sends one request to a backend. When the backend cannot be reached and the request is
// safe to repeat, it marks the backend down (a passive health check) and tries the next one.
func (p *Proxy) RoundTrip(req *http.Request) (*http.Response, error) {
  v := view{b: p.backends, tried: make([]bool, len(p.backends))}
  key := routeKey(req)
  for {
    p.mu.Lock()
    i, ok := p.picker.Pick(v, key)
    p.mu.Unlock()
    if !ok {
      return nil, ErrNoBackend
    }
    b := p.backends[i]
    out := req.Clone(req.Context())
    out.URL.Scheme, out.URL.Host, out.Host = b.URL.Scheme, b.URL.Host, b.URL.Host
    b.inFlight.Add(1)
    start := time.Now()
    resp, err := p.transport.RoundTrip(out)
    if err != nil {
      b.inFlight.Add(-1)
      if !idempotent(req) || req.Context().Err() != nil {
        return nil, fmt.Errorf("backend %s: %w", b.URL.Host, err)
      }
      p.log.Warn("backend unreachable; marking it down and retrying", "backend", b.URL.Host, "err", err)
      b.down.Store(true)
      v.tried[i] = true
      continue
    }
    resp.Body = &trackedBody{ReadCloser: resp.Body, done: func() {
      b.inFlight.Add(-1)
      b.served.Add(1)
      l := time.Since(start).Microseconds()
      old := b.ewma.Load()
      b.ewma.Store(old + (l-old)/8)
    }}
    return resp, nil
  }
}

// Drain stops new requests to backend i and waits until its requests in flight finish. Remove
// the backend, or restart it, after Drain returns.
func (p *Proxy) Drain(ctx context.Context, i int) error {
  b := p.backends[i]
  b.draining.Store(true)
  t := time.NewTicker(5 * time.Millisecond)
  defer t.Stop()
  for b.inFlight.Load() > 0 {
    select {
    case <-ctx.Done():
      return fmt.Errorf("drain %s: %d requests still in flight: %w", b.URL.Host, b.inFlight.Load(), ctx.Err())
    case <-t.C:
    }
  }
  return nil
}