test: add e2e sharding planning step to equally distribute test spec between jobs

.github/workflows/ci.yaml

@@ -363,8 +363,47 @@ jobs:
           if-no-files-found: error
           overwrite: true
 
+  e2e-shard-plan:
+    needs: [ changes ]
+    if: needs.changes.outputs.code == 'true'
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    permissions:
+      contents: read
+      actions: read
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7
+
+      - name: Set up Go
+        uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7
+        with:
+          go-version-file: 'go.mod'
+
+      - name: Go Mod
+        run: go mod download
+
+      - name: Download previous e2e timings
+        run: bash ci/download-e2e-last-run.sh .ci/e2e-timings
+        env:
+          GH_TOKEN: ${{ github.token }}
+
+      - name: Plan e2e shards
+        run: |
+          go run ./tools/e2e-shard-planner \
+            -report-path .ci/e2e-timings/latest \
+            -shards 4 \
+            -output .ci/e2e-shard-plan/e2e-shard-plan.json
+
+      - name: Upload e2e shard plan
+        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7
+        with:
+          name: e2e-shard-plan
+          path: .ci/e2e-shard-plan/e2e-shard-plan.json
+          if-no-files-found: error
+
   e2e-test:
-    needs: [ lint, changes ]
+    needs: [ lint, changes, e2e-shard-plan ]
     if: needs.changes.outputs.code == 'true'
     strategy:
       fail-fast: false
@@ -376,6 +415,12 @@ jobs:
       - name: Checkout code
         uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7
 
+      - name: Download e2e shard plan
+        uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8
+        with:
+          name: e2e-shard-plan
+          path: .ci/e2e-shard-plan
+
       - name: Set up Go
         uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7
         with:
@@ -396,6 +441,7 @@ jobs:
         env:
           E2E_SHARD_INDEX: ${{ matrix.shard }}
           E2E_SHARD_TOTAL: "4"
+          E2E_SHARD_PLAN: ${{ github.workspace }}/.ci/e2e-shard-plan/e2e-shard-plan.json
         run:  make test-e2e
 
       - name: Collect failure diagnostics
@@ -523,7 +569,7 @@ jobs:
     name: All CI checks passed
     runs-on: ubuntu-latest
     timeout-minutes: 5
-    needs: [ changes, lint, bundle, build_and_test, fuzz_specs, helm-test, compat-e2e-test, e2e-test, check-crd-compat ]
+    needs: [ changes, lint, bundle, build_and_test, fuzz_specs, helm-test, compat-e2e-test, e2e-shard-plan, e2e-test, check-crd-compat ]
     if: always()
     steps:
       - name: Determine CI status

ci/download-e2e-last-run.sh

@@ -0,0 +1,62 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+if [ "$#" -ne 1 ]; then
+  echo "usage: $0 <output-dir>" >&2
+  exit 2
+fi
+
+output_dir=$1
+selected_dir="$output_dir/latest"
+
+mkdir -p "$output_dir"
+rm -rf "$selected_dir"
+
+if ! run_ids=$(gh run list \
+  --repo "$GITHUB_REPOSITORY" \
+  --workflow ci.yaml \
+  --branch main \
+  --event push \
+  --status success \
+  --limit 20 \
+  --json databaseId \
+  --jq '.[].databaseId'); then
+  echo "Unable to fetch the previous e2e run; using equal shard weights."
+  exit 0
+fi
+
+while IFS= read -r run_id; do
+  if [ -z "$run_id" ]; then
+    continue
+  fi
+
+  candidate_dir=$(mktemp -d "$output_dir/run.XXXXXX")
+  complete=1
+  for shard in 1 2 3 4; do
+    report_dir="$candidate_dir/e2e-report-shard-$shard"
+    if ! gh run download \
+      "$run_id" \
+      --repo "$GITHUB_REPOSITORY" \
+      --name "e2e-report-shard-$shard" \
+      --dir "$report_dir"; then
+      complete=0
+      break
+    fi
+
+    if ! find "$report_dir" -name ginkgo-report.json -print -quit | grep -q .; then
+      complete=0
+      break
+    fi
+  done
+
+  if [ "$complete" -eq 1 ]; then
+    mv "$candidate_dir" "$selected_dir"
+    echo "Downloaded Ginkgo timings from successful main e2e run $run_id."
+    exit 0
+  fi
+
+  rm -rf "$candidate_dir"
+done <<< "$run_ids"
+
+echo "No successful main e2e run with all Ginkgo reports found; using equal shard weights."

test/e2e/e2e_suite_test.go

