《游戏服务端架构·核心系统与契约》3.3 实例分配、暖池与回收

// file: internal/instancepool/pool.go
// tier: L1
// verify: go test ./internal/instancepool/
package instancepool

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

type Spec struct {
	Region  string
	Version string
	CPU     int
	MemMB   int
}

type Node struct {
	ID       string
	Region   string
	Version  string
	CPUTotal int
	MemTotal int
	CPUUsed  int
	MemUsed  int
	Healthy  bool
}

func (n *Node) fits(s Spec) bool {
	return n.Healthy && n.Region == s.Region && n.Version == s.Version &&
		n.CPUTotal-n.CPUUsed >= s.CPU && n.MemTotal-n.MemUsed >= s.MemMB
}

type Instance struct {
	ID     string
	NodeID string
	Spec   Spec
}

type Lease struct {
	InstanceID string
	SessionID  string
	NodeID     string
	Spec       Spec
	Epoch      uint64
	ExpiresAt  time.Time
}

type Stats struct {
	Healthy   int
	Warm      int
	Active    int
	FreeCPU   int
	FreeMemMB int
}

var (
	ErrEmptySession = errors.New("instancepool: empty session id")
	ErrInvalidSpec  = errors.New("instancepool: invalid spec")
	ErrInvalidNode  = errors.New("instancepool: invalid node")
	ErrNodeExists   = errors.New("instancepool: healthy node already registered")
	ErrDuplicate    = errors.New("instancepool: session already holds a lease")
	ErrNoCapacity   = errors.New("instancepool: no node fits the spec")
	ErrUnknownLease = errors.New("instancepool: unknown lease")
	ErrStaleLease   = errors.New("instancepool: stale lease epoch")
)

type Pool struct {
	mu        sync.Mutex
	now       func() time.Time
	ttl       time.Duration
	nodes     map[string]*Node
	warm      map[Spec][]Instance
	active    map[string]Lease
	bySession map[string]string
	seq       int
	epoch     uint64
}

func New(ttl time.Duration, now func() time.Time) *Pool {
	if now == nil {
		now = time.Now
	}
	return &Pool{now: now, ttl: ttl, nodes: map[string]*Node{},
		warm: map[Spec][]Instance{}, active: map[string]Lease{}, bySession: map[string]string{}}
}

func validateSpec(s Spec) error {
	if s.Region == "" || s.Version == "" || s.CPU <= 0 || s.MemMB <= 0 {
		return ErrInvalidSpec
	}
	return nil
}

