package streaming import ( "context" "sync" "github.com/oklog/ulid/v2" ) type SocketHolder struct { // map of sockets to sockets map[ulid.ULID]*userSockets mu sync.Mutex } func NewSocketHolder() *SocketHolder { return &SocketHolder{ sockets: make(map[ulid.ULID]*userSockets), } } func (sh *SocketHolder) Send(acctID ulid.ULID, et EventType, data any) { userSockets := sh.SocketsFor(acctID) userSockets.mu.Lock() sockets := make([]*Socket, len(userSockets.sockets)) copy(sockets, userSockets.sockets) userSockets.mu.Unlock() for _, s := range sockets { if s.WillAcceptEvent(et) { // the socket might block for a bit, so spin this off into a separate goroutine go func(s *Socket) { s.ch <- Event{Type: et, Data: data} }(s) } } } func (s *SocketHolder) SocketsFor(acct ulid.ULID) *userSockets { s.mu.Lock() defer s.mu.Unlock() us, ok := s.sockets[acct] if !ok { us = &userSockets{} s.sockets[acct] = us } return us } const sessionCountLimit = 50 // no more than 50 concurrent sessions per user type userSockets struct { mu sync.Mutex sockets []*Socket } func (s *userSockets) NewSocket(ctx context.Context, cancel context.CancelFunc) (*Socket, bool) { s.mu.Lock() defer s.mu.Unlock() if len(s.sockets) >= sessionCountLimit { return nil, false } socket := NewSocket(ctx, cancel) s.sockets = append(s.sockets, socket) return socket, true } type Socket struct { ctx context.Context cancel context.CancelFunc ch chan Event types map[EventType]struct{} mu sync.RWMutex } func (s *Socket) WillAcceptEvent(mt EventType) bool { if mt == EventTypeError { return true } s.mu.RLock() _, ok := s.types[mt] s.mu.RUnlock() return ok } func (s *Socket) SetEvent(mt EventType, add bool) { s.mu.Lock() if add { s.types[mt] = struct{}{} } else { delete(s.types, mt) } s.mu.Unlock() } func (s *Socket) Cancel() { s.cancel() } func (s *Socket) Done() <-chan struct{} { return s.ctx.Done() } func (s *Socket) Chan() <-chan Event { return s.ch } func NewSocket(ctx context.Context, cancel context.CancelFunc) *Socket { return &Socket{ ctx: ctx, cancel: cancel, ch: make(chan Event), types: map[EventType]struct{}{ EventTypeEcho: {}, }, } }