《MMO 服务器开发》6.2 心跳与超时检测

// file: internal/ws/server.go
// tier: L2
// verify: go test ./internal/ws/ -run 'TestHeartbeatTimeoutClosesIdleClient|TestServerPingKeepsIdleClientAlive|TestServerEcho' -v

package ws

import (
	"bufio"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

// Close codes. 1001 "going away" is what a gateway sends when it drops a peer
// whose heartbeat went silent.
const (
	CloseGoingAway    uint16 = 1001
	CloseProtocolErr  uint16 = 1002
	ClosePolicyViolat uint16 = 1003
)

// Server is a minimal WebSocket echo/heartbeat gateway. Per connection there is
// one reader goroutine, one watchdog enforcing the read deadline, and one
// optional pinger keeping NAT bindings alive.
type Server struct {
	// ReadTimeout is the longest silence tolerated before the peer is dropped.
	ReadTimeout time.Duration
	// PingInterval, when non-zero, makes the server send periodic pings.
	PingInterval time.Duration
	// MaxPayload caps a single reassembled message; 0 means unlimited.
	MaxPayload uint64
	// OnMessage handles data frames. When nil the server echoes them back.
	OnMessage func(s *Session, opcode byte, payload []byte)
	// RequireClientMask enforces RFC 6455 §5.1 on the read path.
	RequireClientMask bool

	mu       sync.Mutex
	sessions map[*Session]struct{}
	nextID   uint64
}

// NewServer returns a Server with sane defaults.
func NewServer() *Server {
	return &Server{
		MaxPayload:        1 << 20,
		RequireClientMask: true,
		sessions:          make(map[*Session]struct{}),
	}
}

// Session is one accepted WebSocket connection.
type Session struct {
	ID  uint64
	seq uint64 // reconnection sequence handed to the client

	conn   net.Conn
	last   atomic.Int64 // UnixNano of the last inbound frame
	closed atomic.Bool
	wmu    sync.Mutex
}

// LastActive reports the time of the most recent inbound frame.
func (s *Session) LastActive() time.Time { return time.Unix(0, s.last.Load()) }

func (s *Session) touch() { s.last.Store(time.Now().UnixNano()) }

// WriteMessage sends one unfragmented frame. Server-to-client frames are never
// masked, as the RFC requires.
func (s *Session) WriteMessage(opcode byte, payload []byte) error {
	s.wmu.Lock()
	defer s.wmu.Unlock()
	return WriteFrame(s.conn, true, opcode, payload, false)
}

// Close sends a close frame best-effort and then closes the socket. Idempotent,
// so the watchdog and the reader can race without double-closing.
func (s *Session) Close(code uint16, reason string) error {
	if s.closed.Swap(true) {
		return nil
	}
	buf := make([]byte, 2+len(reason))
	buf[0] = byte(code >> 8)
	buf[1] = byte(code)
	copy(buf[2:], reason)
	s.wmu.Lock()
	_ = WriteFrame(s.conn, true, OpClose, buf, false)
	s.wmu.Unlock()
	return s.conn.Close()
}

// SessionCount returns the number of live sessions.
func (s *Server) SessionCount() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.sessions)
}

// PingAll sends one ping to every live session. The table is snapshotted under
// the lock and the writes happen outside it, so one slow socket cannot stall
// the whole sweep.
func (s *Server) PingAll() int {
	s.mu.Lock()
	batch := make([]*Session, 0, len(s.sessions))
	for sess := range s.sessions {
		batch = append(batch, sess)
	}
	s.mu.Unlock()

	sent := 0
	for _, sess := range batch {
		if err := sess.WriteMessage(OpPing, nil); err == nil {
			sent++
		}
	}
	return sent
}

func (s *Server) add(sess *Session) {
	s.mu.Lock()
	s.nextID++
	sess.ID = s.nextID
	s.sessions[sess] = struct{}{}
	s.mu.Unlock()
}

func (s *Server) remove(sess *Session) {
	s.mu.Lock()
	delete(s.sessions, sess)
	s.mu.Unlock()
}

