filter ss events by account

internal/server/sse/sse.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"maps"
 	"net/http"
 	"sync"
 )
@@ -11,18 +12,18 @@ import (
 type (
 	Queue struct {
 		in  chan Event
-		out map[int]chan Event
+		out map[int64]map[int]chan Event
 
 		ctx    context.Context
 		cancel context.CancelFunc
 
-		lock   sync.Mutex
-		prevID int
+		lock sync.Mutex
 	}
 
 	Event struct {
-		Type string
-		Data []byte
+		AccountID int64
+		Type      string
+		Data      []byte
 	}
 )
 
@@ -30,7 +31,7 @@ func NewQueue() *Queue {
 	ctx, cancel := context.WithCancel(context.Background())
 	return &Queue{
 		in:  make(chan Event),
-		out: make(map[int]chan Event),
+		out: make(map[int64]map[int]chan Event),
 
 		ctx:    ctx,
 		cancel: cancel,
@@ -46,9 +47,13 @@ func (q *Queue) Start(ctx context.Context) error {
 			// we're done piping events
 			return nil
 		case e := <-q.in:
-			// share event will all subscribers
+			if e.AccountID == 0 {
+				continue
+			}
+
+			// share event with listeners on the account
 			q.lock.Lock()
-			for _, out := range q.out {
+			for _, out := range q.out[e.AccountID] {
 				out <- e
 			}
 			q.lock.Unlock()
@@ -56,21 +61,34 @@ func (q *Queue) Start(ctx context.Context) error {
 	}
 }
 
-func (q *Queue) Listen(ctx context.Context, fn func(context.Context, *Event) error) error {
+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()
-	q.prevID += 1
-	id := q.prevID
-	q.out[id] = out
+	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, id)
+		delete(q.out[acctID], id)
+		if len(q.out[acctID]) == 0 {
+			delete(q.out, acctID)
+		}
 		q.lock.Unlock()
 	}()
 
@@ -91,7 +109,7 @@ func (q *Queue) Listen(ctx context.Context, fn func(context.Context, *Event) err
 	}
 }
 
-func (q *Queue) Send(ctx context.Context, eventType string, data []byte) error {
+func (q *Queue) Send(ctx context.Context, e Event) error {
 	select {
 	case <-q.ctx.Done():
 		// the queue has closed
@@ -100,10 +118,7 @@ func (q *Queue) Send(ctx context.Context, eventType string, data []byte) error {
 		// sender ran out of time
 		return fmt.Errorf("provided context canceled: %w", ctx.Err())
 		// push the event onto the queue
-	case q.in <- Event{
-		Type: eventType,
-		Data: data,
-	}:
+	case q.in <- e:
 	}
 
 	return nil