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
+}