《游戏服务端编程实践》5.3.2 Gateway 横向扩展策略

// file: ch05/scale/gateway_scale.go
// tier: L1
// verify: go test ./ch05/scale/

package scale

import (
	"errors"
	"sync"
)

var (
	ErrGatewayExists   = errors.New("gateway already registered")
	ErrGatewayNotFound = errors.New("gateway not found")
	ErrNoGateway       = errors.New("no gateway available")
)

// GatewayStat 是注册表对外暴露的只读快照。
type GatewayStat struct {
	ID          string
	Weight      int
	Connections int
	Draining    bool
}

type gatewayState struct {
	weight   int
	current  int // 平滑加权轮询的当前权重
	conns    int
	draining bool
}

// Balancer 只负责“新连接分配到哪个网关”和“实例是否已排空”。
// 玩家存档、会话事实源不在这里；网关之间不需要迁移连接。
type Balancer struct {
	mu          sync.Mutex
	order       []string
	gates       map[string]*gatewayState
	totalWeight int
}

func NewBalancer() *Balancer {
	return &Balancer{gates: make(map[string]*gatewayState)}
}

func (b *Balancer) Register(id string, weight int) error {
	if id == "" {
		return ErrGatewayNotFound
	}
	if weight <= 0 {
		weight = 1
	}
	b.mu.Lock()
	defer b.mu.Unlock()
	if _, ok := b.gates[id]; ok {
		return ErrGatewayExists
	}
	b.gates[id] = &gatewayState{weight: weight}
	b.order = append(b.order, id)
	b.totalWeight += weight
	return nil
}

// Acquire 用平滑加权轮询挑选一个未排空的网关，并把它计入在途连接。
func (b *Balancer) Acquire() (string, error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	best := ""
	bestCurrent := 0
	for _, id := range b.order {
		g := b.gates[id]
		if g.draining {
			continue
		}
		g.current += g.weight
		if best == "" || g.current > bestCurrent {
			best = id
			bestCurrent = g.current
		}
	}
	if best == "" {
		return "", ErrNoGateway
	}
	b.gates[best].current -= b.totalWeight
	b.gates[best].conns++
	return best, nil
}

// Release 在连接关闭后调用；在途计数是判断排空完成的唯一依据。
func (b *Balancer) Release(id string) error {
	b.mu.Lock()
	defer b.mu.Unlock()
	g, ok := b.gates[id]
	if !ok {
		return ErrGatewayNotFound
	}
	if g.conns > 0 {
		g.conns--
	}
	return nil
}

// Drain 停止该实例接收新连接，返回剩余在途连接数与是否已排空。
// 重复调用是幂等的，不会重复扣减总权重。
func (b *Balancer) Drain(id string) (remaining int, drained bool, err error) {
	b.mu.Lock()
	defer b.mu.Unlock()
	g, ok := b.gates[id]
	if !ok {
		return 0, false, ErrGatewayNotFound
	}
	if !g.draining {
		g.draining = true
		g.current = 0
		b.totalWeight -= g.weight
	}
	return g.conns, g.conns == 0, nil
}

func (b *Balancer) Closed(id string) bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	g, ok := b.gates[id]
	return ok && g.draining && g.conns == 0
}

func (b *Balancer) Snapshot() []GatewayStat {
	b.mu.Lock()
	defer b.mu.Unlock()
	out := make([]GatewayStat, 0, len(b.order))
	for _, id := range b.order {
		g := b.gates[id]
		out = append(out, GatewayStat{ID: id, Weight: g.weight, Connections: g.conns, Draining: g.draining})
	}
	return out
}

// file: ch05/scale/gateway_scale_test.go
// tier: L1
// verify: go test ./ch05/scale/

package scale

import (
	"errors"
	"testing"
)

func TestWeightedAllocation(t *testing.T) {
	b := NewBalancer()
	if err := b.Register("gw-a", 2); err != nil {
		t.Fatal(err)
	}
	if err := b.Register("gw-b", 1); err != nil {
		t.Fatal(err)
	}
	counts := map[string]int{}
	for i := 0; i < 300; i++ {
		id, err := b.Acquire()
		if err != nil {
			t.Fatal(err)
		}
		counts[id]++
		if err := b.Release(id); err != nil {
			t.Fatal(err)
		}
	}
	if counts["gw-a"] != 200 || counts["gw-b"] != 100 {
		t.Fatalf("allocation=%v, want gw-a=200 gw-b=100", counts)
	}
}

func TestDrainStopsNewConnectionsAndWaitsForRelease(t *testing.T) {
	b := NewBalancer()
	for _, id := range []string{"gw-a", "gw-b"} {
		if err := b.Register(id, 1); err != nil {
			t.Fatal(err)
		}
	}
	for i := 0; i < 6; i++ {
		if _, err := b.Acquire(); err != nil {
			t.Fatal(err)
		}
	}

	remaining, drained, err := b.Drain("gw-a")
	if err != nil {
		t.Fatal(err)
	}
	if remaining != 3 || drained {
		t.Fatalf("remaining=%d drained=%v, want 3 false", remaining, drained)
	}
	for i := 0; i < 4; i++ {
		id, err := b.Acquire()
		if err != nil {
			t.Fatal(err)
		}
		if id != "gw-b" {
			t.Fatalf("new connection went to draining gateway %q", id)
		}
		if err := b.Release(id); err != nil {
			t.Fatal(err)
		}
	}
	if b.Closed("gw-a") {
		t.Fatal("gw-a closed before in-flight connections released")
	}
	for i := 0; i < 3; i++ {
		if err := b.Release("gw-a"); err != nil {
			t.Fatal(err)
		}
	}
	if !b.Closed("gw-a") {
		t.Fatal("gw-a should be closed after all connections released")
	}
	remaining, drained, err = b.Drain("gw-a")
	if err != nil {
		t.Fatal(err)
	}
	if remaining != 0 || !drained {
		t.Fatalf("idempotent drain remaining=%d drained=%v", remaining, drained)
	}
	stats := b.Snapshot()
	if len(stats) != 2 || stats[0].ID != "gw-a" || !stats[0].Draining || stats[1].Draining {
		t.Fatalf("snapshot=%+v", stats)
	}
}

func TestAllDrainedReturnsNoGateway(t *testing.T) {
	b := NewBalancer()
	if err := b.Register("gw-a", 1); err != nil {
		t.Fatal(err)
	}
	if _, _, err := b.Drain("gw-a"); err != nil {
		t.Fatal(err)
	}
	if _, err := b.Acquire(); !errors.Is(err, ErrNoGateway) {
		t.Fatalf("acquire error=%v, want ErrNoGateway", err)
	}
	if _, _, err := b.Drain("missing"); !errors.Is(err, ErrGatewayNotFound) {
		t.Fatalf("drain error=%v", err)
	}
}

cd .local/checks/gs-practice-lab && gofmt -l ch05/scale && go vet ./ch05/scale/ && go test ./ch05/scale/ -run 'Gateway|Weighted|Drain|AllDrained' -count=1 -v

--- PASS: TestWeightedAllocation (0.00s)
--- PASS: TestDrainStopsNewConnectionsAndWaitsForRelease (0.00s)
--- PASS: TestAllDrainedReturnsNoGateway (0.00s)
ok  	gspractice/ch05/scale	0.242s