Skip to content

Commit

Permalink
Fix error handling in tests file
Browse files Browse the repository at this point in the history
If the errgroup fail (e.g. the writer crash), also stop the deadlineCtx
which will cause program to stop.
Without this, writers will stop and the program will hang.
  • Loading branch information
PierreF committed Mar 8, 2024
1 parent 7ba4d84 commit e2c91b7
Showing 1 changed file with 10 additions and 2 deletions.
12 changes: 10 additions & 2 deletions tests/remote-storage-bench/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -91,16 +91,24 @@ func run(deadlineCtx context.Context) error {
})
}

ctx := context.Background()
group, ctx := errgroup.WithContext(context.Background())

if *runDuration > 0 {
var cancel context.CancelFunc

deadlineCtx, cancel = context.WithTimeout(deadlineCtx, *runDuration)
defer cancel()

go func() {
// If the errgroup context is stopped, stop the deadlineCtx.
// deadlineCtx is used for work generation.
// ctx (errgroup context) is used for writers.
// We use two context, because work generation need to be stopped before writers
<-ctx.Done()
cancel()
}()
}

group, ctx := errgroup.WithContext(ctx)
workChannel := make(chan prompb.WriteRequest, 10)

for _, writer := range writers {
Expand Down

0 comments on commit e2c91b7

Please sign in to comment.