Skip to content

db/seg: collapse compression workers=1 case into batched workers flow #23058

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sudeepdino008 merged 2 commits into from
Aug 10, 2026
Merged

db/seg: collapse compression workers=1 case into batched workers flow #23058

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
dd88959
db/seg: run the single-worker cover phase through the batched pipeline
sudeepdino008 Aug 7, 2026
baae28a
db/seg: cover workers=1 in the backpressure deadlock test
sudeepdino008 Aug 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 14 additions & 20 deletions db/seg/compress_parallel_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,27 +71,21 @@ func readBack(t *testing.T, file string) [][]byte {
return out
}

// assertParallelMatchesSingle compresses corpus at Workers=1 and at each of workers, and
// asserts the decoded words and the on-disk bytes are identical across all of them.
func assertParallelMatchesSingle(t *testing.T, corpus []testWord, workers ...int) {
func assertSameOutputEveryWorkerCount(t *testing.T, corpus []testWord, wantSum uint32, workers ...int) {
t.Helper()
single := compressCorpus(t, corpus, 1)
want := readBack(t, single)
require.Len(t, want, len(corpus))
for i := range corpus {
require.Equalf(t, corpus[i].data, want[i], "single-worker round-trip mismatch at word %d", i)
}
singleSum := checksum(single)
for _, w := range workers {
parallel := compressCorpus(t, corpus, w)
require.Equalf(t, want, readBack(t, parallel), "decoded words differ at Workers=%d", w)
require.Equalf(t, singleSum, checksum(parallel), "output not byte-identical at Workers=%d", w)
file := compressCorpus(t, corpus, w)
got := readBack(t, file)
require.Lenf(t, got, len(corpus), "word count differs at Workers=%d", w)
for i := range corpus {
require.Equalf(t, corpus[i].data, got[i], "round-trip mismatch at word %d, Workers=%d", i, w)
}
require.Equalf(t, wantSum, checksum(file), "output differs from the pre-batching encoder at Workers=%d", w)
}
}

// The parallel cover phase hands each worker a batch of consecutive words. A corpus large
// enough to cross several batch boundaries must round-trip and stay byte-identical to the
// single-worker path.
// The cover phase hands each worker a batch of consecutive words. A corpus large enough to
// cross several batch boundaries must round-trip and keep the original encoding.
func TestCompressParallelBatchingRoundTrip(t *testing.T) {
const n = 4000
corpus := make([]testWord, 0, n+2)
Expand All @@ -100,11 +94,11 @@ func TestCompressParallelBatchingRoundTrip(t *testing.T) {
}
corpus = append(corpus, testWord{data: []byte{}, compressed: true}, testWord{data: []byte("zzz-unique-tail"), compressed: true})

assertParallelMatchesSingle(t, corpus, 2, 4, 8)
assertSameOutputEveryWorkerCount(t, corpus, 1432034479, 1, 2, 4, 8)
}

// Compressed, uncompressed and empty words interleaved exercise the queue's bypass paths
// alongside the batched cover path; output must still match the single-worker path.
// alongside the batched cover path.
func TestCompressParallelBatchingMixedStream(t *testing.T) {
const n = 6000
corpus := make([]testWord, 0, n)
Expand All @@ -119,7 +113,7 @@ func TestCompressParallelBatchingMixedStream(t *testing.T) {
}
}

assertParallelMatchesSingle(t, corpus, 2, 4, 8)
assertSameOutputEveryWorkerCount(t, corpus, 660475778, 1, 2, 4, 8)
}

