Reduce compaction merge fan-in on OOM and keep partial output

cmd/partforge/main.go

@@ -1607,7 +1607,7 @@ func runWorker(ctx context.Context, args []string) error {
 		compactWindow            = fs.Duration("compact-window", defaultCompactWindow, "how long COMPACT_READY artifacts remain eligible for compaction before being promoted to FINISHED and the hard cap for claimed compact merge waits; 0 finalizes as soon as no useful compaction is available")
 		compactMaxArtifacts      = fs.Int("compact-max-artifacts", defaultCompactMaxArtifacts, "maximum single-part artifacts from one destination partition merged in one compaction batch; 1 disables batching")
 		compactMaxBytes          = fs.Uint64("compact-max-bytes", defaultCompactMaxBytes, "maximum summed input bytes_on_disk for one compaction batch; 0 disables the byte cap")
-		compactMaxPartsToMerge   = fs.Int("compact-max-parts-to-merge", rewrite.DefaultCompactMaxPartsToMergeAtOnce, "max_parts_to_merge_at_once for compaction merges; lower values reduce peak merge memory")
+		compactMaxPartsToMerge   = fs.Int("compact-max-parts-to-merge", rewrite.DefaultCompactMaxPartsToMergeAtOnce, "initial max_parts_to_merge_at_once for compaction merges; halved after each memory-limit merge failure down to 2, lower values reduce peak merge memory")
 		metricsAddr              = fs.String("metrics-addr", ":2112", "Prometheus metrics listen address; empty disables PartForge metrics")
 		metricsPath              = fs.String("metrics-path", "/metrics", "HTTP path for PartForge Prometheus metrics")
 		clickHousePrometheusPort = fs.Int("clickhouse-prometheus-port", defaultClickHousePrometheusPort, "port where the local worker ClickHouse exposes native Prometheus metrics")

internal/rewrite/compact_progress.go

@@ -3,6 +3,7 @@ package rewrite
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"log/slog"
@@ -29,9 +30,13 @@ type clickHouseMerge struct {
 	TotalSizeBytesUncompressed uint64   `json:"total_size_bytes_uncompressed"`
 }
 
+// errCompactMergeMemoryExhausted reports that destination merges still exceed
+// the memory limit at the minimum merge fan-in. Compact keeps whatever merged.
+var errCompactMergeMemoryExhausted = errors.New("destination merges exceeded memory limit at minimum max_parts_to_merge_at_once")
+
 func (c Compactor) observeCompactProgress(ctx context.Context, p Processor, item CompactWorkItem, target mergeWaitTarget, inputStats metrics.PartStats) error {
 	lastStateReport := time.Time{}
-	mergeMaxBlockSizeBytes := p.MergeTreeSettings.MergeMaxBlockSizeBytes
+	maxPartsToMerge := c.MaxPartsToMergeAtOnce
 	var handledMemoryLimitFailures uint64
 	for {
 		failedMerges, err := p.destinationFailedMergeSummary(ctx, target)
@@ -48,28 +53,44 @@ func (c Compactor) observeCompactProgress(ctx context.Context, p Processor, item
 			return fmt.Errorf("observe compact merge failures: %w", err)
 		}
 		if failedMerges.MemoryLimitCount > handledMemoryLimitFailures {
-			failures := failedMerges.MemoryLimitCount - handledMemoryLimitFailures
-			next := mergeMaxBlockSizeBytes
-			for range failures {
-				next = max(next/2, minAdaptiveMergeMaxBlockSizeBytes)
-			}
-			if next >= mergeMaxBlockSizeBytes {
-				return fmt.Errorf("destination merge exceeded memory limit with merge_max_block_size_bytes already at minimum %d: %s", mergeMaxBlockSizeBytes, failedMerges.LatestException)
+			// Failures seen in one poll ran at the same fan-in, so step down once per poll.
+			next := max(maxPartsToMerge/2, minAdaptiveMaxPartsToMergeAtOnce)
+			if next >= maxPartsToMerge {
+				if err := p.ClickHouse.Exec(ctx, "SYSTEM STOP MERGES "+target.tableSQL()); err != nil {
+					if ctx.Err() != nil {
+						return nil
+					}
+					return fmt.Errorf("stop destination merges after memory-limit failure at minimum fan-in: %w", err)
+				}
+				slog.Warn(
+					"destination merges exceeded memory limit at minimum fan-in; keeping merged output",
+					"stage", "adjusting_merge_memory",
+					"job_id", item.JobID,
+					"part_id", item.OutputPartID,
+					"destination_table", target.tableSQL(),
+					"memory_limit_failures", failedMerges.MemoryLimitCount,
+					"max_parts_to_merge_at_once", maxPartsToMerge,
+					"latest_exception", failedMerges.LatestException,
+				)
+				return errCompactMergeMemoryExhausted
 			}
 			if err := recoverCompactMergeAfterMemoryFailure(ctx, p, target, next); err != nil {
+				if ctx.Err() != nil {
+					return nil
+				}
 				return err
 			}
 			slog.Warn(
-				"reduced merge block byte limit after memory-limit merge failure",
+				"reduced merge fan-in after memory-limit merge failure",
 				"stage", "adjusting_merge_memory",
 				"job_id", item.JobID,
 				"part_id", item.OutputPartID,
 				"destination_table", target.tableSQL(),
 				"memory_limit_failures", failedMerges.MemoryLimitCount,
-				"previous_merge_max_block_size_bytes", mergeMaxBlockSizeBytes,
-				"merge_max_block_size_bytes", next,
+				"previous_max_parts_to_merge_at_once", maxPartsToMerge,
+				"max_parts_to_merge_at_once", next,
 			)
-			mergeMaxBlockSizeBytes = next
+			maxPartsToMerge = next
 			handledMemoryLimitFailures = failedMerges.MemoryLimitCount
 		}
 		nonMemoryFailures := failedMerges.Count - failedMerges.MemoryLimitCount
@@ -135,18 +156,18 @@ func (c Compactor) observeCompactProgress(ctx context.Context, p Processor, item
 	}
 }
 
-func recoverCompactMergeAfterMemoryFailure(ctx context.Context, p Processor, target mergeWaitTarget, mergeMaxBlockSizeBytes uint64) error {
+func recoverCompactMergeAfterMemoryFailure(ctx context.Context, p Processor, target mergeWaitTarget, maxPartsToMerge int) error {
 	table := target.tableSQL()
 	if err := p.ClickHouse.Exec(ctx, "SYSTEM STOP MERGES "+table); err != nil {
 		return fmt.Errorf("stop destination merges after memory-limit failure: %w", err)
 	}
-	query := "ALTER TABLE " + table + " MODIFY SETTING merge_max_block_size_bytes = " + strconv.FormatUint(mergeMaxBlockSizeBytes, 10)
+	query := "ALTER TABLE " + table + " MODIFY SETTING max_parts_to_merge_at_once = " + strconv.Itoa(maxPartsToMerge)
 	if err := p.ClickHouse.Exec(ctx, query); err != nil {
 		startErr := p.ClickHouse.Exec(ctx, "SYSTEM START MERGES "+table)
 		if startErr != nil {
-			return fmt.Errorf("reduce merge_max_block_size_bytes after memory-limit merge failure: %w; additionally failed to restart merges: %v", err, startErr)
+			return fmt.Errorf("reduce max_parts_to_merge_at_once after memory-limit merge failure: %w; additionally failed to restart merges: %v", err, startErr)
 		}
-		return fmt.Errorf("reduce merge_max_block_size_bytes after memory-limit merge failure: %w", err)
+		return fmt.Errorf("reduce max_parts_to_merge_at_once after memory-limit merge failure: %w", err)
 	}
 	if err := p.ClickHouse.Exec(ctx, "SYSTEM START MERGES "+table); err != nil {
 		return fmt.Errorf("restart destination merges after memory-limit failure: %w", err)

internal/rewrite/compact_progress_test.go

@@ -2,6 +2,7 @@ package rewrite
 
 import (
 	"context"
+	"errors"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -29,7 +30,7 @@ func TestObserveCompactProgressRecoversMemoryLimitedMerge(t *testing.T) {
 			strings.HasPrefix(query, "SYSTEM START MERGES "),
 			strings.HasPrefix(query, "ALTER TABLE "):
 		case strings.Contains(query, "FROM system.part_log"):
-			_, _ = io.WriteString(w, "1\t1\tCode: 241. MEMORY_LIMIT_EXCEEDED\n")
+			_, _ = io.WriteString(w, "3\t3\tCode: 241. MEMORY_LIMIT_EXCEEDED\n")
 		case strings.HasPrefix(query, "SELECT partition_id, count()"):
 			_, _ = io.WriteString(w, "202601\t2\t20\t200\n")
 		case strings.HasPrefix(query, "SELECT partition_id, result_part_name"):
@@ -41,35 +42,39 @@ func TestObserveCompactProgressRecoversMemoryLimitedMerge(t *testing.T) {
 
 	ctx, cancel := context.WithCancel(context.Background())
 	err := (Compactor{
-		ProgressInterval: time.Nanosecond,
+		MaxPartsToMergeAtOnce: 32,
+		ProgressInterval:      time.Nanosecond,
 		ReportProgress: func(context.Context, CompactWorkItem, CompactProgressSnapshot) error {
 			cancel()
 			return nil
 		},
 	}).observeCompactProgress(ctx, Processor{
 		ClickHouse: chhttp.Client{URL: server.URL},
-		MergeTreeSettings: MergeTreeSettings{
-			MergeMaxBlockSizeBytes: 8 * 1024 * 1024,
-		},
 	}, CompactWorkItem{JobID: "job-1", OutputPartID: "compact-1"}, mergeWaitTarget{Database: "db", Table: "events"}, metrics.PartStats{Count: 2})
 	if err != nil {
 		t.Fatal(err)
 	}
-	if !containsQueryWith(queries, "MODIFY SETTING merge_max_block_size_bytes = 4194304") {
-		t.Fatalf("queries = %#v, want merge block bytes halved after memory failure", queries)
+	if !containsQueryWith(queries, "MODIFY SETTING max_parts_to_merge_at_once = 16") {
+		t.Fatalf("queries = %#v, want merge fan-in halved once after memory failures in one poll", queries)
+	}
+	if containsQueryWith(queries, "merge_max_block_size_bytes") {
+		t.Fatalf("queries = %#v, want merge block bytes left unchanged", queries)
 	}
 }
 
-func TestObserveCompactProgressFailsAtMinimumAfterMemoryLimitedMerge(t *testing.T) {
+func TestObserveCompactProgressStopsMergesAtMinimumFanIn(t *testing.T) {
+	var queries []string
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		body, err := io.ReadAll(r.Body)
 		if err != nil {
 			http.Error(w, err.Error(), http.StatusInternalServerError)
 			return
 		}
 		query := string(body)
+		queries = append(queries, query)
 		switch {
-		case query == "SYSTEM FLUSH LOGS":
+		case query == "SYSTEM FLUSH LOGS",
+			strings.HasPrefix(query, "SYSTEM STOP MERGES "):
 		case strings.Contains(query, "FROM system.part_log"):
 			_, _ = io.WriteString(w, "1\t1\tCode: 241. MEMORY_LIMIT_EXCEEDED\n")
 		default:
@@ -78,14 +83,14 @@ func TestObserveCompactProgressFailsAtMinimumAfterMemoryLimitedMerge(t *testing.
 	}))
 	defer server.Close()
 
-	err := (Compactor{}).observeCompactProgress(context.Background(), Processor{
+	err := (Compactor{MaxPartsToMergeAtOnce: minAdaptiveMaxPartsToMergeAtOnce}).observeCompactProgress(context.Background(), Processor{
 		ClickHouse: chhttp.Client{URL: server.URL},
-		MergeTreeSettings: MergeTreeSettings{
-			MergeMaxBlockSizeBytes: minAdaptiveMergeMaxBlockSizeBytes,
-		},
 	}, CompactWorkItem{JobID: "job-1", OutputPartID: "compact-1"}, mergeWaitTarget{Database: "db", Table: "events"}, metrics.PartStats{Count: 2})
-	if err == nil || !strings.Contains(err.Error(), "already at minimum") {
-		t.Fatalf("error = %v, want minimum merge block size failure", err)
+	if !errors.Is(err, errCompactMergeMemoryExhausted) {
+		t.Fatalf("error = %v, want memory exhausted at minimum fan-in", err)
+	}
+	if !containsQueryWith(queries, "SYSTEM STOP MERGES `db`.`events`") {
+		t.Fatalf("queries = %#v, want destination merges stopped", queries)
 	}
 }
 
@@ -131,13 +136,13 @@ func TestRecoverCompactMergeRestartsBackgroundMerges(t *testing.T) {
 		context.Background(),
 		Processor{ClickHouse: chhttp.Client{URL: server.URL}},
 		mergeWaitTarget{Database: "db", Table: "events"},
-		4*1024*1024,
+		16,
 	)
 	if err != nil {
 		t.Fatal(err)
 	}
 	want := "SYSTEM STOP MERGES `db`.`events`\n" +
-		"ALTER TABLE `db`.`events` MODIFY SETTING merge_max_block_size_bytes = 4194304\n" +
+		"ALTER TABLE `db`.`events` MODIFY SETTING max_parts_to_merge_at_once = 16\n" +
 		"SYSTEM START MERGES `db`.`events`"
 	if got := strings.Join(queries, "\n"); got != want {
 		t.Fatalf("queries = %q, want %q", got, want)

internal/rewrite/compactor.go

@@ -224,15 +224,18 @@ func (c Compactor) Compact(ctx context.Context, item CompactWorkItem) (CompactRe
 	}
 	if waitForMerges {
 		observerCtx, cancelObserver := context.WithCancel(phaseCtx)
+		waitCtx, cancelMergeWait := context.WithCancel(c.mergeWaitContext(phaseCtx))
+		defer cancelMergeWait()
 		observerErrCh := make(chan error, 1)
 		go func() {
 			err := c.observeCompactProgress(observerCtx, p, item, target, inputStats)
-			if err != nil {
+			if errors.Is(err, errCompactMergeMemoryExhausted) {
+				cancelMergeWait()
+			} else if err != nil {
 				cancelPhase()
 			}
 			observerErrCh <- err
 		}()
-		waitCtx := c.mergeWaitContext(phaseCtx)
 		cancelWait := func() {}
 		if deadlineActive {
 			waitCtx, cancelWait = context.WithDeadline(waitCtx, c.MergeDeadline)
@@ -249,10 +252,13 @@ func (c Compactor) Compact(ctx context.Context, item CompactWorkItem) (CompactRe
 		}()
 		cancelObserver()
 		observerErr := <-observerErrCh
-		if observerErr != nil {
+		memoryExhausted := errors.Is(observerErr, errCompactMergeMemoryExhausted)
+		if observerErr != nil && !memoryExhausted {
 			return CompactResult{}, observerErr
 		}
-		if err != nil {
+		if memoryExhausted {
+			slog.Warn("compact merges exhausted memory headroom; measuring current output", "stage", "compact_memory_exhausted", "job_id", item.JobID, "part_id", item.OutputPartID, "destination_table", chhttp.TableSQL(item.DestinationDatabase, item.DestinationTable))
+		} else if err != nil {
 			if deadlineActive && errors.Is(err, context.DeadlineExceeded) {
 				slog.Info("compact merge deadline reached; measuring current output", "stage", "compact_window_expired", "job_id", item.JobID, "part_id", item.OutputPartID, "destination_table", chhttp.TableSQL(item.DestinationDatabase, item.DestinationTable), "deadline", c.MergeDeadline)
 			} else if c.shutdownRequested() && errors.Is(err, context.Canceled) {

internal/rewrite/compactor_test.go

@@ -29,6 +29,7 @@ func TestCompactHandlesFragmentedInputStops(t *testing.T) {
 		wantWait             bool
 		checkWait            bool
 		observerError        bool
+		memoryExhausted      bool
 		stop                 string
 		wantError            string
 	}{
@@ -53,6 +54,19 @@ func TestCompactHandlesFragmentedInputStops(t *testing.T) {
 			observerError:        true,
 			wantError:            "observe compact merge failures",
 		},
+		{
+			name:                 "memory exhaustion at minimum fan-in uploads useful output",
+			deadline:             func() time.Time { return time.Time{} },
+			partsAfterFirstQuery: 1,
+			memoryExhausted:      true,
+			wantReduced:          true,
+		},
+		{
+			name:                 "memory exhaustion at minimum fan-in releases unchanged input",
+			deadline:             func() time.Time { return time.Time{} },
+			partsAfterFirstQuery: 2,
+			memoryExhausted:      true,
+		},
 		{
 			name:                 "manual finalize without deadline releases unchanged input",
 			deadline:             func() time.Time { return time.Time{} },
@@ -141,6 +155,8 @@ func TestCompactHandlesFragmentedInputStops(t *testing.T) {
 				case strings.Contains(query, "FROM system.part_log"):
 					if tt.observerError {
 						http.Error(w, "merge observer broke", http.StatusInternalServerError)
+					} else if tt.memoryExhausted {
+						_, _ = io.WriteString(w, "1\t1\tCode: 241. MEMORY_LIMIT_EXCEEDED\n")
 					} else {
 						_, _ = io.WriteString(w, "0\t0\t<none>\n")
 					}
@@ -165,6 +181,10 @@ func TestCompactHandlesFragmentedInputStops(t *testing.T) {
 			}))
 			defer server.Close()
 
+			maxPartsToMerge := DefaultCompactMaxPartsToMergeAtOnce
+			if tt.memoryExhausted {
+				maxPartsToMerge = minAdaptiveMaxPartsToMergeAtOnce
+			}
 			result, err := (Compactor{
 				S3Copy:                s3copy.Copier{Binary: binary},
 				ClickHouse:            chhttp.Client{URL: server.URL},
@@ -174,7 +194,7 @@ func TestCompactHandlesFragmentedInputStops(t *testing.T) {
 				MergeDeadline:         tt.deadline(),
 				MergeStopContext:      manualCtx,
 				ShutdownContext:       shutdownCtx,
-				MaxPartsToMergeAtOnce: DefaultCompactMaxPartsToMergeAtOnce,
+				MaxPartsToMergeAtOnce: maxPartsToMerge,
 				MergeTreeSettings: MergeTreeSettings{
 					MergeMaxBlockSize:        32768,
 					MergeMaxBlockSizeBytes:   10 * 1024 * 1024,

internal/rewrite/processor.go

@@ -35,7 +35,7 @@ const defaultMergePollInterval = time.Second
 const defaultMergeWaitLogInterval = 30 * time.Second
 const maxCompactMergeFailures uint64 = 3
 const minRetryMaxBlockSize = 8192
-const minAdaptiveMergeMaxBlockSizeBytes uint64 = 1024 * 1024
+const minAdaptiveMaxPartsToMergeAtOnce = 2
 
 const (
 	targetMergePartBytes uint64 = 150 * 1024 * 1024 * 1024