Skip to content
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
107 changes: 75 additions & 32 deletions internal/interface/grpc/handlers/broker.go
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,17 @@ type listener[T any] struct {
done chan struct{}
closeDoneMux sync.Once
timeoutTimer *time.Timer
lock *sync.RWMutex
// attached tracks the stream currently consuming ch, if any. Like
// timeoutTimer it is guarded by the owning broker's lock.
attached *attachment
lock *sync.RWMutex
}

// attachment represents a stream's exclusive hold on a listener. It is
// created by attach and released by detach; a newer attach displaces the
// previous holder by closing its displaced channel.
type attachment struct {
displaced chan struct{}
}

func newListener[T any](id string, topics []string) *listener[T] {
Expand Down Expand Up @@ -210,14 +220,52 @@ func (h *broker[T]) removeListener(id string) {
delete(h.listeners, id)
}

func (h *broker[T]) getListenerChannel(id string) (chan T, error) {
h.lock.RLock()
listener, ok := h.listeners[id]
h.lock.RUnlock()
// attach registers the calling stream as the listener's sole consumer and
// cancels any pending removal timeout. If another stream is already attached
// it is displaced: its attachment's displaced channel is closed and ownership
// of the listener's lifecycle moves to the new attachment.
func (h *broker[T]) attach(id string) (*listener[T], *attachment, error) {
h.lock.Lock()
defer h.lock.Unlock()

l, ok := h.listeners[id]
if !ok {
return nil, fmt.Errorf("%w: %s", ErrSubscriptionNotFound, id)
return nil, nil, fmt.Errorf("%w: %s", ErrSubscriptionNotFound, id)
}
if l.timeoutTimer != nil {
l.timeoutTimer.Stop()
l.timeoutTimer = nil
}
if l.attached != nil {
close(l.attached.displaced)
}
return listener.ch, nil
l.attached = &attachment{displaced: make(chan struct{})}
return l, l.attached, nil
}

// release ends att's hold on the listener and settles its lifecycle atomically:
// kept for reconnectWindow if it still has filters, removed otherwise. Returns
// false if att was displaced — the successor owns the lifecycle.
func (h *broker[T]) release(id string, att *attachment, reconnectWindow time.Duration) bool {
h.lock.Lock()
defer h.lock.Unlock()

l, ok := h.listeners[id]
if !ok || l.attached != att {
return false
}
l.attached = nil

l.lock.RLock()
hasFilters := len(l.topics) > 0 || len(l.txFilters) > 0
l.lock.RUnlock()
if reconnectWindow > 0 && hasFilters {
h.scheduleExpiryLocked(l, reconnectWindow)
return true
}
l.closeDone()
delete(h.listeners, id)
return true
}

func (h *broker[T]) getTopics(id string) []string {
Expand Down Expand Up @@ -315,44 +363,39 @@ func compileTxFilters(exprs []string) (map[string]txfilter.Filter, error) {
}

func (h *broker[T]) startTimeout(id string, timeout time.Duration) {
// stop any existing timeout on this listener
h.stopTimeout(id)

h.lock.Lock()
defer h.lock.Unlock()
_, ok := h.listeners[id]

l, ok := h.listeners[id]
if !ok {
return
}
// The timeout reaps a listener no stream is consuming; while one is
// attached it must not be armed (attach cancels it on takeover).
if l.attached != nil {
return
}
h.scheduleExpiryLocked(l, timeout)
}

h.listeners[id].timeoutTimer = time.AfterFunc(timeout, func() {
// scheduleExpiryLocked (re)arms the expiry timer on l; broker lock must be held, and only the current timer may remove the listener.
func (h *broker[T]) scheduleExpiryLocked(l *listener[T], timeout time.Duration) {
if l.timeoutTimer != nil {
l.timeoutTimer.Stop()
}
var timer *time.Timer
timer = time.AfterFunc(timeout, func() {
h.lock.Lock()
defer h.lock.Unlock()

listener, ok := h.listeners[id]
if !ok {
listener, ok := h.listeners[l.id]
if !ok || listener.timeoutTimer != timer {
return
}
if listener.timeoutTimer != nil {
listener.timeoutTimer.Stop()
}
listener.closeDone()
delete(h.listeners, id)
delete(h.listeners, l.id)
})
}

func (h *broker[T]) stopTimeout(id string) {
h.lock.Lock()
defer h.lock.Unlock()

if _, ok := h.listeners[id]; !ok {
return
}

if h.listeners[id].timeoutTimer != nil {
h.listeners[id].timeoutTimer.Stop()
h.listeners[id].timeoutTimer = nil
}
l.timeoutTimer = timer
}

func (h *broker[T]) getListenersCopy() map[string]*listener[T] {
Expand Down
172 changes: 141 additions & 31 deletions internal/interface/grpc/handlers/broker_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package handlers
import (
"fmt"
"sync"
"sync/atomic"
"testing"
"time"

Expand Down Expand Up @@ -234,21 +235,6 @@ func TestBroker(t *testing.T) {
}
})

t.Run("getListenerChannel", func(t *testing.T) {
broker := newBroker[string]()
listener := newListener[string]("test-id", []string{"topic1"})
broker.pushListener(listener)

ch, err := broker.getListenerChannel("test-id")
require.NoError(t, err)
require.Equal(t, listener.ch, ch)

ch, err = broker.getListenerChannel("non-existent")
require.Error(t, err)
require.Nil(t, ch)
require.ErrorIs(t, err, ErrSubscriptionNotFound)
})

t.Run("getTopics", func(t *testing.T) {
broker := newBroker[string]()
topics := []string{"topic1", "topic2", "TOPIC3"}
Expand Down Expand Up @@ -415,20 +401,6 @@ func TestBroker(t *testing.T) {
require.Len(t, listeners, 0)
})

t.Run("stopTimeout", func(t *testing.T) {
broker := newBroker[string]()
listener := newListener[string]("test-id", []string{"topic1"})
broker.pushListener(listener)

broker.startTimeout("test-id", 100*time.Millisecond)
broker.stopTimeout("test-id")

// wait to ensure timeout doesn't trigger
time.Sleep(150 * time.Millisecond)
listeners := broker.getListenersCopy()
require.Len(t, listeners, 1) // should still exist
})

t.Run("concurrent timeout with several listeners", func(t *testing.T) {
const nbListeners = 10
broker := newBroker[string]()
Expand Down Expand Up @@ -459,6 +431,145 @@ func TestBroker(t *testing.T) {
})
})

t.Run("attachment management", func(t *testing.T) {
t.Run("attach unknown id returns not found", func(t *testing.T) {
broker := newBroker[string]()

_, _, err := broker.attach("missing")
require.ErrorIs(t, err, ErrSubscriptionNotFound)
})

t.Run("attach displaces previous attachment", func(t *testing.T) {
broker := newBroker[string]()
listener := newListener[string]("test-id", []string{"topic1"})
broker.pushListener(listener)

_, att1, err := broker.attach("test-id")
require.NoError(t, err)

_, att2, err := broker.attach("test-id")
require.NoError(t, err)

// The first attachment must be displaced, the second must not.
select {
case <-att1.displaced:
case <-time.After(time.Second):
require.Fail(t, "first attachment not displaced by second attach")
}
select {
case <-att2.displaced:
require.Fail(t, "second attachment displaced unexpectedly")
default:
}

// Ownership moved to the second attachment.
require.False(t, broker.release("test-id", att1, time.Hour))
require.True(t, broker.release("test-id", att2, time.Hour))
})

t.Run("release and concurrent attach are atomic", func(t *testing.T) {
// a racing attach either wins (listener survives, successor owns
// it) or gets a clean not-found — never a destroyed listener
for range 200 {
broker := newBroker[string]()
listener := newListener[string]("test-id", nil)
broker.pushListener(listener)

_, att1, err := broker.attach("test-id")
require.NoError(t, err)

attached := make(chan error, 1)
go func() {
_, _, err := broker.attach("test-id")
attached <- err
}()
broker.release("test-id", att1, 0)

if err := <-attached; err != nil {
require.ErrorIs(t, err, ErrSubscriptionNotFound)
require.Empty(t, broker.getListenersCopy())
continue
}
require.Len(t, broker.getListenersCopy(), 1)
select {
case <-listener.done:
t.Fatal("listener destroyed under the attached successor")
default:
}
}
})

t.Run("attach cancels pending timeout", func(t *testing.T) {
broker := newBroker[string]()
listener := newListener[string]("test-id", []string{"topic1"})
broker.pushListener(listener)

broker.startTimeout("test-id", 50*time.Millisecond)
_, _, err := broker.attach("test-id")
require.NoError(t, err)

// wait well past the timeout: the listener must survive because a
// stream attached before it fired
time.Sleep(150 * time.Millisecond)
require.Len(t, broker.getListenersCopy(), 1)
select {
case <-listener.done:
require.Fail(t, "done closed while a stream was attached")
default:
}
})

t.Run("startTimeout is a no-op while attached", func(t *testing.T) {
broker := newBroker[string]()
listener := newListener[string]("test-id", []string{"topic1"})
broker.pushListener(listener)

_, att, err := broker.attach("test-id")
require.NoError(t, err)

broker.startTimeout("test-id", 50*time.Millisecond)
time.Sleep(150 * time.Millisecond)
require.Len(t, broker.getListenersCopy(), 1)

require.True(t, broker.release("test-id", att, 50*time.Millisecond))
time.Sleep(150 * time.Millisecond)
require.Len(t, broker.getListenersCopy(), 0)
})

t.Run("concurrent attach and release on the same id", func(t *testing.T) {
broker := newBroker[string]()
listener := newListener[string]("test-id", []string{"topic1"})
broker.pushListener(listener)

const goroutines = 50
var releaseTrue atomic.Int32
var wg sync.WaitGroup
for range goroutines {
wg.Go(func() {
_, att, err := broker.attach("test-id")
if err != nil {
return
}
broker.release("test-id", att, time.Hour)
if broker.release("test-id", att, time.Hour) {
releaseTrue.Add(1)
}
})
}
wg.Wait()

// only 1 release should return true
require.Zero(t, releaseTrue.Load())
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated

listeners := broker.getListenersCopy()
require.Len(t, listeners, 1)
if att := listeners["test-id"].attached; att != nil {
require.True(t, broker.release("test-id", att, time.Hour))
require.False(t, broker.release("test-id", att, time.Hour))
}
})
})

t.Run("getListenersCopy", func(t *testing.T) {
broker := newBroker[string]()

Expand Down Expand Up @@ -508,8 +619,7 @@ func TestBroker(t *testing.T) {
listener := newListener[string]("test-id", []string{"topic1"})
broker.pushListener(listener)

ch, err := broker.getListenerChannel("test-id")
require.NoError(t, err)
ch := listener.ch

// test sending to channel
go func() {
Expand Down
Loading
Loading