《游戏服务端编程实践》5.3.1 长连接优化与 KeepAlive

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

package scale

import (
	"errors"
	"sort"
	"sync"
	"time"
)

// Heartbeat 维护每条连接最近一次有效心跳的时间。
// TCP keepalive 由内核负责探活，应用层心跳负责确认对端进程仍能处理业务。
type Heartbeat struct {
	mu       sync.Mutex
	interval time.Duration
	timeout  time.Duration
	now      func() time.Time
	lastSeen map[string]time.Time
}

var ErrInvalidHeartbeatConfig = errors.New("heartbeat timeout must be greater than interval")

func NewHeartbeat(interval, timeout time.Duration) (*Heartbeat, error) {
	if interval <= 0 || timeout <= interval {
		return nil, ErrInvalidHeartbeatConfig
	}
	return &Heartbeat{
		interval: interval,
		timeout:  timeout,
		now:      time.Now,
		lastSeen: make(map[string]time.Time),
	}, nil
}

func (h *Heartbeat) Interval() time.Duration { return h.interval }
func (h *Heartbeat) Timeout() time.Duration  { return h.timeout }

// Beat 记录一次有效心跳。同一 id 重复登录时直接覆盖，旧截止时间随之失效。
func (h *Heartbeat) Beat(id string) {
	h.mu.Lock()
	h.lastSeen[id] = h.now()
	h.mu.Unlock()
}

func (h *Heartbeat) LastSeen(id string) (time.Time, bool) {
	h.mu.Lock()
	defer h.mu.Unlock()
	t, ok := h.lastSeen[id]
	return t, ok
}

// Overdue 返回 now 时刻已超过 timeout 未心跳的连接，按 id 排序以便稳定驱逐。
func (h *Heartbeat) Overdue(now time.Time) []string {
	h.mu.Lock()
	defer h.mu.Unlock()
	out := make([]string, 0, len(h.lastSeen))
	for id, t := range h.lastSeen {
		if now.Sub(t) >= h.timeout {
			out = append(out, id)
		}
	}
	sort.Strings(out)
	return out
}

// Evict 在一个锁内完成扫描与删除，避免扫描后到来的心跳被误删。
// 返回被驱逐的连接 id，调用方负责关闭对应的 socket。
func (h *Heartbeat) Evict(now time.Time) []string {
	h.mu.Lock()
	defer h.mu.Unlock()
	var evicted []string
	for id, t := range h.lastSeen {
		if now.Sub(t) >= h.timeout {
			evicted = append(evicted, id)
		}
	}
	sort.Strings(evicted)
	for _, id := range evicted {
		delete(h.lastSeen, id)
	}
	return evicted
}

func (h *Heartbeat) Len() int {
	h.mu.Lock()
	defer h.mu.Unlock()
	return len(h.lastSeen)
}

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

package scale

import (
	"fmt"
	"sync"
	"testing"
	"time"
)

func newTestHeartbeat(t *testing.T) (*Heartbeat, *time.Time) {
	t.Helper()
	h, err := NewHeartbeat(15*time.Second, 45*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	now := time.Unix(1_800_000_000, 0)
	h.now = func() time.Time { return now }
	return h, &now
}

func TestHeartbeatResetsDeadline(t *testing.T) {
	h, now := newTestHeartbeat(t)
	h.Beat("conn-1")
	*now = now.Add(40 * time.Second)
	if got := h.Overdue(*now); len(got) != 0 {
		t.Fatalf("overdue=%v before timeout", got)
	}
	h.Beat("conn-1") // 心跳必须把截止时间推后
	*now = now.Add(40 * time.Second)
	if got := h.Overdue(*now); len(got) != 0 {
		t.Fatalf("overdue=%v after fresh beat", got)
	}
	*now = now.Add(5 * time.Second)
	if got := h.Overdue(*now); len(got) != 1 || got[0] != "conn-1" {
		t.Fatalf("overdue=%v at timeout", got)
	}
}

func TestEvictRemovesOnlyTimedOut(t *testing.T) {
	h, now := newTestHeartbeat(t)
	h.Beat("fresh")
	*now = now.Add(30 * time.Second)
	h.Beat("stale")
	*now = now.Add(20 * time.Second)

	evicted := h.Evict(*now)
	if len(evicted) != 1 || evicted[0] != "fresh" {
		t.Fatalf("evicted=%v, want [fresh]", evicted)
	}
	if _, ok := h.LastSeen("fresh"); ok {
		t.Fatal("fresh connection still tracked after eviction")
	}
	if _, ok := h.LastSeen("stale"); !ok {
		t.Fatal("stale connection should still be tracked")
	}
	if h.Len() != 1 {
		t.Fatalf("tracked=%d, want 1", h.Len())
	}
}

func TestConcurrentBeat(t *testing.T) {
	h, now := newTestHeartbeat(t)
	const workers, rounds = 64, 200
	var wg sync.WaitGroup
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			id := fmt.Sprintf("conn-%d", i%8)
			for j := 0; j < rounds; j++ {
				h.Beat(id)
			}
		}(i)
	}
	wg.Wait()
	if h.Len() != 8 {
		t.Fatalf("tracked=%d, want 8", h.Len())
	}
	if got := h.Evict(*now); len(got) != 0 {
		t.Fatalf("evicted fresh connections: %v", got)
	}
}

func TestHeartbeatRejectsTimeoutNotGreaterThanInterval(t *testing.T) {
	if _, err := NewHeartbeat(30*time.Second, 30*time.Second); err == nil {
		t.Fatal("timeout equal to interval should fail")
	}
	if _, err := NewHeartbeat(0, 45*time.Second); err == nil {
		t.Fatal("zero interval should fail")
	}
}

cd .local/checks/gs-practice-lab && gofmt -l ch05/scale && go vet ./ch05/scale/ && go test ./ch05/scale/ -count=1

ok  	gspractice/ch05/scale	0.311s

cd .local/checks/gs-practice-lab && go test ./ch05/scale/ -race -count=1

ok  	gspractice/ch05/scale	1.516s