From fea1820a7f5d98b18bc08d1381c645a53ac8cb8b Mon Sep 17 00:00:00 2001 From: HeyiSun <31991990+HeyiSun@users.noreply.github.com> Date: Tue, 8 Sep 2026 16:49:44 +0800 Subject: [PATCH 1/3] fix(shared): avoid lost wakeups when canceling Acquire --- .../shared/pkg/utils/resizable_semaphore.go | 8 +- .../pkg/utils/resizable_semaphore_test.go | 96 +++++++++++++++++++ 2 files changed, 102 insertions(+), 2 deletions(-) diff --git a/packages/shared/pkg/utils/resizable_semaphore.go b/packages/shared/pkg/utils/resizable_semaphore.go index db1542335d..a2992d0ae3 100644 --- a/packages/shared/pkg/utils/resizable_semaphore.go +++ b/packages/shared/pkg/utils/resizable_semaphore.go @@ -39,8 +39,12 @@ func (s *AdjustableSemaphore) Acquire(ctx context.Context, n int64) error { return fmt.Errorf("acquiring less than or equal to 0 elements is not supported, got: %d", n) } - // Wake ->cond.Wait when ctx is canceled. - stop := context.AfterFunc(ctx, s.cond.Broadcast) + // Hold mu so cancellation cannot broadcast between the context check and Wait. + stop := context.AfterFunc(ctx, func() { + s.cond.L.Lock() + defer s.cond.L.Unlock() + s.cond.Broadcast() + }) defer stop() // ensure we don’t leak the callback for s.used+n > s.limit { diff --git a/packages/shared/pkg/utils/resizable_semaphore_test.go b/packages/shared/pkg/utils/resizable_semaphore_test.go index 22230be37e..ae6283b912 100644 --- a/packages/shared/pkg/utils/resizable_semaphore_test.go +++ b/packages/shared/pkg/utils/resizable_semaphore_test.go @@ -6,6 +6,7 @@ import ( "runtime" "sync" "testing" + "testing/synctest" "time" "github.com/stretchr/testify/assert" @@ -274,6 +275,101 @@ func TestAcquireRespectsContextCancel(t *testing.T) { } } +// Pause after observing a live context, before Acquire can register its wait. +type semaphoreCancelContext struct { + context.Context + checked chan struct{} + resume chan struct{} + once sync.Once +} + +func (c *semaphoreCancelContext) Err() error { + err := c.Context.Err() + if err == nil { + c.once.Do(func() { + close(c.checked) + <-c.resume + }) + } + + return err +} + +type semaphoreWaitLocker struct { + sync.Locker + waiting chan struct{} + locking chan struct{} + waitOnce sync.Once + lockOnce sync.Once +} + +func (l *semaphoreWaitLocker) Lock() { + l.lockOnce.Do(func() { close(l.locking) }) + l.Locker.Lock() +} + +func (l *semaphoreWaitLocker) Unlock() { + l.waitOnce.Do(func() { close(l.waiting) }) + l.Locker.Unlock() +} + +func TestAcquireCancellationBeforeWait(t *testing.T) { + t.Parallel() + + synctest.Test(t, func(t *testing.T) { + s, err := NewAdjustableSemaphore(1) + require.NoError(t, err) + require.True(t, s.TryAcquire(1)) + + locker := &semaphoreWaitLocker{ + Locker: &s.mu, + waiting: make(chan struct{}), + locking: make(chan struct{}), + } + s.cond.L = locker + + // An existing waiter makes an unlocked cancellation Broadcast observable. + go func() { + s.mu.Lock() + s.cond.Wait() + s.mu.Unlock() + }() + <-locker.waiting + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + paused := &semaphoreCancelContext{ + Context: ctx, + checked: make(chan struct{}), + resume: make(chan struct{}), + } + result := make(chan error, 1) + go func() { result <- s.Acquire(paused, 1) }() + <-paused.checked + cancel() + + // With the fix, the callback tries to lock L. Without it, Broadcast + // wakes the existing waiter, which tries to lock L. Both happen while + // Acquire still holds mu and has not registered its wait. + <-locker.locking + close(paused.resume) + synctest.Wait() + + select { + case err := <-result: + require.ErrorIs(t, err, context.Canceled) + default: + // Rescue a stuck waiter so a regression fails without leaking it. + s.mu.Lock() + s.cond.Broadcast() + s.mu.Unlock() + synctest.Wait() + <-result + t.Fatal("Acquire missed cancellation before entering Wait") + } + }) +} + // ----------------------------------------------------------------------------- // race-detector stress // ----------------------------------------------------------------------------- From 20935e1cd5923704f18b87c44ce0d907048b5e42 Mon Sep 17 00:00:00 2001 From: HeyiSun <31991990+HeyiSun@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:07:04 +0800 Subject: [PATCH 2/3] chore(shared): restore original semaphore comment --- packages/shared/pkg/utils/resizable_semaphore.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/shared/pkg/utils/resizable_semaphore.go b/packages/shared/pkg/utils/resizable_semaphore.go index a2992d0ae3..5170d7aa44 100644 --- a/packages/shared/pkg/utils/resizable_semaphore.go +++ b/packages/shared/pkg/utils/resizable_semaphore.go @@ -39,7 +39,7 @@ func (s *AdjustableSemaphore) Acquire(ctx context.Context, n int64) error { return fmt.Errorf("acquiring less than or equal to 0 elements is not supported, got: %d", n) } - // Hold mu so cancellation cannot broadcast between the context check and Wait. + // Wake ->cond.Wait when ctx is canceled. stop := context.AfterFunc(ctx, func() { s.cond.L.Lock() defer s.cond.L.Unlock() From c15ed0fb0f0f6f8a76208fedd47314f985eeac1b Mon Sep 17 00:00:00 2001 From: HeyiSun <31991990+HeyiSun@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:25:20 +0800 Subject: [PATCH 3/3] refactor(shared): use mu in cancellation callback --- .../shared/pkg/utils/resizable_semaphore.go | 4 +- .../pkg/utils/resizable_semaphore_test.go | 99 +++++-------------- 2 files changed, 27 insertions(+), 76 deletions(-) diff --git a/packages/shared/pkg/utils/resizable_semaphore.go b/packages/shared/pkg/utils/resizable_semaphore.go index 5170d7aa44..38626b839a 100644 --- a/packages/shared/pkg/utils/resizable_semaphore.go +++ b/packages/shared/pkg/utils/resizable_semaphore.go @@ -41,8 +41,8 @@ func (s *AdjustableSemaphore) Acquire(ctx context.Context, n int64) error { // Wake ->cond.Wait when ctx is canceled. stop := context.AfterFunc(ctx, func() { - s.cond.L.Lock() - defer s.cond.L.Unlock() + s.mu.Lock() + defer s.mu.Unlock() s.cond.Broadcast() }) defer stop() // ensure we don’t leak the callback diff --git a/packages/shared/pkg/utils/resizable_semaphore_test.go b/packages/shared/pkg/utils/resizable_semaphore_test.go index ae6283b912..9ddc9837f5 100644 --- a/packages/shared/pkg/utils/resizable_semaphore_test.go +++ b/packages/shared/pkg/utils/resizable_semaphore_test.go @@ -275,97 +275,48 @@ func TestAcquireRespectsContextCancel(t *testing.T) { } } -// Pause after observing a live context, before Acquire can register its wait. type semaphoreCancelContext struct { context.Context - checked chan struct{} - resume chan struct{} - once sync.Once + cancel context.CancelFunc } func (c *semaphoreCancelContext) Err() error { err := c.Context.Err() if err == nil { - c.once.Do(func() { - close(c.checked) - <-c.resume - }) + c.cancel() + runtime.Gosched() } return err } -type semaphoreWaitLocker struct { - sync.Locker - waiting chan struct{} - locking chan struct{} - waitOnce sync.Once - lockOnce sync.Once -} - -func (l *semaphoreWaitLocker) Lock() { - l.lockOnce.Do(func() { close(l.locking) }) - l.Locker.Lock() -} - -func (l *semaphoreWaitLocker) Unlock() { - l.waitOnce.Do(func() { close(l.waiting) }) - l.Locker.Unlock() -} - func TestAcquireCancellationBeforeWait(t *testing.T) { t.Parallel() synctest.Test(t, func(t *testing.T) { - s, err := NewAdjustableSemaphore(1) - require.NoError(t, err) - require.True(t, s.TryAcquire(1)) - - locker := &semaphoreWaitLocker{ - Locker: &s.mu, - waiting: make(chan struct{}), - locking: make(chan struct{}), - } - s.cond.L = locker - - // An existing waiter makes an unlocked cancellation Broadcast observable. - go func() { - s.mu.Lock() - s.cond.Wait() - s.mu.Unlock() - }() - <-locker.waiting - - ctx, cancel := context.WithCancel(t.Context()) - defer cancel() - paused := &semaphoreCancelContext{ - Context: ctx, - checked: make(chan struct{}), - resume: make(chan struct{}), - } - result := make(chan error, 1) - go func() { result <- s.Acquire(paused, 1) }() - <-paused.checked - cancel() - - // With the fix, the callback tries to lock L. Without it, Broadcast - // wakes the existing waiter, which tries to lock L. Both happen while - // Acquire still holds mu and has not registered its wait. - <-locker.locking - close(paused.resume) - synctest.Wait() - - select { - case err := <-result: - require.ErrorIs(t, err, context.Canceled) - default: - // Rescue a stuck waiter so a regression fails without leaking it. - s.mu.Lock() - s.cond.Broadcast() - s.mu.Unlock() + for range 100 { + s, err := NewAdjustableSemaphore(1) + require.NoError(t, err) + require.True(t, s.TryAcquire(1)) + + ctx, cancel := context.WithCancel(t.Context()) + cancelOnCheck := &semaphoreCancelContext{Context: ctx, cancel: cancel} + result := make(chan error, 1) + go func() { result <- s.Acquire(cancelOnCheck, 1) }() synctest.Wait() - <-result - t.Fatal("Acquire missed cancellation before entering Wait") + cancel() + + select { + case err := <-result: + require.ErrorIs(t, err, context.Canceled) + default: + s.mu.Lock() + s.cond.Broadcast() + s.mu.Unlock() + synctest.Wait() + <-result + t.Fatal("Acquire missed cancellation before entering Wait") + } } }) }