136 lines
2.3 KiB
Go
136 lines
2.3 KiB
Go
package sse
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"maps"
|
|
"net/http"
|
|
"sync"
|
|
)
|
|
|
|
type (
|
|
Queue struct {
|
|
in chan Event
|
|
out map[int64]map[int]chan Event
|
|
|
|
ctx context.Context
|
|
cancel context.CancelFunc
|
|
|
|
lock sync.Mutex
|
|
}
|
|
|
|
Event struct {
|
|
AccountID int64
|
|
Type string
|
|
Data []byte
|
|
}
|
|
)
|
|
|
|
func NewQueue() *Queue {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
return &Queue{
|
|
in: make(chan Event),
|
|
out: make(map[int64]map[int]chan Event),
|
|
|
|
ctx: ctx,
|
|
cancel: cancel,
|
|
}
|
|
}
|
|
|
|
func (q *Queue) Start(ctx context.Context) error {
|
|
defer q.cancel()
|
|
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
// we're done piping events
|
|
return nil
|
|
case e := <-q.in:
|
|
if e.AccountID == 0 {
|
|
continue
|
|
}
|
|
|
|
// share event with listeners on the account
|
|
q.lock.Lock()
|
|
for _, out := range q.out[e.AccountID] {
|
|
out <- e
|
|
}
|
|
q.lock.Unlock()
|
|
}
|
|
}
|
|
}
|
|
|
|
func (q *Queue) Listen(ctx context.Context, acctID int64, fn func(context.Context, *Event) error) error {
|
|
// create new out pipe and append it to the queue
|
|
|
|
out := make(chan Event, 1)
|
|
|
|
q.lock.Lock()
|
|
outs := q.out[acctID]
|
|
if outs == nil {
|
|
outs = make(map[int]chan Event)
|
|
q.out[acctID] = outs
|
|
}
|
|
|
|
var maxID int
|
|
for n := range maps.Keys(outs) {
|
|
maxID = max(maxID, n)
|
|
}
|
|
id := maxID + 1
|
|
|
|
outs[id] = out
|
|
q.lock.Unlock()
|
|
|
|
// delete the pipe when done listening
|
|
defer func() {
|
|
q.lock.Lock()
|
|
delete(q.out[acctID], id)
|
|
if len(q.out[acctID]) == 0 {
|
|
delete(q.out, acctID)
|
|
}
|
|
q.lock.Unlock()
|
|
}()
|
|
|
|
for {
|
|
select {
|
|
case <-q.ctx.Done():
|
|
// queue is shut down
|
|
return nil
|
|
case <-ctx.Done():
|
|
// done listening to events
|
|
return nil
|
|
case e := <-out:
|
|
// pass event to the caller
|
|
if err := fn(ctx, &e); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (q *Queue) Send(ctx context.Context, e Event) error {
|
|
select {
|
|
case <-q.ctx.Done():
|
|
// the queue has closed
|
|
return fmt.Errorf("event queue closed: %w", q.ctx.Err())
|
|
case <-ctx.Done():
|
|
// sender ran out of time
|
|
return fmt.Errorf("provided context canceled: %w", ctx.Err())
|
|
// push the event onto the queue
|
|
case q.in <- e:
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (e *Event) Write(w http.ResponseWriter) {
|
|
fmt.Fprintf(
|
|
w,
|
|
"event: %s\ndata: %s\n\n",
|
|
e.Type,
|
|
bytes.ReplaceAll(e.Data, []byte("\n"), []byte(" ")),
|
|
)
|
|
w.(http.Flusher).Flush()
|
|
}
|