《游戏服务端编程实践》2.4.3 实战：实现简易 Echo Server

// file: ch02/echo/server.go
// tier: L1
// verify: go test ./ch02/echo/

package echo

import (
	"context"
	"encoding/binary"
	"io"
	"net"
	"sync"
	"time"
)

const (
	// HeaderSize 是长度前缀的字节数，使用大端序。
	HeaderSize = 4
	// DefaultMaxFrame 是单包负载上限：64 KiB。
	DefaultMaxFrame = 64 << 10
	// DefaultReadTimeout 是单次读操作的超时，防止空闲连接永久占用 goroutine。
	DefaultReadTimeout = 5 * time.Second
)

// Config 是服务器参数。零值字段会被替换为默认值。
type Config struct {
	MaxFrame    uint32
	ReadTimeout time.Duration
}

// Server 是一个带长度前缀的 TCP Echo 服务器。
// 它只负责“收一整包、原样回一整包”，不解析业务协议。
type Server struct {
	ln          net.Listener
	maxFrame    uint32
	readTimeout time.Duration

	wg      sync.WaitGroup
	mu      sync.Mutex
	conns   map[net.Conn]struct{}
	stopped bool

	quit      chan struct{}
	closeOnce sync.Once
	closeErr  error
}

// New 在 addr 上监听并返回服务器。此时还没有开始 Accept。
func New(addr string, cfg Config) (*Server, error) {
	if cfg.MaxFrame == 0 {
		cfg.MaxFrame = DefaultMaxFrame
	}
	if cfg.ReadTimeout <= 0 {
		cfg.ReadTimeout = DefaultReadTimeout
	}
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		return nil, err
	}
	return &Server{
		ln:          ln,
		maxFrame:    cfg.MaxFrame,
		readTimeout: cfg.ReadTimeout,
		conns:       make(map[net.Conn]struct{}),
		quit:        make(chan struct{}),
	}, nil
}

// Addr 返回实际监听地址；测试传入 127.0.0.1:0 时用它取得随机端口。
func (s *Server) Addr() net.Addr { return s.ln.Addr() }

// Serve 接受连接直到 Shutdown 关闭监听器，正常关闭时返回 nil。
func (s *Server) Serve() error {
	for {
		conn, err := s.ln.Accept()
		if err != nil {
			select {
			case <-s.quit:
				return nil
			default:
				return err
			}
		}
		if !s.track(conn) {
			_ = conn.Close()
			return nil
		}
		go func() {
			defer s.wg.Done()
			defer s.forget(conn)
			defer conn.Close()
			s.handle(conn)
		}()
	}
}

// track 在停止接收后拒绝新连接，保证 wg.Add 不会晚于 Shutdown 的 wg.Wait。
func (s *Server) track(conn net.Conn) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.stopped {
		return false
	}
	s.conns[conn] = struct{}{}
	s.wg.Add(1)
	return true
}

func (s *Server) forget(conn net.Conn) {
	s.mu.Lock()
	delete(s.conns, conn)
	s.mu.Unlock()
}

func (s *Server) closeConns() {
	s.mu.Lock()
	for conn := range s.conns {
		_ = conn.Close()
	}
	s.mu.Unlock()
}

// handle 在一个连接上循环：读 4 字节长度、读负载、原样回写。
// 任何协议错误或 I/O 错误都会结束该连接，不影响其他连接。
func (s *Server) handle(conn net.Conn) {
	var head [HeaderSize]byte
	for {
		_ = conn.SetReadDeadline(time.Now().Add(s.readTimeout))
		if _, err := io.ReadFull(conn, head[:]); err != nil {
			return
		}
		n := binary.BigEndian.Uint32(head[:])
		if n == 0 || n > s.maxFrame {
			return
		}
		payload := make([]byte, n)
		if _, err := io.ReadFull(conn, payload); err != nil {
			return
		}
		frame := make([]byte, HeaderSize+len(payload))
		copy(frame, head[:])
		copy(frame[HeaderSize:], payload)
		_ = conn.SetWriteDeadline(time.Now().Add(s.readTimeout))
		if _, err := conn.Write(frame); err != nil {
			return
		}
	}
}

// Shutdown 先停止接受新连接，再在 ctx 预算内等待已有连接处理完。
// 预算耗尽时强制关闭剩余连接，并返回 ctx.Err()。重复调用是安全的。
func (s *Server) Shutdown(ctx context.Context) error {
	s.closeOnce.Do(func() {
		close(s.quit)
		s.closeErr = s.ln.Close()
		s.mu.Lock()
		s.stopped = true
		s.mu.Unlock()
	})
	done := make(chan struct{})
	go func() {
		s.wg.Wait()
		close(done)
	}()
	select {
	case <-done:
		return s.closeErr
	case <-ctx.Done():
		s.closeConns()
		<-done
		return ctx.Err()
	}
}

cd .local/checks/gs-practice-lab
gofmt -l ch02/echo
go vet ./ch02/echo/
go test -race -v ./ch02/echo/

go test -race ./ch02/...
go test -bench=. -benchmem ./ch02/langbench/