github.com/binbinly/pkg@v0.0.11-0.20240321014439-f4fbf666eb0f/transport/ws/manager.go (about)

     1  package ws
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  // Manager 连接管理模块
     8  type Manager struct {
     9  	mu sync.RWMutex
    10  
    11  	connections map[uint64]Connection
    12  }
    13  
    14  // NewManager 创建一个链接管理器
    15  func NewManager() *Manager {
    16  	return &Manager{
    17  		connections: make(map[uint64]Connection),
    18  	}
    19  }
    20  
    21  // Add 添加连接
    22  func (c *Manager) Add(conn Connection) {
    23  	c.mu.Lock()
    24  	defer c.mu.Unlock()
    25  
    26  	c.connections[conn.GetID()] = conn
    27  }
    28  
    29  // Remove 删除连接
    30  func (c *Manager) Remove(conn Connection) {
    31  	c.mu.Lock()
    32  	defer c.mu.Unlock()
    33  
    34  	delete(c.connections, conn.GetID())
    35  }
    36  
    37  // Get 获取连接
    38  func (c *Manager) Get(cid uint64) (Connection, error) {
    39  	c.mu.RLock()
    40  	defer c.mu.RUnlock()
    41  
    42  	if conn, ok := c.connections[cid]; ok {
    43  		return conn, nil
    44  	}
    45  	return nil, ErrConnNotFound
    46  }
    47  
    48  // Len 获取连接数
    49  func (c *Manager) Len() int {
    50  	return len(c.connections)
    51  }
    52  
    53  // Clear 清除并停止所有连接
    54  func (c *Manager) Clear() {
    55  	// 这里不需要单独调用删除连接，也不可以加锁，conn.Stop()里处理了，否则会死锁
    56  	for _, conn := range c.connections {
    57  		conn.Stop()
    58  	}
    59  }
    60  
    61  // Range 遍历所有连接
    62  func (c *Manager) Range(f ConnHandlerFunc) (err error) {
    63  	for cid, conn := range c.connections {
    64  		err = f(cid, conn)
    65  	}
    66  
    67  	return err
    68  }