《MMO 服务器开发》7.2 号段模式

  主段 [cur ───────────────► max]      备用段 [nextCur ───► nextMax]
        ▲                                     ▲
        │ 剩余 < step/2 时触发后台预取 ────────┘
        │
   主段耗尽：备用段就绪 → 直接切换（零等待）
             备用段未就绪 → 同步补取（阻塞一次）

// file: internal/idgen/segment.go
// tier: L1
// verify: go test ./internal/idgen/
package idgen

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/redis/go-redis/v9"
)

// ErrNoSegment 当前段耗尽且预取失败。
var ErrNoSegment = errors.New("idgen: no segment available")

// segmentSource 是号段来源抽象：一次「取号段」返回闭区间 [start, end]。
// 生产上用 Redis 或 DB 实现；测试里可以用内存实现。
type segmentSource interface {
	NextSegment(ctx context.Context, key string, step int64) (start, end int64, err error)
}

// RedisSource 用 INCRBY 原子取号段。多个进程用同一个 key 时，各自拿到的
// 号段天然不相交——这是号段模式不需要协调的根本原因。
type RedisSource struct {
	Client *redis.Client
}

// NextSegment 原子地把 key 推进 step，返回刚划走的那一段。
func (r *RedisSource) NextSegment(ctx context.Context, key string, step int64) (int64, int64, error) {
	end, err := r.Client.IncrBy(ctx, key, step).Result()
	if err != nil {
		return 0, 0, err
	}
	return end - step + 1, end, nil
}

// Segment 是内存双 buffer 的号段分配器：
//
//	主段 [cur, max] 正在发放；剩余不足一半时后台异步预取备用段；
//	主段耗尽瞬间，若备用段已就绪则直接切换，否则同步补取（会阻塞一次）。
type Segment struct {
	mu       sync.Mutex
	src      segmentSource
	key      string
	step     int64
	cur      int64
	max      int64
	nextCur  int64
	nextMax  int64
	nextOK   bool
	loading  bool
	lastErr  error
	loaded   chan struct{} // 一次预取完成的广播信号
	prefetch int64         // 已异步预取的次数（供测试观测）
}

// NewSegment 构造分配器并同步取好第一段。
func NewSegment(ctx context.Context, src segmentSource, key string, step int64) (*Segment, error) {
	if step <= 0 {
		return nil, fmt.Errorf("idgen: step must be positive, got %d", step)
	}
	s := &Segment{src: src, key: key, step: step, loaded: make(chan struct{})}
	cur, max, err := src.NextSegment(ctx, key, step)
	if err != nil {
		return nil, err
	}
	s.cur, s.max = cur, max
	return s, nil
}

// Next 返回下一个全局唯一 ID。并发安全。
func (s *Segment) Next(ctx context.Context) (uint64, error) {
	for {
		s.mu.Lock()
		if s.cur <= s.max {
			id := s.cur
			s.cur++
			if s.remaining() < s.step/2 && !s.nextOK && !s.loading {
				s.loading = true
				s.prefetch++
				// 预取脱离调用方 ctx：它必须活过本次 Next。
				go s.load(context.Background())
			}
			s.mu.Unlock()
			return uint64(id), nil
		}
		// 主段耗尽
		if s.nextOK {
			s.cur, s.max, s.nextOK = s.nextCur, s.nextMax, false
			s.mu.Unlock()
			continue
		}
		if s.loading {
			done := s.loaded
			s.mu.Unlock()
			<-done
			continue
		}
		// 既无备用段又没在加载：同步补取
		s.mu.Unlock()
		if err := s.syncLoad(ctx); err != nil {
			return 0, err
		}
	}
}

func (s *Segment) remaining() int64 { return s.max - s.cur + 1 }

// load 在后台取下一段并写进备用槽。
func (s *Segment) load(ctx context.Context) {
	start, end, err := s.src.NextSegment(ctx, s.key, s.step)
	s.mu.Lock()
	s.loading = false
	if err != nil {
		s.lastErr = err
	} else {
		s.nextCur, s.nextMax, s.nextOK = start, end, true
	}
	old := s.loaded
	s.loaded = make(chan struct{})
	s.mu.Unlock()
	close(old)
}

