《MMO 服务器开发》11.1 请求幂等窗口

// file: internal/idem/window.go
// tier: L1
// verify: go test ./internal/idem/ -run 'TestWindow|TestCountWindow' -v

package idem

import (
	"sync"
	"time"
)

// Window is a sliding dedup window: it maps a uint64 key to an opaque uint64
// value and forgets the pair once the key has been silent for ttl.
//
// Memory is bounded by construction. Keys live in one of n time buckets, and a
// bucket is dropped whole when the window slides past it, so the live set
// always spans exactly ttl: with arrival rate r the entry count stays near
// r*ttl instead of growing with the number of keys ever seen.
type Window struct {
	mu      sync.Mutex
	ttl     time.Duration
	step    time.Duration
	cur     int
	total   int
	evicted int
	buckets []bucket
}

type bucket struct {
	start time.Time
	set   map[uint64]uint64
}

// NewWindow builds a window that remembers keys for ttl across n buckets
// (n >= 2). A larger n makes eviction finer grained (granularity ttl/n) at the
// cost of one extra map lookup per probe.
func NewWindow(ttl time.Duration, n int) *Window {
	if n < 2 {
		n = 2
	}
	w := &Window{
		ttl:     ttl,
		step:    ttl / time.Duration(n),
		buckets: make([]bucket, n),
	}
	for i := range w.buckets {
		w.buckets[i].set = make(map[uint64]uint64)
	}
	return w
}

// Put records key->val at now and returns the value that was live before, if
// any. existed=true is the duplicate signal: the key was already in the window.
func (w *Window) Put(key, val uint64, now time.Time) (prev uint64, existed bool) {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.advance(now)
	for i := range w.buckets {
		if v, ok := w.buckets[i].set[key]; ok {
			w.buckets[i].set[key] = val
			return v, true
		}
	}
	w.buckets[w.cur].set[key] = val
	w.total++
	return 0, false
}

// Lookup reads key without recording anything. It still advances the window,
// so a stale bucket can never answer a fresh probe.
func (w *Window) Lookup(key uint64, now time.Time) (uint64, bool) {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.advance(now)
	for i := range w.buckets {
		if v, ok := w.buckets[i].set[key]; ok {
			return v, true
		}
	}
	return 0, false
}

// Seen is the set-only view for plain dedup.
func (w *Window) Seen(key uint64, now time.Time) bool {
	_, existed := w.Put(key, 0, now)
	return existed
}

// Stats returns the live key count and the number of slots reclaimed so far.
func (w *Window) Stats() (live, evicted int) {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.total, w.evicted
}

func (w *Window) advance(now time.Time) {
	if w.buckets[w.cur].start.IsZero() {
		w.buckets[w.cur].start = now
		return
	}
	// A long stall or a clock jump would otherwise walk the ring once per step;
	// past a whole ttl there is nothing left to keep.
	if now.Sub(w.buckets[w.cur].start) >= w.ttl {
		w.reset(now)
		return
	}
	for now.Sub(w.buckets[w.cur].start) >= w.step {
		next := (w.cur + 1) % len(w.buckets)
		w.drop(next)
		w.buckets[next].start = w.buckets[w.cur].start.Add(w.step)
		w.cur = next
	}
}

func (w *Window) drop(i int) {
	n := len(w.buckets[i].set)
	w.evicted += n
	w.total -= n
	clear(w.buckets[i].set)
}

func (w *Window) reset(now time.Time) {
	for i := range w.buckets {
		w.drop(i)
		w.buckets[i].start = time.Time{}
	}
	w.cur = 0
	w.buckets[0].start = now
}

// snippet: internal/idem/countwindow.go（节选）

// Seen records key and reports whether it was among the last n keys.
func (w *CountWindow) Seen(key uint64) bool {
	w.mu.Lock()
	defer w.mu.Unlock()
	if _, ok := w.set[key]; ok {
		return true
	}
	if w.full {
		delete(w.set, w.ring[w.pos])
	}
	w.ring[w.pos] = key
	w.set[key] = struct{}{}
	w.pos++
	if w.pos == w.n {
		w.pos = 0
		w.full = true
	}
	return false
}

// snippet: internal/idem/window_test.go（节选）

func TestWindowForgetsAfterTTL(t *testing.T) {
	w := NewWindow(100*time.Millisecond, 4)
	w.Seen(7, base)
	if _, ok := w.Lookup(7, base.Add(50*time.Millisecond)); !ok {
		t.Fatal("key vanished before ttl")
	}
	// Past a full ttl the key is gone; a very late repeat looks new again.
	if _, ok := w.Lookup(7, base.Add(500*time.Millisecond)); ok {
		t.Fatal("key survived past ttl")
	}
	if w.Seen(7, base.Add(500*time.Millisecond)) {
		t.Fatal("late repeat still reported as duplicate")
	}
}

func TestWindowMemoryBounded(t *testing.T) {
	// 100k keys, one per millisecond. ttl is 100ms, so at most ~100 keys can be
	// live at once no matter how many have been seen in total.
	const total = 100_000
	w := NewWindow(100*time.Millisecond, 4)
	for i := 0; i < total; i++ {
		w.Seen(uint64(i), base.Add(time.Duration(i)*time.Millisecond))
	}
	live, evicted := w.Stats()
	if live > 110 {
		t.Fatalf("live keys = %d, want <= 110", live)
	}
	if live+evicted != total {
		t.Fatalf("live+evicted = %d, want %d", live+evicted, total)
	}
}