fix(pico): enforce MaxConnections atomically on registration

This commit is contained in:
lc6464 2026-03-24 21:36:51 +08:00
parent 981414d92c
commit 522ef33f69
No known key found for this signature in database
GPG key ID: 53C61B42FEC71D6D

View file

@ -99,10 +99,13 @@ func NewPicoChannel(cfg config.PicoConfig, messageBus *bus.MessageBus) (*PicoCha
}, nil }, nil
} }
// createAndAddConnection generates a unique connID and registers it atomically. // createAndAddConnection checks MaxConnections and registers a connection atomically.
func (c *PicoChannel) createAndAddConnection(conn *websocket.Conn, sessionID string) *picoConn { func (c *PicoChannel) createAndAddConnection(conn *websocket.Conn, sessionID string, maxConns int) (*picoConn, error) {
c.connsMu.Lock() c.connsMu.Lock()
defer c.connsMu.Unlock() defer c.connsMu.Unlock()
if c.connCount >= maxConns {
return nil, channels.ErrSendFailed
}
var connID string var connID string
for { for {
@ -127,7 +130,7 @@ func (c *PicoChannel) createAndAddConnection(conn *websocket.Conn, sessionID str
bySession[pc.id] = pc bySession[pc.id] = pc
c.connCount++ c.connCount++
return pc return pc, nil
} }
// removeConnection deletes a connection from indexes and returns it when found. // removeConnection deletes a connection from indexes and returns it when found.
@ -361,7 +364,16 @@ func (c *PicoChannel) handleWebSocket(w http.ResponseWriter, r *http.Request) {
sessionID = uuid.New().String() sessionID = uuid.New().String()
} }
pc := c.createAndAddConnection(conn, sessionID) pc, err := c.createAndAddConnection(conn, sessionID, maxConns)
if err != nil {
_ = conn.WriteControl(
websocket.CloseMessage,
websocket.FormatCloseMessage(websocket.CloseTryAgainLater, "too many connections"),
time.Now().Add(2*time.Second),
)
_ = conn.Close()
return
}
logger.InfoCF("pico", "WebSocket client connected", map[string]any{ logger.InfoCF("pico", "WebSocket client connected", map[string]any{
"conn_id": pc.id, "conn_id": pc.id,