perf(telemetry): preempt long merges for generation-0 debt and cut span merge allocation

internal/telemetry/compaction_debt.go

@@ -78,6 +78,33 @@ func (p *ParquetStore) setBatchLocked(id string, batch *storedBatch) {
 		close(p.debtChanged)
 		p.debtChanged = nil
 	}
+	if p.compactionDebt >= CompactionDebtLimit && p.debtLimitReached != nil {
+		close(p.debtLimitReached)
+		p.debtLimitReached = nil
+	}
+}
+
+// WaitForCompactionLimit wakes higher-generation work when ingest needs the
+// compactor. Registration and the threshold check share the catalog lock, so
+// debt cannot cross the limit between the check and subscribing. No lock is
+// held while waiting, and caller cancellation releases the waiter.
+func (p *ParquetStore) WaitForCompactionLimit(ctx context.Context) error {
+	p.mu.Lock()
+	if p.compactionDebt >= CompactionDebtLimit {
+		p.mu.Unlock()
+		return nil
+	}
+	if p.debtLimitReached == nil {
+		p.debtLimitReached = make(chan struct{})
+	}
+	reached := p.debtLimitReached
+	p.mu.Unlock()
+	select {
+	case <-reached:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }
 
 // WaitForCompaction releases the catalog lock before sleeping and wakes on a

internal/telemetry/compaction_debt_test.go

@@ -2,11 +2,65 @@ package telemetry
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"testing"
 	"time"
 )
 
