《MMO 服务器开发》7.1 会话表与连接管理

// file: internal/session/table.go
// tier: L1
// verify: go test ./internal/session/

package session

import (
	"errors"
	"sync"
)

// ConnID identifies one logical client connection.
type ConnID uint64

var (
	ErrUnknownConn = errors.New("session: unknown conn")
	ErrConnClosed  = errors.New("session: conn closed")
)

// Conn is the session table's view of a client connection. The table never
// touches a socket: it tracks identity and forwards sends.
type Conn struct {
	ID      ConnID
	Account string

	mu      sync.Mutex
	closed  bool
	send    func(cmd uint16, payload []byte) error
	onClose func(*Conn)
}

// NewConn wraps a transport-level send function. send may be nil for tests.
func NewConn(send func(cmd uint16, payload []byte) error) *Conn {
	return &Conn{send: send}
}

// Send forwards one command to the client unless the conn is already closed.
func (c *Conn) Send(cmd uint16, payload []byte) error {
	c.mu.Lock()
	closed := c.closed
	c.mu.Unlock()
	if closed {
		return ErrConnClosed
	}
	if c.send == nil {
		return nil
	}
	return c.send(cmd, payload)
}

// Close is idempotent and runs onClose exactly once. The old-connection
// disconnect path is expected to call Table.Remove here.
func (c *Conn) Close() error {
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return nil
	}
	c.closed = true
	cb := c.onClose
	c.mu.Unlock()
	if cb != nil {
		cb(c)
	}
	return nil
}

// Closed reports whether Close has run.
func (c *Conn) Closed() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.closed
}

// Table is a bidirectional index: connID -> conn and account -> connID. One
// mutex guards both maps, so a lookup and a rebind can never observe a
// half-updated pair.
type Table struct {
	mu     sync.Mutex
	byConn map[ConnID]*Conn
	byAcc  map[string]ConnID
	nextID ConnID
}

// NewTable returns an empty session table.
func NewTable() *Table {
	return &Table{
		byConn: make(map[ConnID]*Conn),
		byAcc:  make(map[string]ConnID),
	}
}

// Bind registers c as the live connection for account and returns the conn it
// evicted, if any. The evicted conn is removed from the table *inside* the
// lock but is NOT closed here: closing it would re-enter Remove from its
// disconnect handler while we still hold the lock.
func (t *Table) Bind(account string, c *Conn) (evicted *Conn) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.nextID++
	c.ID = t.nextID
	c.Account = account
	if oldID, ok := t.byAcc[account]; ok {
		evicted = t.byConn[oldID]
		delete(t.byConn, oldID)
	}
	t.byConn[c.ID] = c
	t.byAcc[account] = c.ID
	return evicted
}

// Remove detaches connID. It is safe to call from an old connection's
// disconnect path even after the account has been taken over: the account
// index is cleared only when it still points at this conn.
func (t *Table) Remove(connID ConnID) {
	t.mu.Lock()
	defer t.mu.Unlock()
	c, ok := t.byConn[connID]
	if !ok {
		return
	}
	delete(t.byConn, connID)
	if id, ok := t.byAcc[c.Account]; ok && id == connID {
		delete(t.byAcc, c.Account)
	}
}

// Lookup returns the live conn for account.
func (t *Table) Lookup(account string) (*Conn, bool) {
	t.mu.Lock()
	defer t.mu.Unlock()
	id, ok := t.byAcc[account]
	if !ok {
		return nil, false
	}
	c, ok := t.byConn[id]
	return c, ok
}

// Get returns the conn for connID.
func (t *Table) Get(connID ConnID) (*Conn, bool) {
	t.mu.Lock()
	defer t.mu.Unlock()
	c, ok := t.byConn[connID]
	return c, ok
}

// Len reports the number of live connections.
func (t *Table) Len() int {
	t.mu.Lock()
	defer t.mu.Unlock()
	return len(t.byConn)
}

// Snapshot copies all live conns so a fanout can send without holding the lock.
func (t *Table) Snapshot() []*Conn {
	t.mu.Lock()
	defer t.mu.Unlock()
	out := make([]*Conn, 0, len(t.byConn))
	for _, c := range t.byConn {
		out = append(out, c)
	}
	return out
}

// snippet: internal/session/sharded.go（节选：分片绑定与摘除）

const (
	shardBits = 6
	numShards = 1 << shardBits
	shardMask = ConnID(numShards - 1)
	seqBits   = 56
)

type ShardedTable struct {
	shards [numShards]shard
}

type shard struct {
	mu     sync.Mutex
	byConn map[ConnID]*Conn
	byAcc  map[string]ConnID
	seq    ConnID
}

func shardIndex(account string) ConnID {
	h := fnv.New32a()
	_, _ = h.Write([]byte(account))
	return ConnID(h.Sum32()) & shardMask
}

// Bind behaves exactly like Table.Bind, but only locks one shard.
func (t *ShardedTable) Bind(account string, c *Conn) (evicted *Conn) {
	idx := shardIndex(account)
	s := &t.shards[idx]
	s.mu.Lock()
	defer s.mu.Unlock()
	s.seq++
	c.ID = idx<<seqBits | s.seq // 高 6 位 = 分片号
	c.Account = account
	if oldID, ok := s.byAcc[account]; ok {
		evicted = s.byConn[oldID]
		delete(s.byConn, oldID)
	}
	s.byConn[c.ID] = c
	s.byAcc[account] = c.ID
	return evicted
}

// Remove detaches connID from its shard; the shard index rides in the ID.
func (t *ShardedTable) Remove(connID ConnID) {
	idx := (connID >> seqBits) & shardMask
	s := &t.shards[idx]
	s.mu.Lock()
	defer s.mu.Unlock()
	c, ok := s.byConn[connID]
	if !ok {
		return
	}
	delete(s.byConn, connID)
	if id, ok := s.byAcc[c.Account]; ok && id == connID {
		delete(s.byAcc, c.Account)
	}
}

// snippet: internal/session/table_test.go（节选）

// TestOldCloseDoesNotEvictNewSession is the classic takeover race: a second
// login replaces the first, then the first connection finally notices its
// socket died and runs its disconnect path. That path must not delete the new
// session.
func TestOldCloseDoesNotEvictNewSession(t *testing.T) {
	tab := NewTable()

	old := NewConn(nil)
	old.onClose = func(c *Conn) { tab.Remove(c.ID) } // 读循环的清理
	tab.Bind("alice", old)

	fresh := NewConn(nil)
	fresh.onClose = func(c *Conn) { tab.Remove(c.ID) }
	if ev := tab.Bind("alice", fresh); ev != old {
		t.Fatalf("Bind evicted %v, want the old conn", ev)
	}

	_ = old.Close() // 旧 socket 现在才发现自己死了

	if got, ok := tab.Lookup("alice"); !ok || got != fresh {
		t.Fatalf("old Close evicted the new session: got=%v ok=%v", got, ok)
	}
	if _, ok := tab.Get(old.ID); ok {
		t.Fatal("old conn is still indexed after takeover")
	}
}