refactor(pico): replace addConnection with createAndAddConnection for atomic connID generation

This commit is contained in:
lc6464 2026-03-24 21:01:20 +08:00
parent 9eb4ee6d91
commit 0e506585e5
No known key found for this signature in database
GPG key ID: 53C61B42FEC71D6D

View file

@ -99,11 +99,25 @@ func NewPicoChannel(cfg config.PicoConfig, messageBus *bus.MessageBus) (*PicoCha
}, nil }, nil
} }
// addConnection stores a connection in both connID and session indexes. // createAndAddConnection generates a unique connID and registers it atomically.
func (c *PicoChannel) addConnection(pc *picoConn) { func (c *PicoChannel) createAndAddConnection(conn *websocket.Conn, sessionID string) *picoConn {
c.connsMu.Lock() c.connsMu.Lock()
defer c.connsMu.Unlock() defer c.connsMu.Unlock()
var connID string
for {
connID = uuid.New().String()
if _, exists := c.connections[connID]; !exists {
break
}
}
pc := &picoConn{
id: connID,
conn: conn,
sessionID: sessionID,
}
c.connections[pc.id] = pc c.connections[pc.id] = pc
bySession, ok := c.sessionConnections[pc.sessionID] bySession, ok := c.sessionConnections[pc.sessionID]
if !ok { if !ok {
@ -111,6 +125,8 @@ func (c *PicoChannel) addConnection(pc *picoConn) {
c.sessionConnections[pc.sessionID] = bySession c.sessionConnections[pc.sessionID] = bySession
} }
bySession[pc.id] = pc bySession[pc.id] = pc
return pc
} }
// removeConnection deletes a connection from indexes and returns it when found. // removeConnection deletes a connection from indexes and returns it when found.
@ -166,20 +182,6 @@ func (c *PicoChannel) sessionConnectionsSnapshot(sessionID string) []*picoConn {
return conns return conns
} }
// newUniqueConnID generates a connID that is not currently present in the index.
func (c *PicoChannel) newUniqueConnID() string {
for {
id := uuid.New().String()
c.connsMu.RLock()
_, exists := c.connections[id]
c.connsMu.RUnlock()
if !exists {
return id
}
}
}
// 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")
@ -350,13 +352,7 @@ func (c *PicoChannel) handleWebSocket(w http.ResponseWriter, r *http.Request) {
sessionID = uuid.New().String() sessionID = uuid.New().String()
} }
pc := &picoConn{ pc := c.createAndAddConnection(conn, sessionID)
id: c.newUniqueConnID(),
conn: conn,
sessionID: sessionID,
}
c.addConnection(pc)
c.connCount.Add(1) c.connCount.Add(1)
logger.InfoCF("pico", "WebSocket client connected", map[string]any{ logger.InfoCF("pico", "WebSocket client connected", map[string]any{