func (p *Pool) AddNode(n Node) error {
	if n.ID == "" || n.Region == "" || n.Version == "" || n.CPUTotal <= 0 || n.MemTotal <= 0 {
		return ErrInvalidNode
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if old, ok := p.nodes[n.ID]; ok && old.Healthy {
		return ErrNodeExists
	}
	n.CPUUsed, n.MemUsed, n.Healthy = 0, 0, true
	p.nodes[n.ID] = &n
	return nil
}

func (p *Pool) bestNodeLocked(s Spec) *Node {
	var best *Node
	bestCPU, bestMem := 0, 0
	for _, n := range p.nodes {
		if !n.fits(s) {
			continue
		}
		leftCPU, leftMem := n.CPUTotal-n.CPUUsed-s.CPU, n.MemTotal-n.MemUsed-s.MemMB
		if best == nil || leftCPU < bestCPU || (leftCPU == bestCPU &&
			(leftMem < bestMem || (leftMem == bestMem && n.ID < best.ID))) {
			best, bestCPU, bestMem = n, leftCPU, leftMem
		}
	}
	return best
}
func (p *Pool) newInstanceLocked(n *Node, s Spec) Instance {
	n.CPUUsed += s.CPU
	n.MemUsed += s.MemMB
	p.seq++
	return Instance{ID: fmt.Sprintf("inst-%04d", p.seq), NodeID: n.ID, Spec: s}
}
func (p *Pool) warmLocked(s Spec, n int) int {
	created := 0
	for i := 0; i < n; i++ {
		node := p.bestNodeLocked(s)
		if node == nil {
			break
		}
		p.warm[s] = append(p.warm[s], p.newInstanceLocked(node, s))
		created++
	}
	return created
}
func (p *Pool) Warm(s Spec, n int) int {
	if n <= 0 || validateSpec(s) != nil {
		return 0
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.warmLocked(s, n)
}
func (p *Pool) Refill(s Spec, target int) int {
	if target <= 0 || validateSpec(s) != nil {
		return 0
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	if have := len(p.warm[s]); have >= target {
		return 0
	}
	return p.warmLocked(s, target-len(p.warm[s]))
}
func (p *Pool) Acquire(session string, s Spec) (Lease, error) {
	if session == "" {
		return Lease{}, ErrEmptySession
	}
	if err := validateSpec(s); err != nil {
		return Lease{}, err
	}
	p.mu.Lock()
	defer p.mu.Unlock()
	p.reclaimExpiredLocked()
	if _, ok := p.bySession[session]; ok {
		return Lease{}, ErrDuplicate
	}
	var inst Instance
	if list := p.warm[s]; len(list) > 0 {
		inst = list[len(list)-1]
		p.warm[s] = list[:len(list)-1]
	} else {
		node := p.bestNodeLocked(s)
		if node == nil {
			return Lease{}, ErrNoCapacity
		}
		inst = p.newInstanceLocked(node, s)
	}
	p.epoch++
	lease := Lease{InstanceID: inst.ID, SessionID: session, NodeID: inst.NodeID,
		Spec: s, Epoch: p.epoch, ExpiresAt: p.now().Add(p.ttl)}
	p.active[inst.ID] = lease
	p.bySession[session] = inst.ID
	return lease, nil
}

func (p *Pool) Release(instanceID string, epoch uint64) error {
	return p.release(instanceID, epoch, true)
}

func (p *Pool) Evict(instanceID string, epoch uint64) error {
	return p.release(instanceID, epoch, false)
}
func (p *Pool) release(instanceID string, epoch uint64, reuse bool) error {
	p.mu.Lock()
	defer p.mu.Unlock()
	lease, ok := p.active[instanceID]
	if !ok {
		return ErrUnknownLease
	}
	if lease.Epoch != epoch {
		return ErrStaleLease
	}
	delete(p.active, instanceID)
	delete(p.bySession, lease.SessionID)
	p.freeLocked(lease, reuse)
	return nil
}
func (p *Pool) freeLocked(lease Lease, reuse bool) {
	node := p.nodes[lease.NodeID]
	if node == nil || !node.Healthy {
		return // FailNode already zeroed every reservation on this node.
	}
	if reuse {
		p.warm[lease.Spec] = append(p.warm[lease.Spec],
			Instance{ID: lease.InstanceID, NodeID: lease.NodeID, Spec: lease.Spec})
		return
	}
	node.CPUUsed -= lease.Spec.CPU
	node.MemUsed -= lease.Spec.MemMB
}

func (p *Pool) ReclaimExpired() []Lease {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.reclaimExpiredLocked()
}
func (p *Pool) reclaimExpiredLocked() []Lease {
	now := p.now()
	var out []Lease
	for id, lease := range p.active {
		if now.Before(lease.ExpiresAt) {
			continue
		}
		out = append(out, lease)
		delete(p.active, id)
		delete(p.bySession, lease.SessionID)
		p.freeLocked(lease, false)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].InstanceID < out[j].InstanceID })
	return out
}
func (p *Pool) FailNode(nodeID string) []Lease {
	p.mu.Lock()
	defer p.mu.Unlock()
	n, ok := p.nodes[nodeID]
	if !ok {
		return nil
	}
	n.Healthy = false
	var out []Lease
	for id, lease := range p.active {
		if lease.NodeID == nodeID {
			out = append(out, lease)
			delete(p.active, id)
			delete(p.bySession, lease.SessionID)
		}
	}
	for spec, list := range p.warm {
		kept := list[:0]
		for _, inst := range list {
			if inst.NodeID != nodeID {
				kept = append(kept, inst)
			}
		}
		p.warm[spec] = kept
	}
	n.CPUUsed, n.MemUsed = 0, 0
	sort.Slice(out, func(i, j int) bool { return out[i].InstanceID < out[j].InstanceID })
	return out
}
func (p *Pool) Stats() Stats {
	p.mu.Lock()
	defer p.mu.Unlock()
	st := Stats{Active: len(p.active)}
	for _, n := range p.nodes {
		if n.Healthy {
			st.Healthy++
			st.FreeCPU += n.CPUTotal - n.CPUUsed
			st.FreeMemMB += n.MemTotal - n.MemUsed
		}
	}
	for _, list := range p.warm {
		st.Warm += len(list)
	}
	return st
}