@@ -2,8 +2,9 @@ package e2e
 
 import (
 	"context"
+	"encoding/json"
+	"errors"
 	"fmt"
-	"hash/fnv"
 	"os"
 	"path/filepath"
 	"strconv"
@@ -15,6 +16,7 @@ import (
 	"github.com/go-logr/zapr"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
+	"github.com/sethvargo/go-envconfig"
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
@@ -54,9 +56,62 @@ var releases = map[string][]upgrade.ClickHouseVersion{
 	},
 }
 
+type shardingConfig struct {
+	Index    int    `env:"E2E_SHARD_INDEX, default=0"`
+	Total    int    `env:"E2E_SHARD_TOTAL, default=0"`
+	PlanPath string `env:"E2E_SHARD_PLAN"`
+
+	shardAssignments map[string]int
+}
+
+func (c *shardingConfig) Load() error {
+	if err := envconfig.Process(context.Background(), c); err != nil {
+		return fmt.Errorf("load sharding config from env: %w", err)
+	}
+
+	if c.Total <= 1 {
+		return nil
+	}
+
+	if c.PlanPath == "" {
+		return errors.New("sharding plan path must be set if more than 1 shard")
+	}
+
+	if c.Index < 1 || c.Index > c.Total {
+		return fmt.Errorf("invalid shard index %d, should be between 1 and %d", c.Index, c.Total)
+	}
+
+	data, err := os.ReadFile(filepath.Clean(c.PlanPath))
+	if err != nil {
+		return fmt.Errorf("read e2e shard plan: %w", err)
+	}
+
+	if err = json.Unmarshal(data, &c.shardAssignments); err != nil {
+		return fmt.Errorf("decode e2e shard plan: %w", err)
+	}
+
+	return nil
+}
+
+func (c *shardingConfig) Enabled(spec string) (bool, error) {
+	if c.Total <= 1 {
+		return true, nil
+	}
+
+	shard, ok := c.shardAssignments[spec]
+	if !ok {
+		return false, fmt.Errorf("test %q is not assigned in sharding plan", spec)
+	}
+
+	if shard < 1 || shard > c.Total {
+		return false, fmt.Errorf("invalid assignment for %q", spec)
+	}
+
+	return shard == c.Index, nil
+}
+
 var (
-	shardIndex     int
-	shardTotal     int
+	sharding       shardingConfig
 	k8sClient      client.Client
 	config         *rest.Config
 	podDialer      controllerutil.DialContextFunc
@@ -74,6 +129,10 @@ var (
 func TestE2E(t *testing.T) {
 	RegisterFailHandler(Fail)
 
+	if err := sharding.Load(); err != nil {
+		t.Fatalf("failed to load sharding config: %v", err)
+	}
+
 	_, _ = fmt.Fprintf(GinkgoWriter, "Starting clickhouse-operator suite\n")
 
 	RunSpecs(t, "e2e suite")
@@ -91,26 +150,6 @@ var _ = ReportBeforeSuite(func(report Report) {
 	if len(offenders) > 0 {
 		Fail("Ordered containers are not allowed (sharding-incompatible):\n" + strings.Join(offenders, "\n"))
 	}
-
-	indexStr := os.Getenv("E2E_SHARD_INDEX")
-
-	totalStr := os.Getenv("E2E_SHARD_TOTAL")
-	if indexStr == "" && totalStr == "" {
-		return
-	}
-
-	total, err := strconv.Atoi(totalStr)
-	if err != nil || total < 1 {
-		Fail(fmt.Sprintf("invalid E2E_SHARD_TOTAL=%q", totalStr))
-	}
-
-	index, err := strconv.Atoi(indexStr)
-	if err != nil || index < 1 || index > total {
-		Fail(fmt.Sprintf("invalid E2E_SHARD_INDEX=%q (total=%d)", indexStr, total))
-	}
-
-	shardIndex = index
-	shardTotal = total
 })
 
 var _ = BeforeSuite(func(ctx context.Context) {
@@ -193,16 +232,14 @@ var _ = BeforeSuite(func(ctx context.Context) {
 	}
 })
 
-var _ = JustBeforeEach(func() {
-	if shardTotal <= 1 {
-		return
+var _ = BeforeEach(func() {
+	enabled, err := sharding.Enabled(CurrentSpecReport().FullText())
+	if err != nil {
+		Fail(err.Error())
 	}
 
-	h := fnv.New32a()
-
-	_, _ = h.Write([]byte(CurrentSpecReport().FullText()))
-	if int(h.Sum32()%uint32(shardTotal))+1 != shardIndex {
-		Skip(fmt.Sprintf("not in shard %d/%d", shardIndex, shardTotal))
+	if !enabled {
+		Skip(fmt.Sprintf("not in shard %d/%d", sharding.Index, sharding.Total))
 	}
 })