mercury/internal/streaming/sockets.go

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: {},
		},
	}
}