fix(pico): enforce MaxConnections atomically on registration
This commit is contained in:
parent
981414d92c
commit
522ef33f69
1 changed files with 16 additions and 4 deletions
|
|
@ -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,
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue