chore(api-gateway): merge review docs and update dependencies
This commit is contained in:
SpecialX
@@ -1,47 +1,151 @@
|
||||
// Package hub manages the in-memory WebSocket connection pool for push-gateway.
|
||||
//
|
||||
// Hub is the single process-wide registry of all live WebSocket connections,
|
||||
// indexed by userID -> connID -> *Connection. Each Connection owns a send
|
||||
// channel consumed by a dedicated writer goroutine; Hub exposes Register /
|
||||
// Unregister / SendToUser / Broadcast operations.
|
||||
//
|
||||
// The Hub also tracks per-user connection counts to enforce the
|
||||
// MaxConnectionsPerUser limit (default 5) and provides CloseAll for graceful
|
||||
// shutdown. Cross-instance presence and message fanout are delegated to a
|
||||
// RedisClient (see internal/redis), keeping Hub free of any persistence layer.
|
||||
package hub
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// Connection 包装单个 WebSocket 连接及其异步发送通道
|
||||
// ErrTooManyConnections is returned by Register when the per-user connection
|
||||
// cap has been reached. Callers should translate this into a 429 /
|
||||
// PUSH_TOO_MANY_CONNECTIONS response or a close frame with code 1008.
|
||||
var ErrTooManyConnections = errors.New("too many connections per user")
|
||||
|
||||
// ErrHubClosing is returned by Register when the Hub is shutting down and no
|
||||
// new connections are accepted.
|
||||
var ErrHubClosing = errors.New("hub is closing")
|
||||
|
||||
// sendBufferSize is the per-connection send channel capacity. A full channel
|
||||
// triggers message drop with metric increment (see Send).
|
||||
const sendBufferSize = 64
|
||||
|
||||
// Connection wraps a single WebSocket connection together with its asynchronous
|
||||
// send channel and lifecycle metadata.
|
||||
type Connection struct {
|
||||
UserID string
|
||||
conn *websocket.Conn
|
||||
send chan []byte
|
||||
UserID string
|
||||
ConnID string
|
||||
conn *websocket.Conn
|
||||
send chan []byte
|
||||
closeOnce sync.Once
|
||||
closed bool
|
||||
mu sync.Mutex
|
||||
lastPing time.Time
|
||||
}
|
||||
|
||||
// Send 将消息投递到该连接的发送通道（非阻塞，通道满则丢弃）
|
||||
func (c *Connection) Send(message []byte) {
|
||||
// Send delivers message to the connection's writer goroutine. When the send
|
||||
// channel is full the message is dropped (non-blocking) and the caller is
|
||||
// notified via the returned false value so it can increment the
|
||||
// push_gateway_messages_dropped_total{reason="channel_full"} metric.
|
||||
func (c *Connection) Send(message []byte) bool {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if c.closed {
|
||||
return false
|
||||
}
|
||||
select {
|
||||
case c.send <- message:
|
||||
return true
|
||||
default:
|
||||
// 通道满，丢弃消息避免阻塞 hub
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// Outgoing 返回该连接的发送通道，供写协程消费
|
||||
// Outgoing exposes the send channel for the writer goroutine. The channel is
|
||||
// closed by Close, terminating the consumer loop.
|
||||
func (c *Connection) Outgoing() <-chan []byte {
|
||||
return c.send
|
||||
}
|
||||
|
||||
// Hub 管理所有在线 WebSocket 连接，按 userID 索引
|
||||
type Hub struct {
|
||||
mu sync.RWMutex
|
||||
clients map[string]map[*Connection]bool
|
||||
// Close shuts down the connection's send channel exactly once. Subsequent
|
||||
// Send calls become no-ops.
|
||||
func (c *Connection) Close() {
|
||||
c.closeOnce.Do(func() {
|
||||
c.mu.Lock()
|
||||
c.closed = true
|
||||
c.mu.Unlock()
|
||||
close(c.send)
|
||||
})
|
||||
}
|
||||
|
||||
// NewHub 创建 Hub 实例
|
||||
func NewHub() *Hub {
|
||||
// LastPing returns the most recent heartbeat timestamp.
|
||||
func (c *Connection) LastPing() time.Time {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.lastPing
|
||||
}
|
||||
|
||||
// TouchPing updates the last heartbeat timestamp to now.
|
||||
func (c *Connection) TouchPing() {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.lastPing = time.Now()
|
||||
}
|
||||
|
||||
// Conn returns the underlying gorilla websocket.Conn for control-frame
|
||||
// operations (Ping/Pong/Close). Used by the reader goroutine.
|
||||
func (c *Connection) Conn() *websocket.Conn {
|
||||
return c.conn
|
||||
}
|
||||
|
||||
// Hub is the process-wide registry of live WebSocket connections.
|
||||
//
|
||||
// clients maps userID -> connID -> *Connection so a single user may hold
|
||||
// multiple simultaneous connections (multi-device). counters caches the
|
||||
// per-user connection count to make the MaxConnectionsPerUser check O(1).
|
||||
type Hub struct {
|
||||
mu sync.RWMutex
|
||||
clients map[string]map[string]*Connection
|
||||
counters map[string]int
|
||||
closing bool
|
||||
|
||||
// maxConnsPerUser caps the number of simultaneous connections per user.
|
||||
maxConnsPerUser int
|
||||
|
||||
// onRegister / onUnregister are optional hooks used by the Redis presence
|
||||
// layer to maintain the cross-instance online SET. They are invoked outside
|
||||
// the Hub lock to avoid blocking the registry on Redis round-trips.
|
||||
onRegister func(userID, connID string)
|
||||
onUnregister func(userID, connID string)
|
||||
}
|
||||
|
||||
// NewHub returns a Hub with the given per-user connection cap.
|
||||
func NewHub(maxConnsPerUser int) *Hub {
|
||||
if maxConnsPerUser <= 0 {
|
||||
maxConnsPerUser = 5
|
||||
}
|
||||
return &Hub{
|
||||
clients: make(map[string]map[*Connection]bool),
|
||||
clients: make(map[string]map[string]*Connection),
|
||||
counters: make(map[string]int),
|
||||
maxConnsPerUser: maxConnsPerUser,
|
||||
}
|
||||
}
|
||||
|
||||
// ClientCount 返回当前在线连接总数（供 /readyz 探针使用）
|
||||
func (h *Hub) ClientCount() int {
|
||||
// SetPresenceHooks registers callbacks invoked after Register / Unregister.
|
||||
// Used by the Redis presence layer to keep edu:push:online:<userID> in sync.
|
||||
func (h *Hub) SetPresenceHooks(onRegister, onUnregister func(userID, connID string)) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
h.onRegister = onRegister
|
||||
h.onUnregister = onUnregister
|
||||
}
|
||||
|
||||
// ActiveConnections returns the total number of live connections across all
|
||||
// users. Exposed via /healthz and the push_gateway_active_connections gauge.
|
||||
func (h *Hub) ActiveConnections() int {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
count := 0
|
||||
@@ -51,60 +155,162 @@ func (h *Hub) ClientCount() int {
|
||||
return count
|
||||
}
|
||||
|
||||
// Register 注册一个用户连接，返回 Connection 供调用方持有
|
||||
func (h *Hub) Register(userID string, conn *websocket.Conn) *Connection {
|
||||
c := &Connection{
|
||||
UserID: userID,
|
||||
conn: conn,
|
||||
send: make(chan []byte, 64),
|
||||
}
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
if h.clients[userID] == nil {
|
||||
h.clients[userID] = make(map[*Connection]bool)
|
||||
}
|
||||
h.clients[userID][c] = true
|
||||
return c
|
||||
// UserCount returns the number of distinct online users (used by /healthz).
|
||||
func (h *Hub) UserCount() int {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
return len(h.clients)
|
||||
}
|
||||
|
||||
// Unregister 注销一个用户连接并关闭其发送通道
|
||||
func (h *Hub) Unregister(c *Connection) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
conns, ok := h.clients[c.UserID]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if _, exists := conns[c]; exists {
|
||||
delete(conns, c)
|
||||
close(c.send)
|
||||
}
|
||||
if len(conns) == 0 {
|
||||
delete(h.clients, c.UserID)
|
||||
}
|
||||
}
|
||||
|
||||
// SendToUser 向指定用户的所有在线连接推送消息
|
||||
func (h *Hub) SendToUser(userID string, message []byte) error {
|
||||
// HasUser reports whether at least one local connection exists for userID.
|
||||
// Used by the /internal/push path to decide between local delivery and Redis
|
||||
// Pub/Sub fanout.
|
||||
func (h *Hub) HasUser(userID string) bool {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
conns, ok := h.clients[userID]
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
for c := range conns {
|
||||
c.Send(message)
|
||||
}
|
||||
return nil
|
||||
return ok && len(conns) > 0
|
||||
}
|
||||
|
||||
// Broadcast 向所有在线连接广播消息
|
||||
func (h *Hub) Broadcast(message []byte) {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
for _, conns := range h.clients {
|
||||
for c := range conns {
|
||||
c.Send(message)
|
||||
// Register adds a new connection for userID. Returns ErrTooManyConnections
|
||||
// when the per-user cap is exceeded and ErrHubClosing when the Hub is shutting
|
||||
// down. The returned Connection must be Unregistered by the caller on close.
|
||||
func (h *Hub) Register(userID string, conn *websocket.Conn) (*Connection, error) {
|
||||
h.mu.Lock()
|
||||
if h.closing {
|
||||
h.mu.Unlock()
|
||||
return nil, ErrHubClosing
|
||||
}
|
||||
if h.counters[userID] >= h.maxConnsPerUser {
|
||||
h.mu.Unlock()
|
||||
return nil, ErrTooManyConnections
|
||||
}
|
||||
connID := uuid.NewString()
|
||||
c := &Connection{
|
||||
UserID: userID,
|
||||
ConnID: connID,
|
||||
conn: conn,
|
||||
send: make(chan []byte, sendBufferSize),
|
||||
lastPing: time.Now(),
|
||||
}
|
||||
if h.clients[userID] == nil {
|
||||
h.clients[userID] = make(map[string]*Connection)
|
||||
}
|
||||
h.clients[userID][connID] = c
|
||||
h.counters[userID]++
|
||||
onRegister := h.onRegister
|
||||
h.mu.Unlock()
|
||||
|
||||
if onRegister != nil {
|
||||
onRegister(userID, connID)
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
// Unregister removes a connection and closes its send channel. Safe to call
|
||||
// multiple times for the same Connection.
|
||||
func (h *Hub) Unregister(c *Connection) {
|
||||
h.mu.Lock()
|
||||
conns, ok := h.clients[c.UserID]
|
||||
if !ok {
|
||||
h.mu.Unlock()
|
||||
return
|
||||
}
|
||||
if _, exists := conns[c.ConnID]; exists {
|
||||
delete(conns, c.ConnID)
|
||||
h.counters[c.UserID]--
|
||||
if h.counters[c.UserID] <= 0 {
|
||||
delete(h.counters, c.UserID)
|
||||
}
|
||||
if len(conns) == 0 {
|
||||
delete(h.clients, c.UserID)
|
||||
}
|
||||
}
|
||||
onUnregister := h.onUnregister
|
||||
h.mu.Unlock()
|
||||
|
||||
c.Close()
|
||||
if onUnregister != nil {
|
||||
onUnregister(c.UserID, c.ConnID)
|
||||
}
|
||||
}
|
||||
|
||||
// SendToUser delivers message to every local connection of userID. Returns
|
||||
// the number of connections the message was successfully queued to. A zero
|
||||
// return value signals "user not online locally" so the caller can fall back
|
||||
// to Redis Pub/Sub fanout.
|
||||
func (h *Hub) SendToUser(userID string, message []byte) int {
|
||||
h.mu.RLock()
|
||||
conns := h.clients[userID]
|
||||
// Snapshot the slice under the read lock to avoid holding it during Send.
|
||||
snapshot := make([]*Connection, 0, len(conns))
|
||||
for _, c := range conns {
|
||||
snapshot = append(snapshot, c)
|
||||
}
|
||||
h.mu.RUnlock()
|
||||
|
||||
delivered := 0
|
||||
for _, c := range snapshot {
|
||||
if c.Send(message) {
|
||||
delivered++
|
||||
}
|
||||
}
|
||||
return delivered
|
||||
}
|
||||
|
||||
// Broadcast delivers message to every live connection in this instance.
|
||||
// Cross-instance broadcast is handled separately via Redis Pub/Sub.
|
||||
func (h *Hub) Broadcast(message []byte) int {
|
||||
h.mu.RLock()
|
||||
snapshot := make([]*Connection, 0, h.ActiveConnections())
|
||||
for _, conns := range h.clients {
|
||||
for _, c := range conns {
|
||||
snapshot = append(snapshot, c)
|
||||
}
|
||||
}
|
||||
h.mu.RUnlock()
|
||||
|
||||
delivered := 0
|
||||
for _, c := range snapshot {
|
||||
if c.Send(message) {
|
||||
delivered++
|
||||
}
|
||||
}
|
||||
return delivered
|
||||
}
|
||||
|
||||
// CloseAll marks the Hub as closing (rejecting further Register calls) and
|
||||
// sends a close frame to every live connection. Used during graceful shutdown.
|
||||
// The caller should wait for connections to drain after CloseAll returns.
|
||||
func (h *Hub) CloseAll() {
|
||||
h.mu.Lock()
|
||||
h.closing = true
|
||||
all := make([]*Connection, 0, h.ActiveConnections())
|
||||
for _, conns := range h.clients {
|
||||
for _, c := range conns {
|
||||
all = append(all, c)
|
||||
}
|
||||
}
|
||||
h.mu.Unlock()
|
||||
|
||||
for _, c := range all {
|
||||
// Send a graceful close frame (1001 going away) and close the send
|
||||
// channel so the writer goroutine exits.
|
||||
_ = c.conn.WriteControl(
|
||||
websocket.CloseMessage,
|
||||
websocket.FormatCloseMessage(websocket.CloseGoingAway, "server shutting down"),
|
||||
time.Now().Add(2*time.Second),
|
||||
)
|
||||
c.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// ForEachUser iterates over all locally online userIDs. Used by the Redis
|
||||
// presence layer to rebuild the online SET on startup (ISSUE-058).
|
||||
func (h *Hub) ForEachUser(fn func(userID string)) {
|
||||
h.mu.RLock()
|
||||
defer h.mu.RUnlock()
|
||||
for userID := range h.clients {
|
||||
fn(userID)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user