diff --git a/internal/interface/grpc/handlers/broker.go b/internal/interface/grpc/handlers/broker.go index 885530e79..15ba6dd30 100644 --- a/internal/interface/grpc/handlers/broker.go +++ b/internal/interface/grpc/handlers/broker.go @@ -35,7 +35,15 @@ type listener[T any] struct { done chan struct{} closeDoneMux sync.Once timeoutTimer *time.Timer - lock *sync.RWMutex + + attached *attachment + lock *sync.RWMutex +} + +// attachment represents a stream's exclusive hold on a listener. Its displaced +// channel is closed when another stream takes over, telling the old stream to exit. +type attachment struct { + displaced chan struct{} } func newListener[T any](id string, topics []string) *listener[T] { @@ -210,14 +218,49 @@ 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 makes the calling stream the listener's sole consumer, cancelling any +// pending removal timeout and displacing the currently attached stream, if any. +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) + } + l.attached = &attachment{displaced: make(chan struct{})} + return l, l.attached, nil +} + +// release ends att's hold on the listener: kept for reconnectWindow if it still +// has filters, removed otherwise. Returns false if att was displaced. +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 } - return listener.ch, nil + l.closeDone() + delete(h.listeners, id) + return true } func (h *broker[T]) getTopics(id string) []string { @@ -315,44 +358,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] { diff --git a/internal/interface/grpc/handlers/broker_test.go b/internal/interface/grpc/handlers/broker_test.go index ef5b71a1b..a557c8b16 100644 --- a/internal/interface/grpc/handlers/broker_test.go +++ b/internal/interface/grpc/handlers/broker_test.go @@ -3,6 +3,7 @@ package handlers import ( "fmt" "sync" + "sync/atomic" "testing" "time" @@ -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"} @@ -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]() @@ -459,6 +431,146 @@ 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 + } + // release twice + broker.release("test-id", att, time.Hour) + if broker.release("test-id", att, time.Hour) { + releaseTrue.Add(1) + } + }) + } + wg.Wait() + + // every second release should return false + require.Zero(t, releaseTrue.Load()) + + 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]() @@ -508,8 +620,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() { diff --git a/internal/interface/grpc/handlers/indexer.go b/internal/interface/grpc/handlers/indexer.go index d9db763a8..6473ff272 100644 --- a/internal/interface/grpc/handlers/indexer.go +++ b/internal/interface/grpc/handlers/indexer.go @@ -445,16 +445,28 @@ func (h *indexerService) GetSubscription( request *arkv1.GetSubscriptionRequest, stream arkv1.IndexerService_GetSubscriptionServer, ) error { subscriptionId := request.GetSubscriptionId() + reconnectWindow := h.subscriptionTimeoutDuration - var scriptCh chan *arkv1.GetSubscriptionResponse - - if len(subscriptionId) == 0 { + isNew := len(subscriptionId) == 0 + if isNew { // New single-connection flow: create subscription inline. subscriptionId = uuid.NewString() - listener := newListener[*arkv1.GetSubscriptionResponse](subscriptionId, nil) - h.scriptSubsHandler.pushListener(listener) - defer h.scriptSubsHandler.removeListener(subscriptionId) + h.scriptSubsHandler.pushListener( + newListener[*arkv1.GetSubscriptionResponse](subscriptionId, nil), + ) + reconnectWindow = 0 + } + + // Attach as the subscription's sole consumer + // it forces any previous streams attached to the same subscriptinId to be closed + listener, att, err := h.scriptSubsHandler.attach(subscriptionId) + if err != nil { + return subscriptionErr(subscriptionId, err) + } + // On exit, release our hold: the subscription survives for reconnectWindow. + defer h.scriptSubsHandler.release(subscriptionId, att, reconnectWindow) + if isNew { // Apply initial filter, if any, through the same machinery used by // UpdateSubscription. `scripts.remove` is ignored on creation // because the subscription has no scripts to remove yet. @@ -464,8 +476,6 @@ func (h *indexerService) GetSubscription( } } - scriptCh = listener.ch - // Send SubscriptionStartedEvent as first message. startedEvt := &arkv1.GetSubscriptionResponse{ Data: &arkv1.GetSubscriptionResponse_SubscriptionStarted{ @@ -477,27 +487,6 @@ func (h *indexerService) GetSubscription( if err := stream.Send(startedEvt); err != nil { return err } - } else { - // Old flow: subscription_id provided, use existing listener. - h.scriptSubsHandler.stopTimeout(subscriptionId) - defer func() { - // Keep the listener alive on disconnect if either filter type is - // non-empty, so the client can reconnect within the timeout window - // without losing scripts or tx filters. - topics := h.scriptSubsHandler.getTopics(subscriptionId) - txFilters := h.scriptSubsHandler.getTxFilters(subscriptionId) - if len(topics) > 0 || len(txFilters) > 0 { - h.scriptSubsHandler.startTimeout(subscriptionId, h.subscriptionTimeoutDuration) - } else { - h.scriptSubsHandler.removeListener(subscriptionId) - } - }() - - var err error - scriptCh, err = h.scriptSubsHandler.getListenerChannel(subscriptionId) - if err != nil { - return subscriptionErr(subscriptionId, err) - } } // create a Timer that will fire after one heartbeat interval @@ -517,10 +506,40 @@ func (h *indexerService) GetSubscription( } for { + // Two selects are intentional. A single select cannot express + // priority: if both an exit signal and listener.ch are ready, Go + // chooses randomly. That could let a displaced or removed stream + // consume an event intended for its replacement. + // + // The first select is non-blocking. If an exit signal is already + // pending, it returns immediately without draining listener.ch, + // leaving buffered events for the successor. + select { + case <-stream.Context().Done(): + return nil + case <-listener.done: + return nil + case <-att.displaced: + return nil + default: + } + + // The second select blocks waiting for work or shutdown. The exit + // cases are repeated so that a signal arriving while blocked wakes + // the goroutine immediately instead of waiting for the next event or + // heartbeat. select { case <-stream.Context().Done(): return nil - case ev := <-scriptCh: + case <-listener.done: + // Subscription removed (unsubscribed, or reaped after the + // reconnect window expired). + return nil + case <-att.displaced: + // The client reconnected with the same subscription id on a new + // stream; this stream is abandoned. + return nil + case ev := <-listener.ch: if err := stream.Send(ev); err != nil { return err } diff --git a/internal/interface/grpc/handlers/indexer_test.go b/internal/interface/grpc/handlers/indexer_test.go index 82bc82115..002e8729b 100644 --- a/internal/interface/grpc/handlers/indexer_test.go +++ b/internal/interface/grpc/handlers/indexer_test.go @@ -6,6 +6,7 @@ import ( "encoding/hex" stderrors "errors" "fmt" + "sync" "testing" "time" @@ -136,8 +137,7 @@ func TestGetSubscription(t *testing.T) { require.NotEmpty(t, subId) // Push an event via the broker channel. - ch, err := svc.scriptSubsHandler.getListenerChannel(subId) - require.NoError(t, err) + ch := svc.scriptSubsHandler.getListenersCopy()[subId].ch ch <- &arkv1.GetSubscriptionResponse{ Data: &arkv1.GetSubscriptionResponse_Event{ @@ -177,8 +177,7 @@ func TestGetSubscription(t *testing.T) { require.NotEmpty(t, subId) // Listener should be present while the stream is open. - _, err := svc.scriptSubsHandler.getListenerChannel(subId) - require.NoError(t, err) + require.Contains(t, svc.scriptSubsHandler.getListenersCopy(), subId) cancel() @@ -190,8 +189,7 @@ func TestGetSubscription(t *testing.T) { } // After handler returns, listener must be removed (defer removeListener). - _, err = svc.scriptSubsHandler.getListenerChannel(subId) - require.Error(t, err) + require.NotContains(t, svc.scriptSubsHandler.getListenersCopy(), subId) }) t.Run("new flow invalid scripts returns error", func(t *testing.T) { @@ -327,8 +325,7 @@ func TestGetSubscription(t *testing.T) { ) // Push an event via the broker channel. - ch, err := svc.scriptSubsHandler.getListenerChannel(subId) - require.NoError(t, err) + ch := svc.scriptSubsHandler.getListenersCopy()[subId].ch ch <- &arkv1.GetSubscriptionResponse{ Data: &arkv1.GetSubscriptionResponse_Event{ @@ -375,8 +372,7 @@ func TestGetSubscription(t *testing.T) { subId := msg.GetSubscriptionStarted().GetSubscriptionId() require.NotEmpty(t, subId) - ch, err := svc.scriptSubsHandler.getListenerChannel(subId) - require.NoError(t, err) + ch := svc.scriptSubsHandler.getListenersCopy()[subId].ch // Send an event before the heartbeat fires. ch <- &arkv1.GetSubscriptionResponse{ @@ -417,8 +413,7 @@ func TestGetSubscription(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) stream := newMockGetSubscriptionServer(ctx) - ch, err := svc.scriptSubsHandler.getListenerChannel(subId) - require.NoError(t, err) + ch := svc.scriptSubsHandler.getListenersCopy()[subId].ch errCh := make(chan error, 1) go func() { @@ -445,13 +440,12 @@ func TestGetSubscription(t *testing.T) { } // Listener should still exist (timeout not yet expired). - _, err = svc.scriptSubsHandler.getListenerChannel(subId) - require.NoError(t, err) + require.Contains(t, svc.scriptSubsHandler.getListenersCopy(), subId) // After the timeout fires, listener should be cleaned up. require.Eventually(t, func() bool { - _, err := svc.scriptSubsHandler.getListenerChannel(subId) - return err != nil + _, ok := svc.scriptSubsHandler.getListenersCopy()[subId] + return !ok }, 2*time.Second, 50*time.Millisecond) }) @@ -499,8 +493,7 @@ func TestGetSubscription(t *testing.T) { } // Listener should be removed immediately (no scripts → no timeout). - _, err = svc.scriptSubsHandler.getListenerChannel(subId) - require.Error(t, err) + require.NotContains(t, svc.scriptSubsHandler.getListenersCopy(), subId) }) t.Run("old flow existing subscription_id works", func(t *testing.T) { @@ -524,8 +517,7 @@ func TestGetSubscription(t *testing.T) { // Grab the channel before starting GetSubscription; it is the same // channel the handler will read from. - ch, err := svc.scriptSubsHandler.getListenerChannel(subId) - require.NoError(t, err) + ch := svc.scriptSubsHandler.getListenersCopy()[subId].ch errCh := make(chan error, 1) go func() { @@ -556,6 +548,333 @@ func TestGetSubscription(t *testing.T) { t.Fatal("GetSubscription did not return") } }) + + t.Run("old flow reconnect displaces previous stream", func(t *testing.T) { + t.Parallel() + svc := newTestIndexerService(t) + + subResp, err := svc.SubscribeForScripts(context.Background(), + &arkv1.SubscribeForScriptsRequest{Scripts: []string{testScript1}}, + ) + require.NoError(t, err) + subId := subResp.GetSubscriptionId() + + // First stream. Its context is never cancelled, simulating a client + // that vanished behind an intermediary that keeps the connection open, + // so the server never observes the disconnect. + ctx1, cancel1 := context.WithCancel(context.Background()) + defer cancel1() + stream1 := newMockGetSubscriptionServer(ctx1) + + errCh1 := make(chan error, 1) + go func() { + errCh1 <- svc.GetSubscription( + &arkv1.GetSubscriptionRequest{SubscriptionId: subId}, + stream1, + ) + }() + + // Prove stream1 is attached before reconnecting: it must consume an + // event from the listener channel. + ch := svc.scriptSubsHandler.getListenersCopy()[subId].ch + ch <- &arkv1.GetSubscriptionResponse{ + Data: &arkv1.GetSubscriptionResponse_Event{ + Event: &arkv1.IndexerSubscriptionEvent{Txid: "warmup"}, + }, + } + stream1.recv(t, time.Second) + + // The client reconnects with the same subscription id on a new stream. + ctx2, cancel2 := context.WithCancel(context.Background()) + defer cancel2() + stream2 := newMockGetSubscriptionServer(ctx2) + + errCh2 := make(chan error, 1) + go func() { + errCh2 <- svc.GetSubscription( + &arkv1.GetSubscriptionRequest{SubscriptionId: subId}, + stream2, + ) + }() + + // The reconnect must terminate the previous stream; otherwise it stays + // attached forever, competing for events and leaking its goroutines. + select { + case err := <-errCh1: + require.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("previous stream still running after reconnect with same subscription id") + } + + // Events must now reach the new stream only. + ch <- &arkv1.GetSubscriptionResponse{ + Data: &arkv1.GetSubscriptionResponse_Event{ + Event: &arkv1.IndexerSubscriptionEvent{Txid: "after-reconnect"}, + }, + } + got := stream2.recv(t, time.Second) + require.Equal(t, "after-reconnect", got.GetEvent().GetTxid()) + + cancel2() + select { + case err := <-errCh2: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("GetSubscription did not return") + } + }) + + t.Run("old flow displaced stream does not consume buffered events", func(t *testing.T) { + t.Parallel() + // a displaced stream must leave buffered events to its successor; the + // select picks randomly among ready cases so repeat to make it decisive + for range 20 { + svc := newTestIndexerService(t) + svc.heartbeat = 10 * time.Second // keep heartbeats out of the ordering + + subResp, err := svc.SubscribeForScripts(context.Background(), + &arkv1.SubscribeForScriptsRequest{Scripts: []string{testScript1}}, + ) + require.NoError(t, err) + subId := subResp.GetSubscriptionId() + + ch := svc.scriptSubsHandler.getListenersCopy()[subId].ch + + pred := newGatedSubscriptionServer(context.Background()) + errCh := make(chan error, 1) + go func() { + errCh <- svc.GetSubscription( + &arkv1.GetSubscriptionRequest{SubscriptionId: subId}, + pred, + ) + }() + + // The predecessor consumes this event and parks in its first Send, + // so it is out of the select loop while we set up the race. + ch <- &arkv1.GetSubscriptionResponse{ + Data: &arkv1.GetSubscriptionResponse_Event{ + Event: &arkv1.IndexerSubscriptionEvent{Txid: "warmup"}, + }, + } + select { + case <-pred.entered: + case <-time.After(time.Second): + t.Fatal("predecessor did not reach its first Send") + } + + // Buffer an event, then displace the predecessor while it is still + // parked in Send. attach here stands in for a reconnecting client. + ch <- &arkv1.GetSubscriptionResponse{ + Data: &arkv1.GetSubscriptionResponse_Event{ + Event: &arkv1.IndexerSubscriptionEvent{Txid: "for-successor"}, + }, + } + if _, _, err := svc.scriptSubsHandler.attach(subId); err != nil { + t.Fatalf("attach (simulated reconnect) failed: %v", err) + } + + // Release the predecessor. It resumes displaced and must return + // without touching the buffered event. + close(pred.release) + select { + case err := <-errCh: + require.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("displaced predecessor did not return") + } + + require.Len(t, ch, 1, + "displaced stream consumed an event meant for the successor") + } + }) + + t.Run("old flow stream ends when subscription removed", func(t *testing.T) { + t.Parallel() + svc := newTestIndexerService(t) + + subResp, err := svc.SubscribeForScripts(context.Background(), + &arkv1.SubscribeForScriptsRequest{Scripts: []string{testScript1}}, + ) + require.NoError(t, err) + subId := subResp.GetSubscriptionId() + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + stream := newMockGetSubscriptionServer(ctx) + + errCh := make(chan error, 1) + go func() { + errCh <- svc.GetSubscription( + &arkv1.GetSubscriptionRequest{SubscriptionId: subId}, + stream, + ) + }() + + // Prove the stream is attached before removing the subscription. + ch := svc.scriptSubsHandler.getListenersCopy()[subId].ch + ch <- &arkv1.GetSubscriptionResponse{ + Data: &arkv1.GetSubscriptionResponse_Event{ + Event: &arkv1.IndexerSubscriptionEvent{Txid: "warmup"}, + }, + } + stream.recv(t, time.Second) + + // Unsubscribing from all scripts removes the listener; the stream must + // terminate rather than keep heartbeating with no listener behind it. + _, err = svc.UnsubscribeForScripts(context.Background(), + &arkv1.UnsubscribeForScriptsRequest{SubscriptionId: subId}, + ) + require.NoError(t, err) + + select { + case err := <-errCh: + require.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("stream still running after its subscription was removed") + } + }) + + t.Run("old flow displaced stream leaves subscription to successor", func(t *testing.T) { + t.Parallel() + svc := newTestIndexerService(t) + svc.subscriptionTimeoutDuration = 200 * time.Millisecond + + subResp, err := svc.SubscribeForScripts(context.Background(), + &arkv1.SubscribeForScriptsRequest{Scripts: []string{testScript1}}, + ) + require.NoError(t, err) + subId := subResp.GetSubscriptionId() + + ctx1, cancel1 := context.WithCancel(context.Background()) + defer cancel1() + stream1 := newMockGetSubscriptionServer(ctx1) + + errCh1 := make(chan error, 1) + go func() { + errCh1 <- svc.GetSubscription( + &arkv1.GetSubscriptionRequest{SubscriptionId: subId}, + stream1, + ) + }() + + ch := svc.scriptSubsHandler.getListenersCopy()[subId].ch + ch <- &arkv1.GetSubscriptionResponse{ + Data: &arkv1.GetSubscriptionResponse_Event{ + Event: &arkv1.IndexerSubscriptionEvent{Txid: "warmup"}, + }, + } + stream1.recv(t, time.Second) + + ctx2, cancel2 := context.WithCancel(context.Background()) + defer cancel2() + stream2 := newMockGetSubscriptionServer(ctx2) + + errCh2 := make(chan error, 1) + go func() { + errCh2 <- svc.GetSubscription( + &arkv1.GetSubscriptionRequest{SubscriptionId: subId}, + stream2, + ) + }() + + select { + case err := <-errCh1: + require.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("previous stream still running after reconnect with same subscription id") + } + + // The displaced stream's exit must not arm the reconnect timeout while + // the successor is attached: well past the timeout, the subscription + // must still exist and serve the new stream. + time.Sleep(500 * time.Millisecond) + require.Contains(t, svc.scriptSubsHandler.getListenersCopy(), subId, + "subscription reaped while a live stream was attached") + + ch <- &arkv1.GetSubscriptionResponse{ + Data: &arkv1.GetSubscriptionResponse_Event{ + Event: &arkv1.IndexerSubscriptionEvent{Txid: "still-alive"}, + }, + } + got := stream2.recv(t, time.Second) + require.Equal(t, "still-alive", got.GetEvent().GetTxid()) + + cancel2() + select { + case err := <-errCh2: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("GetSubscription did not return") + } + }) + + t.Run("new flow stream displaced by reconnect with announced id", func(t *testing.T) { + t.Parallel() + svc := newTestIndexerService(t) + + // New flow announces the subscription id in SubscriptionStartedEvent. + ctx1, cancel1 := context.WithCancel(context.Background()) + defer cancel1() + stream1 := newMockGetSubscriptionServer(ctx1) + + errCh1 := make(chan error, 1) + go func() { + errCh1 <- svc.GetSubscription( + &arkv1.GetSubscriptionRequest{ + Filter: scriptsAddFilter(testScript1), + }, + stream1, + ) + }() + + msg := stream1.recv(t, time.Second) + subId := msg.GetSubscriptionStarted().GetSubscriptionId() + require.NotEmpty(t, subId) + + // The client reconnects with the announced id while the server still + // considers the first stream alive. + ctx2, cancel2 := context.WithCancel(context.Background()) + defer cancel2() + stream2 := newMockGetSubscriptionServer(ctx2) + + errCh2 := make(chan error, 1) + go func() { + errCh2 <- svc.GetSubscription( + &arkv1.GetSubscriptionRequest{SubscriptionId: subId}, + stream2, + ) + }() + + select { + case err := <-errCh1: + require.NoError(t, err) + case <-time.After(2 * time.Second): + t.Fatal("previous stream still running after reconnect with same subscription id") + } + + // The displaced stream's exit must not remove the listener out from + // under the stream that took over. + listeners := svc.scriptSubsHandler.getListenersCopy() + require.Contains(t, listeners, subId, + "listener removed by displaced stream while successor attached") + ch := listeners[subId].ch + + ch <- &arkv1.GetSubscriptionResponse{ + Data: &arkv1.GetSubscriptionResponse_Event{ + Event: &arkv1.IndexerSubscriptionEvent{Txid: "after-takeover"}, + }, + } + got := stream2.recv(t, time.Second) + require.Equal(t, "after-takeover", got.GetEvent().GetTxid()) + + cancel2() + select { + case err := <-errCh2: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("GetSubscription did not return") + } + }) } func TestUpdateSubscription(t *testing.T) { @@ -1260,8 +1579,7 @@ func TestTxFilter(t *testing.T) { } // Listener should still be present (timeout-scheduled), not removed. - _, err := svc.scriptSubsHandler.getListenerChannel("sub-old-tx") - require.NoError(t, err) + require.Contains(t, svc.scriptSubsHandler.getListenersCopy(), "sub-old-tx") require.ElementsMatch( t, []string{hasExtension}, svc.scriptSubsHandler.getTxFilters("sub-old-tx"), @@ -1391,8 +1709,7 @@ func TestTxFilter(t *testing.T) { require.ElementsMatch( t, []string{hasExtension}, svc.scriptSubsHandler.getTxFilters("sub-keep"), ) - _, err = svc.scriptSubsHandler.getListenerChannel("sub-keep") - require.NoError(t, err) + require.Contains(t, svc.scriptSubsHandler.getListenersCopy(), "sub-keep") }) t.Run("UnsubscribeForScripts removes listener without tx filters", func(t *testing.T) { @@ -1408,8 +1725,7 @@ func TestTxFilter(t *testing.T) { ) require.NoError(t, err) - _, err = svc.scriptSubsHandler.getListenerChannel("sub-drop") - require.ErrorIs(t, err, ErrSubscriptionNotFound) + require.NotContains(t, svc.scriptSubsHandler.getListenersCopy(), "sub-drop") }) t.Run("UnsubscribeForScripts unknown subscription returns NotFound", func(t *testing.T) { @@ -1607,6 +1923,42 @@ func (m *mockGetSubscriptionServer) recv(t *testing.T, timeout time.Duration) *a } } +// gatedSubscriptionServer blocks in its first Send until release is closed, +// signalling entered when it gets there. It lets a test park the handler out +// of its select loop to set up a precise ordering before it resumes. +type gatedSubscriptionServer struct { + ctx context.Context + sendCh chan *arkv1.GetSubscriptionResponse + entered chan struct{} + release chan struct{} + once sync.Once +} + +func newGatedSubscriptionServer(ctx context.Context) *gatedSubscriptionServer { + return &gatedSubscriptionServer{ + ctx: ctx, + sendCh: make(chan *arkv1.GetSubscriptionResponse, 100), + entered: make(chan struct{}), + release: make(chan struct{}), + } +} + +func (m *gatedSubscriptionServer) Send(resp *arkv1.GetSubscriptionResponse) error { + m.once.Do(func() { + close(m.entered) + <-m.release + }) + m.sendCh <- resp + return nil +} + +func (m *gatedSubscriptionServer) Context() context.Context { return m.ctx } +func (m *gatedSubscriptionServer) SetHeader(metadata.MD) error { return nil } +func (m *gatedSubscriptionServer) SendHeader(metadata.MD) error { return nil } +func (m *gatedSubscriptionServer) SetTrailer(metadata.MD) {} +func (m *gatedSubscriptionServer) SendMsg(any) error { return nil } +func (m *gatedSubscriptionServer) RecvMsg(any) error { return nil } + // mockAppIndexer is a minimal application.IndexerService used to assert the // gRPC handler's GetVtxoChain wiring. Only GetVtxoChain is implemented; any // other method would panic via the embedded nil interface (none are called). diff --git a/internal/interface/grpc/service.go b/internal/interface/grpc/service.go index b08fd18a9..1856d1cdf 100644 --- a/internal/interface/grpc/service.go +++ b/internal/interface/grpc/service.go @@ -29,6 +29,7 @@ import ( "google.golang.org/grpc/credentials" "google.golang.org/grpc/credentials/insecure" grpchealth "google.golang.org/grpc/health/grpc_health_v1" + "google.golang.org/grpc/keepalive" "google.golang.org/grpc/status" ) @@ -328,6 +329,11 @@ func (s *service) newServer(tlsConfig *tls.Config, withPprof, withChannelz bool) s.macaroonSvc, s.readinessSvc, getVersionGuard, getDigestGuard, ), grpc.StatsHandler(otelHandler), + // ping clients, close dead ones + grpc.KeepaliveParams(keepalive.ServerParameters{ + Time: 30 * time.Second, + Timeout: 20 * time.Second, + }), } creds := insecure.NewCredentials() if !s.config.insecure() {