diff --git a/db/etl/buffers.go b/db/etl/buffers.go index e7bf1750c23..88056d9b550 100644 --- a/db/etl/buffers.go +++ b/db/etl/buffers.go @@ -26,6 +26,7 @@ import ( "sort" "strconv" "sync" + "unsafe" "github.com/c2h5oh/datasize" @@ -44,7 +45,9 @@ const ( // BufIOSize - 128 pages | default is 1 page | increasing over `64 * 4096` doesn't show speedup on SSD/NVMe, but show speedup in cloud drives BufIOSize = 128 * 4096 - entryLocSize = 16 // sizeof(entryLoc): insertionOrder(4) + offset(4) + keyLen(4) + valLen(4) + entryLocSize = 4 // sizeof(entryLoc) + entryLocAlign = 4 // what a chunk's length has to be a multiple of, so the index lands aligned + entryHeaderSize = 4 // valLen, in front of the entry's bytes in its chunk ) // writeSortedEntries writes buffer entries to w in varint-length-prefixed format. @@ -69,7 +72,8 @@ func putValLen(dst []byte, valLen int32) { binary.NativeEndian.PutUint32(dst, uint32(valLen)) //nolint:gosec } -func writeSortedEntries(w io.Writer, entries []sortableBufferEntry, numBuf []byte) error { +func writeSortedEntries(w io.Writer, entries []sortableBufferEntry) error { + var numBuf [valLenSize]byte for _, entry := range entries { keyLen, valLen := int32(len(entry.key)), int32(len(entry.value)) //nolint:gosec if entry.key == nil { @@ -78,14 +82,14 @@ func writeSortedEntries(w io.Writer, entries []sortableBufferEntry, numBuf []byt if entry.value == nil { valLen = -1 } - putKeyLen(numBuf, keyLen) + putKeyLen(numBuf[:], keyLen) if _, err := w.Write(numBuf[:keyLenSize]); err != nil { return err } if _, err := w.Write(entry.key); err != nil { return err } - putValLen(numBuf, valLen) + putValLen(numBuf[:], valLen) if _, err := w.Write(numBuf[:valLenSize]); err != nil { return err } @@ -121,22 +125,17 @@ var ( ) const ( - // sortableBuffer stores key/value bytes in chunks of a power-of-two size, so - // entryLoc.offset can pack the chunk index with the offset inside the chunk - // and splitting the two is a shift and a mask. 1MB is also the least a - // collector can hold once it takes a chunk at all. - // A key must fit a chunk, since Sort slices it straight out of one. - maxKeyLen = 4096 - + // Each chunk carries the index of its own entries, so an entry addresses + // only bytes inside its chunk. 1MB is also the least a collector holds + // once it takes a chunk at all. dataChunkBits = 20 dataChunkSize = 1 << dataChunkBits // 1MB - // The chunk index takes what is left of a positive int32, so one buffer - // addresses at most maxDataChunks*dataChunkSize bytes (~2GB); nextChunk - // panics past that. NewSortableBuffer's MaxInt32 bound on optimalSize - // does not fully rule this out, since Put can grow the buffer past - // optimalSize before CheckFlushSize is checked. - maxDataChunks = math.MaxInt32>>dataChunkBits + 1 + // What is left of entryLoc's uint32 once the offset has dataChunkBits, + // less one for the bias. Sort slices a key out of its chunk, so only a + // value may outgrow one. Collect turns a longer key into an error before + // it reaches Put, which panics. + maxKeyLen = 1<<(32-dataChunkBits) - 2 ) // dataChunks are shared by all sortableBuffer instances: a buffer takes chunks as @@ -146,26 +145,21 @@ var dataChunks = sync.Pool{New: func() any { return &c }} -func getDataChunk() []byte { return *dataChunks.Get().(*[]byte) } - -// isPooledChunk reports whether c came from the pool. An oversized entry gets a -// private chunk instead, and handing that one back would let a later -// getDataChunk give an unrelated buffer a chunk of the wrong size. -func isPooledChunk(c []byte) bool { return len(c) == dataChunkSize } - -func putDataChunk(c []byte) { - if !isPooledChunk(c) { - return +func getDataChunk() *[]byte { return dataChunks.Get().(*[]byte) } +func putDataChunk(ref *[]byte) { + if ref != nil { + dataChunks.Put(ref) } - dataChunks.Put(&c) } type Buffer interface { // Put does copy `k` and `v` Put(k, v []byte) - // Next returns the entries in order, one goroutine at a time. The slices - // point into the buffer's own storage and must not be modified. Sort puts - // the cursor back at the first entry, so Sorting again re-reads. + // Next returns the entries in key order, one goroutine at a time. A + // buffer fills, then Sorts, then is read: Put after Sort, and Next or + // Write before it, both panic. Sort puts the cursor at the first entry, + // so Sorting again is how a buffer is read twice. The slices point into + // the buffer's own storage and must not be modified. Next() (k, v []byte, ok bool) Len() int Reset() @@ -176,6 +170,15 @@ type Buffer interface { CheckFlushSize() bool } +// panicIfUnsorted guards every read - Next and Write both - since a buffer +// read after a Put and before a Sort would hand back the previous run, which +// duplicates rows silently. +func panicIfUnsorted(unsorted bool) { + if unsorted { + panic("etl: buffer read before Sort") + } +} + type sortableBufferEntry struct { key []byte value []byte @@ -187,18 +190,17 @@ var ( _ Buffer = &oldestEntrySortableBuffer{} ) -// entryLoc stores the location of a key/value pair inside sortableBuffer. -// offset packs the chunk index and the offset inside that chunk: -// idx<>dataChunkBits) - 1 } //nolint:gosec +func (e entryLoc) offset() int32 { return int32(e) & (dataChunkSize - 1) } func NewSortableBuffer(bufferOptimalSize datasize.ByteSize) *sortableBuffer { if bufferOptimalSize.Bytes() > math.MaxInt32 { @@ -206,124 +208,191 @@ func NewSortableBuffer(bufferOptimalSize datasize.ByteSize) *sortableBuffer { } return &sortableBuffer{ optimalSize: int(bufferOptimalSize.Bytes()), + sortedN: -1, } } +// dataChunk holds entry bytes growing up from the front and the index of those +// entries growing down from the back, and is full when the two meet. So the +// index costs no allocation and a buffer's footprint is the chunks it holds. +type dataChunk struct { + buf []byte + // The pool stores *[]byte, so holding its own pointer lets Reset give the + // chunk back without boxing a slice header. nil for a chunk sized to a + // single oversized entry, which must never enter the pool. + ref *[]byte + end int32 // data grows up to here; only the chunk being filled moves it + entTop int32 // the index grows down to here +} + +func keyOf(buf []byte, e entryLoc) []byte { + if kLen := e.keyLen(); kLen > 0 { + off := e.offset() + entryHeaderSize + return buf[off : off+kLen] + } + return nil +} + +// entries views the chunk's index as entryLoc. Sort leaves it in key order; +// before that it runs newest-first, because it grows downward. +func (c *dataChunk) entries() []entryLoc { + n := (len(c.buf) - int(c.entTop)) / entryLocSize + if n == 0 { + return nil + } + // Aligned: chunkSizeFor only returns multiples of entryLocAlign, and + // entTop starts at len(buf) and moves by entryLocSize. + return unsafe.Slice((*entryLoc)(unsafe.Pointer(&c.buf[c.entTop])), n) +} + type sortableBuffer struct { - entries []entryLoc - at int // Next's cursor into entries + // A copy of the chunk being filled - always the last of chunks - so Put + // reaches its bytes without indexing the slice. completeCurrentChunk + // puts it back, and has to run before anything reads that chunk. + currentChunk dataChunk + n int + + chunks []dataChunk + + // Sort orders each chunk on its own, so reading in key order is a k-way + // merge over the chunks. + mrg merger + sortedN int // n as of the last Sort; -1 while unsorted - // Write's length scratch. w.Write takes an io.Writer, so a local array - // escapes and costs an allocation on every Write. - numBuf [valLenSize]byte - // chunks hold the key/value bytes. Growing by chunk instead of by one big - // slice keeps Put from re-allocating and copying everything collected so - // far. All chunks are dataChunkSize, except the private chunk an entry - // larger than that gets. cur is the chunk being filled. - chunks [][]byte - cur []byte - curBase int32 // packed location of cur's first byte: curIdx<= maxDataChunks { - panic(fmt.Sprintf("etl: sortableBuffer exceeded %d chunks", maxDataChunks)) +// chunkSizeFor returns the size of a chunk able to hold an entry of n bytes +// and its index slot, or 0 if no chunk can. Bounding n first matters: the +// rounding below wraps for n near math.MaxInt, and a wrapped size would come +// back small enough to pass for a pooled chunk. +func chunkSizeFor(n int) int { + if n > math.MaxInt32-entryLocSize-entryLocAlign { + return 0 } - if n > dataChunkSize { - b.cur = make([]byte, n) + if size := n + entryLocSize; size > dataChunkSize { + return (size + entryLocAlign - 1) &^ (entryLocAlign - 1) + } + return dataChunkSize +} + +// nextChunk starts a chunk able to hold an entry of n bytes and its index +// slot. An entry never straddles chunks, so Next hands out direct references. +func (b *sortableBuffer) nextChunk(n int) { + size := chunkSizeFor(n) + if size == 0 { + panic(fmt.Sprintf("etl: no chunk can hold an entry of %d bytes", n)) + } + var buf []byte + var ref *[]byte + if size == dataChunkSize { + ref = getDataChunk() + buf = *ref } else { - b.cur = getDataChunk() + buf = make([]byte, size) } - b.chunks = append(b.chunks, b.cur) - b.curBase = int32(len(b.chunks)-1) << dataChunkBits //nolint:gosec - b.curOff = 0 - b.chunkBytes += len(b.cur) + b.completeCurrentChunk() + b.currentChunk = dataChunk{buf: buf, ref: ref, entTop: int32(len(buf))} //nolint:gosec + b.chunks = append(b.chunks, b.currentChunk) + b.chunkBytes += len(buf) } -// entryData returns e's bytes: the key, immediately followed by the value. -func (b *sortableBuffer) entryData(e *entryLoc) []byte { - return b.chunks[e.offset>>dataChunkBits][e.offset&(dataChunkSize-1):] +// completeCurrentChunk puts the filled copy back where the readers look. It +// runs when the chunk has no room left, and when Sort ends the filling phase. +func (b *sortableBuffer) completeCurrentChunk() { + if len(b.chunks) == 0 { + return + } + b.chunks[len(b.chunks)-1] = b.currentChunk } // Put adds key and value to the buffer. These slices will not be accessed later, // so no copying is necessary func (b *sortableBuffer) Put(k, v []byte) { - e := entryLoc{ - keyLen: int32(len(k)), //nolint:gosec - valLen: int32(len(v)), //nolint:gosec - } - if k == nil { - e.keyLen = -1 - } - if v == nil { - e.valLen = -1 - } - // An entry with no bytes still takes one, so that no two entries share an - // offset - the sort orders duplicate keys by it. - n := max(len(k)+len(v), 1) - off := b.curOff - if int(off)+n > len(b.cur) { - b.nextChunk(n) - off = 0 - } - data := b.cur[off:] - copy(data, k) - copy(data[len(k):], v) - e.offset = b.curBase | off - b.curOff = off + int32(n) //nolint:gosec - b.entries = append(b.entries, e) -} - -// Size counts the stored bytes, the tails wasted by the chunks already filled, -// and entryLocSize bytes of metadata per entry. -func (b *sortableBuffer) Size() int { - return b.chunkBytes - (len(b.cur) - int(b.curOff)) + len(b.entries)*entryLocSize + kLen, vLen := len(k), len(v) + n := entryHeaderSize + kLen + vLen + off := int(b.currentChunk.end) + // One test for all three, so the fast path holds no call and Put keeps + // its arguments in registers. + if kLen > maxKeyLen || off+n+entryLocSize > int(b.currentChunk.entTop) || b.sortedN == b.n { + b.putSlow(k, v) + return + } + // As on main, a nil is stored apart from an empty: both lengths go to -1. + keyLen, valLen := int32(-1), int32(-1) + if k != nil { + keyLen = int32(kLen) //nolint:gosec + } + if v != nil { + valLen = int32(vLen) //nolint:gosec + } + // Capacity included, so the compiler can prove each copy's destination + // length and drop the min against the source. Worth 10% of Put. + data := b.currentChunk.buf[off : off+n : off+n] + binary.NativeEndian.PutUint32(data, uint32(valLen)) //nolint:gosec + b.currentChunk.entTop -= entryLocSize + binary.NativeEndian.PutUint32(b.currentChunk.buf[b.currentChunk.entTop:], uint32(makeEntryLoc(keyLen, int32(off)))) //nolint:gosec + b.currentChunk.end = int32(off + n) //nolint:gosec + b.n++ + copy(data[entryHeaderSize:entryHeaderSize+kLen], k) + copy(data[entryHeaderSize+kLen:], v) +} + +// putSlow handles what Put's single guard rejects: a Put after Sort, a key +// too long to index, and an entry the current chunk has no room for. +// nextChunk always leaves room, so the retry cannot come back here. +// +//go:noinline +func (b *sortableBuffer) putSlow(k, v []byte) { + if b.sortedN == b.n { + panic("etl: Put after Sort") + } + if len(k) > maxKeyLen { + panic(fmt.Sprintf("etl: key of %d bytes exceeds %d", len(k), maxKeyLen)) + } + b.nextChunk(entryHeaderSize + len(k) + len(v)) + b.Put(k, v) } -func (b *sortableBuffer) Len() int { - return len(b.entries) +// Size counts every chunk taken, less what is free in the one being filled. +// The entry index lives inside the chunks, so it is counted. +func (b *sortableBuffer) Size() int { + return b.chunkBytes - int(b.currentChunk.entTop-b.currentChunk.end) } +func (b *sortableBuffer) Len() int { return b.n } + +// Next returns the entry the read cursor sits on and moves it along. The +// buffer carries the merge state, so no two goroutines may read at once. func (b *sortableBuffer) Next() ([]byte, []byte, bool) { - if b.at >= len(b.entries) { + panicIfUnsorted(b.sortedN != b.n) + buf, e, ok := b.mrg.next() + if !ok { return nil, nil, false } - e := &b.entries[b.at] - b.at++ - kLen, vLen := int(e.keyLen), int(e.valLen) + data := buf[e.offset():] + kLen := e.keyLen() + vLen := int32(binary.NativeEndian.Uint32(data)) //nolint:gosec + data = data[entryHeaderSize:] var key, val []byte - if kLen == 0 { - key = []byte{} - } - if vLen == 0 { - val = []byte{} - } - if kLen <= 0 && vLen <= 0 { - return key, val, true - } - data := b.entryData(e) - if kLen > 0 { + if kLen >= 0 { key = data[:kLen:kLen] data = data[kLen:] } - if vLen > 0 { + if vLen >= 0 { val = data[:vLen:vLen] } return key, val, true } -// Prealloc sizes the entries slice. predictDataSize only reserves room in the -// chunks slice for the chunk pointers; the chunks themselves are still taken -// one at a time, which is what keeps an idle buffer from holding its peak. -func (b *sortableBuffer) Prealloc(predictKeysAmount, predictDataSize int) Buffer { - if cap(b.entries) < predictKeysAmount { - b.entries = make([]entryLoc, 0, predictKeysAmount) - } +// Prealloc only reserves room for the chunk headers. The chunks come from +// their pool one at a time, so an idle buffer holds nothing. +func (b *sortableBuffer) Prealloc(_, predictDataSize int) Buffer { if n := predictDataSize/dataChunkSize + 1; cap(b.chunks) < n { b.chunks = slices.Grow(b.chunks, n) } @@ -331,41 +400,58 @@ func (b *sortableBuffer) Prealloc(predictKeysAmount, predictDataSize int) Buffer } func (b *sortableBuffer) Reset() { - b.at = 0 - b.entries = b.entries[:0] - for i, c := range b.chunks { - putDataChunk(c) - b.chunks[i] = nil - } + // The cursors and currentChunk alias the chunks, so drop them before the + // pool hands those chunks to another buffer. + b.mrg.release() + b.currentChunk = dataChunk{} + for i := range b.chunks { + putDataChunk(b.chunks[i].ref) + } + clear(b.chunks) b.chunks = b.chunks[:0] - b.cur, b.curBase, b.curOff = nil, 0, 0 - b.chunkBytes = 0 + b.n, b.chunkBytes = 0, 0 + b.sortedN = -1 } + func (b *sortableBuffer) SizeLimit() int { return b.optimalSize } + +// Sort orders each chunk on its own, so it stays inside 1MB however large the +// buffer is; reading the buffer back merges the runs. It also puts the read +// cursor at the first entry, so Sorting again is how a buffer is read twice. func (b *sortableBuffer) Sort() { - b.at = 0 - chunks := b.chunks - // Key extraction stays inside cmp: pdqsortCmpFunc calls the comparator - // indirectly, so a separate closure never inlines and costs a call per key. - cmp := func(x, y entryLoc) int { - var xk, yk []byte - if x.keyLen > 0 { - off := x.offset & (dataChunkSize - 1) - xk = chunks[x.offset>>dataChunkBits][off : off+x.keyLen] - } - if y.keyLen > 0 { - off := y.offset & (dataChunkSize - 1) - yk = chunks[y.offset>>dataChunkBits][off : off+y.keyLen] + if b.sortedN != b.n { + b.completeCurrentChunk() + for i := range b.chunks { + b.chunks[i].sort() } - if c := bytes.Compare(xk, yk); c != 0 { - return c - } - return int(x.offset - y.offset) // StableSort: offsets rise with insertion order + b.sortedN = b.n } - if slices.IsSortedFunc(b.entries, cmp) { + b.mrg.rewind(b.chunks) +} + +// sort orders the chunk's index by the keys it holds. +func (c *dataChunk) sort() { + ents := c.entries() + if len(ents) < 2 { return } - slices.SortFunc(b.entries, cmp) + buf := c.buf + cmp := func(x, y entryLoc) int { + if r := bytes.Compare(keyOf(buf, x), keyOf(buf, y)); r != 0 { + return r + } + return int(x.offset() - y.offset()) // StableSort: offsets rise with insertion order + } + // The index grows downward, so ascending keys arrive reversed. pdqsort + // spots that too, but only after sampling for a pivot. Equal keys leave + // the offsets descending, which the byte compare alone already accepts. + for j := 1; j < len(ents); j++ { + if bytes.Compare(keyOf(buf, ents[j-1]), keyOf(buf, ents[j])) < 0 { + slices.SortFunc(ents, cmp) + return + } + } + slices.Reverse(ents) } func (b *sortableBuffer) CheckFlushSize() bool { @@ -373,37 +459,39 @@ func (b *sortableBuffer) CheckFlushSize() bool { } func (b *sortableBuffer) Write(w io.Writer) error { + panicIfUnsorted(b.sortedN != b.n) numBuf := b.numBuf[:] - for i := range b.entries { - e := &b.entries[i] - kLen, vLen := int(e.keyLen), int(e.valLen) - var data []byte - if kLen > 0 || vLen > 0 { - data = b.entryData(e) + for { + k, v, ok := b.Next() + if !ok { + return nil + } + keyLen, valLen := int32(len(k)), int32(len(v)) //nolint:gosec + if k == nil { + keyLen = -1 + } + if v == nil { + valLen = -1 } - // write key - putKeyLen(numBuf, e.keyLen) + putKeyLen(numBuf, keyLen) if _, err := w.Write(numBuf[:keyLenSize]); err != nil { return err } - if kLen > 0 { - if _, err := w.Write(data[:kLen]); err != nil { + if len(k) > 0 { + if _, err := w.Write(k); err != nil { return err } - data = data[kLen:] } - // write value - putValLen(numBuf, e.valLen) + putValLen(numBuf, valLen) if _, err := w.Write(numBuf[:valLenSize]); err != nil { return err } - if vLen > 0 { - if _, err := w.Write(data[:vLen]); err != nil { + if len(v) > 0 { + if _, err := w.Write(v); err != nil { return err } } } - return nil } func NewAppendBuffer(bufferOptimalSize datasize.ByteSize) *appendSortableBuffer { @@ -417,7 +505,8 @@ func NewAppendBuffer(bufferOptimalSize datasize.ByteSize) *appendSortableBuffer type appendSortableBuffer struct { entries map[string][]byte sortedBuf []sortableBufferEntry - at int // Next's cursor into sortedBuf + at int + unsorted bool // sortedBuf does not hold what entries does size int optimalSize int } @@ -429,6 +518,7 @@ func (b *appendSortableBuffer) Put(k, v []byte) { } b.size += len(v) b.entries[string(k)] = append(stored, v...) + b.unsorted = true } func (b *appendSortableBuffer) Size() int { return b.size } @@ -439,7 +529,11 @@ func (b *appendSortableBuffer) Len() int { } func (b *appendSortableBuffer) Sort() { - b.sortedBuf, b.at = b.sortedBuf[:0], 0 + if !b.unsorted { + b.at = 0 // already flattened; Sort still positions the cursor + return + } + b.sortedBuf, b.at, b.unsorted = b.sortedBuf[:0], 0, false if cap(b.sortedBuf) < len(b.entries) { b.sortedBuf = make([]sortableBufferEntry, 0, len(b.entries)) } @@ -458,6 +552,7 @@ func (b *appendSortableBuffer) Swap(i, j int) { } func (b *appendSortableBuffer) Next() ([]byte, []byte, bool) { + panicIfUnsorted(b.unsorted) if b.at >= len(b.sortedBuf) { return nil, nil, false } @@ -467,13 +562,17 @@ func (b *appendSortableBuffer) Next() ([]byte, []byte, bool) { } func (b *appendSortableBuffer) Reset() { - b.sortedBuf, b.at = nil, 0 + b.sortedBuf = nil + b.at, b.unsorted = 0, false b.entries = make(map[string][]byte) b.size = 0 } func (b *appendSortableBuffer) Prealloc(predictKeysAmount, predictDataSize int) Buffer { b.entries = make(map[string][]byte, predictKeysAmount) // maps have no cap(), always recreate + // The new map holds nothing, so neither does the run flattened out of the + // old one, and neither does what Size reports. + b.sortedBuf, b.at, b.unsorted, b.size = b.sortedBuf[:0], 0, false, 0 if cap(b.sortedBuf) < predictKeysAmount { b.sortedBuf = make([]sortableBufferEntry, 0, predictKeysAmount) } @@ -481,8 +580,8 @@ func (b *appendSortableBuffer) Prealloc(predictKeysAmount, predictDataSize int) } func (b *appendSortableBuffer) Write(w io.Writer) error { - var numBuf [valLenSize]byte - return writeSortedEntries(w, b.sortedBuf, numBuf[:]) + panicIfUnsorted(b.unsorted) + return writeSortedEntries(w, b.sortedBuf) } func (b *appendSortableBuffer) CheckFlushSize() bool { @@ -500,7 +599,8 @@ func NewOldestEntryBuffer(bufferOptimalSize datasize.ByteSize) *oldestEntrySorta type oldestEntrySortableBuffer struct { entries map[string][]byte sortedBuf []sortableBufferEntry - at int // Next's cursor into sortedBuf + at int + unsorted bool // sortedBuf does not hold what entries does size int optimalSize int } @@ -514,6 +614,7 @@ func (b *oldestEntrySortableBuffer) Put(k, v []byte) { b.size += len(k)*2 + len(v) b.entries[string(k)] = bytes.Clone(v) + b.unsorted = true } func (b *oldestEntrySortableBuffer) Size() int { return b.size } @@ -524,7 +625,11 @@ func (b *oldestEntrySortableBuffer) Len() int { } func (b *oldestEntrySortableBuffer) Sort() { - b.sortedBuf, b.at = b.sortedBuf[:0], 0 + if !b.unsorted { + b.at = 0 // already flattened; Sort still positions the cursor + return + } + b.sortedBuf, b.at, b.unsorted = b.sortedBuf[:0], 0, false if cap(b.sortedBuf) < len(b.entries) { b.sortedBuf = make([]sortableBufferEntry, 0, len(b.entries)) } @@ -543,6 +648,7 @@ func (b *oldestEntrySortableBuffer) Swap(i, j int) { } func (b *oldestEntrySortableBuffer) Next() ([]byte, []byte, bool) { + panicIfUnsorted(b.unsorted) if b.at >= len(b.sortedBuf) { return nil, nil, false } @@ -552,13 +658,17 @@ func (b *oldestEntrySortableBuffer) Next() ([]byte, []byte, bool) { } func (b *oldestEntrySortableBuffer) Reset() { - b.sortedBuf, b.at = nil, 0 + b.sortedBuf = nil + b.at, b.unsorted = 0, false b.entries = make(map[string][]byte) b.size = 0 } func (b *oldestEntrySortableBuffer) Prealloc(predictKeysAmount, predictDataSize int) Buffer { b.entries = make(map[string][]byte, predictKeysAmount) // maps have no cap(), always recreate + // The new map holds nothing, so neither does the run flattened out of the + // old one, and neither does what Size reports. + b.sortedBuf, b.at, b.unsorted, b.size = b.sortedBuf[:0], 0, false, 0 if cap(b.sortedBuf) < predictKeysAmount { b.sortedBuf = make([]sortableBufferEntry, 0, predictKeysAmount) } @@ -566,8 +676,8 @@ func (b *oldestEntrySortableBuffer) Prealloc(predictKeysAmount, predictDataSize } func (b *oldestEntrySortableBuffer) Write(w io.Writer) error { - var numBuf [valLenSize]byte - return writeSortedEntries(w, b.sortedBuf, numBuf[:]) + panicIfUnsorted(b.unsorted) + return writeSortedEntries(w, b.sortedBuf) } func (b *oldestEntrySortableBuffer) CheckFlushSize() bool { diff --git a/db/etl/collector.go b/db/etl/collector.go index 264a1c91ffc..6157b93521e 100644 --- a/db/etl/collector.go +++ b/db/etl/collector.go @@ -93,15 +93,9 @@ func (c *Collector) SortAndFlushInBackground(v bool) *Collector { return c } -// errKeyTooLong is out of line so that building it does not cost the caller -// its inlining budget on a path that never runs. -// -//go:noinline -func errKeyTooLong(logPrefix string, n int) error { - return fmt.Errorf("%s: key of %d bytes exceeds %d", logPrefix, n, maxKeyLen) -} - func (c *Collector) extractNextFunc(originalK, k []byte, v []byte) error { + // sortableBuffer.Put panics past this, and Collect sits under enough + // error-returning callers that a long key should fail the stage instead. if len(k) > maxKeyLen { return errKeyTooLong(c.logPrefix, len(k)) } @@ -117,6 +111,14 @@ func (c *Collector) extractNextFunc(originalK, k []byte, v []byte) error { } // Collect does copy `k` and `v` +// errKeyTooLong is out of line so that building it does not cost the caller +// its inlining budget on a path that never runs. +// +//go:noinline +func errKeyTooLong(logPrefix string, n int) error { + return fmt.Errorf("%s: key of %d bytes exceeds %d", logPrefix, n, maxKeyLen) +} + func (c *Collector) Collect(k, v []byte) error { return c.extractNextFunc(k, k, v) } diff --git a/db/etl/etl_test.go b/db/etl/etl_test.go index cdbc8b21a52..75ece88ebaa 100644 --- a/db/etl/etl_test.go +++ b/db/etl/etl_test.go @@ -25,6 +25,7 @@ import ( "errors" "fmt" "io" + "math" "os" "os/exec" "path/filepath" @@ -34,6 +35,8 @@ import ( "sync/atomic" "testing" + "unsafe" + "github.com/c2h5oh/datasize" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -131,11 +134,12 @@ func TestWriteAndReadBufferEntry(t *testing.T) { entries := make([]sortableBufferEntry, 100) for i := range entries { - entries[i].key = fmt.Appendf(nil, "key-%d", i) + entries[i].key = fmt.Appendf(nil, "key-%03d", i) entries[i].value = fmt.Appendf(nil, "value-%d", i) b.Put(entries[i].key, entries[i].value) } + b.Sort() // a buffer fills, then Sorts, then is read if err := b.Write(buffer); err != nil { t.Error(err) } @@ -495,11 +499,12 @@ func TestReuseCollectorAfterLoad(t *testing.T) { require.Equal(t, 1, see) c.Close() - // buffer state resets for reuse: entries keep their cap, chunks are cleared + // buffer state resets for reuse: chunks go back to the pool and take the + // entry index with them, the chunk header slice keeps its cap require.Empty(t, buf.chunks) - require.Empty(t, buf.entries) + require.Zero(t, buf.Len()) require.Zero(t, buf.Size()) - require.NotZero(t, cap(buf.entries)) + require.NotZero(t, cap(buf.chunks)) // teset that no data visible see = 0 @@ -1269,8 +1274,8 @@ func BenchmarkSortableBufferPutOnly(b *testing.B) { } // BenchmarkSortableBufferRead reads a sorted buffer end to end. The sort runs -// once before the loop, so an iteration is the read itself - it re-reads one -// buffer, where a collector reads one only once. +// once before the loop, so an iteration is the cursor rewind plus the read - +// and it re-reads one buffer, where a collector reads one only once. func BenchmarkSortableBufferRead(b *testing.B) { const keyLen = 32 const valLen = 64 @@ -1302,10 +1307,10 @@ func BenchmarkSortableBufferRead(b *testing.B) { binary.BigEndian.PutUint64(val, uint64(i)) buf.Put(key, val) } - buf.Sort() + buf.Sort() // once: an iteration is the rewind plus the read b.ResetTimer() for b.Loop() { - buf.at = 0 // rewind; Sort would re-check sortedness first + buf.Sort() for _, _, ok := buf.Next(); ok; _, _, ok = buf.Next() { } } @@ -1447,7 +1452,7 @@ func BenchmarkMemoryDataProviderNext(b *testing.B) { b.Run(name+"/Next", func(b *testing.B) { b.ReportAllocs() for i := 0; i < b.N; i++ { - buf.at = 0 + buf.Sort() p := &memoryDataProvider{buffer: buf} for { _, _, err := p.Next() @@ -1464,7 +1469,7 @@ func BenchmarkMemoryDataProviderNext(b *testing.B) { b.Run(name+"/Buffer", func(b *testing.B) { b.ReportAllocs() for i := 0; i < b.N; i++ { - buf.at = 0 + buf.Sort() for _, _, ok := buf.Next(); ok; _, _, ok = buf.Next() { } } @@ -1603,9 +1608,10 @@ func TestSortableBufferChunks(t *testing.T) { require.Equal(t, entries, buf.Len()) require.Greater(t, len(buf.chunks), 1, "data must be split into chunks") for i, c := range buf.chunks { - require.Equal(t, dataChunkSize, cap(c), "chunk %d", i) + require.Equal(t, dataChunkSize, cap(c.buf), "chunk %d", i) } + buf.Sort() for i, e := range drainBuffer(buf) { binary.BigEndian.PutUint64(key, uint64(i)) require.Equal(t, key, e.key, "entry %d", i) @@ -1648,12 +1654,14 @@ func TestSortableBufferOversizedEntry(t *testing.T) { buf.Put([]byte{0x02}, big) buf.Put([]byte{0x03}, []byte("after")) + buf.Sort() want := drainBuffer(buf) require.Equal(t, []byte{0x02}, want[1].key) require.Equal(t, big, want[1].value) require.Equal(t, []byte{0x03}, want[2].key) require.Equal(t, []byte("after"), want[2].value) + buf.Sort() // drainBuffer consumed the cursor; Write needs it back w := bytes.NewBuffer(nil) require.NoError(t, buf.Write(w)) m := &mmapBytesReader{data: w.Bytes()} @@ -1684,30 +1692,58 @@ func TestSortableBufferResetReleasesChunks(t *testing.T) { require.Zero(t, buf.Len()) buf.Put([]byte{0x01}, []byte("reused")) - got2 := drainBuffer(buf) - require.Equal(t, []byte{0x01}, got2[0].key) - require.Equal(t, []byte("reused"), got2[0].value) + buf.Sort() + got := drainBuffer(buf) + require.Equal(t, []byte{0x01}, got[0].key) + require.Equal(t, []byte("reused"), got[0].value) } -// TestPutDataChunkRejectsOversized: an entry's private chunk (bigger than -// dataChunkSize) must never enter the shared pool — a later getDataChunk -// handing it out under a normal chunk index would corrupt an unrelated buffer. -func TestPutDataChunkRejectsOversized(t *testing.T) { +// TestChunkSizeFor: a chunk's entTop is an int32, so the size has to fit one. +// n is bounded before the rounding, which would otherwise wrap and come back +// small enough to pass for a pooled chunk. +func TestChunkSizeFor(t *testing.T) { + const biggest = math.MaxInt32 - entryLocSize - entryLocAlign for _, tc := range []struct { - name string - length int - pooled bool + name string + n int + size int }{ - {"short", dataChunkSize - 1, false}, - {"exact", dataChunkSize, true}, - {"oversized", dataChunkSize + 7, false}, + {"empty", 0, dataChunkSize}, + {"fills a chunk", dataChunkSize - entryLocSize, dataChunkSize}, + {"one byte over", dataChunkSize - entryLocSize + 1, dataChunkSize + entryLocSize}, + {"largest that fits", biggest, (biggest + entryLocSize + entryLocAlign - 1) &^ (entryLocAlign - 1)}, + {"one past", biggest + 1, 0}, + {"two gigabytes", 1 << 31, 0}, + {"maxint", math.MaxInt, 0}, } { t.Run(tc.name, func(t *testing.T) { - require.Equal(t, tc.pooled, isPooledChunk(make([]byte, tc.length))) + got := chunkSizeFor(tc.n) + require.Equal(t, tc.size, got) + require.LessOrEqual(t, got, math.MaxInt32, "a chunk must fit entTop") + // entries() views the chunk tail as []entryLoc through + // unsafe.Slice, so every size it can hand out must be aligned. + require.Zero(t, got%entryLocAlign, "chunk size must fit whole index slots") }) } } +// TestPutDataChunkRejectsOversized: an entry's private chunk must never enter +// the shared pool - a later getDataChunk handing it out as a normal chunk +// would corrupt an unrelated buffer. Only a chunk the pool gave out carries a +// ref, so only those can go back. +func TestPutDataChunkRejectsOversized(t *testing.T) { + buf := NewSortableBuffer(256 * datasize.MB) + defer buf.Reset() + buf.Put([]byte{0x01}, []byte("small")) + buf.Put([]byte{0x02}, bytes.Repeat([]byte{0xCD}, dataChunkSize+7)) + require.Len(t, buf.chunks, 2) + + require.NotNil(t, buf.chunks[0].ref, "a pooled chunk goes back") + require.Equal(t, dataChunkSize, len(buf.chunks[0].buf)) + require.Nil(t, buf.chunks[1].ref, "an oversized chunk has no ref, so it cannot") + require.Greater(t, len(buf.chunks[1].buf), dataChunkSize) +} + // disposeProbe records whether the collector still owned its data chunks when // the provider was disposed. type disposeProbe struct { @@ -1736,9 +1772,8 @@ func TestCloseDisposesProvidersBeforeBuffer(t *testing.T) { require.True(t, probe.sawOwnChunks, "buffer was recycled before its providers were disposed") } -// TestSortableBufferAllEmptyEntries: entries whose key and value are both -// zero-length keep insertion order too, which they only can if each still has -// an offset of its own. nil and empty stay distinguishable. +// TestSortableBufferAllEmptyEntries: zero-length keys and values keep +// insertion order too, and nil stays distinguishable from empty. func TestSortableBufferAllEmptyEntries(t *testing.T) { buf := NewSortableBuffer(256 * 1024) @@ -1750,10 +1785,12 @@ func TestSortableBufferAllEmptyEntries(t *testing.T) { buf.Put([]byte{}, nil) seen := map[int32]bool{} - for i := range buf.entries { - off := buf.entries[i].offset - require.False(t, seen[off], "entry %d reuses offset %d, so Sort cannot order it", i, off) - seen[off] = true + for i := range buf.chunks { + for _, e := range buf.chunks[i].entries() { + require.False(t, seen[e.offset()], + "Sort orders equal keys by offset, so every entry needs one of its own") + seen[e.offset()] = true + } } buf.Sort() @@ -1769,16 +1806,66 @@ func TestSortableBufferAllEmptyEntries(t *testing.T) { assert.Equal(t, []byte("last"), entries[4].value) } +// TestSortableBufferChunkBoundary: an entry never straddles a chunk. The sizes +// stop the fill at a different offset in the last chunk each time. +func TestSortableBufferChunkBoundary(t *testing.T) { + const keyLen = 8 + // entryHeaderSize+keyLen+valLen divides dataChunkSize for 4, 20, 52, 116 and + // 4084, so those land an entry's last byte exactly on a chunk boundary. + for _, valLen := range []int{0, 1, 4, 7, 20, 52, 63, 64, 116, 4084, 4095, 4096} { + t.Run(fmt.Sprintf("val%d", valLen), func(t *testing.T) { + buf := NewSortableBuffer(64 * 1024 * 1024) + entrySize := entryHeaderSize + keyLen + valLen + count := 2*dataChunkSize/entrySize + 3 + + key := make([]byte, keyLen) + for i := range count { // descending, so Sort has real work to do + binary.BigEndian.PutUint64(key, uint64(count-1-i)) //nolint:gosec + buf.Put(key, bytes.Repeat([]byte{byte(i)}, valLen)) + } + require.Greater(t, len(buf.chunks), 1, "test must cross a chunk boundary") + + buf.Sort() + require.Equal(t, count, buf.Len()) + for i, e := range drainBuffer(buf) { + binary.BigEndian.PutUint64(key, uint64(i)) //nolint:gosec + require.Equal(t, key, e.key, "entry %d", i) + require.Equal(t, bytes.Repeat([]byte{byte(count - 1 - i)}, valLen), e.value, "entry %d", i) + } + }) + } +} + +// TestSortableBufferRejectsOversizedKey: Sort reads keys straight out of a +// chunk, so a key must fit one. Values may still be any size. +func TestSortableBufferRejectsOversizedKey(t *testing.T) { + buf := NewSortableBuffer(256 * datasize.MB) + + require.Panics(t, func() { buf.Put(make([]byte, maxKeyLen+1), []byte("v")) }) + + // maxKeyLen is what keeps keyLen inside entryLoc, so read the edge back. + buf.Put(make([]byte, maxKeyLen), []byte("v")) + buf.Put([]byte{0x01}, make([]byte, dataChunkSize+7)) + buf.Put(nil, nil) + require.Equal(t, 3, buf.Len()) + + // Sorted: the nil key, then the all-zero key of maxKeyLen, then 0x01. + buf.Sort() + got := drainBuffer(buf) + require.Len(t, got[1].key, maxKeyLen) + require.Nil(t, got[0].key) +} + // TestSortableBufferStableSortAcrossChunks: duplicate keys spread over several -// data chunks are the case the offset tie-break has to get right, since the -// packed offset carries the chunk index in its high bits. +// chunks are the case the per-chunk sort plus merge can reorder. func TestSortableBufferStableSortAcrossChunks(t *testing.T) { buf := NewSortableBuffer(256 * datasize.MB) dupKey := []byte{0x05, 0x05, 0x05, 0x05, 0x05, 0x05, 0x05, 0x05} pad := make([]byte, 4096) // few entries per chunk, so the dups spread out val := make([]byte, 8) - const dups = 1200 + const dups = 2000 + for i := range dups { binary.BigEndian.PutUint64(val, uint64(i)) //nolint:gosec buf.Put(dupKey, val) @@ -1789,6 +1876,7 @@ func TestSortableBufferStableSortAcrossChunks(t *testing.T) { require.Greater(t, len(buf.chunks), 4, "dups must spread over several chunks") buf.Sort() + seq := 0 for i, e := range drainBuffer(buf) { if !bytes.Equal(e.key, dupKey) { @@ -1800,18 +1888,6 @@ func TestSortableBufferStableSortAcrossChunks(t *testing.T) { require.Equal(t, dups, seq) } -// TestCollectRejectsOversizedKey: Sort slices a key straight out of a chunk, -// so a key has to fit one. Collect sits under Load and the stage loop, which -// return errors, so it fails the stage rather than the process. -func TestCollectRejectsOversizedKey(t *testing.T) { - c := NewCollector(t.Name(), t.TempDir(), NewSortableBuffer(1*datasize.MB), log.New()) - defer c.Close() - require.NoError(t, c.Collect(make([]byte, maxKeyLen), []byte("v"))) - err := c.Collect(make([]byte, maxKeyLen+1), []byte("v")) - require.Error(t, err) - require.Contains(t, err.Error(), "exceeds") -} - // BenchmarkSortableBufferPutOnlyCold fills a fresh buffer without Prealloc. // Only the first iteration misses the chunk pool - Reset hands the chunks // straight back - so this measures a new buffer against a warm pool, not a @@ -1883,8 +1959,194 @@ func BenchmarkSortableBufferWrite(b *testing.B) { } } -// TestSortableBufferReadIsAllocFree: reading a sorted buffer must not allocate -// per entry - the slices point into the chunks the buffer already holds. +// TestSortableBufferMergesChunks: ascending keys leave the chunks ordered end +// to end, descending keys interleave them. Both must read back in key order. +func TestSortableBufferMergesChunks(t *testing.T) { + const count = 40_000 // several chunks at 4+8+64 bytes an entry + for _, ascending := range []bool{true, false} { + t.Run(fmt.Sprintf("ascending%v", ascending), func(t *testing.T) { + buf := NewSortableBuffer(256 * datasize.MB) + key := make([]byte, 8) + val := make([]byte, 64) + for i := range count { + n := uint64(i) //nolint:gosec + if !ascending { + n = uint64(count - 1 - i) //nolint:gosec + } + binary.BigEndian.PutUint64(key, n) + binary.BigEndian.PutUint64(val, n) + buf.Put(key, val) + } + require.Greater(t, len(buf.chunks), 2, "must cross chunk boundaries") + + buf.Sort() + require.Equal(t, ascending, buf.mrg.concat, "ascending keys must skip the heap") + + for i, e := range drainBuffer(buf) { + binary.BigEndian.PutUint64(key, uint64(i)) //nolint:gosec + require.Equal(t, key, e.key, "entry %d", i) + require.Equal(t, uint64(i), binary.BigEndian.Uint64(e.value), "entry %d", i) //nolint:gosec + } + // a second pass must restart cleanly + buf.Sort() + binary.BigEndian.PutUint64(key, 0) + k, _, ok := buf.Next() + require.True(t, ok) + require.Equal(t, key, k) + }) + } +} + +func TestSortableBufferPutAfterSort(t *testing.T) { + buf := NewSortableBuffer(1 * datasize.MB) + defer buf.Reset() + for _, k := range []byte{5, 3, 9, 1} { + buf.Put([]byte{k}, []byte{k}) + } + buf.Sort() + // Sort permutes each chunk's index, so a later Put has no insertion order + // left to fall back on. Filling and reading are separate phases. + require.Panics(t, func() { buf.Put([]byte{7}, []byte{7}) }) +} + +// TestBufferNextBeforeSort: reading a buffer that was never sorted must be +// loud. Returning the previous run instead would duplicate rows into whatever +// the load feeds. +func TestBufferNextBeforeSort(t *testing.T) { + for _, bt := range allBufferTypes { + t.Run(bt.name, func(t *testing.T) { + buf := bt.new() + buf.Put([]byte{3}, []byte("c")) + require.Panics(t, func() { buf.Next() }, "Next before any Sort") + require.Panics(t, func() { _ = buf.Write(io.Discard) }, "Write before any Sort") + + buf.Sort() + require.Len(t, drainBuffer(buf), 1) + }) + } +} + +// TestSortableBufferWriteAfterPartialRead: Write drives the same cursor Next +// does, so a Sort has to come between them - otherwise a flush would silently +// drop the entries already read and mergeSortFiles would panic on the short +// file. +func TestSortableBufferWriteAfterPartialRead(t *testing.T) { + buf := NewSortableBuffer(1 * datasize.MB) + defer buf.Reset() + for i := range 10 { + buf.Put([]byte{byte(i)}, []byte{byte(i)}) + } + buf.Sort() + buf.Next() + buf.Next() + + buf.Sort() // what sortAndFlush does before it writes + w := bytes.NewBuffer(nil) + require.NoError(t, buf.Write(w)) + m := &mmapBytesReader{data: w.Bytes()} + for i := range 10 { + k, err := readKeyField(m) + require.NoError(t, err, "entry %d", i) + require.Equal(t, []byte{byte(i)}, k) + _, err = readValField(m) + require.NoError(t, err) + } + _, err := readKeyField(m) + require.Equal(t, io.EOF, err) +} + +// TestChunksInOrderAcrossEmptyChunk: an empty chunk must not hide the pair on +// either side of it. Skipping both comparisons let concat mode read chunks +// that were out of order, which emits unsorted entries. +func TestChunksInOrderAcrossEmptyChunk(t *testing.T) { + buf := NewSortableBuffer(1 * datasize.MB) + defer buf.Reset() + buf.Put([]byte{9}, []byte("hi")) + buf.nextChunk(0) // an empty chunk between two that are out of order + buf.nextChunk(entryHeaderSize + 2) + buf.Put([]byte{1}, []byte("lo")) + buf.Sort() + + require.False(t, buf.mrg.concat, "chunk 0 sorts after chunk 2, so the heap is needed") + var got []byte + for _, e := range drainBuffer(buf) { + got = append(got, e.key[0]) + } + require.Equal(t, []byte{1, 9}, got) +} + +// TestOversizedChunkEntryIndex: entries() views a chunk's tail as []entryLoc +// through unsafe.Slice, and a private chunk is sized by chunkSizeFor rather +// than being a round 1MB. Run under -race for checkptr's alignment check. +func TestOversizedChunkEntryIndex(t *testing.T) { + buf := NewSortableBuffer(256 * datasize.MB) + defer buf.Reset() + + big := bytes.Repeat([]byte{0xCD}, dataChunkSize+7) + buf.Put([]byte{0x02}, big) // takes a chunk of its own + buf.Put([]byte{0x01}, []byte("nx")) // and the next entry starts a fresh one + + c := &buf.chunks[0] + require.Nil(t, c.ref, "chunk 0 must be the private one") + require.Zero(t, uintptr(unsafe.Pointer(&c.buf[c.entTop]))%entryLocSize) + + ents := c.entries() + require.Len(t, ents, 1) + require.Equal(t, []byte{0x02}, keyOf(c.buf, ents[0])) + + buf.Sort() + got := drainBuffer(buf) + require.Equal(t, []byte{0x01}, got[0].key) + require.Equal(t, big, got[1].value) +} + +// TestMergerMatchesReferenceSort: the merge is a heap over per-chunk runs, so +// build one from random keys spread over many chunks and check it against a +// plain sort of the same pairs, duplicates and insertion order included. +func TestMergerMatchesReferenceSort(t *testing.T) { + buf := NewSortableBuffer(256 * datasize.MB) + defer buf.Reset() + + const count = 60_000 + pad := make([]byte, 200) // few entries per chunk, so the runs interleave + want := make([]sortableBufferEntry, 0, count) + key := make([]byte, 8) + for i := range count { + binary.BigEndian.PutUint64(key, uint64(i)*6364136223846793005%1000) //nolint:gosec + val := make([]byte, len(pad)) + binary.BigEndian.PutUint64(val, uint64(i)) //nolint:gosec + want = append(want, sortableBufferEntry{key: bytes.Clone(key), value: val}) + buf.Put(key, val) + } + require.Greater(t, len(buf.chunks), 8, "keys must spread over many chunks") + + slices.SortStableFunc(want, func(a, b sortableBufferEntry) int { + return bytes.Compare(a.key, b.key) + }) + + buf.Sort() + require.False(t, buf.mrg.concat, "random keys must go through the heap") + + // rewind heapifies with the same sift the reads use, so check the + // invariant directly - a broken heapify need not show up in the order. + m := &buf.mrg + require.Greater(t, len(m.heap), 8, "a heap of a few elements proves little") + for i := 1; i < len(m.heap); i++ { + require.False(t, m.less(m.heap[i], m.heap[(i-1)/2]), + "heap[%d] sorts before its parent %d", i, (i-1)/2) + } + + got := drainBuffer(buf) + require.Len(t, got, count) + for i := range want { + require.Equal(t, want[i].key, got[i].key, "entry %d", i) + require.Equal(t, want[i].value, got[i].value, "entry %d value (insertion order)", i) + } +} + +// TestSortableBufferReadIsAllocFree: reading a sorted buffer walks a merge, so +// it must not allocate per entry - the whole point of holding the index inside +// the chunks. Sort's own state is allocated once per buffer, so warm it first. func TestSortableBufferReadIsAllocFree(t *testing.T) { const count = 50_000 buf := NewSortableBuffer(256 * datasize.MB) @@ -1900,6 +2162,7 @@ func TestSortableBufferReadIsAllocFree(t *testing.T) { } require.Greater(t, len(buf.chunks), 2, "the read must cross chunks") buf.Sort() + require.False(t, buf.mrg.concat, "random keys must go through the heap") got := 0 n := testing.AllocsPerRun(3, func() { @@ -1912,3 +2175,93 @@ func TestSortableBufferReadIsAllocFree(t *testing.T) { require.Equal(t, count, got) require.Zero(t, n, "Sort and a full read must not allocate") } + +// TestMapBufferSortIsIdempotent: sortAndFlush calls Sort and then Write, and +// Write sorts too. Flattening the map and re-sorting it on the second call +// would do every flush's work twice. +func TestMapBufferSortIsIdempotent(t *testing.T) { + for _, bt := range allBufferTypes { + if bt.name == "sortable" { + continue // its Sort only rewinds; the map-backed ones rebuild + } + t.Run(bt.name, func(t *testing.T) { + buf := bt.new() + for _, k := range []string{"c", "a", "b"} { + buf.Put([]byte(k), []byte(k)) + } + buf.Sort() + + // A rebuild would drop this, since it comes from the map. + switch b := buf.(type) { + case *appendSortableBuffer: + b.sortedBuf[0].value = []byte("marker") + case *oldestEntrySortableBuffer: + b.sortedBuf[0].value = []byte("marker") + } + buf.Sort() + + _, v, ok := buf.Next() + require.True(t, ok) + require.Equal(t, "marker", string(v), "Sort re-flattened an unchanged map") + }) + } +} + +// TestMapBufferPreallocClearsState: Prealloc replaces the entry map, so the run +// flattened out of the old one, the read cursor and Size must go with it. +func TestMapBufferPreallocClearsState(t *testing.T) { + for _, bt := range allBufferTypes { + if bt.name == "sortable" { + continue // its Prealloc only reserves chunk headers, it drops nothing + } + t.Run(bt.name, func(t *testing.T) { + buf := bt.new() + for _, k := range []string{"a", "b", "c", "d"} { + buf.Put([]byte(k), []byte(k)) + } + buf.Sort() + buf.Prealloc(2, 2) // cap(sortedBuf) still covers it, so it survives + + require.Zero(t, buf.Len()) + require.Empty(t, drainBuffer(buf), "read the wiped map, not the old run") + if s, ok := buf.(interface{ Size() int }); ok { + require.Zero(t, s.Size()) + } + }) + } +} + +// TestBufferSortPositionsCursor: a buffer fills, Sorts, then is read, and +// Sorting again is the only way to read it twice. +func TestBufferSortPositionsCursor(t *testing.T) { + for _, bt := range allBufferTypes { + t.Run(bt.name, func(t *testing.T) { + buf := bt.new() + for _, k := range []byte{3, 1, 2} { + buf.Put([]byte{k}, []byte{k}) + } + require.Panics(t, func() { buf.Next() }, "read before Sort") + + buf.Sort() + k, _, ok := buf.Next() + require.True(t, ok) + require.Equal(t, []byte{1}, k) + + buf.Sort() // already ordered, but the cursor goes back + k, _, ok = buf.Next() + require.True(t, ok) + require.Equal(t, []byte{1}, k) + }) + } +} + +// TestCollectRejectsOversizedKey: Put panics past maxKeyLen, but Collect sits +// under Load and the stage loop, which return errors. +func TestCollectRejectsOversizedKey(t *testing.T) { + c := NewCollector(t.Name(), t.TempDir(), NewSortableBuffer(1*datasize.MB), log.New()) + defer c.Close() + require.NoError(t, c.Collect(make([]byte, maxKeyLen), []byte("v"))) + err := c.Collect(make([]byte, maxKeyLen+1), []byte("v")) + require.Error(t, err) + require.Contains(t, err.Error(), "exceeds") +} diff --git a/db/etl/heap.go b/db/etl/heap.go index f1dd08d5292..9cfd84fe662 100644 --- a/db/etl/heap.go +++ b/db/etl/heap.go @@ -18,6 +18,7 @@ package etl import ( "bytes" + "slices" ) type HeapElem struct { @@ -120,3 +121,141 @@ func down(h *Heap, i0, n int) bool { } return i > i0 } + +// ------ the merge over a sortableBuffer's sorted chunks + +// merger walks already-sorted chunks in key order, a cursor per chunk under a +// heap of chunk ids. +type merger struct { + heap []int32 // chunk ids, ordered by their cursor's key + cur []cursor // by chunk id + + // Chunks already in order end to end, which ascending keys produce, are + // read straight through instead of merged. + concat bool + chunk int // chunk the straight-through cursor sits in +} + +type cursor struct { + ents []entryLoc + buf []byte + at int32 + key []byte +} + +// rewind puts the cursor on the first entry in key order. +func (m *merger) rewind(chunks []dataChunk) { + clear(m.cur) // a shorter run would leave the old cursors pinning their chunks + m.cur = slices.Grow(m.cur[:0], len(chunks))[:len(chunks)] + for i := range chunks { + m.cur[i] = cursor{ents: chunks[i].entries(), buf: chunks[i].buf} + } + m.chunk = 0 + + if m.concat = m.chunksInOrder(); m.concat { + return + } + m.heap = m.heap[:0] + for i := range m.cur { + if len(m.cur[i].ents) == 0 { + continue + } + m.load(int32(i)) //nolint:gosec + m.heap = append(m.heap, int32(i)) + } + for i := len(m.heap)/2 - 1; i >= 0; i-- { + m.siftRoot(i) + } +} + +// next returns the entry the cursor sits on and moves it to the next in key +// order. +func (m *merger) next() ([]byte, entryLoc, bool) { + if m.concat { + for ; m.chunk < len(m.cur); m.chunk++ { + c := &m.cur[m.chunk] + if int(c.at) < len(c.ents) { + e := c.ents[c.at] + c.at++ + return c.buf, e, true + } + } + return nil, 0, false + } + if len(m.heap) == 0 { + return nil, 0, false + } + id := m.heap[0] + c := &m.cur[id] + buf, e := c.buf, c.ents[c.at] + c.at++ + if int(c.at) == len(c.ents) { + last := len(m.heap) - 1 + m.heap[0] = m.heap[last] + m.heap = m.heap[:last] + } else { + m.load(id) + } + if len(m.heap) > 0 { + m.siftRoot(0) + } + return buf, e, true +} + +func (m *merger) release() { + clear(m.cur) + m.cur, m.heap = m.cur[:0], m.heap[:0] + m.chunk, m.concat = 0, false +} + +func (m *merger) load(id int32) { + c := &m.cur[id] + c.key = keyOf(c.buf, c.ents[c.at]) +} + +// chunksInOrder reports whether every chunk's last key comes before the next +// chunk's first. A tie keeps the earlier chunk, which is insertion order. +func (m *merger) chunksInOrder() bool { + prev := -1 // last chunk holding anything, so an empty one does not hide a pair + for i := range m.cur { + cur := &m.cur[i] + if len(cur.ents) == 0 { + continue + } + if prev >= 0 { + p := &m.cur[prev] + if bytes.Compare(keyOf(p.buf, p.ents[len(p.ents)-1]), keyOf(cur.buf, cur.ents[0])) > 0 { + return false + } + } + prev = i + } + return true +} + +// less orders two cursors by the key they sit on. Chunks fill in insertion +// order, so the lower id wins a tie and equal keys keep the order they went in. +func (m *merger) less(x, y int32) bool { + if r := bytes.Compare(m.cur[x].key, m.cur[y].key); r != 0 { + return r < 0 + } + return x < y +} + +// siftRoot restores the heap under i, whose element changed. +func (m *merger) siftRoot(i int) { + for { + s, l, r := i, 2*i+1, 2*i+2 + if l < len(m.heap) && m.less(m.heap[l], m.heap[s]) { + s = l + } + if r < len(m.heap) && m.less(m.heap[r], m.heap[s]) { + s = r + } + if s == i { + return + } + m.heap[i], m.heap[s] = m.heap[s], m.heap[i] + i = s + } +} diff --git a/db/etl/read_bench_test.go b/db/etl/read_bench_test.go index ef2630fcbc6..8ed366bb1bb 100644 --- a/db/etl/read_bench_test.go +++ b/db/etl/read_bench_test.go @@ -204,12 +204,13 @@ func benchBufioU16(b *testing.B, fname string, bufSize int) { defer f.Close() r := bufio.NewReaderSize(f, bufSize) + lenBuf := make([]byte, 2) var buf []byte for { - if buf, err = readFieldBufioU16(r, buf); err != nil { + if buf, err = readFieldBufioU16(r, lenBuf, buf); err != nil { break } - if buf, err = readFieldBufioU16(r, buf); err != nil { + if buf, err = readFieldBufioU16(r, lenBuf, buf); err != nil { break } } @@ -224,20 +225,22 @@ func benchBufioU32(b *testing.B, fname string, bufSize int) { defer f.Close() r := bufio.NewReaderSize(f, bufSize) + lenBuf := make([]byte, 4) var buf []byte for { - if buf, err = readFieldBufioU32(r, buf); err != nil { + if buf, err = readFieldBufioU32(r, lenBuf, buf); err != nil { break } - if buf, err = readFieldBufioU32(r, buf); err != nil { + if buf, err = readFieldBufioU32(r, lenBuf, buf); err != nil { break } } } -func readFieldBufioU16(r *bufio.Reader, buf []byte) ([]byte, error) { - var lenBuf [2]byte - if _, err := io.ReadFull(r, lenBuf[:]); err != nil { +func readFieldBufioU16(r *bufio.Reader, lenBuf, buf []byte) ([]byte, error) { + // lenBuf comes from the caller: io.ReadFull takes an io.Reader, so a local + // array escapes and costs an allocation on every field read. + if _, err := io.ReadFull(r, lenBuf); err != nil { return buf, err } n := int(*(*uint16)(unsafe.Pointer(&lenBuf[0]))) @@ -255,9 +258,10 @@ func readFieldBufioU16(r *bufio.Reader, buf []byte) ([]byte, error) { return buf, nil } -func readFieldBufioU32(r *bufio.Reader, buf []byte) ([]byte, error) { - var lenBuf [4]byte - if _, err := io.ReadFull(r, lenBuf[:]); err != nil { +func readFieldBufioU32(r *bufio.Reader, lenBuf, buf []byte) ([]byte, error) { + // lenBuf comes from the caller: io.ReadFull takes an io.Reader, so a local + // array escapes and costs an allocation on every field read. + if _, err := io.ReadFull(r, lenBuf); err != nil { return buf, err } n := int(*(*uint32)(unsafe.Pointer(&lenBuf[0])))