From 7f6207fae2386d747347199936d56a105cbdf557 Mon Sep 17 00:00:00 2001 From: Loi Nguyen Date: Fri, 31 Jul 2026 00:06:00 +0700 Subject: [PATCH] lib/batcher: prevent commits racing shutdown - fixes #9687 --- lib/batcher/batcher.go | 6 +++ lib/batcher/batcher_test.go | 89 +++++++++++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/lib/batcher/batcher.go b/lib/batcher/batcher.go index 1ac7f589d..ace41f952 100644 --- a/lib/batcher/batcher.go +++ b/lib/batcher/batcher.go @@ -47,6 +47,7 @@ type Batcher[Item, Result any] struct { in chan request[Item, Result] // incoming items to batch closed chan struct{} // close to indicate batcher shut down atexit atexit.FnHandle // atexit handle + admitMu sync.Mutex // serializes Commit admission with shutdown shutOnce sync.Once // make sure we shutdown once only wg sync.WaitGroup // wait for shutdown } @@ -239,6 +240,7 @@ func (b *Batcher[Item, Result]) Shutdown() { b.shutOnce.Do(func() { atexit.Unregister(b.atexit) fs.Infof(b.f, "Committing uploads - please wait...") + b.admitMu.Lock() // show that batcher is shutting down close(b.closed) // quit the commitLoop by sending a quitRequest message @@ -247,6 +249,7 @@ func (b *Batcher[Item, Result]) Shutdown() { // cause write to closed channel in Commit when we are // exiting due to a signal. b.in <- request[Item, Result]{quit: true} + b.admitMu.Unlock() b.wg.Wait() }) } @@ -261,8 +264,10 @@ func (b *Batcher[Item, Result]) Shutdown() { // This should not be called if batching is off - check first with // IsBatching. func (b *Batcher[Item, Result]) Commit(ctx context.Context, name string, item Item) (entry Result, err error) { + b.admitMu.Lock() select { case <-b.closed: + b.admitMu.Unlock() return entry, fserrors.FatalError(errors.New("batcher is shutting down")) default: } @@ -273,6 +278,7 @@ func (b *Batcher[Item, Result]) Commit(ctx context.Context, name string, item It name: name, result: resp, } + b.admitMu.Unlock() // If running async then don't wait for the result if b.async { return entry, nil diff --git a/lib/batcher/batcher_test.go b/lib/batcher/batcher_test.go index 3391bd81b..72873f12c 100644 --- a/lib/batcher/batcher_test.go +++ b/lib/batcher/batcher_test.go @@ -19,6 +19,21 @@ type ( Item string ) +// blockingStringer pauses Commit in its debug log after the shutdown check. +type blockingStringer struct { + started chan struct{} + release chan struct{} + once sync.Once +} + +func (b *blockingStringer) String() string { + b.once.Do(func() { + close(b.started) + <-b.release + }) + return "batcher test" +} + func TestBatcherNew(t *testing.T) { ctx := context.Background() ci := fs.GetConfig(ctx) @@ -221,6 +236,80 @@ func TestBatcherCommitShutdown(t *testing.T) { assert.Equal(t, 10, totalSize) } +func TestBatcherCommitRacingShutdown(t *testing.T) { + for _, mode := range []string{"sync", "async"} { + t.Run(mode, func(t *testing.T) { + ctx := context.Background() + ci := fs.GetConfig(ctx) + oldLogLevel := ci.LogLevel + ci.LogLevel = fs.LogLevelDebug + defer func() { ci.LogLevel = oldLogLevel }() + + committed := make(chan struct{}) + commitBatch := func(ctx context.Context, items []Item, results []Result, errors []error) error { + close(committed) + for i := range items { + results[i] = Result(items[i]) + } + return nil + } + blocker := &blockingStringer{ + started: make(chan struct{}), + release: make(chan struct{}), + } + b, err := New[Item, Result](ctx, blocker, commitBatch, Options{ + Mode: mode, + Size: 1, + Timeout: time.Hour, + MaxBatchSize: 1000, + }) + require.NoError(t, err) + + commitDone := make(chan error, 1) + go func() { + _, err := b.Commit(ctx, "item", Item("item")) + commitDone <- err + }() + select { + case <-blocker.started: + case <-time.After(time.Second): + t.Fatal("commit did not reach the admission point") + } + + ci.LogLevel = oldLogLevel + shutdownDone := make(chan struct{}) + go func() { + b.Shutdown() + close(shutdownDone) + }() + + // Give Shutdown a chance to contend with the blocked admission. + select { + case <-b.closed: + case <-time.After(100 * time.Millisecond): + } + close(blocker.release) + + select { + case err := <-commitDone: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("commit hung while racing shutdown") + } + select { + case <-shutdownDone: + case <-time.After(time.Second): + t.Fatal("shutdown hung while racing commit") + } + select { + case <-committed: + case <-time.After(time.Second): + t.Fatal("accepted commit was dropped during shutdown") + } + }) + } +} + func TestBatcherCommitAsync(t *testing.T) { ctx := context.Background()