《后端架构设计·演进与治理》3.2 客户端负载均衡：算法、局部性与重试放大

// file: internal/platform/discovery/balancer.go
// tier: L1
// verify: go test ./internal/platform/discovery/ -count=1
package discovery

import "errors"
import "fmt"
import "math"
import "math/rand"
import "sort"
import "sync"

// Balancer selects endpoints with power-of-two choices (P2C) and prefers the
// caller's zone when it has at least two endpoints. In-flight counts make a
// slow instance lose comparisons before its health check fails.
type Balancer struct {
	mu        sync.Mutex
	endpoints []Endpoint
	localZone string
	inflight  map[string]int
	rnd       *rand.Rand
}

// NewBalancer builds a balancer. A fixed seed makes tests reproducible;
// production callers should pass time.Now().UnixNano().
func NewBalancer(localZone string, seed int64) *Balancer {
	return &Balancer{
		localZone: localZone,
		inflight:  make(map[string]int),
		rnd:       rand.New(rand.NewSource(seed)),
	}
}

// Update replaces the endpoint set with one registry snapshot.
func (b *Balancer) Update(s Snapshot) {
	b.mu.Lock()
	defer b.mu.Unlock()
	eps := make([]Endpoint, 0, len(s.Endpoints))
	for _, ep := range s.Endpoints {
		if ep.Addr != "" && ep.Weight > 0 {
			eps = append(eps, ep)
		}
	}
	sort.Slice(eps, func(i, j int) bool { return eps[i].InstanceID < eps[j].InstanceID })
	b.endpoints = eps
}

// Pick returns one endpoint and a release function. The caller must invoke
// release after the request finishes; forgetting it biases P2C forever.
func (b *Balancer) Pick() (Endpoint, func(), error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	cands := b.candidatesLocked()
	if len(cands) == 0 {
		return Endpoint{}, nil, errors.New("discovery: no ready endpoint")
	}
	chosen := 0
	if len(cands) > 1 {
		i := b.rnd.Intn(len(cands))
		j := b.rnd.Intn(len(cands) - 1)
		if j >= i {
			j++
		}
		chosen = i
		if lessLoaded(cands[j], cands[i], b.inflight) {
			chosen = j
		}
	}
	ep := cands[chosen]
	b.inflight[ep.InstanceID]++
	released := false
	release := func() {
		b.mu.Lock()
		defer b.mu.Unlock()
		if !released {
			released = true
			if b.inflight[ep.InstanceID] > 0 {
				b.inflight[ep.InstanceID]--
			}
		}
	}
	return ep, release, nil
}

func (b *Balancer) candidatesLocked() []Endpoint {
	if b.localZone == "" {
		return b.endpoints
	}
	local := make([]Endpoint, 0, len(b.endpoints))
	for _, ep := range b.endpoints {
		if ep.Zone == b.localZone {
			local = append(local, ep)
		}
	}
	// With fewer than two local endpoints, P2C has nothing to compare, so
	// include remote zones and let the comparison prefer the idle local one.
	if len(local) >= 2 {
		return local
	}
	return b.endpoints
}

// lessLoaded compares in-flight per unit weight without floating point.
func lessLoaded(a, b Endpoint, inflight map[string]int) bool {
	return inflight[a.InstanceID]*b.Weight < inflight[b.InstanceID]*a.Weight
}

const maxRetryTokens = 10

// RetryBudget is a token bucket refilled by a fixed ratio of every request.
// It caps retries so an outage cannot turn one client call into an unbounded
// retry storm. A production variant may add a small per-second floor.
type RetryBudget struct {
	mu                sync.Mutex
	ratio             float64
	tokens            float64
	requests, retries int
}

// NewRetryBudget validates ratio in (0, 1].
func NewRetryBudget(ratio float64) (*RetryBudget, error) {
	if ratio <= 0 || ratio > 1 {
		return nil, fmt.Errorf("discovery: retry ratio must be in (0,1], got %v", ratio)
	}
	return &RetryBudget{ratio: ratio}, nil
}

// RecordRequest adds one request's share to the retry budget.
func (rb *RetryBudget) RecordRequest() {
	rb.mu.Lock()
	defer rb.mu.Unlock()
	rb.requests++
	rb.tokens += rb.ratio
	if rb.tokens > maxRetryTokens {
		rb.tokens = maxRetryTokens
	}
}

// AllowRetry consumes one token. False means the caller must fail now.
func (rb *RetryBudget) AllowRetry() bool {
	rb.mu.Lock()
	defer rb.mu.Unlock()
	if rb.tokens < 1 {
		return false
	}
	rb.tokens--
	rb.retries++
	return true
}

// Amplification is offered load divided by original requests.
func Amplification(requests, retries int) float64 {
	if requests <= 0 {
		return 0
	}
	return float64(requests+retries) / float64(requests)
}