// A partial (unflushed) batch of compressible words followed by more than queueLimit
Expand All @@ -140,7 +134,7 @@ func TestCompressParallelBatchingBackpressureNoDeadlock(t *testing.T) {
done := make(chan struct{})
go func() {
defer close(done)
assertParallelMatchesSingle(t, corpus, 4)
assertSameOutputEveryWorkerCount(t, corpus, 2163753720, 1, 4)
}()
select {
case <-done:
Expand Down
186 changes: 72 additions & 114 deletions db/seg/parallel_compress.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -356,11 +356,10 @@ func compressWithPatternCandidates(ctx context.Context, trace bool, cfg Cfg, log
if lvl < log.LvlTrace {
logger.Log(lvl, fmt.Sprintf("[%s] dictionary file parsed", logPrefix), "entries", len(code2pattern))
}
// we pass consecutive words so that AC mather's prefix-resume functionality
// can process words faster (the words are in sorted order);
// so we send a batch of coverBatchSize consecutive words to each worker
// Consecutive words per batch keep the AC matcher's prefix-resume working (words
// arrive sorted); a small input shrinks the batch so every worker still gets a share.
coverBatchSize := 512
if n := int(uncompressedFile.count); cfg.Workers > 1 && n < coverBatchSize*cfg.Workers {
if n := int(uncompressedFile.count); n < coverBatchSize*cfg.Workers {
coverBatchSize = max(1, n/cfg.Workers)
}
ch := make(chan []*CompressionWord, cfg.Workers*4)
Expand All @@ -377,31 +376,19 @@ func compressWithPatternCandidates(ctx context.Context, trace bool, cfg Cfg, log
heap.Init(&compressionQueue)
queueLimit := 128 * 1024

// For the case of workers == 1
var output = make([]byte, 0, 256)
var uncovered = make([]int, 256)
var patterns = make([]int, 0, 256)
var cells = make([]DynamicCell, 0, 256)
mf3 := patricia.NewACMatcher(ac)

var posMaps []*posCounter
posMaps := make([]*posCounter, 0, 1+cfg.Workers)
uncompPosMap := &posCounter{} // For the uncompressed words
posMaps = append(posMaps, uncompPosMap)
var wg sync.WaitGroup
if cfg.Workers > 1 {
for i := 0; i < cfg.Workers; i++ {
posMap := &posCounter{}
posMaps = append(posMaps, posMap)
wg.Go(func() {
coverWordsByPatternsWorker(trace, ch, out, ac, inputSize, outputSize, posMap)
})
}
}
var curBatch []*CompressionWord // consecutive words accumulating for the next batch
var freeList []*CompressionWord // written words available for reuse
if cfg.Workers > 1 {
curBatch = make([]*CompressionWord, 0, coverBatchSize)
}
for range cfg.Workers {
posMap := &posCounter{}
posMaps = append(posMaps, posMap)
wg.Go(func() {
coverWordsByPatternsWorker(trace, ch, out, ac, inputSize, outputSize, posMap)
})
}
curBatch := make([]*CompressionWord, 0, coverBatchSize) // consecutive words accumulating for the next batch
var freeList []*CompressionWord // written words available for reuse
t := time.Now()

var err error
Expand Down Expand Up @@ -435,108 +422,79 @@ func compressWithPatternCandidates(ctx context.Context, trace bool, cfg Cfg, log
}
}

if cfg.Workers > 1 {
// take processed batches in non-blocking way and push their words to the queue
outer:
for {
// take processed batches in non-blocking way and push their words to the queue
outer:
for {
select {
case batch := <-out:
for _, w := range batch {
heap.Push(&compressionQueue, w)
}
default:
break outer
}
}
// queue[0].order is never below outCount, so > means the next word to write is missing:
// nothing can be written, so wait for results instead of reading more input.
for compressionQueue.Len() >= queueLimit && compressionQueue[0].order > outCount {
if len(curBatch) > 0 {
// The missing word may sit in curBatch, and only we can send it — so offer it
// while waiting, else <-out waits forever. Rare: AddWord followed by a long run
// of AddUncompressedWord (see TestCompressParallelBatchingBackpressureNoDeadlock).
select {
case ch <- curBatch:
curBatch = make([]*CompressionWord, 0, coverBatchSize)
case batch := <-out:
for _, w := range batch {
heap.Push(&compressionQueue, w)
}
default:
break outer
}
}
// queue[0].order is never below outCount, so > means the next word to write is
// missing: nothing can be written, so wait for results instead of reading more input.
for compressionQueue.Len() >= queueLimit && compressionQueue[0].order > outCount {
if len(curBatch) > 0 {
// The missing word may sit in curBatch, and only we can send it — so offer it
// while waiting, else <-out waits forever. Rare: AddWord followed by a long run
// of AddUncompressedWord (see TestCompressParallelBatchingBackpressureNoDeadlock).
select {
case ch <- curBatch:
curBatch = make([]*CompressionWord, 0, coverBatchSize)
case batch := <-out:
for _, w := range batch {
heap.Push(&compressionQueue, w)
}
}
continue
}
batch := <-out
for _, w := range batch {
heap.Push(&compressionQueue, w)
}
}
// Write any in-order words at the top of the queue, recycling them onto freeList
for compressionQueue.Len() > 0 && compressionQueue[0].order == outCount {
w := heap.Pop(&compressionQueue).(*CompressionWord)
outCount++
if _, e := intermediateW.Write(w.word); e != nil {
return e
}
freeList = append(freeList, w)
}
var compW *CompressionWord
if k := len(freeList); k > 0 {
compW = freeList[k-1]
freeList = freeList[:k-1]
} else {
compW = &CompressionWord{}
continue
}
compW.order = inCount
switch {
case len(v) == 0:
// Empty word, cannot be compressed
compW.word = append(compW.word[:0], 0)
uncompPosMap.add(1)
uncompPosMap.add(0)
heap.Push(&compressionQueue, compW) // Push to the queue directly, bypassing compression
case compression:
compW.word = append(compW.word[:0], v...)
curBatch = append(curBatch, compW)
if len(curBatch) >= coverBatchSize {
ch <- curBatch // Send for compression
curBatch = make([]*CompressionWord, 0, coverBatchSize)
}
default:
// Prepend word with encoding of length + zero byte, which indicates no patterns to be found in this word
wordLen := uint64(len(v))
n := binary.PutUvarint(numBuf[:], wordLen)
uncompPosMap.add(wordLen + 1)
uncompPosMap.add(0)
compW.word = append(append(append(compW.word[:0], numBuf[:n]...), 0), v...)
heap.Push(&compressionQueue, compW) // Push to the queue directly, bypassing compression
batch := <-out
for _, w := range batch {
heap.Push(&compressionQueue, w)
}
} else {
}
// Write any in-order words at the top of the queue, recycling them onto freeList
for compressionQueue.Len() > 0 && compressionQueue[0].order == outCount {
w := heap.Pop(&compressionQueue).(*CompressionWord)
outCount++
wordLen := uint64(len(v))
n := binary.PutUvarint(numBuf[:], wordLen)
if _, e := intermediateW.Write(numBuf[:n]); e != nil {
if _, e := intermediateW.Write(w.word); e != nil {
return e
}
if wordLen > 0 {
if compression {
output, patterns, uncovered, cells = coverWordByPatterns(trace, v, mf3, output[:0], uncovered, patterns, cells, uncompPosMap)
if _, e := intermediateW.Write(output); e != nil {
return e
}
outputSize.Add(uint64(len(output)))
} else {
if e := intermediateW.WriteByte(0); e != nil {
return e
}
if _, e := intermediateW.Write(v); e != nil {
return e
}
outputSize.Add(1 + uint64(len(v)))
}
freeList = append(freeList, w)
}
var compW *CompressionWord
if k := len(freeList); k > 0 {
compW = freeList[k-1]
freeList = freeList[:k-1]
} else {
compW = &CompressionWord{}
}
compW.order = inCount
switch {
case len(v) == 0:
// Empty word, cannot be compressed
compW.word = append(compW.word[:0], 0)
uncompPosMap.add(1)
uncompPosMap.add(0)
heap.Push(&compressionQueue, compW) // Push to the queue directly, bypassing compression
case compression:
compW.word = append(compW.word[:0], v...)
curBatch = append(curBatch, compW)
if len(curBatch) >= coverBatchSize {
ch <- curBatch // Send for compression
curBatch = make([]*CompressionWord, 0, coverBatchSize)
}
inputSize.Add(1 + wordLen)
default:
// Prepend word with encoding of length + zero byte, which indicates no patterns to be found in this word
wordLen := uint64(len(v))
n := binary.PutUvarint(numBuf[:], wordLen)
uncompPosMap.add(wordLen + 1)
uncompPosMap.add(0)
compW.word = append(append(append(compW.word[:0], numBuf[:n]...), 0), v...)
heap.Push(&compressionQueue, compW) // Push to the queue directly, bypassing compression
}
inCount++
if len(v) == 0 {
Expand Down
Loading