git.r.bdr.sh - rbdr/mobius/blob - hotline/client_manager.go

   1 package hotline
   2
   3 import (
   4         "cmp"
   5         "encoding/binary"
   6         "github.com/stretchr/testify/mock"
   7         "slices"
   8         "sync"
   9         "sync/atomic"
  10 )
  11
  12 type ClientID [2]byte
  13
  14 type ClientManager interface {
  15         List() []*ClientConn // Returns list of sorted clients
  16         Get(id ClientID) *ClientConn
  17         Add(cc *ClientConn)
  18         Delete(id ClientID)
  19 }
  20
  21 type MockClientMgr struct {
  22         mock.Mock
  23 }
  24
  25 func (m *MockClientMgr) List() []*ClientConn {
  26         args := m.Called()
  27
  28         return args.Get(0).([]*ClientConn)
  29 }
  30
  31 func (m *MockClientMgr) Get(id ClientID) *ClientConn {
  32         args := m.Called(id)
  33
  34         return args.Get(0).(*ClientConn)
  35 }
  36
  37 func (m *MockClientMgr) Add(cc *ClientConn) {
  38         m.Called(cc)
  39 }
  40 func (m *MockClientMgr) Delete(id ClientID) {
  41         m.Called(id)
  42 }
  43
  44 type MemClientMgr struct {
  45         clients map[ClientID]*ClientConn
  46
  47         mu           sync.Mutex
  48         nextClientID atomic.Uint32
  49 }
  50
  51 func NewMemClientMgr() *MemClientMgr {
  52         return &MemClientMgr{
  53                 clients: make(map[ClientID]*ClientConn),
  54         }
  55 }
  56
  57 // List returns slice of sorted clients.
  58 func (cm *MemClientMgr) List() []*ClientConn {
  59         cm.mu.Lock()
  60         defer cm.mu.Unlock()
  61
  62         var clients []*ClientConn
  63         for _, client := range cm.clients {
  64                 clients = append(clients, client)
  65         }
  66
  67         slices.SortFunc(clients, func(a, b *ClientConn) int {
  68                 return cmp.Compare(
  69                         binary.BigEndian.Uint16(a.ID[:]),
  70                         binary.BigEndian.Uint16(b.ID[:]),
  71                 )
  72         })
  73
  74         return clients
  75 }
  76
  77 func (cm *MemClientMgr) Get(id ClientID) *ClientConn {
  78         cm.mu.Lock()
  79         defer cm.mu.Unlock()
  80
  81         return cm.clients[id]
  82 }
  83
  84 func (cm *MemClientMgr) Add(cc *ClientConn) {
  85         cm.mu.Lock()
  86         defer cm.mu.Unlock()
  87
  88         cm.nextClientID.Add(1)
  89         binary.BigEndian.PutUint16(cc.ID[:], uint16(cm.nextClientID.Load()))
  90
  91         cm.clients[cc.ID] = cc
  92 }
  93
  94 func (cm *MemClientMgr) Delete(id ClientID) {
  95         cm.mu.Lock()
  96         defer cm.mu.Unlock()
  97
  98         delete(cm.clients, id)
  99 }