diff --git a/compactor.go b/compactor.go index 2bdbda7..21155f7 100644 --- a/compactor.go +++ b/compactor.go @@ -2,6 +2,7 @@ package ltx import ( "context" + "errors" "fmt" "io" "sync/atomic" @@ -54,7 +55,9 @@ func NewCompactor(w io.Writer, rdrs []io.Reader) (*Compactor, error) { c := &Compactor{enc: enc} c.inputs = make([]*compactorInput, len(rdrs)) for i := range c.inputs { - c.inputs[i] = &compactorInput{dec: NewDecoder(rdrs[i])} + dec := NewDecoder(rdrs[i]) + dec.SetRetainPageIndex(false) // inputs are streamed; never build a database-sized map + c.inputs[i] = &compactorInput{dec: dec} } return c, nil } @@ -74,8 +77,32 @@ func (c *Compactor) Status() CompactorStatus { } } -// Compact merges the input readers into a single LTX writer. -func (c *Compactor) Compact(ctx context.Context) error { +// SetSpillDir enables spilling the output page index to a temp file in dir +// once it grows past the encoder's spill threshold. See Encoder.SetSpillDir. +func (c *Compactor) SetSpillDir(dir string) { c.enc.SetSpillDir(dir) } + +// SetSpillThreshold sets the number of in-memory page index entries that +// triggers a spill of the output index. See Encoder.SetSpillThreshold. +func (c *Compactor) SetSpillThreshold(n int) { c.enc.SetSpillThreshold(n) } + +// Cleanup removes any spill file left by an abandoned compaction and aborts +// the output encoder if it was not closed successfully. It is safe to call +// after Compact returns, successfully or not. See Encoder.Cleanup. +func (c *Compactor) Cleanup() error { return c.enc.Cleanup() } + +// Compact merges the input readers into a single LTX writer. A failed +// compaction cannot be resumed, so its output spill file, if any, is removed +// before returning. +func (c *Compactor) Compact(ctx context.Context) (err error) { + defer func() { + if err == nil { + return + } + if cerr := c.enc.Cleanup(); cerr != nil { + err = errors.Join(err, fmt.Errorf("cleanup spill: %w", cerr)) + } + }() + if len(c.inputs) == 0 { return fmt.Errorf("at least one input reader required") } diff --git a/compactor_test.go b/compactor_test.go index 95e1174..065f052 100644 --- a/compactor_test.go +++ b/compactor_test.go @@ -382,3 +382,104 @@ func TestCompactor_Compact(t *testing.T) { } }) } + +func TestCompactor_Spill(t *testing.T) { + page := func(pgno uint32) ltx.PageSpec { + return ltx.PageSpec{Header: ltx.PageHeader{Pgno: pgno}, Data: bytes.Repeat([]byte{byte(pgno)}, 512)} + } + inputs := []*ltx.FileSpec{ + { + Header: ltx.Header{Version: ltx.Version, PageSize: 512, Commit: 6, MinTXID: 2, MaxTXID: 2, Timestamp: 1000, PreApplyChecksum: ltx.ChecksumFlag | 1}, + Pages: []ltx.PageSpec{page(1), page(2), page(3), page(4), page(5), page(6)}, + Trailer: ltx.Trailer{PostApplyChecksum: ltx.ChecksumFlag | 2}, + }, + { + Header: ltx.Header{Version: ltx.Version, PageSize: 512, Commit: 6, MinTXID: 3, MaxTXID: 3, Timestamp: 2000, PreApplyChecksum: ltx.ChecksumFlag | 2}, + Pages: []ltx.PageSpec{page(2), page(5)}, + Trailer: ltx.Trailer{PostApplyChecksum: ltx.ChecksumFlag | 3}, + }, + } + readers := func() []io.Reader { + rdrs := make([]io.Reader, len(inputs)) + for i, input := range inputs { + var buf bytes.Buffer + writeFileSpec(t, &buf, input) + rdrs[i] = &buf + } + return rdrs + } + + var want bytes.Buffer + c, err := ltx.NewCompactor(&want, readers()) + if err != nil { + t.Fatal(err) + } + if err := c.Compact(context.Background()); err != nil { + t.Fatal(err) + } + + t.Run("OutputIdenticalWithSpill", func(t *testing.T) { + dir := t.TempDir() + var got bytes.Buffer + c, err := ltx.NewCompactor(&got, readers()) + if err != nil { + t.Fatal(err) + } + c.SetSpillDir(dir) + c.SetSpillThreshold(2) + if err := c.Compact(context.Background()); err != nil { + t.Fatal(err) + } + if !bytes.Equal(got.Bytes(), want.Bytes()) { + t.Fatal("spilled compaction output differs from in-memory output") + } + if names := dirEntries(t, dir); len(names) != 0 { + t.Fatalf("spill file left behind: %v", names) + } + if err := c.Cleanup(); err != nil { + t.Fatalf("Cleanup()=%v", err) + } + }) + + t.Run("FailedCompactionRemovesSpill", func(t *testing.T) { + dir := t.TempDir() + var entriesAtFailure []string + w := &callCountingWriter{failOnCall: 11, onFail: func() { entriesAtFailure = dirEntries(t, dir) }} + c, err := ltx.NewCompactor(w, readers()) + if err != nil { + t.Fatal(err) + } + c.SetSpillDir(dir) + c.SetSpillThreshold(2) + if err := c.Compact(context.Background()); !errors.Is(err, errInjected) { + t.Fatalf("Compact()=%v, want %v", err, errInjected) + } + if len(entriesAtFailure) != 1 { + t.Fatalf("expected one spill file at the time of failure, got %v", entriesAtFailure) + } + if names := dirEntries(t, dir); len(names) != 0 { + t.Fatalf("spill file left behind after failed compaction: %v", names) + } + if err := c.Cleanup(); err != nil { + t.Fatalf("Cleanup()=%v", err) + } + }) +} + +type callCountingWriter struct { + calls int + failOnCall int + onFail func() +} + +func (w *callCountingWriter) Write(p []byte) (int, error) { + w.calls++ + if w.calls >= w.failOnCall { + if w.onFail != nil { + w.onFail() + w.onFail = nil + } + return 0, errInjected + } + return len(p), nil +} diff --git a/decoder.go b/decoder.go index 8419ba7..88a03e8 100644 --- a/decoder.go +++ b/decoder.go @@ -1,8 +1,10 @@ package ltx import ( + "bufio" "bytes" "encoding/binary" + "errors" "fmt" "hash" "hash/crc64" @@ -28,22 +30,41 @@ type Decoder struct { pageIndex map[uint32]PageIndexElem state string + // retainPageIndex controls whether Close materializes the page index map. + // Callers that only stream pages (e.g. compaction inputs) turn it off so a + // database-sized map is never built. + retainPageIndex bool + chksum Checksum hash hash.Hash64 pageN int // pages read n int64 // bytes read + + // pageSeq is a running hash of the decoded page numbers in order; Close + // compares it with the same hash over the page index so the index must + // name exactly the decoded pages, without retaining them. + pageSeq hash.Hash64 } // NewDecoder returns a new instance of Decoder. func NewDecoder(r io.Reader) *Decoder { return &Decoder{ - r: r, - zr: lz4.NewReader(r), - state: stateHeader, - hash: crc64.New(crc64.MakeTable(crc64.ISO)), + r: r, + zr: lz4.NewReader(r), + state: stateHeader, + hash: crc64.New(crc64.MakeTable(crc64.ISO)), + pageSeq: crc64.New(crc64.MakeTable(crc64.ISO)), + retainPageIndex: true, } } +// SetRetainPageIndex controls whether Close builds the in-memory page index +// returned by PageIndex. It defaults to true. When false, Close still reads, +// validates, and checksums the index but discards the entries, so memory no +// longer scales with the number of pages in the file; PageIndex returns nil. +// Must be called before Close. +func (dec *Decoder) SetRetainPageIndex(retain bool) { dec.retainPageIndex = retain } + // N returns the number of bytes read. func (dec *Decoder) N() int64 { return dec.n } @@ -79,28 +100,57 @@ func (dec *Decoder) Close() error { return fmt.Errorf("cannot close, expected %s", dec.state) } - // Slurp the remaining data in to memory so we can use the ByteReader interface. - remainingBytes, err := io.ReadAll(dec.r) - if err != nil { - return fmt.Errorf("read all: %w", err) - } - remaining := bytes.NewReader(remainingBytes) - - // Write everything but the file checksum to the hash. - dec.writeToHash(remainingBytes[:len(remainingBytes)-ChecksumSize]) - - // Read page index. - if dec.pageIndex, err = DecodePageIndex(remaining, 0, dec.header.MinTXID, dec.header.MaxTXID); err != nil { + // Stream the page index straight from the reader, hashing bytes as they + // are consumed, instead of slurping the tail of the file into memory. + // The index is only materialized when retention is requested. + br := bufio.NewReader(dec.r) + var index map[uint32]PageIndexElem + if dec.retainPageIndex { + index = make(map[uint32]PageIndexElem) + } + // Index entries must describe one frame per decoded page, in ascending + // page order, with non-overlapping frames after the header. (Offsets are + // file positions of compressed frames, which the decoder does not track, + // so exact page-block bounds are not checked here.) + v := pageIndexValidator{commit: dec.header.Commit, nextOffset: HeaderSize, seq: crc64.New(crc64.MakeTable(crc64.ISO))} + if err := dec.streamPageIndex(br, func(pgno uint32, offset, size int64) error { + if err := v.check(pgno, offset, size); err != nil { + return err + } + if index != nil { + index[pgno] = PageIndexElem{ + MinTXID: dec.header.MinTXID, + MaxTXID: dec.header.MaxTXID, + Offset: offset, + Size: size, + } + } + return nil + }); err != nil { return fmt.Errorf("read page index: %w", err) } + if v.n != dec.pageN { + return fmt.Errorf("page index has %d entries but %d pages were decoded", v.n, dec.pageN) + } + if v.seq.Sum64() != dec.pageSeq.Sum64() { + return errors.New("page index does not match the decoded page numbers") + } + dec.pageIndex = index - // Read trailer. + // Read trailer. Everything except the trailing file checksum is hashed. b := make([]byte, TrailerSize) - if _, err := io.ReadFull(remaining, b); err != nil { - return err - } else if err := dec.trailer.UnmarshalBinary(b); err != nil { + if _, err := io.ReadFull(br, b); err != nil { + return fmt.Errorf("read trailer: %w", err) + } + dec.writeToHash(b[:TrailerChecksumOffset]) + if err := dec.trailer.UnmarshalBinary(b); err != nil { return fmt.Errorf("unmarshal trailer: %w", err) } + if _, err := br.ReadByte(); err == nil { + return errors.New("unexpected data after trailer") + } else if !errors.Is(err, io.EOF) { + return fmt.Errorf("read after trailer: %w", err) + } // TODO: Ensure last read page is equal to the commit for snapshot LTX files @@ -217,6 +267,7 @@ func (dec *Decoder) DecodePage(hdr *PageHeader, data []byte) error { dec.writeToHash(data) dec.pageN++ + hashPgno(dec.pageSeq, hdr.Pgno) // Calculate checksum while decoding snapshots if tracking checksums. if dec.header.IsSnapshot() && !dec.header.NoChecksum() { @@ -366,41 +417,158 @@ func DecodePageData(b []byte) (hdr PageHeader, data []byte, err error) { return hdr, data, err } -// DecodePageIndex decodes the page index from r. -func DecodePageIndex(r io.ByteReader, level int, minTXID, maxTXID TXID) (map[uint32]PageIndexElem, error) { - pageIndex := make(map[uint32]PageIndexElem) +// streamPageIndex reads the page index section (records, end marker, and +// size field) from br, hashing every byte consumed and validating that page +// numbers ascend and that the size field matches the bytes read. fn is called +// for each record in file order. +func (dec *Decoder) streamPageIndex(br *bufio.Reader, fn func(pgno uint32, offset, size int64) error) error { + return parsePageIndex(br, dec.writeToHash, fn) +} +// pageIndexValidator checks that index entries are structurally plausible +// without retaining them: page numbers within the commit size, and frames +// that start after the header, do not overlap, and have a positive size. +type pageIndexValidator struct { + commit uint32 + nextOffset int64 // earliest offset the next frame may start at + n int + seq hash.Hash64 // running hash of page numbers, compared with Decoder.pageSeq +} + +func (v *pageIndexValidator) check(pgno uint32, offset, size int64) error { + if pgno > v.commit { + return fmt.Errorf("page index pgno %d exceeds commit %d", pgno, v.commit) + } + if offset < v.nextOffset { + return fmt.Errorf("page index pgno %d offset %d overlaps previous frame ending at %d", pgno, offset, v.nextOffset) + } + if size <= PageHeaderSize || offset > math.MaxInt64-size { + return fmt.Errorf("page index pgno %d has invalid frame size %d", pgno, size) + } + v.nextOffset = offset + size + v.n++ + hashPgno(v.seq, pgno) + return nil +} + +// hashPgno feeds pgno into a page-sequence hash. +func hashPgno(h hash.Hash64, pgno uint32) { + var b [4]byte + binary.BigEndian.PutUint32(b[:], pgno) + _, _ = h.Write(b[:]) +} + +// parsePageIndex reads page index records from br until the end marker, +// then the size field, validating that page numbers ascend and that the size +// field equals the bytes consumed. observe, if non-nil, receives every byte +// consumed (for checksumming); fn receives every record. +func parsePageIndex(br io.ByteReader, observe func([]byte), fn func(pgno uint32, offset, size int64) error) error { + var scratch [3 * binary.MaxVarintLen64]byte + var consumed int64 + var prevPgno uint32 for { - pgno, err := binary.ReadUvarint(r) + buf := scratch[:0] + pgno, err := readUvarintInto(br, &buf) if err != nil { - return nil, fmt.Errorf("read page index pgno: %w", err) - } else if pgno == 0 { - break // End when we hit the end marker. + return fmt.Errorf("read page index pgno: %w", err) } + if pgno == 0 { + if observe != nil { + observe(buf) + } + consumed += int64(len(buf)) + break // end marker + } + if pgno > math.MaxUint32 { + return fmt.Errorf("page index pgno %d out of range", pgno) + } + if uint32(pgno) <= prevPgno { + return fmt.Errorf("page index out of order: %d after %d", pgno, prevPgno) + } + offset, err := readUvarintInto(br, &buf) + if err != nil { + return fmt.Errorf("read page index offset: %w", err) + } + size, err := readUvarintInto(br, &buf) + if err != nil { + return fmt.Errorf("read page index size: %w", err) + } + if offset > math.MaxInt64 || size > math.MaxInt64 { + return fmt.Errorf("page index pgno %d offset/size out of range", pgno) + } + if observe != nil { + observe(buf) + } + consumed += int64(len(buf)) + prevPgno = uint32(pgno) + if err := fn(uint32(pgno), int64(offset), int64(size)); err != nil { + return err + } + } - offset, err := binary.ReadUvarint(r) + var sizeBuf [8]byte + for i := range sizeBuf { + b, err := br.ReadByte() if err != nil { - return nil, fmt.Errorf("read page index offset: %w", err) + if errors.Is(err, io.EOF) { + err = io.ErrUnexpectedEOF + } + return fmt.Errorf("read page index size: %w", err) } - size, err := binary.ReadUvarint(r) + sizeBuf[i] = b + } + if observe != nil { + observe(sizeBuf[:]) + } + if indexSize := binary.BigEndian.Uint64(sizeBuf[:]); indexSize != uint64(consumed) { + return fmt.Errorf("page index size mismatch: field=%d read=%d", indexSize, consumed) + } + return nil +} + +// readUvarintInto reads a uvarint from br, appending the consumed bytes to *buf. +func readUvarintInto(br io.ByteReader, buf *[]byte) (uint64, error) { + var x uint64 + var s uint + for i := 0; i < binary.MaxVarintLen64; i++ { + b, err := br.ReadByte() if err != nil { - return nil, fmt.Errorf("read page index size: %w", err) + if errors.Is(err, io.EOF) && i > 0 { + err = io.ErrUnexpectedEOF + } + return 0, err + } + *buf = append(*buf, b) + if b < 0x80 { + if i == binary.MaxVarintLen64-1 && b > 1 { + return 0, errors.New("uvarint overflows 64 bits") + } + return x | uint64(b)< enc.header.Commit { @@ -239,42 +343,162 @@ func (enc *Encoder) EncodePage(hdr PageHeader, data []byte) (err error) { writeData := enc.compressBuf[:n] - // Write page header. + // Write page header. Once any bytes of a frame have been written the + // file cannot be completed, so write failures abort the encoder. b, err := hdr.MarshalBinary() if err != nil { return fmt.Errorf("marshal: %w", err) } else if _, err := enc.write(b); err != nil { - return fmt.Errorf("write page header: %w", err) + return enc.abort(fmt.Errorf("write page header: %w", err)) } // Write data size (4 bytes, big-endian). sizeBuf := make([]byte, 4) binary.BigEndian.PutUint32(sizeBuf, uint32(len(writeData))) if _, err := enc.write(sizeBuf); err != nil { - return fmt.Errorf("write data size: %w", err) + return enc.abort(fmt.Errorf("write data size: %w", err)) } // Write page data (compressed or uncompressed). - if _, err := enc.w.Write(writeData); err != nil { - return fmt.Errorf("write page data: %w", err) + if n, err := enc.w.Write(writeData); err != nil { + return enc.abort(fmt.Errorf("write page data: %w", err)) + } else if n < len(writeData) { + return enc.abort(fmt.Errorf("write page data: %w", io.ErrShortWrite)) } _, _ = enc.hash.Write(data) // hash the uncompressed data enc.n += int64(len(writeData)) enc.pagesWritten++ enc.prevPgno = hdr.Pgno - enc.index.append(pageIndexEntry{ + entry := pageIndexEntry{ pgno: hdr.Pgno, offset: offset, size: enc.n - offset, - }) + } + // The page bytes are already in the output, so an index failure here + // leaves a file that cannot be completed: abort the encoder. + if enc.spill != nil { + if err := enc.spill.append(entry); err != nil { + return enc.abort(fmt.Errorf("spill page index: %w", err)) + } + return nil + } + enc.index.append(entry) + if enc.spillDir != "" && enc.index.n >= enc.spillThreshold { + if err := enc.startSpill(); err != nil { + return enc.abort(fmt.Errorf("spill page index: %w", err)) + } + } + + return nil +} + +// startSpill moves the in-memory page index to a temp file in spillDir and +// releases the in-memory chunks. +func (enc *Encoder) startSpill() error { + spill, err := newIndexSpill(enc.spillDir) + if err != nil { + return err + } + // Reference the spill before migrating so a failure here (which aborts + // the encoder) still leaves Cleanup able to remove, or retry removing, + // the temp file. + enc.spill = spill + for _, chunk := range enc.index.chunks { + for _, elem := range chunk { + if err := spill.append(elem); err != nil { + return err + } + } + } + enc.index = pageIndex{} + return nil +} + +// appendPageIndexRecord appends elem's on-disk varint tuple to b. +func appendPageIndexRecord(b []byte, elem pageIndexEntry) []byte { + b = binary.AppendUvarint(b, uint64(elem.pgno)) + b = binary.AppendUvarint(b, uint64(elem.offset)) + b = binary.AppendUvarint(b, uint64(elem.size)) + return b +} + +// indexSpill is a buffered temp file holding page index records in their +// final on-disk format. +type indexSpill struct { + f *os.File + w *bufio.Writer + buf []byte + path string // retained after a failed removal so it can be retried +} + +func newIndexSpill(dir string) (*indexSpill, error) { + f, err := os.CreateTemp(dir, ".ltx-page-index-*.tmp") + if err != nil { + return nil, err + } + return &indexSpill{f: f, w: bufio.NewWriterSize(f, 256<<10), buf: make([]byte, 0, 3*binary.MaxVarintLen64)}, nil +} + +func (s *indexSpill) append(elem pageIndexEntry) error { + s.buf = appendPageIndexRecord(s.buf[:0], elem) + _, err := s.w.Write(s.buf) + return err +} + +// copyTo flushes the spill file and streams its records through enc.write so +// they are hashed and counted exactly as in-memory records are. +func (s *indexSpill) copyTo(enc *Encoder) error { + if err := s.w.Flush(); err != nil { + return err + } + if _, err := s.f.Seek(0, io.SeekStart); err != nil { + return err + } + buf := make([]byte, 256<<10) + for { + n, err := s.f.Read(buf) + if n > 0 { + if _, werr := enc.write(buf[:n]); werr != nil { + return werr + } + } + if err == io.EOF { + return nil + } else if err != nil { + return err + } + } +} +// remove closes and deletes the spill file. The handle is closed before the +// file is removed so the sequence is valid on Windows too. Safe to call more +// than once; a failed removal keeps the path so it can be retried. +func (s *indexSpill) remove() error { + if s.f != nil { + s.path = s.f.Name() + err := s.f.Close() + s.f = nil + if err != nil { + return err + } + } + if s.path == "" { + return nil + } + if err := os.Remove(s.path); err != nil && !errors.Is(err, os.ErrNotExist) { + return err + } + s.path = "" return nil } // write to the uncompressed writer & add to the checksum. -func (enc *Encoder) write(b []byte) (n int, err error) { - n, err = enc.w.Write(b) +func (enc *Encoder) write(b []byte) (int, error) { + n, err := enc.w.Write(b) + if err == nil && n < len(b) { + err = io.ErrShortWrite + } enc.writeToHash(b[:n]) return n, err } @@ -309,9 +533,11 @@ const ( // and no copy of the whole index when it grows. type pageIndex struct { chunks [][]pageIndexEntry + n int // total entries } func (idx *pageIndex) append(e pageIndexEntry) { + idx.n++ n := len(idx.chunks) if n == 0 || len(idx.chunks[n-1]) == cap(idx.chunks[n-1]) { size := pageIndexMinChunk diff --git a/encoder_spill_test.go b/encoder_spill_test.go new file mode 100644 index 0000000..cab30f0 --- /dev/null +++ b/encoder_spill_test.go @@ -0,0 +1,467 @@ +package ltx_test + +import ( + "bytes" + "errors" + "os" + "path/filepath" + "runtime" + "testing" + + "github.com/superfly/ltx" +) + +var errInjected = errors.New("injected write failure") + +func encodeWithSpill(t *testing.T, spillDir string, threshold int, pgnos []uint32, commit uint32, closeIt bool) (*ltx.Encoder, []byte) { + t.Helper() + var buf bytes.Buffer + enc, err := ltx.NewEncoder(&buf) + if err != nil { + t.Fatal(err) + } + if spillDir != "" { + enc.SetSpillDir(spillDir) + } + if threshold > 0 { + enc.SetSpillThreshold(threshold) + } + if err := enc.EncodeHeader(ltx.Header{ + Version: ltx.Version, + PageSize: 512, + Commit: commit, + MinTXID: 2, + MaxTXID: 2, + Timestamp: 1000, + PreApplyChecksum: ltx.ChecksumFlag | 1, + }); err != nil { + t.Fatal(err) + } + page := make([]byte, 512) + for _, pgno := range pgnos { + page[0] = byte(pgno) + if err := enc.EncodePage(ltx.PageHeader{Pgno: pgno}, page); err != nil { + t.Fatalf("pgno %d: %v", pgno, err) + } + } + enc.SetPostApplyChecksum(ltx.ChecksumFlag | 2) + if closeIt { + if err := enc.Close(); err != nil { + t.Fatal(err) + } + } + return enc, buf.Bytes() +} + +func dirEntries(t *testing.T, dir string) []string { + t.Helper() + entries, err := os.ReadDir(dir) + if err != nil { + t.Fatal(err) + } + names := make([]string, 0, len(entries)) + for _, e := range entries { + names = append(names, e.Name()) + } + return names +} + +func requireDirectoryRemovalFailure(t *testing.T) { + t.Helper() + if runtime.GOOS == "windows" { + t.Skip("Windows does not enforce Unix directory write permissions") + } + if os.Getuid() == 0 { + t.Skip("root can remove files from read-only directories") + } +} + +func TestEncoder_SpillFile(t *testing.T) { + pgnos := make([]uint32, 0, 1000) + for pgno := uint32(3); pgno <= 3000; pgno += 3 { + pgnos = append(pgnos, pgno) + } + + t.Run("OutputIdenticalWithAndWithoutSpill", func(t *testing.T) { + dir := t.TempDir() + _, plain := encodeWithSpill(t, "", 0, pgnos, 4000, true) + _, spilled := encodeWithSpill(t, dir, 8, pgnos, 4000, true) + if !bytes.Equal(plain, spilled) { + t.Fatal("spilled output differs from in-memory output") + } + if names := dirEntries(t, dir); len(names) != 0 { + t.Fatalf("spill file left behind after Close: %v", names) + } + dec := ltx.NewDecoder(bytes.NewReader(spilled)) + if err := dec.Verify(); err != nil { + t.Fatal(err) + } + if got, want := len(dec.PageIndex()), len(pgnos); got != want { + t.Fatalf("decoded index len=%d, want %d", got, want) + } + }) + + t.Run("BelowThresholdNeverTouchesDisk", func(t *testing.T) { + dir := t.TempDir() + encodeWithSpill(t, dir, len(pgnos)+1, pgnos, 4000, true) + if names := dirEntries(t, dir); len(names) != 0 { + t.Fatalf("unexpected spill file: %v", names) + } + }) + + t.Run("NoSpillDirNeverSpills", func(t *testing.T) { + enc, _ := encodeWithSpill(t, "", 8, pgnos, 4000, true) + if enc.Spilled() { + t.Fatal("encoder spilled without a spill dir") + } + }) + + t.Run("CleanupWithoutClose", func(t *testing.T) { + dir := t.TempDir() + enc, _ := encodeWithSpill(t, dir, 8, pgnos, 4000, false) + if !enc.Spilled() { + t.Fatal("expected encoder to have spilled") + } + if names := dirEntries(t, dir); len(names) != 1 { + t.Fatalf("expected one spill file while open, got %v", names) + } + if err := enc.Cleanup(); err != nil { + t.Fatal(err) + } + if err := enc.Cleanup(); err != nil { // idempotent + t.Fatal(err) + } + if names := dirEntries(t, dir); len(names) != 0 { + t.Fatalf("spill file left behind after Cleanup: %v", names) + } + }) + + t.Run("UnwritableSpillDir", func(t *testing.T) { + notADir := filepath.Join(t.TempDir(), "file") + if err := os.WriteFile(notADir, []byte("x"), 0o600); err != nil { + t.Fatal(err) + } + var buf bytes.Buffer + enc, err := ltx.NewEncoder(&buf) + if err != nil { + t.Fatal(err) + } + enc.SetSpillDir(notADir) + enc.SetSpillThreshold(2) + if err := enc.EncodeHeader(ltx.Header{Version: ltx.Version, PageSize: 512, Commit: 4, MinTXID: 2, MaxTXID: 2, Timestamp: 1000, PreApplyChecksum: ltx.ChecksumFlag | 1}); err != nil { + t.Fatal(err) + } + page := make([]byte, 512) + if err := enc.EncodePage(ltx.PageHeader{Pgno: 1}, page); err != nil { + t.Fatal(err) + } + if err := enc.EncodePage(ltx.PageHeader{Pgno: 2}, page); err == nil { + t.Fatal("expected spill error for unwritable spill dir") + } + }) +} + +func TestEncoder_SpillAcrossChunkBoundaries(t *testing.T) { + const n = 2*65536 + 259 + pgnos := make([]uint32, 0, n) + for pgno := uint32(2); pgno <= n+1; pgno++ { + pgnos = append(pgnos, pgno) + } + dir := t.TempDir() + _, plain := encodeWithSpill(t, "", 0, pgnos, n+1, true) + _, spilled := encodeWithSpill(t, dir, 70000, pgnos, n+1, true) + if !bytes.Equal(plain, spilled) { + t.Fatal("spilled output differs from in-memory output across chunk boundaries") + } +} + +// failingWriter fails every write once failAfter bytes have been accepted. +type failingWriter struct { + buf bytes.Buffer + failAfter int + failed bool + err error +} + +func (w *failingWriter) Write(p []byte) (int, error) { + if w.failed || w.buf.Len()+len(p) > w.failAfter { + w.failed = true + return 0, w.err + } + return w.buf.Write(p) +} + +func newSpilledEncoder(t *testing.T, w interface{ Write([]byte) (int, error) }, dir string, pages int) *ltx.Encoder { + t.Helper() + enc, err := ltx.NewEncoder(w) + if err != nil { + t.Fatal(err) + } + enc.SetSpillDir(dir) + enc.SetSpillThreshold(2) + if err := enc.EncodeHeader(ltx.Header{Version: ltx.Version, PageSize: 512, Commit: uint32(pages + 1), MinTXID: 2, MaxTXID: 2, Timestamp: 1000, PreApplyChecksum: ltx.ChecksumFlag | 1}); err != nil { + t.Fatal(err) + } + page := make([]byte, 512) + for i := 0; i < pages; i++ { + if err := enc.EncodePage(ltx.PageHeader{Pgno: uint32(i + 2)}, page); err != nil { + t.Fatal(err) + } + } + enc.SetPostApplyChecksum(ltx.ChecksumFlag | 2) + if !enc.Spilled() { + t.Fatal("expected encoder to have spilled") + } + return enc +} + +func TestEncoder_SpillLifecycle(t *testing.T) { + t.Run("CleanupBeforeCloseAborts", func(t *testing.T) { + dir := t.TempDir() + var buf bytes.Buffer + enc := newSpilledEncoder(t, &buf, dir, 4) + if err := enc.Cleanup(); err != nil { + t.Fatal(err) + } + if err := enc.Close(); err != ltx.ErrEncoderAborted { + t.Fatalf("Close()=%v, want ErrEncoderAborted", err) + } + if err := enc.EncodePage(ltx.PageHeader{Pgno: 9}, make([]byte, 512)); err != ltx.ErrEncoderAborted { + t.Fatalf("EncodePage()=%v, want ErrEncoderAborted", err) + } + if names := dirEntries(t, dir); len(names) != 0 { + t.Fatalf("spill file left behind: %v", names) + } + }) + + t.Run("FirstCloseWriteFailureRemovesSpill", func(t *testing.T) { + dir := t.TempDir() + w := &failingWriter{err: errInjected, failAfter: 1 << 30} + enc, err := ltx.NewEncoder(w) + if err != nil { + t.Fatal(err) + } + enc.SetSpillDir(dir) + enc.SetSpillThreshold(2) + if err := enc.EncodeHeader(ltx.Header{Version: ltx.Version, PageSize: 512, Commit: 8, MinTXID: 2, MaxTXID: 2, Timestamp: 1000, PreApplyChecksum: ltx.ChecksumFlag | 1}); err != nil { + t.Fatal(err) + } + page := make([]byte, 512) + for pgno := uint32(2); pgno <= 5; pgno++ { + if err := enc.EncodePage(ltx.PageHeader{Pgno: pgno}, page); err != nil { + t.Fatal(err) + } + } + w.failAfter = w.buf.Len() // the end-of-pages marker is the first write to fail + enc.SetPostApplyChecksum(ltx.ChecksumFlag | 2) + if err := enc.Close(); err == nil { + t.Fatal("expected Close to fail") + } + if names := dirEntries(t, dir); len(names) != 0 { + t.Fatalf("spill file leaked after failed Close: %v", names) + } + if err := enc.Close(); err != ltx.ErrEncoderAborted { + t.Fatalf("second Close()=%v, want ErrEncoderAborted", err) + } + }) + + t.Run("FailedSpilledCloseIsNotRetryable", func(t *testing.T) { + dir := t.TempDir() + w := &failingWriter{err: errInjected, failAfter: 1 << 30} + enc, err := ltx.NewEncoder(w) + if err != nil { + t.Fatal(err) + } + enc.SetSpillDir(dir) + enc.SetSpillThreshold(2) + if err := enc.EncodeHeader(ltx.Header{Version: ltx.Version, PageSize: 512, Commit: 8, MinTXID: 2, MaxTXID: 2, Timestamp: 1000, PreApplyChecksum: ltx.ChecksumFlag | 1}); err != nil { + t.Fatal(err) + } + page := make([]byte, 512) + for pgno := uint32(2); pgno <= 5; pgno++ { + if err := enc.EncodePage(ltx.PageHeader{Pgno: pgno}, page); err != nil { + t.Fatal(err) + } + } + w.failAfter = w.buf.Len() + ltx.PageHeaderSize + 1 // marker succeeds, index copy fails + enc.SetPostApplyChecksum(ltx.ChecksumFlag | 2) + if err := enc.Close(); err == nil { + t.Fatal("expected Close to fail during index copy") + } + w.failed, w.failAfter = false, 1<<30 + if err := enc.Close(); err != ltx.ErrEncoderAborted { + t.Fatalf("retried Close()=%v, want ErrEncoderAborted", err) + } + if names := dirEntries(t, dir); len(names) != 0 { + t.Fatalf("spill file leaked: %v", names) + } + }) + + t.Run("CleanupReportsRemovalFailure", func(t *testing.T) { + requireDirectoryRemovalFailure(t) + dir := t.TempDir() + var buf bytes.Buffer + enc := newSpilledEncoder(t, &buf, dir, 4) + if err := os.Chmod(dir, 0o500); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = os.Chmod(dir, 0o700) }) + if err := enc.Cleanup(); err == nil { + t.Fatal("expected Cleanup to report the removal failure") + } + if err := os.Chmod(dir, 0o700); err != nil { + t.Fatal(err) + } + if err := enc.Cleanup(); err != nil { + t.Fatalf("retried Cleanup()=%v", err) + } + if names := dirEntries(t, dir); len(names) != 0 { + t.Fatalf("spill file left behind after retry: %v", names) + } + }) + + t.Run("CloseRetriesFailedSpillRemoval", func(t *testing.T) { + requireDirectoryRemovalFailure(t) + dir := t.TempDir() + var buf bytes.Buffer + enc := newSpilledEncoder(t, &buf, dir, 4) + if err := os.Chmod(dir, 0o500); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = os.Chmod(dir, 0o700) }) + if err := enc.Close(); err == nil { + t.Fatal("expected Close to report the spill removal failure") + } + // The file itself is complete and must verify. + if err := ltx.NewDecoder(bytes.NewReader(buf.Bytes())).Verify(); err != nil { + t.Fatal(err) + } + if err := enc.Close(); err == nil { + t.Fatal("expected the second Close to keep reporting the pending removal") + } + if err := os.Chmod(dir, 0o700); err != nil { + t.Fatal(err) + } + if err := enc.Close(); err != nil { + t.Fatalf("Close() after fixing permissions = %v", err) + } + if names := dirEntries(t, dir); len(names) != 0 { + t.Fatalf("spill file left behind: %v", names) + } + if len(buf.Bytes()) == 0 || enc.Spilled() { + t.Fatal("expected spill reference cleared after successful removal") + } + }) + + t.Run("FailedCloseRetriesFailedSpillRemoval", func(t *testing.T) { + requireDirectoryRemovalFailure(t) + dir := t.TempDir() + w := &failingWriter{err: errInjected, failAfter: 1 << 30} + enc := newSpilledEncoder(t, w, dir, 4) + w.failAfter = w.buf.Len() + if err := os.Chmod(dir, 0o500); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = os.Chmod(dir, 0o700) }) + err := enc.Close() + if !errors.Is(err, errInjected) { + t.Fatalf("Close()=%v, want injected write failure", err) + } + var pathErr *os.PathError + if !errors.As(err, &pathErr) { + t.Fatalf("Close()=%v, want spill cleanup failure", err) + } + if names := dirEntries(t, dir); len(names) != 1 || !enc.Spilled() { + t.Fatalf("expected pending spill after failed removal: %v", names) + } + if err := os.Chmod(dir, 0o700); err != nil { + t.Fatal(err) + } + if err := enc.Close(); err != ltx.ErrEncoderAborted { + t.Fatalf("retried Close()=%v, want ErrEncoderAborted", err) + } + if names := dirEntries(t, dir); len(names) != 0 || enc.Spilled() { + t.Fatalf("spill file left behind after retry: %v", names) + } + }) + + t.Run("NoChecksumAndDeletionOutputMatch", func(t *testing.T) { + encode := func(dir string, commit uint32, pages int, flags uint32) []byte { + var buf bytes.Buffer + enc, err := ltx.NewEncoder(&buf) + if err != nil { + t.Fatal(err) + } + if dir != "" { + enc.SetSpillDir(dir) + enc.SetSpillThreshold(1) + } + hdr := ltx.Header{Version: ltx.Version, Flags: flags, PageSize: 512, Commit: commit, MinTXID: 1, MaxTXID: 1, Timestamp: 1000} + if err := enc.EncodeHeader(hdr); err != nil { + t.Fatal(err) + } + page := make([]byte, 512) + for i := 0; i < pages; i++ { + if err := enc.EncodePage(ltx.PageHeader{Pgno: uint32(i + 1)}, page); err != nil { + t.Fatal(err) + } + } + switch { + case commit == 0: + enc.SetPostApplyChecksum(ltx.ChecksumFlag) // deletion files carry the empty checksum + case flags<x.HeaderFlagNoChecksum == 0: + enc.SetPostApplyChecksum(ltx.ChecksumFlag | 2) + } + if err := enc.Close(); err != nil { + t.Fatal(err) + } + return buf.Bytes() + } + if a, b := encode("", 3, 3, ltx.HeaderFlagNoChecksum), encode(t.TempDir(), 3, 3, ltx.HeaderFlagNoChecksum); !bytes.Equal(a, b) { + t.Fatal("NoChecksum output differs with spill") + } + if a, b := encode("", 0, 0, 0), encode(t.TempDir(), 0, 0, 0); !bytes.Equal(a, b) { + t.Fatal("deletion output differs with spill") + } + }) +} + +func TestEncoder_PageWriteFailureAborts(t *testing.T) { + w := &failingWriter{err: errInjected, failAfter: 1 << 30} + enc, err := ltx.NewEncoder(w) + if err != nil { + t.Fatal(err) + } + if err := enc.EncodeHeader(ltx.Header{Version: ltx.Version, PageSize: 512, Commit: 4, MinTXID: 2, MaxTXID: 2, Timestamp: 1000, PreApplyChecksum: ltx.ChecksumFlag | 1}); err != nil { + t.Fatal(err) + } + page := make([]byte, 512) + if err := enc.EncodePage(ltx.PageHeader{Pgno: 1}, page); err != nil { + t.Fatal(err) + } + w.failAfter = w.buf.Len() + ltx.PageHeaderSize + 2 // fail inside the second page's size field + if err := enc.EncodePage(ltx.PageHeader{Pgno: 2}, page); err == nil { + t.Fatal("expected page write to fail") + } + w.failed, w.failAfter = false, 1<<30 + enc.SetPostApplyChecksum(ltx.ChecksumFlag | 2) + if err := enc.Close(); err != ltx.ErrEncoderAborted { + t.Fatalf("Close() after failed page write = %v, want ErrEncoderAborted", err) + } +} + +func TestEncoder_HeaderWriteFailureAborts(t *testing.T) { + w := &failingWriter{err: errInjected, failAfter: 10} + enc, err := ltx.NewEncoder(w) + if err != nil { + t.Fatal(err) + } + hdr := ltx.Header{Version: ltx.Version, PageSize: 512, Commit: 4, MinTXID: 2, MaxTXID: 2, Timestamp: 1000, PreApplyChecksum: ltx.ChecksumFlag | 1} + if err := enc.EncodeHeader(hdr); err == nil { + t.Fatal("expected header write to fail") + } + w.failed, w.failAfter = false, 1<<30 + if err := enc.EncodeHeader(hdr); err != ltx.ErrEncoderAborted { + t.Fatalf("retried EncodeHeader()=%v, want ErrEncoderAborted", err) + } +} diff --git a/ltx.go b/ltx.go index bd819d2..ee4882e 100644 --- a/ltx.go +++ b/ltx.go @@ -46,6 +46,11 @@ var ( ErrDecoderClosed = errors.New("ltx decoder closed") ErrEncoderClosed = errors.New("ltx encoder closed") + // ErrEncoderAborted is returned by an Encoder whose file cannot be + // completed: Cleanup ran before a successful Close, or an earlier write + // or Close failed. + ErrEncoderAborted = errors.New("ltx encoder aborted") + ErrNoChecksum = errors.New("no file checksum") ErrInvalidChecksumFormat = errors.New("invalid file checksum format") ErrChecksumMismatch = errors.New("file checksum mismatch") @@ -60,6 +65,11 @@ const ( statePage = "page" stateClose = "close" stateClosed = "closed" + // stateAborted is terminal: the encoder was cleaned up before a + // successful Close, or a Close failed part-way. Further writes and + // closes are refused so a structurally incomplete file cannot be + // finished with a valid checksum. + stateAborted = "aborted" ) // Pos represents the transactional position of a database.