diff --git a/bytes_test.go b/bytes_test.go index 57de86a..27fb471 100644 --- a/bytes_test.go +++ b/bytes_test.go @@ -8,6 +8,7 @@ import ( "encoding/json" "math/rand" "testing" + "time" "github.com/lanrat/extsort" ) @@ -233,3 +234,35 @@ func TestRandom1M(t *testing.T) { } } } + +func TestDeadLockContextCancel(t *testing.T) { + inputChan := make(chan extsort.SortType, 2) + config := extsort.DefaultConfig() + config.ChunkSize = 2 + config.SortedChanBuffSize = 2 + // for simplicity, set ChanBuffSize to zero. the deadlock can happen with any value. + // see https://github.com/lanrat/extsort/issues/7 for details. + config.ChanBuffSize = 0 + lessFunc := func(a, b extsort.SortType) bool { + time.Sleep(300 * time.Millisecond) // emulate long operation + return false + } + sort, _, _ := extsort.New(inputChan, fromBytesForTest, lessFunc, config) + ctx, cf := context.WithCancel(context.Background()) + defer cf() + waitCh := make(chan struct{}) + go func() { + defer close(waitCh) + sort.Sort(ctx) + }() + inputChan <- val{Key: 1, Order: 1} + inputChan <- val{Key: 2, Order: 2} + close(inputChan) + // cancel the context. the sort.Sort should now be waiting inside lessFunc. + cf() + select { + case <- waitCh: + case <-time.After(time.Second): + t.Fatal("deadlock") + } +} diff --git a/sort_sorttype.go b/sort_sorttype.go index 45f56aa..3a8e498 100644 --- a/sort_sorttype.go +++ b/sort_sorttype.go @@ -176,7 +176,11 @@ func (s *SortTypeSorter) sortChunks() error { // sort sort.Sort(b) // save - s.saveChunkChan <- b + select { + case s.saveChunkChan <- b: + case <-s.buildSortCtx.Done(): + return s.buildSortCtx.Err() + } } else { return nil }