封装的方法类

connPoolx/more/connPoolx.go

@@ -0,0 +1,115 @@
+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)
+		}
+	}
+}

connPoolx/simple/connPoolx.go

@@ -0,0 +1,57 @@
+package connPoolx
+
+import (
+	"errors"
+	"sync"
+)
+
+// 连接池
+
+// 核心思想：连接放在这里面统一管理，所有的新建与删除需要外部操作
+
+// 适用场景：每一个key对应一个val
+
+type connPool struct {
+	pool sync.Map
+}
+
+type PoolValue interface {
+	Close() error
+}
+
+func NewConnPool() *connPool {
+	return &connPool{}
+}
+
+// 同名将覆盖
+func (c *connPool) Store(key string, value PoolValue) {
+	c.pool.Store(key, value)
+}
+
+func (c *connPool) Load(key string) (value PoolValue, ok bool) {
+	val, ok := c.pool.Load(key)
+	if !ok {
+		return nil, false
+	}
+	value, ok = val.(PoolValue)
+	if !ok {
+		return nil, false
+	}
+	return value, true
+}
+
+func (c *connPool) Close(key string) error {
+	val, ok := c.pool.LoadAndDelete(key)
+	if ok {
+		value, ok := val.(PoolValue)
+		if !ok {
+			return errors.New("解析异常")
+		}
+		return value.Close()
+	}
+	return nil
+}
+
+func (c *connPool) Range(f func(key any, value any) bool) {
+	c.pool.Range(f)
+}