func (s *Segment) syncLoad(ctx context.Context) error {
	start, end, err := s.src.NextSegment(ctx, s.key, s.step)
	if err != nil {
		return err
	}
	s.mu.Lock()
	s.cur, s.max = start, end
	s.mu.Unlock()
	return nil
}

// PrefetchCount 返回异步预取次数，供测试断言双 buffer 真的生效。
func (s *Segment) PrefetchCount() int64 {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.prefetch
}

// WaitPrefetch 阻塞到当前正在进行的预取结束（或超时），供测试用。
func (s *Segment) WaitPrefetch(d time.Duration) {
	s.mu.Lock()
	if !s.loading {
		s.mu.Unlock()
		return
	}
	done := s.loaded
	s.mu.Unlock()
	select {
	case <-done:
	case <-time.After(d):
	}
}

// file: internal/idgen/segment_redis_test.go
// tier: L2
// verify: go test ./internal/idgen/ -run TestSegmentRedisTwoInstancesDisjoint -v
package idgen

import (
	"context"
	"fmt"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/redis/go-redis/v9"
)

func TestSegmentRedisTwoInstancesDisjoint(t *testing.T) {
	ctx := context.Background()
	addr := os.Getenv("AEGIS_REDIS_ADDR")
	if addr == "" {
		addr = "127.0.0.1:26379"
	}
	c := redis.NewClient(&redis.Options{Addr: addr})
	defer c.Close()
	if err := c.Ping(ctx).Err(); err != nil {
		t.Skipf("Redis 不可达（%s），跳过 L2：%v", addr, err)
	}
	key := fmt.Sprintf("aegis:test:ch07:seg:%d", os.Getpid())
	// 后台预取可能在测试收尾后仍 INCRBY 一次，故等它落定后再删、并补删一次。
	defer func() {
		time.Sleep(150 * time.Millisecond)
		c.Del(ctx, key)
		time.Sleep(150 * time.Millisecond)
		c.Del(ctx, key)
	}()

	src := &RedisSource{Client: c}
	const step = 1000
	a, err := NewSegment(ctx, src, key, step)
	if err != nil {
		t.Fatalf("new a: %v", err)
	}
	b, err := NewSegment(ctx, src, key, step)
	if err != nil {
		t.Fatalf("new b: %v", err)
	}

	const per = 2500
	seen := make(map[uint64]struct{}, per*2)
	var mu sync.Mutex
	consume := func(seg *Segment, tag string) {
		local := make([]uint64, 0, per)
		for i := 0; i < per; i++ {
			id, err := seg.Next(ctx)
			if err != nil {
				t.Errorf("%s next: %v", tag, err)
				return
			}
			local = append(local, id)
		}
		mu.Lock()
		for _, id := range local {
			seen[id] = struct{}{}
		}
		mu.Unlock()
	}
	var wg sync.WaitGroup
	wg.Add(2)
	go func() { defer wg.Done(); consume(a, "a") }()
	go func() { defer wg.Done(); consume(b, "b") }()
	wg.Wait()

	if len(seen) != per*2 {
		t.Fatalf("两实例应共产出 %d 个互斥 ID，实际 %d", per*2, len(seen))
	}
	// 号段模式下 Redis 只被 INCRBY 了 ceil(5000/1000)=5 次
	got, err := c.Get(ctx, key).Int64()
	if err != nil {
		t.Fatalf("get key: %v", err)
	}
	if got%step != 0 || got < per*2 {
		t.Fatalf("Redis 计数应为 %d 的整数倍且不小于 %d，实际 %d", step, per*2, got)
	}
	t.Logf("两实例共 %d 个 ID，Redis INCRBY 次数 = %d（号段步长 %d）", per*2, got/step, step)
}

两实例共 5000 个 ID，Redis INCRBY 次数 = 6（号段步长 1000）