// LayeredAmplification models retries at several hops: (1+ratio)^layers.
func LayeredAmplification(ratio float64, layers int) float64 {
	if ratio <= 0 || layers <= 0 {
		return 1
	}
	return math.Pow(1+ratio, float64(layers))
}

// file: internal/platform/discovery/balancer_test.go
// tier: L1
// verify: go test ./internal/platform/discovery/ -count=1
package discovery

import "math"
import "testing"

func TestBalancerPrefersLocalZone(t *testing.T) {
	b := NewBalancer("az-a", 11)
	b.Update(Snapshot{Version: 7, Endpoints: []Endpoint{
		{"a-1", "10.0.0.1:8080", "az-a", 1},
		{"a-2", "10.0.0.2:8080", "az-a", 1},
		{"b-1", "10.0.1.1:8080", "az-b", 1},
		{"b-2", "10.0.1.2:8080", "az-b", 1},
	}})
	for i := 0; i < 100; i++ {
		ep, release, err := b.Pick()
		if err != nil {
			t.Fatalf("Pick: %v", err)
		}
		if ep.Zone != "az-a" {
			t.Fatalf("pick %d zone=%s, want az-a", i, ep.Zone)
		}
		release()
	}
}

func TestBalancerP2CAvoidsLoadedEndpoint(t *testing.T) {
	b := NewBalancer("", 5)
	b.Update(Snapshot{Version: 1, Endpoints: []Endpoint{
		{"a", "10.0.0.1:8080", "", 1},
		{"b", "10.0.0.2:8080", "", 1},
	}})
	first, releaseFirst, _ := b.Pick()
	second, releaseSecond, err := b.Pick()
	if err != nil {
		t.Fatalf("second Pick: %v", err)
	}
	if second.InstanceID == first.InstanceID {
		t.Fatalf("second pick=%s, want the idle instance", second.InstanceID)
	}
	releaseFirst()
	releaseSecond()
}

func TestBalancerFallsBackWhenSingleLocal(t *testing.T) {
	b := NewBalancer("az-a", 3)
	b.Update(Snapshot{Version: 2, Endpoints: []Endpoint{
		{"a-1", "10.0.0.1:8080", "az-a", 1},
		{"b-1", "10.0.1.1:8080", "az-b", 1},
		{"b-2", "10.0.1.2:8080", "az-b", 1},
		{"b-3", "10.0.1.3:8080", "az-b", 1},
	}})
	remote := 0
	for i := 0; i < 100; i++ {
		ep, release, err := b.Pick()
		if err != nil {
			t.Fatalf("Pick: %v", err)
		}
		if ep.Zone == "az-b" {
			remote++
		}
		release()
	}
	if remote == 0 {
		t.Fatal("single local endpoint blocked all remote spillover")
	}
}

func TestRetryBudgetRatio(t *testing.T) {
	rb, err := NewRetryBudget(0.25)
	if err != nil {
		t.Fatalf("NewRetryBudget: %v", err)
	}
	allowed := 0
	for i := 0; i < 100; i++ {
		rb.RecordRequest()
		if rb.AllowRetry() {
			allowed++
		}
	}
	if allowed != 25 {
		t.Fatalf("allowed=%d, want 25", allowed)
	}
	if got := Amplification(100, allowed); math.Abs(got-1.25) > 1e-9 {
		t.Fatalf("amplification=%v, want 1.25", got)
	}
}

func TestLayeredAmplification(t *testing.T) {
	if got := LayeredAmplification(0.2, 3); math.Abs(got-1.728) > 1e-9 {
		t.Fatalf("3 layers at 20%%=%v, want 1.728", got)
	}
	if got := LayeredAmplification(0, 3); got != 1 {
		t.Fatalf("zero ratio=%v, want 1", got)
	}
}

$ cd .local/checks/be-arch-gov-ch03
$ gofmt -l . && go vet ./... && go test ./internal/platform/discovery/ -count=1 -v -run '^(TestBalancer|TestRetryBudget|TestLayered)'
=== RUN   TestBalancerPrefersLocalZone
--- PASS: TestBalancerPrefersLocalZone (0.00s)
=== RUN   TestBalancerP2CAvoidsLoadedEndpoint
--- PASS: TestBalancerP2CAvoidsLoadedEndpoint (0.00s)
=== RUN   TestBalancerFallsBackWhenSingleLocal
--- PASS: TestBalancerFallsBackWhenSingleLocal (0.00s)
=== RUN   TestRetryBudgetRatio
--- PASS: TestRetryBudgetRatio (0.00s)
=== RUN   TestLayeredAmplification
--- PASS: TestLayeredAmplification (0.00s)
PASS
ok  	archlab/internal/platform/discovery	0.297s