// ServeListener accepts connections until the listener is closed.
func (s *Server) ServeListener(ln net.Listener) error {
	for {
		c, err := ln.Accept()
		if err != nil {
			return err
		}
		go func() { _ = s.Serve(c) }()
	}
}

// Serve runs the full lifecycle for one accepted TCP connection: handshake,
// frame loop, keepalive, and timeout enforcement.
func (s *Server) Serve(conn net.Conn) error {
	br := bufio.NewReader(conn)
	if _, err := Handshake(br, conn); err != nil {
		_ = conn.Close()
		return err
	}

	sess := &Session{conn: conn}
	sess.touch()
	s.add(sess)

	defer func() {
		sess.closed.Store(true)
		_ = conn.Close()
		s.remove(sess)
	}()

	done := make(chan struct{})
	go func() {
		defer close(done)
		for {
			op, payload, err := ReadMessageMasked(br, s.MaxPayload, s.RequireClientMask)
			if err != nil {
				return
			}
			sess.touch()
			switch op {
			case OpPing:
				_ = sess.WriteMessage(OpPong, payload)
			case OpPong:
				// Liveness only; touch() already recorded it.
			case OpClose:
				_ = sess.WriteMessage(OpClose, payload)
				return
			default:
				if s.OnMessage != nil {
					s.OnMessage(sess, op, payload)
				} else {
					_ = sess.WriteMessage(op, payload)
				}
			}
		}
	}()

	if s.PingInterval > 0 {
		go func() {
			t := time.NewTicker(s.PingInterval)
			defer t.Stop()
			for {
				select {
				case <-done:
					return
				case <-t.C:
					_ = sess.WriteMessage(OpPing, nil)
				}
			}
		}()
	}

	deadline := s.ReadTimeout
	if deadline <= 0 {
		deadline = 60 * time.Second
	}
	tick := deadline / 8
	if tick < time.Millisecond {
		tick = time.Millisecond
	}
	t := time.NewTicker(tick)
	defer t.Stop()
	for {
		select {
		case <-done:
			return nil
		case <-t.C:
			if time.Since(sess.LastActive()) > deadline {
				_ = sess.Close(CloseGoingAway, "heartbeat timeout")
				return nil
			}
		}
	}
}

// snippet: internal/ws/server_test.go（节选）

// TestHeartbeatTimeoutClosesIdleClient is the core L2 assertion: a client that
// never sends anything is dropped and its session is reclaimed.
func TestHeartbeatTimeoutClosesIdleClient(t *testing.T) {
	const timeout = 300 * time.Millisecond
	s := NewServer()
	s.ReadTimeout = timeout
	s.PingInterval = 0 // no server pings: pure silence must trip the watchdog
	addr := startServer(t, s)

	c, br := dialWS(t, addr)
	defer c.Close()

	if got := s.SessionCount(); got != 1 {
		t.Fatalf("session count after handshake = %d, want 1", got)
	}

	start := time.Now()
	op, _, err := ReadMessage(br, 0) // blocks until the server's close frame
	elapsed := time.Since(start)
	if err != nil {
		t.Fatalf("expected a close frame, got error: %v", err)
	}
	if op != OpClose {
		t.Fatalf("op = %#x, want close", op)
	}
	if elapsed < timeout {
		t.Fatalf("closed after %v, before the %v deadline", elapsed, timeout)
	}
	if elapsed > timeout+500*time.Millisecond {
		t.Fatalf("closed after %v, far past the %v deadline", elapsed, timeout)
	}

	deadline := time.Now().Add(2 * time.Second)
	for s.SessionCount() != 0 && time.Now().Before(deadline) {
		time.Sleep(5 * time.Millisecond)
	}
	if got := s.SessionCount(); got != 0 {
		t.Fatalf("session leaked: count = %d", got)
	}
	t.Logf("timeout precision: configured %v, observed %v, delta %v",
		timeout, elapsed.Round(time.Millisecond), (elapsed - timeout).Round(time.Millisecond))
}