eod commit

util/cache/conn_cache.go

@@ -0,0 +1,88 @@
+package cache
+
+import (
+	"github.com/gorilla/websocket"
+)
+
+func NewConnectionCache() *ConnectionCache {
+	c := &ConnectionCache{}
+	c.state.Store(&connectionMappings{
+		id2c: make(map[string]*conn),
+	})
+
+	return c
+}
+
+func (c *ConnectionCache) Set(id string, connection *websocket.Conn, meta *ConnectionMetaData) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	current := c.state.Load()
+	next := current.clone()
+
+	next.id2c[id] = &conn{
+		connection: connection,
+		meta:       meta,
+	}
+
+	c.state.Store(next)
+}
+
+func (c *ConnectionCache) GetById(id string) (*websocket.Conn, *ConnectionMetaData, bool) {
+	m := c.state.Load()
+	if conn, ok := m.id2c[id]; ok {
+		return conn.connection, conn.meta, true
+	}
+
+	return nil, nil, false
+}
+
+func (c *ConnectionCache) RemoveById(id string) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	current := c.state.Load()
+	if _, ok := current.id2c[id]; !ok {
+		return
+	}
+
+	next := current.clone()
+	delete(next.id2c, id)
+
+	c.state.Store(next)
+}
+
+func (c *ConnectionCache) Range(fn func(id string, connection *websocket.Conn, meta *ConnectionMetaData)) {
+	m := c.state.Load()
+	for id, conn := range m.id2c {
+		fn(id, conn.connection, conn.meta)
+	}
+}
+
+func (c *ConnectionCache) GetAllConnections() []*websocket.Conn {
+	m := c.state.Load()
+	conns := make([]*websocket.Conn, 0, len(m.id2c))
+	for _, conn := range m.id2c {
+		conns = append(conns, conn.connection)
+	}
+	return conns
+}
+
+func (c *ConnectionCache) Clear() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	c.state.Store(&connectionMappings{
+		id2c: make(map[string]*conn),
+	})
+}
+
+func (cm *connectionMappings) clone() *connectionMappings {
+	newCM := &connectionMappings{
+		id2c: make(map[string]*conn, len(cm.id2c)),
+	}
+	for k, v := range cm.id2c {
+		newCM.id2c[k] = v
+	}
+	return newCM
+}

util/cache/structs.go

@@ -3,6 +3,8 @@ package cache
 import (
 	"sync"
 	"sync/atomic"
+
+	"github.com/gorilla/websocket"
 )
 
 type Cache struct {
@@ -10,7 +12,26 @@ type Cache struct {
 	state atomic.Pointer[mappings]
 }
 
+type ConnectionCache struct {
+	mu    sync.Mutex
+	state atomic.Pointer[connectionMappings]
+}
+
 type mappings struct {
 	s2c map[string]string
 	c2s map[string]string
 }
+
+type connectionMappings struct {
+	id2c map[string]*conn
+}
+
+type ConnectionMetaData struct {
+	ConnectionType string // "mod" or "bot"
+	ID             string // server_id or bot_id for logging
+}
+
+type conn struct {
+	connection *websocket.Conn
+	meta       *ConnectionMetaData
+}