feat: add WithID/WithLabel options and monotonic counter for subscription IDs

This commit is contained in:
Jay
2026-05-17 22:37:57 -04:00
parent 715dfa17b0
commit 5bbfd9523b
2 changed files with 183 additions and 54 deletions
+68 -30
View File
@@ -2,13 +2,12 @@ package prism
import (
"context"
"crypto/rand"
"encoding/base32"
"fmt"
"git.wisehodl.dev/jay/go-mana-component"
"git.wisehodl.dev/jay/go-roots-ws"
"log/slog"
"sync"
"sync/atomic"
"time"
)
@@ -30,6 +29,7 @@ type ReqClosed struct {
type RequestManager struct {
reqs map[string]*request
counter atomic.Uint64
envoy *Envoy
events <-chan OutboundPoolEvent
@@ -43,6 +43,29 @@ type RequestManager struct {
logger *slog.Logger
}
// ----------------------------------------------------------------------------
// Options
// ----------------------------------------------------------------------------
type RequestOption func(*requestOptions)
type requestOptions struct {
id string
label string
}
// WithID sets an explicit subscription ID. Returns an error from Stream or
// Query if the ID is already in use.
func WithID(id string) RequestOption {
return func(o *requestOptions) { o.id = id }
}
// WithLabel sets the prefix for the generated subscription ID. The default
// prefix is "req". The counter is shared across all labels.
func WithLabel(label string) RequestOption {
return func(o *requestOptions) { o.label = label }
}
type request struct {
id string
filters [][]byte
@@ -58,21 +81,6 @@ type request struct {
closedOnce sync.Once
}
// ----------------------------------------------------------------------------
// Helpers
// ----------------------------------------------------------------------------
var encoder = base32.StdEncoding.WithPadding(base32.NoPadding)
func generateID() string {
b := make([]byte, 5)
if _, err := rand.Read(b); err != nil {
panic(fmt.Sprintf("generateID: %v", err))
}
return encoder.EncodeToString(b)
}
// ----------------------------------------------------------------------------
// Request Manager
// ----------------------------------------------------------------------------
@@ -107,16 +115,32 @@ func NewRequestManager(e *Envoy) *RequestManager {
func (m *RequestManager) Stream(
filters [][]byte,
) (string, <-chan ReqEvent, <-chan ReqClosed) {
id, events, closed := m.newStream(filters, false)
return id, events, closed
opts ...RequestOption,
) (string, <-chan ReqEvent, <-chan ReqClosed, error) {
id, events, closed, err := m.newStream(filters, false, opts...)
return id, events, closed, err
}
func (m *RequestManager) newStream(
filters [][]byte,
isQuery bool,
) (string, <-chan ReqEvent, <-chan ReqClosed) {
id := generateID()
opts ...RequestOption,
) (string, <-chan ReqEvent, <-chan ReqClosed, error) {
var o requestOptions
for _, opt := range opts {
opt(&o)
}
var id string
if o.id != "" {
id = o.id
} else {
label := o.label
if label == "" {
label = "req"
}
id = fmt.Sprintf("%s:%d", label, m.counter.Add(1))
}
buffer := make(chan ReqEvent, 64)
closed := make(chan ReqClosed, 1)
@@ -145,24 +169,38 @@ func (m *RequestManager) newStream(
}
m.mu.Lock()
if _, exists := m.reqs[id]; exists {
m.mu.Unlock()
close(buffer)
if !isQuery {
// drain bufferedPipe goroutine
for range events {
}
}
return "", nil, nil, fmt.Errorf("Stream: id %q already in use", id)
}
m.reqs[id] = req
if m.envoy.IsConnected() {
m.activate(req)
}
m.mu.Unlock()
return id, events, closed
return id, events, closed, nil
}
func (m *RequestManager) Query(
filters [][]byte,
timeout time.Duration,
) ([]ReqEvent, *ReqClosed) {
opts ...RequestOption,
) ([]ReqEvent, *ReqClosed, error) {
if !m.envoy.IsConnected() {
return nil, nil
return nil, nil, nil
}
id, eventsCh, closedCh := m.newStream(filters, true)
id, eventsCh, closedCh, err := m.newStream(filters, true, opts...)
if err != nil {
return nil, nil, err
}
ctx, cancel := context.WithTimeout(m.ctx, timeout)
defer cancel()
@@ -172,17 +210,17 @@ func (m *RequestManager) Query(
select {
case ev, ok := <-eventsCh:
if !ok {
return result, nil
return result, nil, nil
}
result = append(result, ev)
case cl, ok := <-closedCh:
if !ok {
return result, nil
return result, nil, nil
}
return result, &cl
return result, &cl, nil
case <-ctx.Done():
m.Cancel(id)
return result, nil
return result, nil, nil
}
}
}
+114 -23
View File
@@ -7,6 +7,86 @@ import (
"time"
)
func TestRequestManager_Options(t *testing.T) {
t.Run("default id uses req label and monotonic counter", func(t *testing.T) {
_, envoy := newMockEnvoy(t)
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
filters := [][]byte{[]byte(`{}`)}
idA, _, _, err := m.Stream(filters)
assert.NoError(t, err)
idB, _, _, err := m.Stream(filters)
assert.NoError(t, err)
assert.Equal(t, "req:1", idA)
assert.Equal(t, "req:2", idB)
})
t.Run("WithLabel sets prefix", func(t *testing.T) {
_, envoy := newMockEnvoy(t)
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
filters := [][]byte{[]byte(`{}`)}
idA, _, _, err := m.Stream(filters, WithLabel("feed"))
assert.NoError(t, err)
idB, _, _, err := m.Stream(filters, WithLabel("profile"))
assert.NoError(t, err)
assert.Equal(t, "feed:1", idA)
assert.Equal(t, "profile:2", idB)
})
t.Run("WithID uses caller id", func(t *testing.T) {
_, envoy := newMockEnvoy(t)
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
filters := [][]byte{[]byte(`{}`)}
id, _, _, err := m.Stream(filters, WithID("my-custom-id"))
assert.NoError(t, err)
assert.Equal(t, "my-custom-id", id)
})
t.Run("WithID wins over WithLabel", func(t *testing.T) {
_, envoy := newMockEnvoy(t)
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
filters := [][]byte{[]byte(`{}`)}
id, _, _, err := m.Stream(filters, WithLabel("feed"), WithID("explicit"))
assert.NoError(t, err)
assert.Equal(t, "explicit", id)
})
t.Run("WithID returns error on collision", func(t *testing.T) {
_, envoy := newMockEnvoy(t)
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
filters := [][]byte{[]byte(`{}`)}
_, _, _, err := m.Stream(filters, WithID("dup"))
assert.NoError(t, err)
_, _, _, err = m.Stream(filters, WithID("dup"))
assert.Error(t, err)
})
t.Run("WithID does not advance counter", func(t *testing.T) {
_, envoy := newMockEnvoy(t)
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
filters := [][]byte{[]byte(`{}`)}
_, _, _, err := m.Stream(filters, WithID("explicit"))
assert.NoError(t, err)
id, _, _, err := m.Stream(filters)
assert.NoError(t, err)
assert.Equal(t, "req:1", id)
})
}
func TestRequestManager_Stream(t *testing.T) {
t.Run("sends req when connected", func(t *testing.T) {
p, envoy := newMockEnvoy(t)
@@ -17,7 +97,8 @@ func TestRequestManager_Stream(t *testing.T) {
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
filters := [][]byte{[]byte(`{}`)}
id, events, closed := m.Stream(filters)
id, events, closed, err := m.Stream(filters)
assert.NoError(t, err)
assert.NotEmpty(t, id)
assert.NotNil(t, events)
@@ -42,7 +123,8 @@ func TestRequestManager_Stream(t *testing.T) {
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
filters := [][]byte{[]byte(`{}`)}
id, events, closed := m.Stream(filters)
id, events, closed, err := m.Stream(filters)
assert.NoError(t, err)
assert.NotEmpty(t, id)
assert.NotNil(t, events)
@@ -66,7 +148,8 @@ func TestRequestManager_Stream(t *testing.T) {
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
filters := [][]byte{[]byte(`{}`)}
id, events, _ := m.Stream(filters)
id, events, _, err := m.Stream(filters)
assert.NoError(t, err)
// drain the REQ send
Eventually(t, func() bool {
@@ -110,7 +193,8 @@ func TestRequestManager_Stream(t *testing.T) {
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
filters := [][]byte{[]byte(`{}`)}
id, events, closed := m.Stream(filters)
id, events, closed, err := m.Stream(filters)
assert.NoError(t, err)
// drain the REQ send
Eventually(t, func() bool {
@@ -162,7 +246,8 @@ func TestRequestManager_Stream(t *testing.T) {
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
filters := [][]byte{[]byte(`{}`)}
id, events, closed := m.Stream(filters)
id, events, closed, err := m.Stream(filters)
assert.NoError(t, err)
// drain the REQ send
Eventually(t, func() bool {
@@ -221,7 +306,8 @@ func TestRequestManager_Cancel(t *testing.T) {
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
filters := [][]byte{[]byte(`{}`)}
id, events, _ := m.Stream(filters)
id, events, _, streamErr := m.Stream(filters)
assert.NoError(t, streamErr)
// drain the REQ send
Eventually(t, func() bool {
@@ -269,7 +355,8 @@ func TestRequestManager_Cancel(t *testing.T) {
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
filters := [][]byte{[]byte(`{}`)}
id, events, _ := m.Stream(filters)
id, events, _, streamErr := m.Stream(filters)
assert.NoError(t, streamErr)
err := m.Cancel(id)
assert.NoError(t, err)
@@ -327,7 +414,8 @@ func TestRequestManager_Query(t *testing.T) {
p.receive(envelope.EncloseEOSE(subID))
}()
events, closed := m.Query(filters, TestTimeout)
events, closed, err := m.Query(filters, TestTimeout)
assert.NoError(t, err)
assert.Len(t, events, 3)
assert.Nil(t, closed)
@@ -364,7 +452,8 @@ func TestRequestManager_Query(t *testing.T) {
p.receive(envelope.EncloseClosed(subID, reason))
}()
events, closed := m.Query(filters, TestTimeout)
events, closed, err := m.Query(filters, TestTimeout)
assert.NoError(t, err)
assert.Empty(t, events)
if assert.NotNil(t, closed) {
@@ -398,8 +487,9 @@ func TestRequestManager_Query(t *testing.T) {
}()
start := time.Now()
events, closed := m.Query(filters, queryTimeout)
events, closed, err := m.Query(filters, queryTimeout)
elapsed := time.Since(start)
assert.NoError(t, err)
assert.GreaterOrEqual(t, elapsed, queryTimeout)
assert.Len(t, events, 2)
@@ -413,7 +503,8 @@ func TestRequestManager_Query(t *testing.T) {
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
events, closed := m.Query([][]byte{[]byte(`{}`)}, TestTimeout)
events, closed, err := m.Query([][]byte{[]byte(`{}`)}, TestTimeout)
assert.NoError(t, err)
assert.Nil(t, events)
assert.Nil(t, closed)
@@ -429,8 +520,8 @@ func TestRequestManager_Reconnect(t *testing.T) {
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
filters := [][]byte{[]byte(`{}`)}
idA, _, _ := m.Stream(filters)
idB, _, _ := m.Stream(filters)
idA, _, _, _ := m.Stream(filters)
idB, _, _, _ := m.Stream(filters)
// drain both REQ sends
for range 2 {
@@ -463,8 +554,8 @@ func TestRequestManager_Reconnect(t *testing.T) {
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
filters := [][]byte{[]byte(`{}`)}
idA, eventsA, closedA := m.Stream(filters)
idB, eventsB, closedB := m.Stream(filters)
idA, eventsA, closedA, _ := m.Stream(filters)
idB, eventsB, closedB, _ := m.Stream(filters)
for range 2 {
Eventually(t, func() bool {
@@ -539,8 +630,8 @@ func TestRequestManager_Reconnect(t *testing.T) {
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
filters := [][]byte{[]byte(`{}`)}
idA, _, _ := m.Stream(filters)
idB, _, _ := m.Stream(filters)
idA, _, _, _ := m.Stream(filters)
idB, _, _, _ := m.Stream(filters)
for range 2 {
Eventually(t, func() bool {
@@ -599,7 +690,7 @@ func TestRequestManager_Reconnect(t *testing.T) {
m := NewRequestManager(envoy)
t.Cleanup(func() { m.Close() })
filters := [][]byte{[]byte(`{}`)}
id, events, _ := m.Stream(filters)
id, events, _, _ := m.Stream(filters)
Eventually(t, func() bool {
select {
@@ -650,9 +741,9 @@ func TestRequestManager_Close(t *testing.T) {
m := NewRequestManager(envoy)
filters := [][]byte{[]byte(`{}`)}
m.Stream(filters)
m.Stream(filters)
m.Stream(filters)
_, _, _, _ = m.Stream(filters)
_, _, _, _ = m.Stream(filters)
_, _, _, _ = m.Stream(filters)
// drain all three REQ sends
for range 3 {
@@ -699,8 +790,8 @@ func TestRequestManager_Close(t *testing.T) {
m := NewRequestManager(envoy)
filters := [][]byte{[]byte(`{}`)}
_, eventsA, _ := m.Stream(filters)
_, eventsB, _ := m.Stream(filters)
_, eventsA, _, _ := m.Stream(filters)
_, eventsB, _, _ := m.Stream(filters)
for range 2 {
Eventually(t, func() bool {