《游戏服务端架构·品类、状态与拓扑》11.2 有状态服务的边界

// file: internal/stateful/stateful.go
// tier: L1
// verify: go test ./internal/stateful/
package stateful

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

var (
	ErrEmptyPartition = errors.New("stateful: partition must not be blank")
	ErrEmptyOwner     = errors.New("stateful: owner must not be blank")
	ErrNoLease        = errors.New("stateful: partition has no lease")
	ErrLeaseExpired   = errors.New("stateful: lease expired; take over with a new epoch")
	ErrNotOwner       = errors.New("stateful: owner or epoch does not match the lease")
	ErrDualPrimary    = errors.New("stateful: another owner holds a live lease")
	ErrDraining       = errors.New("stateful: partition is draining; new owners are rejected")
)

// Lease is the right to serve one partition. Epoch increases on every
// takeover, so a paused old owner cannot write after the new owner starts.
type Lease struct {
	Partition string
	Owner     string
	Epoch     uint64
	ExpiresAt time.Time
}

// Registry is the lease store's contract. In production the same operations
// run as conditional writes (compare owner and epoch) in etcd, Redis or a
// database; the map here only makes the rules executable in tests.
type Registry struct {
	mu       sync.Mutex
	now      func() time.Time
	ttl      time.Duration
	leases   map[string]Lease
	draining map[string]bool
}

func NewRegistry(ttl time.Duration, now func() time.Time) *Registry {
	if now == nil {
		now = time.Now
	}
	return &Registry{
		now: now, ttl: ttl,
		leases:   make(map[string]Lease),
		draining: make(map[string]bool),
	}
}

// Acquire takes a partition for owner. A live lease held by the same owner is
// renewed; a live lease held by anyone else is a dual-primary attempt. An
// expired lease is taken over with a higher epoch, unless the partition is
// draining, in which case the migration must finish first.
func (r *Registry) Acquire(partition, owner string) (Lease, error) {
	if strings.TrimSpace(partition) == "" {
		return Lease{}, ErrEmptyPartition
	}
	if strings.TrimSpace(owner) == "" {
		return Lease{}, ErrEmptyOwner
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	now := r.now()
	if old, ok := r.leases[partition]; ok && now.Before(old.ExpiresAt) {
		if old.Owner != owner {
			if r.draining[partition] {
				return Lease{}, ErrDraining
			}
			return Lease{}, ErrDualPrimary
		}
		old.ExpiresAt = now.Add(r.ttl)
		r.leases[partition] = old
		return old, nil
	}
	if r.draining[partition] {
		return Lease{}, ErrDraining
	}
	lease := Lease{Partition: partition, Owner: owner, Epoch: 1, ExpiresAt: now.Add(r.ttl)}
	if old, ok := r.leases[partition]; ok {
		lease.Epoch = old.Epoch + 1
	}
	r.leases[partition] = lease
	return lease, nil
}

// Renew extends a live lease. The owner and epoch must both match: matching
// only the owner would let a paused process resume after a takeover.
func (r *Registry) Renew(partition, owner string, epoch uint64) (Lease, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	lease, ok := r.leases[partition]
	if !ok {
		return Lease{}, ErrNoLease
	}
	if !r.now().Before(lease.ExpiresAt) {
		return Lease{}, ErrLeaseExpired
	}
	if lease.Owner != owner {
		return Lease{}, ErrDualPrimary
	}
	if lease.Epoch != epoch {
		return Lease{}, ErrNotOwner
	}
	lease.ExpiresAt = r.now().Add(r.ttl)
	r.leases[partition] = lease
	return lease, nil
}

// Drain marks a partition for migration. Renew still works so the current
// owner can reach a checkpoint; Acquire for a new owner is rejected until the
// current owner releases.
func (r *Registry) Drain(partition string) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	if _, ok := r.leases[partition]; !ok {
		return ErrNoLease
	}
	r.draining[partition] = true
	return nil
}

// Release drops the lease and clears the drain mark. The epoch check rejects a
// stale owner that was replaced while it was paused.
func (r *Registry) Release(partition, owner string, epoch uint64) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	lease, ok := r.leases[partition]
	if !ok {
		return ErrNoLease
	}
	if lease.Owner != owner || lease.Epoch != epoch {
		return ErrNotOwner
	}
	delete(r.leases, partition)
	delete(r.draining, partition)
	return nil
}

// Current returns a copy for status pages and for the drain controller.
func (r *Registry) Current(partition string) (Lease, bool) {
	r.mu.Lock()
	defer r.mu.Unlock()
	lease, ok := r.leases[partition]
	return lease, ok
}

// Partitions lists lease keys in sorted order so snapshots are reproducible.
func (r *Registry) Partitions() []string {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := make([]string, 0, len(r.leases))
	for partition := range r.leases {
		out = append(out, partition)
	}
	sort.Strings(out)
	return out
}

$ cd .local/checks/gsarch-lab
$ gofmt -l internal/stateful
$ go vet ./internal/stateful/
$ go test -count=1 ./internal/stateful/
ok  	gsarchlab/internal/stateful	0.245s