fix(store): stop a killed compaction from reselecting the group that killed it

internal/telemetry/store/compaction.go

@@ -27,6 +27,10 @@ const (
 	// byte ceiling exists to prevent.
 	assumedCompactionBytesPerRow = 256
 	minCompactionInputs          = 8
+
+	// minCompactionMergeInputs is the smallest group worth merging, and the
+	// floor the interrupted-pass halving stops at.
+	minCompactionMergeInputs = 2
 )
 
 var parquetSignals = [...]string{"spans", "logs", "metrics"}
@@ -36,6 +40,63 @@ type compactionMarker struct {
 	Inputs []string                `json:"inputs"`
 }
 
+// compactionAttemptFile records that a merge was started. The completion
+// marker is written only after PrepareReplacement, so a process killed during
+// the merge -- which is the likely way a large merge ends -- leaves no trace
+// that it ever ran. The next start reselects the same group by the same rules
+// and dies identically, and nothing in the system observes the loop. This file
+// is what a restart reads to know the last attempt was too big.
+const compactionAttemptFile = "COMPACTION-ATTEMPT.json"
+
+type compactionAttempt struct {
+	Inputs int `json:"inputs"`
+}
+
+// compactionBatchCap halves the ceiling after an interrupted pass, so repeated
+// kills walk the group down instead of retrying the same one. It never goes
+// below a pair: selection always admits two however large they are, so a lower
+// floor would stop compaction rather than shrink it. A merge that still dies
+// at two inputs is a single file too large to merge, which is a different
+// problem and not one a smaller group can solve.
+func compactionBatchCap(maxBatches, lastAttempt int) int {
+	if lastAttempt <= 0 {
+		return maxBatches
+	}
+	capped := min(maxBatches, lastAttempt) / 2
+	return max(capped, minCompactionMergeInputs)
+}
+
+// readCompactionAttempt reports the size of an interrupted pass, or zero when
+// there was none. A missing, damaged or nonsensical file reads as zero: losing
+// this record costs one oversized merge, while treating it as fatal would cost
+// the ability to compact at all.
+func readCompactionAttempt(root string) int {
+	data, err := os.ReadFile(filepath.Join(root, compactionAttemptFile))
+	if err != nil {
+		return 0
+	}
+	var attempt compactionAttempt
+	if err := json.Unmarshal(data, &attempt); err != nil || attempt.Inputs <= 0 {
+		return 0
+	}
+	return attempt.Inputs
+}
+
+func writeCompactionAttempt(root string, inputs int) error {
+	data, err := json.Marshal(compactionAttempt{Inputs: inputs})
+	if err != nil {
+		return err
+	}
+	return writeDurableFile(filepath.Join(root, compactionAttemptFile), data)
+}
+
+func clearCompactionAttempt(root string) error {
+	if err := os.Remove(filepath.Join(root, compactionAttemptFile)); err != nil && !errors.Is(err, os.ErrNotExist) {
+		return err
+	}
+	return nil
+}
+
 type compactionKey struct {
 	day        int64
 	generation uint32
@@ -63,8 +124,10 @@ func (r *Repository) CompactParquet(ctx context.Context, publisher ParquetPublis
 			return 0, fmt.Errorf("recover pending Parquet compaction: %w", err)
 		}
 	}
