feat: add checksum

checksum_merge_test.go

@@ -0,0 +1,103 @@
+package extsort
+
+import (
+	"bufio"
+	"cmp"
+	"context"
+	"errors"
+	"io"
+	"slices"
+	"testing"
+)
+
+type checksumFailureReader struct {
+	data []byte
+	err  error
+}
+
+func (r *checksumFailureReader) Read(p []byte) (int, error) {
+	if len(r.data) == 0 {
+		return 0, r.err
+	}
+	n := copy(p, r.data)
+	r.data = r.data[n:]
+	if len(r.data) == 0 {
+		return n, r.err
+	}
+	return n, nil
+}
+
+type checksumFailureTempReader struct {
+	readers []*bufio.Reader
+}
+
+func (r *checksumFailureTempReader) Close() error {
+	return nil
+}
+
+func (r *checksumFailureTempReader) Size() int {
+	return len(r.readers)
+}
+
+func (r *checksumFailureTempReader) Read(i int) *bufio.Reader {
+	return r.readers[i]
+}
+
+func TestChecksumFailureStopsMerge(t *testing.T) {
+	for _, test := range []struct {
+		name          string
+		numWorkers    int
+		bufferSize    int
+		secondSection []byte
+		want          []int
+		exact         bool
+	}{
+		{
+			name:          "single threaded",
+			numWorkers:    2,
+			bufferSize:    10,
+			secondSection: []byte{1, 2, 1, 4},
+			want:          []int{1, 2, 3},
+			exact:         true,
+		},
+		{
+			name:          "parallel",
+			numWorkers:    1,
+			secondSection: []byte{1, 2, 1, 4},
+			want:          []int{1, 2, 3},
+		},
+		{name: "first block/single threaded", numWorkers: 2, bufferSize: 10, exact: true},
+		{name: "first block/parallel", numWorkers: 1, exact: true},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			checksumErr := errors.New("temporary section 1 block 0 checksum mismatch")
+			sorter := newSorter[int](
+				nil,
+				func(data []byte) (int, error) { return int(data[0]), nil },
+				func(value int) ([]byte, error) { return []byte{byte(value)}, nil },
+				cmp.Compare[int],
+				&Config{ChunkSize: 2, NumWorkers: test.numWorkers, SortedChanBuffSize: test.bufferSize},
+			)
+			sorter.tempReader = &checksumFailureTempReader{readers: []*bufio.Reader{
+				bufio.NewReader(&checksumFailureReader{data: []byte{1, 1, 1, 3, 1, 5}, err: io.EOF}),
+				bufio.NewReader(&checksumFailureReader{data: test.secondSection, err: checksumErr}),
+			}}
+
+			go sorter.mergeNChunks(context.Background())
+
+			var got []int
+			for value := range sorter.mergeChunkChan {
+				got = append(got, value)
+			}
+			if len(got) > len(test.want) || !slices.Equal(got, test.want[:len(got)]) {
+				t.Fatalf("output before checksum failure = %v, want a prefix of %v", got, test.want)
+			}
+			if test.exact && !slices.Equal(got, test.want) {
+				t.Fatalf("output before checksum failure = %v, want %v", got, test.want)
+			}
+			if err := <-sorter.mergeErrChan; !errors.Is(err, checksumErr) {
+				t.Fatalf("merge error = %v, want %v", err, checksumErr)
+			}
+		})
+	}
+}

config.go

@@ -36,6 +36,10 @@ type Config struct {
 	//
 	// Default: "" (intelligent selection).
 	TempFilesDir string
+
+	// Checksum verifies temporary file blocks before records are returned by the
+	// merge phase. Default: false.
+	Checksum bool
 }
 
 // DefaultConfig returns a Config with sensible default values optimized for

ordered_test.go

@@ -3,6 +3,7 @@ package extsort_test
 import (
 	"context"
 	"reflect"
+	"slices"
 	"testing"
 
 	"github.com/lanrat/extsort"
@@ -240,3 +241,32 @@ done5:
 		t.Errorf("Expected empty result, got %v", result)
 	}
 }
+
+func TestOrderedWithChecksum(t *testing.T) {
+	input := make(chan int, 6)
+	for _, value := range []int{6, 1, 5, 2, 4, 3} {
+		input <- value
+	}
+	close(input)
+
+	config := extsort.DefaultConfig()
+	config.ChunkSize = 2
+	config.TempFilesDir = t.TempDir()
+	config.Checksum = true
+
+	sorter, output, errChan := extsort.Ordered(input, config)
+	sorter.Sort(context.Background())
+
+	var got []int
+	for value := range output {
+		got = append(got, value)
+	}
+	if err := <-errChan; err != nil {
+		t.Fatal(err)
+	}
+
+	want := []int{1, 2, 3, 4, 5, 6}
+	if !slices.Equal(got, want) {
+		t.Fatalf("sorted values = %v, want %v", got, want)
+	}
+}

sort_generic.go

