104 lines
1.9 KiB
Go
104 lines
1.9 KiB
Go
|
package streaming
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"sync"
|
||
|
|
|
||
|
|
"github.com/oklog/ulid/v2"
|
||
|
|
)
|
||
|
|
|
||
|
|
var SocketHolder socketHolder
|
||
|
|
|
||
|
|
type socketHolder struct {
|
||
|
|
// map of sockets to
|
||
|
|
sockets map[ulid.ULID]*userSockets
|
||
|
|
mu sync.Mutex
|
||
|
|
}
|
||
|
|
|
||
|
|
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()
|
||
|
|
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 newSocket(ctx context.Context, cancel context.CancelFunc) *socket {
|
||
|
|
return &socket{
|
||
|
|
ctx: ctx,
|
||
|
|
cancel: cancel,
|
||
|
|
ch: make(chan Event),
|
||
|
|
types: make(map[EventType]struct{}),
|
||
|
|
}
|
||
|
|
}
|