lib/batcher: prevent commits racing shutdown - fixes #9687
This commit is contained in:
committed by
Nick Craig-Wood
parent
5dd34275dc
commit
7f6207fae2
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user