fix: stop dropping the bytes between the chunk boundary and the read size

engine/chunk/chunk.go

@@ -7,6 +7,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"io"
 	"sync"
 	"unicode"
 
@@ -134,28 +135,30 @@ func (c *Chunk) GetFileThreshold() int64 {
 
 // ReadChunk reads the next chunk of data from file
 func (c *Chunk) ReadChunk(reader *bufio.Reader, totalLines int) (string, error) {
-	// borrow a []bytes from the pool and seed it with raw data from file (up to chunk size + peek size)
-	rawData, ok := c.GetPeekedBuf()
-	if !ok {
-		return "", fmt.Errorf("expected *bytes.Buffer, got %T", rawData)
-	}
-	defer c.PutPeekedBuf(rawData)
-	n, err := reader.Read(*rawData)
-
-	var chunkStr string
-	// "Callers should always process the n > 0 bytes returned before considering the error err."
-	// https://pkg.go.dev/io#Reader
-	if n > 0 {
-		// only check the filetype at the start of file
-		if totalLines == 0 && ShouldSkipFile((*rawData)[:n]) {
-			return "", fmt.Errorf("skipping file: %w", ErrUnsupportedFileType)
-		}
+	// look ahead without consuming (up to chunk size + peek size): whatever sits past the
+	// chunk boundary has to stay in the reader, otherwise the next call never sees it
+	rawData, err := reader.Peek(c.size + c.maxPeekSize)
+	if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, bufio.ErrBufferFull) {
+		return "", err
+	}
+	if len(rawData) == 0 {
+		return "", io.EOF
+	}
 
-		chunkStr, err = c.generateChunk((*rawData)[:n])
+	// only check the filetype at the start of file
+	if totalLines == 0 && ShouldSkipFile(rawData) {
+		return "", fmt.Errorf("skipping file: %w", ErrUnsupportedFileType)
 	}
+
+	chunkStr, err := c.generateChunk(rawData)
 	if err != nil {
 		return "", err
 	}
+
+	// consume exactly what the chunk covers
+	if _, err := reader.Discard(len(chunkStr)); err != nil {
+		return "", err
+	}
 	return chunkStr, nil
 }
 

engine/chunk/chunk_test.go

@@ -160,3 +160,47 @@ func TestGenerateChunk(t *testing.T) {
 		})
 	}
 }
+
+func TestReadChunkKeepsEveryByte(t *testing.T) {
+	// Arrange
+	testCases := []struct {
+		name  string
+		input string
+	}{
+		{
+			name:  "boundary inside the peek window",
+			input: "abc\ndef\n\n\n\n\nghi\njkl",
+		},
+		{
+			name:  "boundary right after the chunk size",
+			input: "0123456789\n\nSECRET\n",
+		},
+		{
+			name:  "no boundary at all",
+			input: "abc\ndef\nghi\njkl\nmno\npqr\nstu\nvwx\nyz",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			c := New(WithSize(chunkSize), WithMaxPeekSize(maxPeekSize), WithSmallFileThreshold(smallFileThreshold))
+			reader := bufio.NewReaderSize(strings.NewReader(tc.input), chunkSize+maxPeekSize)
+
+			// Act
+			var scanned strings.Builder
+			totalLines := 0
+			for {
+				chunkStr, err := c.ReadChunk(reader, totalLines)
+				if err == io.EOF {
+					break
+				}
+				require.NoError(t, err)
+				totalLines += strings.Count(chunkStr, "\n")
+				scanned.WriteString(chunkStr)
+			}
+
+			// Assert
+			require.Equal(t, tc.input, scanned.String())
+		})
+	}
+}