89 lines
1.7 KiB
Go
89 lines
1.7 KiB
Go
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
|
|
}
|