WIP add some sync stuff

util/multiqueue/multiqueue.go

@@ -3,8 +3,9 @@ package multiqueue
 import (
 	"context"
 	"errors"
-	"github.com/cheggaaa/mb/v3"
 	"sync"
+
+	"github.com/cheggaaa/mb/v3"
 )
 
 var (
@@ -25,6 +26,7 @@ type HandleFunc[T any] func(msg T)
 type MultiQueue[T any] interface {
 	Add(ctx context.Context, threadId string, msg T) (err error)
 	CloseThread(threadId string) (err error)
+	ThreadIds() []string
 	Close() (err error)
 }
 
@@ -36,6 +38,16 @@ type multiQueue[T any] struct {
 	closed       bool
 }
 
+func (m *multiQueue[T]) ThreadIds() []string {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	ids := make([]string, 0, len(m.threads))
+	for id := range m.threads {
+		ids = append(ids, id)
+	}
+	return ids
+}
+
 func (m *multiQueue[T]) Add(ctx context.Context, threadId string, msg T) (err error) {
 	m.mu.Lock()
 	if m.closed {

util/multiqueue/queue.go

@@ -0,0 +1,35 @@
+package multiqueue
+
+import (
+	"context"
+
+	"github.com/cheggaaa/mb/v3"
+)
+
+type QueueHandler[T any] func(id string, msg T, q *mb.MB[T]) error
+
+type Queue[T any] struct {
+	id      string
+	q       *mb.MB[T]
+	handler QueueHandler[T]
+}
+
+func NewQueue[T any](id string, size int, h QueueHandler[T]) *Queue[T] {
+	return &Queue[T]{
+		id:      id,
+		q:       mb.New[T](size),
+		handler: h,
+	}
+}
+
+func (q *Queue[T]) TryAdd(msg T) error {
+	return q.handler(q.id, msg, q.q)
+}
+
+func (q *Queue[T]) WaitOne(ctx context.Context) (T, error) {
+	return q.q.WaitOne(ctx)
+}
+
+func (q *Queue[T]) Close() error {
+	return q.q.Close()
+}

util/multiqueue/queueprovider.go

@@ -0,0 +1,47 @@
+package multiqueue
+
+import (
+	"sync"
+)
+
+type QueueProvider[T any] interface {
+	GetQueue(id string) *Queue[T]
+	RemoveQueue(id string) error
+}
+
+type queueProvider[T any] struct {
+	queues  map[string]*Queue[T]
+	mx      sync.Mutex
+	size    int
+	handler QueueHandler[T]
+}
+
+func NewQueueProvider[T any](size int, handler QueueHandler[T]) QueueProvider[T] {
+	return &queueProvider[T]{
+		queues:  make(map[string]*Queue[T]),
+		size:    size,
+		handler: handler,
+	}
+}
+
+func (p *queueProvider[T]) GetQueue(id string) *Queue[T] {
+	p.mx.Lock()
+	defer p.mx.Unlock()
+	q, ok := p.queues[id]
+	if !ok {
+		q = NewQueue(id, p.size, p.handler)
+		p.queues[id] = q
+	}
+	return q
+}
+
+func (p *queueProvider[T]) RemoveQueue(id string) error {
+	p.mx.Lock()
+	defer p.mx.Unlock()
+	q, ok := p.queues[id]
+	if !ok {
+		return nil
+	}
+	delete(p.queues, id)
+	return q.Close()
+}