@@ -5,6 +5,7 @@ import (
 	"bufio"
 	"context"
 	"encoding/binary"
+	"errors"
 	"io"
 	"slices"
 	"sync"
@@ -162,7 +163,11 @@ func (s *GenericSorter[E]) initMemoryPools() *memoryPools {
 func Generic[E any](input <-chan E, fromBytes FromBytesGeneric[E], toBytes ToBytesGeneric[E], compareFunc CompareGeneric[E], config *Config) (*GenericSorter[E], <-chan E, <-chan error) {
 	var err error
 	s := newSorter(input, fromBytes, toBytes, compareFunc, config)
-	s.tempWriter, err = tempfile.New(s.config.TempFilesDir, true)
+	if s.config.Checksum {
+		s.tempWriter, err = tempfile.NewChecksummed(s.config.TempFilesDir, true)
+	} else {
+		s.tempWriter, err = tempfile.New(s.config.TempFilesDir, true)
+	}
 	if err != nil {
 		s.mergeErrChan <- err
 		close(s.mergeErrChan)
@@ -194,9 +199,12 @@ func MockGeneric[E any](input <-chan E, fromBytes FromBytesGeneric[E], toBytes T
 // Merge uses the same context and runs in a goroutine after Sort returns().
 // for example, if calling sort in an errGroup, you must pass the group's parent context into sort.
 func (s *GenericSorter[E]) Sort(ctx context.Context) {
+	sortCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	var buildSortErrGroup, saveErrGroup *errgroup.Group
-	buildSortErrGroup, s.buildSortCtx = errgroup.WithContext(ctx)
-	saveErrGroup, s.saveCtx = errgroup.WithContext(ctx)
+	buildSortErrGroup, s.buildSortCtx = errgroup.WithContext(sortCtx)
+	saveErrGroup, s.saveCtx = errgroup.WithContext(sortCtx)
 
 	//start creating chunks
 	buildSortErrGroup.Go(s.buildChunks)
@@ -207,22 +215,34 @@ func (s *GenericSorter[E]) Sort(ctx context.Context) {
 	}
 
 	// Start the save worker that will handle single-chunk optimization
-	saveErrGroup.Go(s.saveChunksOptimized)
+	saveErrGroup.Go(func() error {
+		err := s.saveChunksOptimized()
+		if err != nil {
+			cancel()
+		}
+		return err
+	})
 
-	err := buildSortErrGroup.Wait()
-	if err != nil {
-		s.mergeErrChan <- err
-		close(s.mergeErrChan)
-		close(s.mergeChunkChan)
-		return
+	buildErr := buildSortErrGroup.Wait()
+	if buildErr != nil {
+		cancel()
 	}
 
 	// Close saveChunkChan to signal end of chunks
 	close(s.saveChunkChan)
 
 	// Wait for save worker to complete
-	err = saveErrGroup.Wait()
+	saveErr := saveErrGroup.Wait()
+	err := buildErr
+	if saveErr != nil && (err == nil || errors.Is(err, context.Canceled)) {
+		err = saveErr
+	}
 	if err != nil {
+		if s.tempReader != nil {
+			_ = s.tempReader.Close()
+		} else {
+			_ = s.tempWriter.Close()
+		}
 		s.mergeErrChan <- err
 		close(s.mergeErrChan)
 		close(s.mergeChunkChan)
@@ -241,6 +261,20 @@ func (s *GenericSorter[E]) Sort(ctx context.Context) {
 	go s.mergeNChunks(ctx)
 }
 
+func (s *GenericSorter[E]) Next(ctx context.Context) (value E, ok bool, err error) {
+	select {
+	case value, ok = <-s.mergeChunkChan:
+		if ok {
+			return value, true, nil
+		}
+	case <-ctx.Done():
+		return value, false, ctx.Err()
+	}
+
+	err, _ = <-s.mergeErrChan
+	return value, false, err
+}
+
 // buildChunks reads data from the input chan to builds chunks and pushes them to chunkChan
 func (s *GenericSorter[E]) buildChunks() error {
 	defer close(s.chunkChan) // if this is not called on error, causes a deadlock
@@ -457,6 +491,7 @@ func (s *GenericSorter[E]) saveChunk(b *genericChunk[E]) error {
 // mergeNChunks runs asynchronously in the background feeding data to getNext
 // sends errors to s.mergeErrorChan. Uses parallel merging for better performance.
 func (s *GenericSorter[E]) mergeNChunks(ctx context.Context) {
+	defer close(s.mergeErrChan)
 	defer close(s.mergeChunkChan)
 	defer func() {
 		if s.tempReader != nil {
@@ -470,8 +505,6 @@ func (s *GenericSorter[E]) mergeNChunks(ctx context.Context) {
 			}
 		}
 	}()
-	// Always ensure error channel is closed
-	defer close(s.mergeErrChan)
 
 	if s.tempReader == nil {
 		return
@@ -504,13 +537,13 @@ func (s *GenericSorter[E]) mergeNChunksSingleThreaded(ctx context.Context) {
 			reader:    s.tempReader.Read(i),
 		}
 		_, ok, err := merge.getNext() // start the merge by preloading the values
-		if err == io.EOF || !ok {
-			continue
-		}
 		if err != nil {
 			s.mergeErrChan <- err
 			return
 		}
+		if !ok {
+			continue
+		}
 		pq.Push(merge)
 	}
 
@@ -640,12 +673,12 @@ func (s *GenericSorter[E]) mergeWorkerSimple(ctx context.Context, startChunk, en
 			reader:    s.tempReader.Read(i),
 		}
 		_, ok, err := merge.getNext()
-		if err == io.EOF || !ok {
-			continue
-		}
 		if err != nil {
 			return err
 		}
+		if !ok {
+			continue
+		}
 		pq.Push(merge)
 	}