refactor(pico): make connCount a regular int guarded by connsMu

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

View file

@ -59,7 +59,7 @@ type PicoChannel struct {
connections map[string]*picoConn // connID -> *picoConn connections map[string]*picoConn // connID -> *picoConn
sessionConnections map[string]map[string]*picoConn // sessionID -> connID -> *picoConn sessionConnections map[string]map[string]*picoConn // sessionID -> connID -> *picoConn
connsMu sync.RWMutex connsMu sync.RWMutex
connCount atomic.Int32 connCount int
ctx context.Context ctx context.Context
cancel context.CancelFunc cancel context.CancelFunc
} }
@ -125,7 +125,7 @@ func (c *PicoChannel) createAndAddConnection(conn *websocket.Conn, sessionID str
c.sessionConnections[pc.sessionID] = bySession c.sessionConnections[pc.sessionID] = bySession
} }
bySession[pc.id] = pc bySession[pc.id] = pc
c.connCount.Add(1) c.connCount++
return pc return pc
} }
@ -147,7 +147,7 @@ func (c *PicoChannel) removeConnection(connID string) *picoConn {
delete(c.sessionConnections, pc.sessionID) delete(c.sessionConnections, pc.sessionID)
} }
} }
c.connCount.Add(-1) c.connCount--
return pc return pc
} }
@ -163,7 +163,7 @@ func (c *PicoChannel) takeAllConnections() []*picoConn {
} }
clear(c.connections) clear(c.connections)
clear(c.sessionConnections) clear(c.sessionConnections)
c.connCount.Store(0) c.connCount = 0
return all return all
} }
@ -185,6 +185,13 @@ func (c *PicoChannel) sessionConnectionsSnapshot(sessionID string) []*picoConn {
return conns return conns
} }
// currentConnCount returns a lock-protected snapshot of active connection count.
func (c *PicoChannel) currentConnCount() int {
c.connsMu.RLock()
defer c.connsMu.RUnlock()
return c.connCount
}
// Start implements Channel. // Start implements Channel.
func (c *PicoChannel) Start(ctx context.Context) error { func (c *PicoChannel) Start(ctx context.Context) error {
logger.InfoC("pico", "Starting Pico Protocol channel") logger.InfoC("pico", "Starting Pico Protocol channel")
@ -329,7 +336,7 @@ func (c *PicoChannel) handleWebSocket(w http.ResponseWriter, r *http.Request) {
if maxConns <= 0 { if maxConns <= 0 {
maxConns = 100 maxConns = 100
} }
if int(c.connCount.Load()) >= maxConns { if c.currentConnCount() >= maxConns {
http.Error(w, "too many connections", http.StatusServiceUnavailable) http.Error(w, "too many connections", http.StatusServiceUnavailable)
return return
} }