-	selected := selectCompactionBatches(r.Parquet.BatchMetadata(), maxBatches)
-	if len(selected) < 2 {
+	// A pass that was killed mid-merge left a record of how much it tried to
+	// take. Take less.
+	selected := selectCompactionBatches(r.Parquet.BatchMetadata(), compactionBatchCap(maxBatches, readCompactionAttempt(r.root)))
+	if len(selected) < minCompactionMergeInputs {
 		return 0, nil
 	}
 	output := telemetry.BatchMetadata{
@@ -122,6 +185,12 @@ func (r *Repository) CompactParquet(ctx context.Context, publisher ParquetPublis
 		}
 		plans = append(plans, mergePlan{signal: signal, inputs: inputs, output: filepath.Join(stage, signal+".parquet")})
 	}
+	// Written before the merge, on purpose: the completion marker below is
+	// only reached if the merge returns, and the failure this guards against
+	// is the merge not returning.
+	if err := writeCompactionAttempt(r.root, len(selected)); err != nil {
+		return 0, err
+	}
 	group, mergeCtx := errgroup.WithContext(ctx)
 	for _, plan := range plans {
 		group.Go(func() error {
@@ -157,11 +226,19 @@ func (r *Repository) CompactParquet(ctx context.Context, publisher ParquetPublis
 	if err := r.completeCompaction(ctx, marker, publisher.PublishParquet); err != nil {
 		return 0, err
 	}
+	if err := clearCompactionAttempt(r.root); err != nil {
+		return 0, err
+	}
 	return len(selected), nil
 }
 
 func selectCompactionBatches(batches []telemetry.BatchMetadata, maxBatches int) []telemetry.BatchMetadata {
-	if maxBatches < minCompactionInputs {
+	// minCompactionInputs is a "worth the effort" policy and belongs to the
+	// caller, which applies it to the ceiling it was asked for. Here the floor
+	// is only what a merge needs: after an interrupted pass the ceiling is
+	// deliberately halved below that policy, and refusing it would stop
+	// compaction at exactly the moment it most needs to make smaller progress.
+	if maxBatches < minCompactionMergeInputs {
 		return nil
 	}
 	groups := make(map[compactionKey][]telemetry.BatchMetadata)

internal/telemetry/store/compaction_attempt_test.go

@@ -0,0 +1,121 @@
+package store
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+)
+
+// A merge that dies takes the process with it and leaves nothing behind saying
+// so: the completion marker is written after PrepareReplacement, so an
+// out-of-memory kill mid-merge is indistinguishable from never having started.
+// The next start then reselects the same group by the same rules and dies the
+// same way. Recording the attempt before the merge is what turns that loop
+// into a descent.
+func TestCompactionBatchCapHalvesAfterAnInterruptedPass(t *testing.T) {
+	tests := []struct {
+		name        string
+		maxBatches  int
+		lastAttempt int
+		want        int
+	}{
+		{name: "no previous attempt leaves the ceiling alone", maxBatches: 128, lastAttempt: 0, want: 128},
+		{name: "an interrupted pass halves it", maxBatches: 128, lastAttempt: 128, want: 64},
+		{name: "halving compounds across restarts", maxBatches: 128, lastAttempt: 64, want: 32},
+		{name: "never below a pair, which always merges", maxBatches: 128, lastAttempt: 2, want: 2},
+		{name: "an attempt larger than the ceiling still halves from the ceiling", maxBatches: 16, lastAttempt: 128, want: 8},
+		{name: "a nonsensical attempt does not raise the ceiling", maxBatches: 16, lastAttempt: -5, want: 16},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			if got := compactionBatchCap(test.maxBatches, test.lastAttempt); got != test.want {
+				t.Errorf("compactionBatchCap(%d, %d) = %d, want %d",
+					test.maxBatches, test.lastAttempt, got, test.want)
+			}
+		})
+	}
+}
+
+// The record has to survive a kill -9, so it is a file rather than process
+// state, and a missing or unreadable one must read as "no attempt" rather than
+// stopping compaction: a corrupt byte here should cost a larger merge, not the
+// ability to compact at all.
+func TestCompactionAttemptRoundTripsAndToleratesDamage(t *testing.T) {
+	root := t.TempDir()
+
+	if got := readCompactionAttempt(root); got != 0 {
+		t.Errorf("readCompactionAttempt on a clean root = %d, want 0", got)
+	}
+
+	if err := writeCompactionAttempt(root, 32); err != nil {
+		t.Fatalf("writeCompactionAttempt: %v", err)
+	}
+	if got := readCompactionAttempt(root); got != 32 {
+		t.Errorf("readCompactionAttempt after write = %d, want 32", got)
+	}
+
+	if err := os.WriteFile(filepath.Join(root, compactionAttemptFile), []byte("{not json"), 0o600); err != nil {
+		t.Fatal(err)
+	}
+	if got := readCompactionAttempt(root); got != 0 {
+		t.Errorf("readCompactionAttempt on damaged file = %d, want 0", got)
+	}
+
+	if err := clearCompactionAttempt(root); err != nil {
+		t.Fatalf("clearCompactionAttempt: %v", err)
+	}
+	if got := readCompactionAttempt(root); got != 0 {
+		t.Errorf("readCompactionAttempt after clear = %d, want 0", got)
+	}
+	if err := clearCompactionAttempt(root); err != nil {
+		t.Errorf("clearCompactionAttempt must be idempotent, got %v", err)
+	}
+}
+
+// End to end: a record left by a pass that never returned must shrink the next
+// one, and a pass that does return must clear it so the group can grow back.
+// Without the second half the first kill would permanently halve compaction.
+func TestCompactParquetShrinksAfterAnInterruptedPassAndRecovers(t *testing.T) {
+	dir := t.TempDir()
+	repository, err := Open(dir)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer repository.Close()
+	for i := range minCompactionInputs {
+		batch := testBatch()
+		batch.ID = fmt.Sprintf("batch-%d", i)
+		batch.Spans[0].SpanID = fmt.Sprintf("span-%d", i)
+		batch.Spans[0].StartUnixNanos = int64(100 + i)
+		if err := repository.Commit(context.Background(), batch); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	// Stand in for a process killed mid-merge: the record is on disk, the
+	// completion marker never got written.
+	if err := writeCompactionAttempt(dir, minCompactionInputs); err != nil {
+		t.Fatal(err)
+	}
+
+	compactor := &testParquetCompactor{}
+	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
+	defer cancel()
+	compacted, err := repository.CompactParquet(ctx, compactor, minCompactionInputs)
+	if err != nil {
+		t.Fatalf("CompactParquet after an interrupted pass: %v", err)
+	}
+	if compacted == 0 {
+		t.Fatal("compacted nothing; an interrupted pass must shrink the group, not stop compaction")
+	}
+	if compacted >= minCompactionInputs {
+		t.Errorf("compacted %d inputs, want fewer than the %d the interrupted pass tried",
+			compacted, minCompactionInputs)
+	}
+	if left := readCompactionAttempt(dir); left != 0 {
+		t.Errorf("attempt record still reads %d after a completed pass; the group could never grow back", left)
+	}
+}