perf(pico): implement O(1) session lookup for pico connections
- Replace `sync.Map` with `connections` and `sessionConnections`. - Add `addConnection`, `removeConnection`, `sessionConnectionsSnapshot`, and `takeAllConnections` with `connsMu` for concurrency. - `broadcastToSession` now dispatches directly to `sessionConnections`. - Add `newUniqueConnID` to avoid UUID collision/overwrites. - Ensure `Stop` and `readLoop` use the new helpers for safe cleanup and correct `connCount` updates.
This commit is contained in:
parent
f2f6987f00
commit
9eb4ee6d91
1 changed files with 115 additions and 39 deletions
|
|
@ -56,7 +56,9 @@ type PicoChannel struct {
|
||||||
*channels.BaseChannel
|
*channels.BaseChannel
|
||||||
config config.PicoConfig
|
config config.PicoConfig
|
||||||
upgrader websocket.Upgrader
|
upgrader websocket.Upgrader
|
||||||
connections sync.Map // connID → *picoConn
|
connections map[string]*picoConn // connID -> *picoConn
|
||||||
|
sessionConnections map[string]map[string]*picoConn // sessionID -> connID -> *picoConn
|
||||||
|
connsMu sync.RWMutex
|
||||||
connCount atomic.Int32
|
connCount atomic.Int32
|
||||||
ctx context.Context
|
ctx context.Context
|
||||||
cancel context.CancelFunc
|
cancel context.CancelFunc
|
||||||
|
|
@ -92,9 +94,92 @@ func NewPicoChannel(cfg config.PicoConfig, messageBus *bus.MessageBus) (*PicoCha
|
||||||
ReadBufferSize: 1024,
|
ReadBufferSize: 1024,
|
||||||
WriteBufferSize: 1024,
|
WriteBufferSize: 1024,
|
||||||
},
|
},
|
||||||
|
connections: make(map[string]*picoConn),
|
||||||
|
sessionConnections: make(map[string]map[string]*picoConn),
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// addConnection stores a connection in both connID and session indexes.
|
||||||
|
func (c *PicoChannel) addConnection(pc *picoConn) {
|
||||||
|
c.connsMu.Lock()
|
||||||
|
defer c.connsMu.Unlock()
|
||||||
|
|
||||||
|
c.connections[pc.id] = pc
|
||||||
|
bySession, ok := c.sessionConnections[pc.sessionID]
|
||||||
|
if !ok {
|
||||||
|
bySession = make(map[string]*picoConn)
|
||||||
|
c.sessionConnections[pc.sessionID] = bySession
|
||||||
|
}
|
||||||
|
bySession[pc.id] = pc
|
||||||
|
}
|
||||||
|
|
||||||
|
// removeConnection deletes a connection from indexes and returns it when found.
|
||||||
|
func (c *PicoChannel) removeConnection(connID string) *picoConn {
|
||||||
|
c.connsMu.Lock()
|
||||||
|
defer c.connsMu.Unlock()
|
||||||
|
|
||||||
|
pc, ok := c.connections[connID]
|
||||||
|
if !ok {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
delete(c.connections, connID)
|
||||||
|
if bySession, ok := c.sessionConnections[pc.sessionID]; ok {
|
||||||
|
delete(bySession, connID)
|
||||||
|
if len(bySession) == 0 {
|
||||||
|
delete(c.sessionConnections, pc.sessionID)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return pc
|
||||||
|
}
|
||||||
|
|
||||||
|
// takeAllConnections snapshots and clears all connection indexes.
|
||||||
|
func (c *PicoChannel) takeAllConnections() []*picoConn {
|
||||||
|
c.connsMu.Lock()
|
||||||
|
defer c.connsMu.Unlock()
|
||||||
|
|
||||||
|
all := make([]*picoConn, 0, len(c.connections))
|
||||||
|
for connID, pc := range c.connections {
|
||||||
|
all = append(all, pc)
|
||||||
|
delete(c.connections, connID)
|
||||||
|
}
|
||||||
|
clear(c.sessionConnections)
|
||||||
|
|
||||||
|
return all
|
||||||
|
}
|
||||||
|
|
||||||
|
// sessionConnectionsSnapshot returns all active connections for a session.
|
||||||
|
func (c *PicoChannel) sessionConnectionsSnapshot(sessionID string) []*picoConn {
|
||||||
|
c.connsMu.RLock()
|
||||||
|
defer c.connsMu.RUnlock()
|
||||||
|
|
||||||
|
bySession, ok := c.sessionConnections[sessionID]
|
||||||
|
if !ok || len(bySession) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
conns := make([]*picoConn, 0, len(bySession))
|
||||||
|
for _, pc := range bySession {
|
||||||
|
conns = append(conns, pc)
|
||||||
|
}
|
||||||
|
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")
|
||||||
|
|
@ -110,13 +195,10 @@ func (c *PicoChannel) Stop(ctx context.Context) error {
|
||||||
c.SetRunning(false)
|
c.SetRunning(false)
|
||||||
|
|
||||||
// Close all connections
|
// Close all connections
|
||||||
c.connections.Range(func(key, value any) bool {
|
for _, pc := range c.takeAllConnections() {
|
||||||
if pc, ok := value.(*picoConn); ok {
|
|
||||||
pc.close()
|
pc.close()
|
||||||
}
|
}
|
||||||
c.connections.Delete(key)
|
c.connCount.Store(0)
|
||||||
return true
|
|
||||||
})
|
|
||||||
|
|
||||||
if c.cancel != nil {
|
if c.cancel != nil {
|
||||||
c.cancel()
|
c.cancel()
|
||||||
|
|
@ -133,8 +215,8 @@ func (c *PicoChannel) WebhookPath() string { return "/pico/" }
|
||||||
func (c *PicoChannel) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
func (c *PicoChannel) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||||
path := strings.TrimPrefix(r.URL.Path, "/pico")
|
path := strings.TrimPrefix(r.URL.Path, "/pico")
|
||||||
|
|
||||||
switch {
|
switch path {
|
||||||
case path == "/ws" || path == "/ws/":
|
case "/ws", "/ws/":
|
||||||
c.handleWebSocket(w, r)
|
c.handleWebSocket(w, r)
|
||||||
default:
|
default:
|
||||||
http.NotFound(w, r)
|
http.NotFound(w, r)
|
||||||
|
|
@ -208,12 +290,7 @@ func (c *PicoChannel) broadcastToSession(chatID string, msg PicoMessage) error {
|
||||||
msg.SessionID = sessionID
|
msg.SessionID = sessionID
|
||||||
|
|
||||||
var sent bool
|
var sent bool
|
||||||
c.connections.Range(func(key, value any) bool {
|
for _, pc := range c.sessionConnectionsSnapshot(sessionID) {
|
||||||
pc, ok := value.(*picoConn)
|
|
||||||
if !ok {
|
|
||||||
return true
|
|
||||||
}
|
|
||||||
if pc.sessionID == sessionID {
|
|
||||||
if err := pc.writeJSON(msg); err != nil {
|
if err := pc.writeJSON(msg); err != nil {
|
||||||
logger.DebugCF("pico", "Write to connection failed", map[string]any{
|
logger.DebugCF("pico", "Write to connection failed", map[string]any{
|
||||||
"conn_id": pc.id,
|
"conn_id": pc.id,
|
||||||
|
|
@ -223,8 +300,6 @@ func (c *PicoChannel) broadcastToSession(chatID string, msg PicoMessage) error {
|
||||||
sent = true
|
sent = true
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return true
|
|
||||||
})
|
|
||||||
|
|
||||||
if !sent {
|
if !sent {
|
||||||
return fmt.Errorf("no active connections for session %s: %w", sessionID, channels.ErrSendFailed)
|
return fmt.Errorf("no active connections for session %s: %w", sessionID, channels.ErrSendFailed)
|
||||||
|
|
@ -276,12 +351,12 @@ func (c *PicoChannel) handleWebSocket(w http.ResponseWriter, r *http.Request) {
|
||||||
}
|
}
|
||||||
|
|
||||||
pc := &picoConn{
|
pc := &picoConn{
|
||||||
id: uuid.New().String(),
|
id: c.newUniqueConnID(),
|
||||||
conn: conn,
|
conn: conn,
|
||||||
sessionID: sessionID,
|
sessionID: sessionID,
|
||||||
}
|
}
|
||||||
|
|
||||||
c.connections.Store(pc.id, pc)
|
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{
|
||||||
|
|
@ -341,12 +416,13 @@ func (c *PicoChannel) matchedSubprotocol(r *http.Request) string {
|
||||||
func (c *PicoChannel) readLoop(pc *picoConn) {
|
func (c *PicoChannel) readLoop(pc *picoConn) {
|
||||||
defer func() {
|
defer func() {
|
||||||
pc.close()
|
pc.close()
|
||||||
c.connections.Delete(pc.id)
|
if removed := c.removeConnection(pc.id); removed != nil {
|
||||||
c.connCount.Add(-1)
|
c.connCount.Add(-1)
|
||||||
logger.InfoCF("pico", "WebSocket client disconnected", map[string]any{
|
logger.InfoCF("pico", "WebSocket client disconnected", map[string]any{
|
||||||
"conn_id": pc.id,
|
"conn_id": removed.id,
|
||||||
"session_id": pc.sessionID,
|
"session_id": removed.sessionID,
|
||||||
})
|
})
|
||||||
|
}
|
||||||
}()
|
}()
|
||||||
|
|
||||||
readTimeout := time.Duration(c.config.ReadTimeout) * time.Second
|
readTimeout := time.Duration(c.config.ReadTimeout) * time.Second
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue