《后端架构设计·分布式与数据一致性》5.2 再平衡：迁移、双写与客户端路由

// file: internal/platform/migration/migration.go
// tier: L2
// verify: go test ./internal/platform/migration/ -count=1
package migration

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

// ErrInvalidState is returned when a step is not allowed in the current state.
var ErrInvalidState = errors.New("invalid migration state")

// ErrIncomplete is returned when cutover runs before every key is copied.
var ErrIncomplete = errors.New("migration incomplete")

// Record is a versioned value. Version is the only conflict resolver, so a
// stale snapshot never overwrites a newer target value.
type Record struct {
	Value   string
	Version int64
}

// Store is the minimal shard interface used by the state machine.
type Store interface {
	Get(key string) (Record, bool)
	Put(key string, rec Record)
	ConditionalPut(key string, rec Record) bool
	Keys() []string
}

// MemoryStore is an in-process Store for tests and local drills.
type MemoryStore struct {
	mu   sync.Mutex
	data map[string]Record
}

// NewMemoryStore returns an empty store.
func NewMemoryStore() *MemoryStore { return &MemoryStore{data: make(map[string]Record)} }

// Get returns the stored record.
func (s *MemoryStore) Get(key string) (Record, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	rec, ok := s.data[key]
	return rec, ok
}

// Put writes unconditionally.
func (s *MemoryStore) Put(key string, rec Record) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.data[key] = rec
}

// ConditionalPut writes only when rec is newer than the current value.
func (s *MemoryStore) ConditionalPut(key string, rec Record) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	if old, ok := s.data[key]; ok && old.Version >= rec.Version {
		return false
	}
	s.data[key] = rec
	return true
}

// Keys returns all keys in sorted order.
func (s *MemoryStore) Keys() []string {
	s.mu.Lock()
	defer s.mu.Unlock()
	keys := make([]string, 0, len(s.data))
	for k := range s.data {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	return keys
}

// State is the migration state exposed to routing and operators.
type State string

const (
	StateIdle       State = "idle"
	StateDualWrite  State = "dual-write"
	StateBackfill   State = "backfill"
	StateDone       State = "done"
	StateRolledBack State = "rolled-back"
)

// Migration copies one shard to another with dual writes. Every step is
// idempotent: replaying a batch or resuming after a crash converges.
type Migration struct {
	Source Store
	Target Store

	state   State
	keys    []string
	journal map[string]bool
	order   []string
}

// NewMigration returns a migration in StateIdle.
func NewMigration(source, target Store) *Migration {
	return &Migration{Source: source, Target: target, state: StateIdle, journal: make(map[string]bool)}
}

// State returns the current state.
func (m *Migration) State() State { return m.state }

// Begin snapshots the source key list and turns on dual writes.
func (m *Migration) Begin() error {
	if m.state != StateIdle {
		return fmt.Errorf("migration: begin in state %s: %w", m.state, ErrInvalidState)
	}
	m.keys = m.Source.Keys()
	m.state = StateDualWrite
	return nil
}

// Write sends writes to both shards during migration, target after cutover,
// and source after rollback.
func (m *Migration) Write(key string, rec Record) error {
	switch m.state {
	case StateDualWrite, StateBackfill:
		m.Source.Put(key, rec)
		m.Target.Put(key, rec)
	case StateDone:
		m.Target.Put(key, rec)
	case StateRolledBack:
		m.Source.Put(key, rec)
	default:
		return fmt.Errorf("migration: write in state %s: %w", m.state, ErrInvalidState)
	}
	return nil
}

// Read reads the source until cutover and the target afterwards.
func (m *Migration) Read(key string) (Record, bool) {
	if m.state == StateDone {
		return m.Target.Get(key)
	}
	return m.Source.Get(key)
}

// CopyBatch copies at most limit not-yet-journaled keys. Replay is safe: the
// journal skips done keys and ConditionalPut rejects stale versions.
func (m *Migration) CopyBatch(limit int) (copied, remaining int, err error) {
	if m.state != StateDualWrite && m.state != StateBackfill {
		return 0, m.remaining(), fmt.Errorf("migration: copy in state %s: %w", m.state, ErrInvalidState)
	}
	if limit <= 0 {
		return 0, m.remaining(), fmt.Errorf("migration: batch limit must be positive, got %d", limit)
	}
	m.state = StateBackfill
	for _, key := range m.keys {
		if copied >= limit {
			break
		}
		if m.journal[key] {
			continue
		}
		if rec, ok := m.Source.Get(key); ok {
			m.Target.ConditionalPut(key, rec)
		}
		m.journal[key] = true
		m.order = append(m.order, key)
		copied++
	}
	return copied, m.remaining(), nil
}

// Cutover requires every snapshot key to be copied, then switches reads to
// the target. It moves no data, so it is fast and retryable.
func (m *Migration) Cutover() error {
	if m.state != StateBackfill {
		return fmt.Errorf("migration: cutover in state %s: %w", m.state, ErrInvalidState)
	}
	if left := m.remaining(); left > 0 {
		return fmt.Errorf("migration: %d keys remain: %w", left, ErrIncomplete)
	}
	m.state = StateDone
	return nil
}

// Rollback abandons the migration before cutover and keeps the source
// authoritative.
func (m *Migration) Rollback() error {
	switch m.state {
	case StateDualWrite, StateBackfill:
		m.state = StateRolledBack
		return nil
	default:
		return fmt.Errorf("migration: rollback in state %s: %w", m.state, ErrInvalidState)
	}
}

// SnapshotKeys returns the key list captured by Begin.
func (m *Migration) SnapshotKeys() []string { return append([]string(nil), m.keys...) }

// CopiedKeys returns journaled keys in copy order.
func (m *Migration) CopiedKeys() []string { return append([]string(nil), m.order...) }

// Remaining reports how many snapshot keys are not journaled.
func (m *Migration) Remaining() int { return m.remaining() }

func (m *Migration) remaining() int {
	n := 0
	for _, key := range m.keys {
		if !m.journal[key] {
			n++
		}
	}
	return n
}

// Resume rebuilds a worker from a snapshot key list and a copied-key journal.
func Resume(source, target Store, keys, copied []string) *Migration {
	m := NewMigration(source, target)
	m.state = StateBackfill
	m.keys = append([]string(nil), keys...)
	m.order = append([]string(nil), copied...)
	for _, key := range copied {
		m.journal[key] = true
	}
	return m
}

$ go test ./internal/platform/migration/ -v -count=1
=== RUN   TestMigrationResumesAfterCrash
    migration_test.go:54: resumed: first_copied=30 remaining_after_first=70 resumed_batches=4 resumed_copied=70
--- PASS: TestMigrationResumesAfterCrash (0.00s)
=== RUN   TestDualWriteKeepsTargetFresh
    migration_test.go:101: dual write survived backfill: value=paid version=2
--- PASS: TestDualWriteKeepsTargetFresh (0.00s)
=== RUN   TestStaleCopyDoesNotOverwriteNewerTarget
    migration_test.go:115: stale copy rejected: value=newer version=2
--- PASS: TestStaleCopyDoesNotOverwriteNewerTarget (0.00s)
=== RUN   TestCutoverRequiresCompleteBackfill
--- PASS: TestCutoverRequiresCompleteBackfill (0.00s)
=== RUN   TestRollbackReturnsWritesToSource
    migration_test.go:164: rollback state=rolled-back source_value=after-rollback
--- PASS: TestRollbackReturnsWritesToSource (0.00s)
PASS
ok  	archlab/internal/platform/migration	0.265s