+func TestCompactionLimitNotificationDoesNotMissCrossings(t *testing.T) {
+	p := &ParquetStore{batches: make(map[string]*storedBatch)}
+	for i := range CompactionDebtLimit + CompactionBatchLimit - 2 {
+		p.installBatch(&storedBatch{metadata: BatchMetadata{ID: fmt.Sprint(i), Spans: 1, MaxIngestedNanos: 1}})
+	}
+	for cycle := range 2 {
+		ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second)
+		done := make(chan error, 2)
+		for range 2 {
+			go func() { done <- p.WaitForCompactionLimit(ctx) }()
+		}
+		until := time.Now().Add(time.Second)
+		for {
+			p.mu.RLock()
+			registered := p.debtLimitReached != nil
+			p.mu.RUnlock()
+			if registered {
+				break
+			}
+			if time.Now().After(until) {
+				cancel()
+				t.Fatal("limit waiter not registered")
+			}
+			time.Sleep(time.Millisecond)
+		}
+		p.installBatch(&storedBatch{metadata: BatchMetadata{ID: "crossing", Spans: 1, MaxIngestedNanos: 1}})
+		for range 2 {
+			if err := <-done; err != nil {
+				cancel()
+				t.Fatal(err)
+			}
+		}
+		// Arriving after the crossing also observes it, without waiting for
+		// another publication. Retiring it rearms the next crossing.
+		if err := p.WaitForCompactionLimit(ctx); err != nil {
+			cancel()
+			t.Fatal(err)
+		}
+		cancel()
+		p.mu.Lock()
+		p.setBatchLocked("crossing", nil)
+		p.mu.Unlock()
+		if p.CompactionDebt() != CompactionDebtLimit-1 {
+			t.Fatalf("cycle %d debt=%d", cycle, p.CompactionDebt())
+		}
+	}
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+	if err := p.WaitForCompactionLimit(ctx); !errors.Is(err, context.Canceled) {
+		t.Fatal(err)
+	}
+}
+
 func TestGenerationZeroInventoryTracksCatalogChanges(t *testing.T) {
 	p := &ParquetStore{batches: make(map[string]*storedBatch)}
 	add := func(id string, generation int) {

internal/telemetry/compaction_index_test.go

@@ -118,7 +118,9 @@ type cancelMergeContext struct {
 
 func (c *cancelMergeContext) Err() error {
 	c.checks++
-	if c.checks == 3 {
+	// Initial admission, input open and key priming now check cancellation too.
+	// Cancel at the next row-group boundary, after the first group was written.
+	if c.checks == 5 {
 		c.cancel()
 	}
 	return c.Context.Err()

internal/telemetry/parquet.go

@@ -94,6 +94,7 @@ type ParquetStore struct {
 	generation0ByDay   map[int64]int
 	compactionDebt     int
 	debtChanged        chan struct{}
+	debtLimitReached   chan struct{}
 }
 
 type ParquetStats struct {

internal/telemetry/parquet_merge.go

@@ -1,6 +1,7 @@
 package telemetry
 
 import (
+	"bytes"
 	"container/heap"
 	"context"
 	"errors"
@@ -37,6 +38,9 @@ func mergeTypedParquet[T any](ctx context.Context, inputs []string, output strin
 	schema := telemetrySchema[T]()
 	var groups []parquet.RowGroup
 	for _, input := range inputs {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
 		f, e := os.Open(input)
 		if e != nil {
 			return e
@@ -96,18 +100,25 @@ func mergeTypedParquet[T any](ctx context.Context, inputs []string, output strin
 			}
 		}()
 		for i, group := range groups {
-			c := &mergeKeyCursor{input: i, reader: parquet.NewGenericRowGroupReader[mergeSpanKey](group)}
+			c := &mergeKeyCursor{input: i, remaining: group.NumRows()}
+			for k, name := range [...]string{"_trace_hash", "start_unix_nano", "span_id"} {
+				leaf, _ := schema.Lookup(name)
+				c.readers[k] = parquet.NewColumnChunkValueReader(group.ColumnChunks()[leaf.ColumnIndex])
+			}
 			// Keep empty/exhausted cursors in a separate lifetime list below.
 			keys = append(keys, c)
 		}
 		allKeys := append(mergeKeyHeap(nil), keys...)
 		defer func() {
 			for _, c := range allKeys {
-				err = errors.Join(err, c.reader.Close())
+				err = errors.Join(err, c.close())
 			}
 		}()
 		keys = keys[:0]
 		for _, c := range allKeys {
+			if err := ctx.Err(); err != nil {
+				return err
+			}
 			has, e := c.next()
 			if e != nil {
 				return e
@@ -285,42 +296,60 @@ func (c *mergeColumn) write(ctx context.Context, writer *parquet.ColumnWriter, r
 	return writer.Flush()
 }
 
+// Key values borrow their column reader's current page. Each reader advances
+// only after its current key has been consumed; no per-span string or row copy
+// is needed, and the three columns may have independent page boundaries.
 type mergeSpanKey struct {
-	Hash   uint64 `parquet:"_trace_hash"`
-	Start  int64  `parquet:"start_unix_nano"`
-	SpanID string `parquet:"span_id"`
+	Hash   uint64
+	Start  int64
+	SpanID []byte
 }
 type mergeKeyCursor struct {
-	input  int
-	reader *parquet.GenericReader[mergeSpanKey]
-	buffer [256]mergeSpanKey
-	pos, n int
-	eof    bool
+	input     int
+	readers   [3]parquet.ColumnChunkValueReader
+	buffer    [3][256]parquet.Value
+	pos, n    [3]int
+	remaining int64
+	key       mergeSpanKey
 }
 
-func (c *mergeKeyCursor) current() mergeSpanKey { return c.buffer[c.pos-1] }
-func (c *mergeKeyCursor) next() (bool, error) {
-	if c.pos == c.n {
-		if c.eof {
-			return false, nil
+func (c *mergeKeyCursor) close() (err error) {
+	for i, r := range c.readers {
+		if r != nil {
+			err = errors.Join(err, r.Close())
+			c.readers[i] = nil
 		}
-		clear(c.buffer[:])
-		n, e := c.reader.Read(c.buffer[:])
-		c.pos, c.n = 0, n
-		if e != nil {
-			if e != io.EOF {
+	}
+	clear(c.buffer[:])
+	c.key = mergeSpanKey{}
+	return err
+}
+func (c *mergeKeyCursor) current() mergeSpanKey { return c.key }
+func (c *mergeKeyCursor) next() (bool, error) {
+	if c.remaining == 0 {
+		return false, c.close()
+	}
+	var values [3]parquet.Value
+	for k, r := range c.readers {
+		if c.pos[k] == c.n[k] {
+			clear(c.buffer[k][:])
+			n, e := r.ReadValues(c.buffer[k][:])
+			c.pos[k], c.n[k] = 0, n
+			if e != nil && e != io.EOF {
 				return false, e
 			}
-			c.eof = true
-		}
-		if n == 0 {
-			if c.eof {
-				return false, nil
+			if n == 0 {
+				if e == io.EOF {
+					return false, io.ErrUnexpectedEOF
+				}
+				return false, io.ErrNoProgress
 			}
-			return false, io.ErrNoProgress
 		}
+		values[k] = c.buffer[k][c.pos[k]]
+		c.pos[k]++
 	}
-	c.pos++
+	c.key = mergeSpanKey{Hash: values[0].Uint64(), Start: values[1].Int64(), SpanID: values[2].ByteArray()}
+	c.remaining--
 	return true, nil
 }
 
@@ -335,8 +364,8 @@ func (h mergeKeyHeap) Less(i, j int) bool {
 	if a.Start != b.Start {
 		return a.Start < b.Start
 	}
-	if a.SpanID != b.SpanID {
-		return a.SpanID < b.SpanID
+	if order := bytes.Compare(a.SpanID, b.SpanID); order != 0 {
+		return order < 0
 	}
 	return h[i].input < h[j].input
 }

internal/telemetry/store/backpressure_test.go

@@ -120,7 +120,7 @@ func TestCompactionPressureDrainsGenerationZeroBeforeHigherGenerations(t *testin
 			}
 			debt := telemetry.CompactionDebtLimit
 			for pass := 0; debt >= telemetry.CompactionDebtLimit/2; pass++ {
-				selected := selectCompactionBatches(batches, 128, debt)
+				selected := selectCompactionBatchesWithHigherCap(batches, 128, 128, debt)
 				if len(selected) < 2 {
 					t.Fatal("positive debt has no selection")
 				}
@@ -146,7 +146,7 @@ func TestCompactionPressureDrainsGenerationZeroBeforeHigherGenerations(t *testin
 				batches = next
 			}
 			if strings.HasPrefix(mode, "higher-") {
-				selected := selectCompactionBatches(batches, 128, debt)
+				selected := selectCompactionBatchesWithHigherCap(batches, 128, 128, debt)
 				if len(selected) < 2 || selected[0].Generation == 0 {
 					t.Fatalf("higher generations did not resume at debt=%d: %v", debt, selected)
 				}
@@ -160,14 +160,14 @@ func TestCompactionPressureDrainsGenerationZeroBeforeHigherGenerations(t *testin
 		small = append(small, telemetry.BatchMetadata{ID: fmt.Sprint(i), Spans: 1, MaxIngestedNanos: 1})
 	}
 	small = append(small, telemetry.BatchMetadata{ID: "today", Spans: 1, MaxIngestedNanos: int64(24*time.Hour) + 1})
-	if len(selectCompactionBatches(small, 128, 0)) != 8 {
+	if len(selectCompactionBatchesWithHigherCap(small, 128, 128, 0)) != 8 {
 		t.Fatal("finished eight-input group stranded")
 	}
 	var terminal []telemetry.BatchMetadata
 	for i := 0; i < 1024; i++ {
 		terminal = append(terminal, telemetry.BatchMetadata{ID: fmt.Sprint(i), Spans: maxCompactionRows, MaxIngestedNanos: 1})
 	}
-	if len(selectCompactionBatches(terminal, 128, 0)) != 0 {
+	if len(selectCompactionBatchesWithHigherCap(terminal, 128, 128, 0)) != 0 {
 		t.Fatal("terminal counterexample changed")
 	}
 }