《后端架构设计·分布式与数据一致性》10.1 分布式 ID：Snowflake、号段与时钟回拨

// file: internal/platform/idgen/snowflake.go
// tier: L1
// verify: go test ./idgen/ -run TestSnowflake -count=1
package idgen

import (
	"errors"
	"fmt"
	"sync"
	"time"
)

const (
	epochMillis int64 = 1704067200000 // 2024-01-01T00:00:00Z
	nodeBits          = 10
	stepBits          = 12
	maxNode     int64 = (1 << nodeBits) - 1
	maxStep     int64 = (1 << stepBits) - 1
	maxBackward       = 5 * time.Millisecond
)

var (
	// ErrClockBackwards means the wall clock jumped back too far to absorb.
	ErrClockBackwards = errors.New("idgen: clock moved backwards")
	// ErrInvalidNode means the node id does not fit in nodeBits.
	ErrInvalidNode = errors.New("idgen: node id out of range")
)

// Clock is the wall clock used by Snowflake, injectable for tests.
type Clock interface {
	Now() time.Time
	Sleep(time.Duration)
}

// SystemClock is the production clock.
type SystemClock struct{}

// Now returns the current wall time.
func (SystemClock) Now() time.Time { return time.Now() }

// Sleep pauses the caller for d.
func (SystemClock) Sleep(d time.Duration) { time.Sleep(d) }

// Snowflake packs millis|node|step into one int64.
type Snowflake struct {
	mu     sync.Mutex
	node   int64
	clock  Clock
	lastMS int64
	step   int64
}

// NewSnowflake validates node and returns a generator.
func NewSnowflake(node int64, clock Clock) (*Snowflake, error) {
	if node < 0 || node > maxNode {
		return nil, fmt.Errorf("snowflake: node %d: %w", node, ErrInvalidNode)
	}
	return &Snowflake{node: node, clock: clock, lastMS: -1}, nil
}

// Next returns the next id. A backwards jump up to maxBackward is absorbed by
// waiting; a larger jump fails rather than risking duplicate ids.
func (s *Snowflake) Next() (int64, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	nowMS := s.clock.Now().UnixMilli()
	if nowMS < s.lastMS {
		drift := s.lastMS - nowMS
		if drift > maxBackward.Milliseconds() {
			return 0, fmt.Errorf("snowflake: now=%d last=%d drift=%dms: %w", nowMS, s.lastMS, drift, ErrClockBackwards)
		}
		s.clock.Sleep(time.Duration(drift) * time.Millisecond)
		nowMS = s.lastMS
	}
	if nowMS == s.lastMS {
		s.step++
		if s.step > maxStep {
			for nowMS <= s.lastMS {
				s.clock.Sleep(time.Millisecond)
				nowMS = s.clock.Now().UnixMilli()
			}
			s.step = 0
		}
	} else {
		s.step = 0
	}
	s.lastMS = nowMS
	return ((nowMS - epochMillis) << (nodeBits + stepBits)) | (s.node << stepBits) | s.step, nil
}

// file: internal/platform/idgen/segment.go
// tier: L1
// verify: go test ./idgen/ -run TestSegment -count=1
package idgen

import (
	"errors"
	"fmt"
	"sync"
)

const prefetchAt int64 = 200

// ErrStoreUnavailable means the segment store could not be reached.
var ErrStoreUnavailable = errors.New("idgen: segment store unavailable")

// Segment is a half-open range [Start, End).
type Segment struct {
	Start int64
	End   int64
}

// SegmentStore hands out monotonically increasing segments. The production
// implementation is one UPDATE ... RETURNING transaction per call.
type SegmentStore interface {
	NextSegment(name string, step int64) (Segment, error)
}

// MemorySegmentStore is a deterministic in-process SegmentStore.
type MemorySegmentStore struct {
	mu    sync.Mutex
	next  int64
	fail  bool
	calls int
}

// NewMemorySegmentStore starts allocating at start.
func NewMemorySegmentStore(start int64) *MemorySegmentStore {
	return &MemorySegmentStore{next: start}
}

// NextSegment returns the next range or fails when injected.
func (s *MemorySegmentStore) NextSegment(_ string, step int64) (Segment, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.calls++
	if s.fail {
		return Segment{}, ErrStoreUnavailable
	}
	seg := Segment{Start: s.next, End: s.next + step}
	s.next = seg.End
	return seg, nil
}

// SetFail toggles the injected store failure.
func (s *MemorySegmentStore) SetFail(fail bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.fail = fail
}

// Calls reports how many segment requests were made, including failures.
func (s *MemorySegmentStore) Calls() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.calls
}

// SegmentAllocator serves ids from a current segment and keeps one prefetched
// segment as a buffer.
type SegmentAllocator struct {
	mu      sync.Mutex
	store   SegmentStore
	name    string
	step    int64
	cur     *Segment
	next    *Segment
	curID   int64
	fetches int
}

// NewSegmentAllocator returns an allocator with no buffered segment.
func NewSegmentAllocator(store SegmentStore, name string, step int64) *SegmentAllocator {
	return &SegmentAllocator{store: store, name: name, step: step}
}

// Next returns the next id, fetching or switching segments as needed.
func (a *SegmentAllocator) Next() (int64, error) {
	a.mu.Lock()
	defer a.mu.Unlock()

	if a.cur == nil || a.curID >= a.cur.End {
		if a.next != nil {
			a.cur, a.next = a.next, nil
		} else {
			seg, err := a.store.NextSegment(a.name, a.step)
			if err != nil {
				return 0, fmt.Errorf("idgen: allocate %q: %w", a.name, err)
			}
			a.fetches++
			a.cur = &seg
		}
		a.curID = a.cur.Start
	}
	id := a.curID
	a.curID++
	if a.next == nil && a.cur.End-a.curID <= prefetchAt {
		// Prefetch failure is recorded by the store and retried on the next
		// call; the current segment keeps serving ids in the meantime.
		if seg, err := a.store.NextSegment(a.name, a.step); err == nil {
			a.fetches++
			a.next = &seg
		}
	}
	return id, nil
}

// Fetches reports how many segment ranges were obtained.
func (a *SegmentAllocator) Fetches() int {
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.fetches
}

$ go test ./idgen/ -v -count=1
=== RUN   TestSnowflakeLayout
--- PASS: TestSnowflakeLayout (0.00s)
=== RUN   TestSnowflakeInvalidNode
--- PASS: TestSnowflakeInvalidNode (0.00s)
=== RUN   TestSnowflakeMonotonic
--- PASS: TestSnowflakeMonotonic (0.00s)
=== RUN   TestSnowflakeConcurrentUnique
--- PASS: TestSnowflakeConcurrentUnique (0.01s)
=== RUN   TestSnowflakeRollbackAbsorbed
--- PASS: TestSnowflakeRollbackAbsorbed (0.00s)
=== RUN   TestSnowflakeLargeRollback
--- PASS: TestSnowflakeLargeRollback (0.00s)
=== RUN   TestSnowflakeSequenceExhaustion
--- PASS: TestSnowflakeSequenceExhaustion (0.00s)
=== RUN   TestSegmentUniqueAndMonotonic
--- PASS: TestSegmentUniqueAndMonotonic (0.00s)
=== RUN   TestSegmentRecoversAfterStoreFailure
--- PASS: TestSegmentRecoversAfterStoreFailure (0.00s)
=== RUN   TestSegmentConcurrentUnique
--- PASS: TestSegmentConcurrentUnique (0.00s)
PASS
ok  	archlab/idgen	0.295s
$ go test ./idgen/ -race -count=1
ok  	archlab/idgen	1.354s