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

@@ -162,7 +162,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)
@@ -457,6 +461,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 +475,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 +507,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 +643,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)
 	}
 

tempfile/checksum.go

@@ -0,0 +1,128 @@
+package tempfile
+
+import (
+	"bufio"
+	"fmt"
+	"hash/crc32"
+	"io"
+)
+
+const checksumBlockSize = 1 << 16
+
+var checksumTable = crc32.MakeTable(crc32.Castagnoli)
+
+type checksumWriter struct {
+	blockSize     int
+	blockChecksum uint32
+	checksums     []uint32
+}
+
+func (w *checksumWriter) write(dst *bufio.Writer, p []byte) (int, error) {
+	written := 0
+	for len(p) > 0 {
+		blockRemaining := checksumBlockSize - w.blockSize
+		if blockRemaining > len(p) {
+			blockRemaining = len(p)
+		}
+		chunk := p[:blockRemaining]
+		n, err := dst.Write(chunk)
+		if n > 0 {
+			w.blockChecksum = crc32.Update(w.blockChecksum, checksumTable, chunk[:n])
+			w.blockSize += n
+		}
+		written += n
+		if w.blockSize == checksumBlockSize {
+			w.finishBlock()
+		}
+		if err != nil {
+			return written, err
+		}
+		if n != len(chunk) {
+			return written, io.ErrShortWrite
+		}
+		p = p[n:]
+	}
+	return written, nil
+}
+
+func (w *checksumWriter) writeString(dst *bufio.Writer, s string) (int, error) {
+	buf := make([]byte, min(len(s), checksumBlockSize))
+	written := 0
+	for len(s) > 0 {
+		n := copy(buf, s)
+		m, err := w.write(dst, buf[:n])
+		written += m
+		if err != nil {
+			return written, err
+		}
+		s = s[m:]
+	}
+	return written, nil
+}
+
+func (w *checksumWriter) finishBlock() {
+	if w.blockSize == 0 {
+		return
+	}
+	w.checksums = append(w.checksums, w.blockChecksum)
+	w.blockSize = 0
+	w.blockChecksum = 0
+}
+
+func (w *checksumWriter) finishSection() []uint32 {
+	w.finishBlock()
+	checksums := w.checksums
+	w.checksums = nil
+	return checksums
+}
+
+func newChecksummedReader(reader io.Reader, section int, checksums []uint32, size int64) *bufio.Reader {
+	return bufio.NewReaderSize(&checksummedReader{
+		reader:    reader,
+		section:   section,
+		checksums: checksums,
+		remaining: size,
+	}, 16)
+}
+
+type checksummedReader struct {
+	reader    io.Reader
+	section   int
+	checksums []uint32
+	nextBlock int
+	remaining int64
+	data      []byte
+	offset    int
+}
+
+func (r *checksummedReader) Read(p []byte) (int, error) {
+	if len(p) == 0 {
+		return 0, nil
+	}
+	if r.offset == len(r.data) {
+		if r.nextBlock == len(r.checksums) {
+			return 0, io.EOF
+		}
+		blockSize := checksumBlockSize
+		if int64(blockSize) > r.remaining {
+			blockSize = int(r.remaining)
+		}
+		if cap(r.data) < blockSize {
+			r.data = make([]byte, blockSize)
+		} else {
+			r.data = r.data[:blockSize]
+		}
+		if _, err := io.ReadFull(r.reader, r.data); err != nil {
+			return 0, fmt.Errorf("read temporary section %d block %d: %w", r.section, r.nextBlock, err)
+		}
+		if crc32.Checksum(r.data, checksumTable) != r.checksums[r.nextBlock] {
+			return 0, fmt.Errorf("temporary section %d block %d checksum mismatch", r.section, r.nextBlock)
+		}
+		r.nextBlock++
+		r.remaining -= int64(blockSize)
+		r.offset = 0
+	}
+	n := copy(p, r.data[r.offset:])
+	r.offset += n
+	return n, nil
+}