diff --git a/engine/chunk/chunk.go b/engine/chunk/chunk.go index 4011d7d3..c705a3bc 100644 --- a/engine/chunk/chunk.go +++ b/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 } diff --git a/engine/chunk/chunk_test.go b/engine/chunk/chunk_test.go index aaa17e5a..22cffa02 100644 --- a/engine/chunk/chunk_test.go +++ b/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()) + }) + } +}