《MMO 服务器开发》11.2 多进程编排与验收

// file: internal/battlecluster/orchestrator.go
// tier: L2
// verify: go test -run TestOrchestrator ./internal/battlecluster/
package battlecluster

import (
	"bufio"
	"fmt"
	"io"
	"net"
	"os/exec"
	"strings"
	"sync"
	"time"
)

// NodeSpec describes one process to run. ArgsFor, when set, computes the
// arguments at spawn time from the addresses of already-started nodes, so a
// gateway can be pointed at the hub without the caller knowing its address.
type NodeSpec struct {
	Name    string
	Bin     string
	Args    []string
	ArgsFor func(get func(name string) string) []string
}

// managed is a running node plus the state the orchestrator tracks for it.
type managed struct {
	spec     NodeSpec
	cmd      *exec.Cmd
	addr     string
	restarts int
}

// Orchestrator brings the cluster up in dependency order, gates each step on a
// passing health check, restarts crashed nodes, and shuts down in reverse order.
type Orchestrator struct {
	out         io.Writer
	order       []string
	specs       map[string]NodeSpec
	maxRestarts int

	mu    sync.Mutex
	procs map[string]*managed
	down  bool
}

// New builds an orchestrator; nodes start in the order given.
func New(specs []NodeSpec, out io.Writer) *Orchestrator {
	o := &Orchestrator{out: out, specs: map[string]NodeSpec{}, procs: map[string]*managed{}, maxRestarts: 5}
	for _, s := range specs {
		o.order = append(o.order, s.Name)
		o.specs[s.Name] = s
	}
	return o
}

// StartAll starts every node in order. Each node must answer a heartbeat before
// the next is started, so no node comes up before its dependency.
func (o *Orchestrator) StartAll() error {
	for _, name := range o.order {
		m, err := o.spawn(o.specs[name])
		if err != nil {
			return fmt.Errorf("start %s: %w", name, err)
		}
		o.mu.Lock()
		o.procs[name] = m
		o.mu.Unlock()
		fmt.Fprintf(o.out, "UP %s %s\n", name, m.addr)
		go o.supervise(name)
	}
	return nil
}

// spawn runs one process, reads its advertised address, and blocks until it
// answers a heartbeat (or the readiness budget runs out).
func (o *Orchestrator) spawn(spec NodeSpec) (*managed, error) {
	args := spec.Args
	if spec.ArgsFor != nil {
		args = spec.ArgsFor(o.Addr)
	}
	cmd := exec.Command(spec.Bin, args...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}
	cmd.Stderr = o.out
	if err := cmd.Start(); err != nil {
		return nil, err
	}

	addrCh := make(chan string, 1)
	go func() {
		sc := bufio.NewScanner(stdout)
		for sc.Scan() {
			if f := strings.Fields(sc.Text()); len(f) >= 2 && f[0] == "LISTEN" {
				addrCh <- f[1]
				return
			}
		}
	}()

	var addr string
	select {
	case addr = <-addrCh:
	case <-time.After(5 * time.Second):
		_ = cmd.Process.Kill()
		return nil, fmt.Errorf("%s never advertised a LISTEN address", spec.Name)
	}
	if err := o.waitHealthy(addr, 5*time.Second); err != nil {
		_ = cmd.Process.Kill()
		return nil, err
	}
	return &managed{spec: spec, cmd: cmd, addr: addr}, nil
}

// waitHealthy polls the node's heartbeat until it answers or the budget ends.
func (o *Orchestrator) waitHealthy(addr string, budget time.Duration) error {
	deadline := time.Now().Add(budget)
	for time.Now().Before(deadline) {
		if h, err := Probe(addr); err == nil && h.Healthy {
			return nil
		}
		time.Sleep(20 * time.Millisecond)
	}
	return fmt.Errorf("node at %s never became healthy", addr)
}

// supervise waits for its node to exit and restarts it unless a shutdown is in
// progress. Restarts are bounded so a crash loop cannot spin forever.
func (o *Orchestrator) supervise(name string) {
	for {
		o.mu.Lock()
		m := o.procs[name]
		o.mu.Unlock()
		_ = m.cmd.Wait()

		o.mu.Lock()
		if o.down {
			o.mu.Unlock()
			return
		}
		m.restarts++
		over, spec := m.restarts > o.maxRestarts, m.spec
		o.mu.Unlock()
		if over {
			fmt.Fprintf(o.out, "GIVEUP %s\n", name)
			return
		}
		nm, err := o.spawn(spec)
		if err != nil {
			fmt.Fprintf(o.out, "RESTART_FAIL %s %v\n", name, err)
			time.Sleep(100 * time.Millisecond)
			continue
		}
		o.mu.Lock()
		o.procs[name] = nm
		o.mu.Unlock()
		fmt.Fprintf(o.out, "RESTART %s %s\n", name, nm.addr)
	}
}

// Addr returns a node's advertised address.
func (o *Orchestrator) Addr(name string) string {
	o.mu.Lock()
	defer o.mu.Unlock()
	if m := o.procs[name]; m != nil {
		return m.addr
	}
	return ""
}

// Shutdown stops nodes in reverse startup order: battles first, hub last. Each
// node gets a drain notice before it is stopped.
func (o *Orchestrator) Shutdown() {
	o.mu.Lock()
	o.down = true
	o.mu.Unlock()
	for i := len(o.order) - 1; i >= 0; i-- {
		name := o.order[i]
		o.mu.Lock()
		m := o.procs[name]
		o.mu.Unlock()
		if m == nil {
			continue
		}
		drain(m.addr)
		_ = m.cmd.Process.Kill()
		fmt.Fprintf(o.out, "DOWN %s\n", name)
	}
}

// drain tells a node to stop accepting new work (best effort).
func drain(addr string) {
	c, err := net.Dial("tcp", addr)
	if err != nil {
		return
	}
	peer := NewPeer(c, nil, nil)
	_ = peer.Push(CmdDrainNotice, nil)
	peer.Close()
}

// snippet: cmd/orchestrator/main.go 里构造节点清单的一段
specs := []battlecluster.NodeSpec{{Name: "hub", Bin: *hubBin}}
for i := 1; i <= *gateways; i++ {
	n := strconv.Itoa(i)
	specs = append(specs, battlecluster.NodeSpec{
		Name: "gateway-" + n,
		Bin:  *gwBin,
		ArgsFor: func(get func(string) string) []string {
			return []string{"-hub", get("hub"), "-server", n}
		},
	})
}

UP hub 127.0.0.1:56854
UP gateway-1 127.0.0.1:56857
UP gateway-2 127.0.0.1:56860
UP battle-0 127.0.0.1:56863
CLUSTER READY
（……运行中，杀掉某个网关会打印 RESTART gateway-1 127.0.0.1:57xxx……）
DOWN battle-0
DOWN gateway-2
DOWN gateway-1
DOWN hub

=== RUN   TestOrchestrator
--- PASS: TestOrchestrator (4.59s)
PASS
ok  	aegis/internal/battlecluster	4.6s