116 lines
2.0 KiB
Go
116 lines
2.0 KiB
Go
package more
|
|
|
|
import (
|
|
"sync"
|
|
|
|
uuid "github.com/satori/go.uuid"
|
|
)
|
|
|
|
// 同一个key对应多个连接
|
|
// 2023年6月29日13:45:25
|
|
|
|
// 场景：
|
|
// 1. 同时支持多个用户在线
|
|
// 1. 一个用户对应多个连接
|
|
// 2. 在其中一个用户连接中发送消息，同用户其他连接也能收到
|
|
|
|
type connPool struct {
|
|
pool map[string]map[string]PoolValue
|
|
idx map[string]string // 索引
|
|
mu sync.RWMutex
|
|
}
|
|
|
|
type PoolValue interface {
|
|
Close() error
|
|
}
|
|
|
|
// 初始化连接池
|
|
func NewConnPool() *connPool {
|
|
return &connPool{
|
|
pool: make(map[string]map[string]PoolValue),
|
|
idx: make(map[string]string),
|
|
mu: sync.RWMutex{},
|
|
}
|
|
}
|
|
|
|
// 添加连接, 返回connId
|
|
func (c *connPool) Store(key string, value PoolValue) string {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
val, _ := c.pool[key]
|
|
|
|
connId := uuid.NewV4().String()
|
|
val[connId] = value
|
|
c.pool[key] = val
|
|
c.idx[connId] = key
|
|
return connId
|
|
}
|
|
|
|
// 获取一个用户的所有连接
|
|
func (c *connPool) LoadByKey(key string) (conns map[string]PoolValue, ok bool) {
|
|
c.mu.RLock()
|
|
defer c.mu.RUnlock()
|
|
|
|
valMap, ok := c.pool[key]
|
|
|
|
return valMap, ok
|
|
}
|
|
|
|
// 根据连接ID获取连接
|
|
func (c *connPool) LoadByConnId(connId string) (*PoolValue, bool) {
|
|
c.mu.RLock()
|
|
defer c.mu.Unlock()
|
|
|
|
key, ok := c.idx[connId]
|
|
if !ok {
|
|
return nil, false
|
|
}
|
|
|
|
valMap, ok := c.pool[key]
|
|
if !ok {
|
|
return nil, false
|
|
}
|
|
|
|
val, ok := valMap[connId]
|
|
if !ok {
|
|
return nil, false
|
|
}
|
|
return &val, true
|
|
}
|
|
|
|
// 关闭连接
|
|
func (c *connPool) Close(connId string) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
|
|
key, ok := c.idx[connId]
|
|
if !ok {
|
|
return
|
|
}
|
|
valMap, ok := c.pool[key]
|
|
if !ok {
|
|
return
|
|
}
|
|
val, ok := valMap[connId]
|
|
if !ok {
|
|
return
|
|
}
|
|
delete(valMap, connId)
|
|
|
|
// 异步关闭
|
|
go val.Close()
|
|
}
|
|
|
|
// 遍历所有的连接
|
|
func (c *connPool) Range(f func(key string, connId string, values PoolValue)) {
|
|
c.mu.RLock()
|
|
defer c.mu.Unlock()
|
|
|
|
for key, val := range c.pool {
|
|
val := val
|
|
for keyPoo, valPoo := range val {
|
|
go f(key, keyPoo, valPoo)
|
|
}
|
|
}
|
|
}
|