《MMO 服务器开发》4.3 连接排空（draining）

// file: internal/gateway/drain.go
// tier: L1
// verify: go test ./internal/gateway/ -run TestDrain
package gateway

import (
	"sync"
	"time"
)

// State 是网关生命周期的三态。
type State int

const (
	StateAccepting State = iota // 正常接客
	StateDraining               // 排空中：不再接受新连接，等旧连接自然结束
	StateClosed                 // 已关闭
)

func (s State) String() string {
	switch s {
	case StateAccepting:
		return "accepting"
	case StateDraining:
		return "draining"
	default:
		return "closed"
	}
}

// DrainResult 是一次排空的结果。TimedOut 为真表示有连接在超时前没走，
// 只能兜底断开。
type DrainResult struct {
	Drained   bool
	TimedOut  bool
	Remaining int
	Waited    time.Duration
}

// Drainer 是排空状态机：accepting → draining → closed。它只负责两件事——
// 决定「还要不要接新连接」，以及「旧连接都走完了没有」。真正的广播与断开
// 由 onDrain 回调和调用方完成，这样状态机本身可以在纯逻辑下测。
type Drainer struct {
	mu          sync.Mutex
	state       State
	live        int
	empty       chan struct{} // live 归零且已进入 draining 时关闭
	emptyClosed bool
	onDrain     func() // 进入 draining 时触发一次，用于广播排空通知
}

// NewDrainer 构造一个处于 accepting 态的状态机。onDrain 在进入 draining 时
// 被调用一次，用来向所有在线连接广播 CmdDrainNotice 与 CmdReconnectTicket。
func NewDrainer(onDrain func()) *Drainer {
	return &Drainer{state: StateAccepting, empty: make(chan struct{}), onDrain: onDrain}
}

// AddConn 尝试登记一条新连接。只有 accepting 态才接受；draining 之后返回
// false，调用方应立即关掉这条连接——这就是「排空后新连接被拒」。
func (d *Drainer) AddConn() bool {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.state != StateAccepting {
		return false
	}
	d.live++
	return true
}

// RemoveConn 注销一条连接。若正处于 draining 且这是最后一条，关闭 empty，
// 唤醒 WaitDrained。
func (d *Drainer) RemoveConn() {
	d.mu.Lock()
	defer d.mu.Unlock()
	if d.live > 0 {
		d.live--
	}
	if d.state == StateDraining && d.live == 0 && !d.emptyClosed {
		close(d.empty)
		d.emptyClosed = true
	}
}

// TryAccept 供 accept 循环快速判断是否还应继续接客。
func (d *Drainer) TryAccept() bool {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.state == StateAccepting
}

// BeginDrain 进入 draining 态并触发一次广播。重复调用返回 false。
func (d *Drainer) BeginDrain() bool {
	d.mu.Lock()
	if d.state != StateAccepting {
		d.mu.Unlock()
		return false
	}
	d.state = StateDraining
	if d.live == 0 && !d.emptyClosed {
		close(d.empty)
		d.emptyClosed = true
	}
	cb := d.onDrain
	d.mu.Unlock()
	if cb != nil {
		cb()
	}
	return true
}

// WaitDrained 阻塞等待旧连接自然结束；超时则兜底，返回剩余连接数。
func (d *Drainer) WaitDrained(timeout time.Duration) DrainResult {
	t0 := time.Now()
	select {
	case <-d.empty:
		d.mu.Lock()
		d.state = StateClosed
		d.mu.Unlock()
		return DrainResult{Drained: true, Waited: time.Since(t0)}
	case <-time.After(timeout):
		d.mu.Lock()
		rem := d.live
		d.state = StateClosed
		d.mu.Unlock()
		return DrainResult{TimedOut: true, Remaining: rem, Waited: time.Since(t0)}
	}
}

// Drain 是 BeginDrain + WaitDrained 的合并入口。
func (d *Drainer) Drain(timeout time.Duration) DrainResult {
	d.BeginDrain()
	return d.WaitDrained(timeout)
}

// State 返回当前状态。
func (d *Drainer) State() State {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.state
}

// Live 返回当前在线连接数。
func (d *Drainer) Live() int {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.live
}

// file: cmd/gwd/main.go
// tier: L2
// verify: go build ./cmd/gwd/ && go test ./internal/gateway/ -run TestL2Drain
package main

import (
	"flag"
	"log"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"aegis/internal/gateway"
)

const (
	cmdHeartbeatReq  = 0x0001
	cmdHeartbeatResp = 0x0002
	cmdLoginReq      = 0x0010
	cmdLoginResp     = 0x0011
)

// server 是可排空的网关：accepting 态接客，收到 SIGTERM 后关闭监听、广播
// 排空通知与重连票、等旧连接自然结束，超时兜底退出。
type server struct {
	mu    sync.Mutex
	conns map[net.Conn]struct{}
	drain *gateway.Drainer
	key   []byte
}

func (s *server) serve(c net.Conn) {
	if !s.drain.AddConn() { // draining 之后进来的连接一律拒绝
		_ = c.Close()
		return
	}
	s.mu.Lock()
	s.conns[c] = struct{}{}
	s.mu.Unlock()
	defer func() {
		s.mu.Lock()
		delete(s.conns, c)
		s.mu.Unlock()
		s.drain.RemoveConn()
		_ = c.Close()
	}()
	for {
		f, err := gateway.ReadFrame(c)
		if err != nil {
			return
		}
		switch f.Cmd {
		case cmdHeartbeatReq:
			_ = gateway.WriteFrame(c, gateway.Frame{Cmd: cmdHeartbeatResp, Seq: f.Seq})
		case cmdLoginReq:
			_ = gateway.WriteFrame(c, gateway.Frame{Cmd: cmdLoginResp, Seq: f.Seq})
		}
	}
}

// broadcastDrain 在进入 draining 时触发：先发 CmdDrainNotice 0x1001，再发
// CmdReconnectTicket 0x1002，客户端据此换网关重连。
func (s *server) broadcastDrain() {
	tok := gateway.IssueTicket(s.key, 0, "gateway-next", time.Now().Add(time.Minute).Unix())
	s.mu.Lock()
	defer s.mu.Unlock()
	for c := range s.conns {
		_ = gateway.WriteFrame(c, gateway.Frame{Cmd: gateway.CmdDrainNotice})
		_ = gateway.WriteFrame(c, gateway.Frame{Cmd: gateway.CmdReconnectTicket, Payload: []byte(tok)})
	}
}

func main() {
	listen := flag.String("listen", "127.0.0.1:32211", "listen addr")
	timeout := flag.Duration("drain-timeout", 5*time.Second, "drain timeout")
	flag.Parse()

	ln, err := net.Listen("tcp", *listen)
	if err != nil {
		log.Fatalf("listen: %v", err)
	}
	s := &server{conns: make(map[net.Conn]struct{}), key: []byte("aegis-dist-ch04-drain")}
	s.drain = gateway.NewDrainer(s.broadcastDrain)

	go func() {
		for {
			c, err := ln.Accept()
			if err != nil {
				return
			}
			go s.serve(c)
		}
	}()
	log.Printf("gwd listening on %s", *listen)

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGTERM, syscall.SIGINT)
	<-sig
	log.Printf("draining: stop accepting, live=%d", s.drain.Live())
	_ = ln.Close() // 1) 停止接受新连接
	s.drain.BeginDrain()
	res := s.drain.WaitDrained(*timeout)
	log.Printf("drain done: drained=%v timedout=%v remaining=%d waited=%s",
		res.Drained, res.TimedOut, res.Remaining, res.Waited)
}