diff --git a/pkg/store/store_adapter.go b/pkg/store/store_adapter.go index e5724e565b..a07da8430a 100644 --- a/pkg/store/store_adapter.go +++ b/pkg/store/store_adapter.go @@ -104,10 +104,37 @@ func (hs *heightSub) Wait(ctx context.Context, height uint64) error { case <-ch: return nil case <-ctx.Done(): + hs.removeWaiter(height, ch) return ctx.Err() } } +// removeWaiter unregisters a specific waiter from the requested height. +func (hs *heightSub) removeWaiter(height uint64, target chan struct{}) { + hs.heightMu.Lock() + defer hs.heightMu.Unlock() + + chs, ok := hs.heightChs[height] + if !ok { + return + } + + for i, ch := range chs { + if ch != target { + continue + } + copy(chs[i:], chs[i+1:]) + chs[len(chs)-1] = nil + chs = chs[:len(chs)-1] + if len(chs) == 0 { + delete(hs.heightChs, height) + } else { + hs.heightChs[height] = chs + } + return + } +} + // notifyUpTo notifies all waiters for heights <= h. func (hs *heightSub) notifyUpTo(h uint64) { hs.heightMu.Lock() diff --git a/pkg/store/store_adapter_test.go b/pkg/store/store_adapter_test.go index e88fd76f12..65844388f0 100644 --- a/pkg/store/store_adapter_test.go +++ b/pkg/store/store_adapter_test.go @@ -4,6 +4,7 @@ import ( "context" "sync" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -311,6 +312,97 @@ func TestPendingCache_ConcurrentAccess(t *testing.T) { assert.GreaterOrEqual(t, cache.len(), 0) } +func TestHeightSubWaitCancellationRemovesWaiter(t *testing.T) { + t.Parallel() + + hs := newHeightSub(1) + for range 100 { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + require.ErrorIs(t, hs.Wait(ctx, 1_000), context.Canceled) + } + + hs.heightMu.Lock() + defer hs.heightMu.Unlock() + assert.Empty(t, hs.heightChs) +} + +func TestHeightSubWaitCancellationPreservesOtherWaiters(t *testing.T) { + t.Parallel() + + hs := newHeightSub(1) + ctx1, cancel1 := context.WithCancel(context.Background()) + ctx2 := context.Background() + waitDone := make(chan error, 2) + go func() { + waitDone <- hs.Wait(ctx1, 1_000) + }() + go func() { + waitDone <- hs.Wait(ctx2, 1_000) + }() + + require.Eventually(t, func() bool { + hs.heightMu.Lock() + defer hs.heightMu.Unlock() + return len(hs.heightChs[1_000]) == 2 + }, time.Second, time.Millisecond) + + cancel1() + require.Eventually(t, func() bool { + hs.heightMu.Lock() + defer hs.heightMu.Unlock() + return len(hs.heightChs[1_000]) == 1 + }, time.Second, time.Millisecond) + + hs.SetHeight(1_000) + results := []error{<-waitDone, <-waitDone} + assert.Contains(t, results, context.Canceled) + assert.Contains(t, results, nil) +} + +func TestHeightSubWaitCancellationAndSetHeightConcurrent(t *testing.T) { + t.Parallel() + + for range 100 { + hs := newHeightSub(1) + ctx, cancel := context.WithCancel(context.Background()) + waitDone := make(chan error, 1) + go func() { + waitDone <- hs.Wait(ctx, 1_000) + }() + + // Ensure the waiter is registered before racing cancellation and notification. + require.Eventually(t, func() bool { + hs.heightMu.Lock() + defer hs.heightMu.Unlock() + return len(hs.heightChs[1_000]) == 1 + }, time.Second, time.Millisecond) + + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + cancel() + }() + go func() { + defer wg.Done() + hs.SetHeight(1_000) + }() + wg.Wait() + + err := <-waitDone + if err != nil { + assert.ErrorIs(t, err, context.Canceled) + } + + cancel() + hs.heightMu.Lock() + assert.Empty(t, hs.heightChs) + hs.heightMu.Unlock() + } +} + // TestStoreAdapter_Backpressure tests that Append blocks when cache is full func TestStoreAdapter_Backpressure(t *testing.T) { t.Parallel()