tidb,tiproxy: expose max surge on tidb/tiproxygroup spec (#6797)

YangKeao · web-flow · commit b38569cf249e · 2026-04-02T12:27:27.000Z
diff --git a/api/core/v1alpha1/tidb_types.go b/api/core/v1alpha1/tidb_types.go
@@ -138,6 +138,13 @@ type TiDBGroupSpec struct {
 	// +optional
 	MinReadySeconds *int64 `json:"minReadySeconds,omitempty"`
 
+	// MaxSurge specifies the maximum number of additional instances that may be created during rolling restart.
+	// It only takes effect when the update requires restarting TiDB instances.
+	// Defaults to 1.
+	// +kubebuilder:validation:Minimum=1
+	// +optional
+	MaxSurge *int32 `json:"maxSurge,omitempty"`
+
 	// Template is the instance template
 	Template TiDBTemplate `json:"template"`
 }
diff --git a/api/core/v1alpha1/tiproxy_types.go b/api/core/v1alpha1/tiproxy_types.go
@@ -123,6 +123,13 @@ type TiProxyGroupSpec struct {
 	// +optional
 	MinReadySeconds *int64 `json:"minReadySeconds,omitempty"`
 
+	// MaxSurge specifies the maximum number of additional instances that may be created during rolling restart.
+	// It only takes effect when the update requires restarting TiProxy instances.
+	// Defaults to 1.
+	// +kubebuilder:validation:Minimum=1
+	// +optional
+	MaxSurge *int32 `json:"maxSurge,omitempty"`
+
 	Template TiProxyTemplate `json:"template"`
 }
 
diff --git a/api/core/v1alpha1/zz_generated.deepcopy.go b/api/core/v1alpha1/zz_generated.deepcopy.go
diff --git a/manifests/crd/core.pingcap.com_tidbgroups.yaml b/manifests/crd/core.pingcap.com_tidbgroups.yaml
@@ -110,6 +110,14 @@ spec:
                   - IndependentKVEngineWorker
                   type: string
                 type: array
+              maxSurge:
+                description: |-
+                  MaxSurge specifies the maximum number of additional instances that may be created during rolling restart.
+                  It only takes effect when the update requires restarting TiDB instances.
+                  Defaults to 1.
+                format: int32
+                minimum: 1
+                type: integer
               minReadySeconds:
                 description: MinReadySeconds specifies the minimum number of seconds
                   for which a newly created pod be ready without any of its containers
diff --git a/manifests/crd/core.pingcap.com_tiproxygroups.yaml b/manifests/crd/core.pingcap.com_tiproxygroups.yaml
@@ -110,6 +110,14 @@ spec:
                   - IndependentKVEngineWorker
                   type: string
                 type: array
+              maxSurge:
+                description: |-
+                  MaxSurge specifies the maximum number of additional instances that may be created during rolling restart.
+                  It only takes effect when the update requires restarting TiProxy instances.
+                  Defaults to 1.
+                format: int32
+                minimum: 1
+                type: integer
               minReadySeconds:
                 description: MinReadySeconds specifies the minimum number of seconds
                   for which a newly created pod be ready without any of its containers
diff --git a/pkg/controllers/tidbgroup/tasks/updater.go b/pkg/controllers/tidbgroup/tasks/updater.go
@@ -89,6 +89,9 @@ func TaskUpdater(state *ReconcileContext, c client.Client, af tracker.AllocateFa
 		noUpdate := false
 		if needRestart {
 			maxSurge, maxUnavailable = 1, 0
+			if dbg.Spec.MaxSurge != nil {
+				maxSurge = int(*dbg.Spec.MaxSurge)
+			}
 			noUpdate = true
 		}
 
diff --git a/pkg/controllers/tidbgroup/tasks/updater_test.go b/pkg/controllers/tidbgroup/tasks/updater_test.go
@@ -254,6 +254,28 @@ func TestTaskUpdater(t *testing.T) {
 			expectedStatus:  task.SWait,
 			expectedTiDBNum: 3,
 		},
+		{
+			desc: "rolling restart honors maxSurge",
+			state: &ReconcileContext{
+				State: &state{
+					dbg: fake.FakeObj("aaa", func(obj *v1alpha1.TiDBGroup) *v1alpha1.TiDBGroup {
+						obj.Spec.Replicas = ptr.To[int32](2)
+						obj.Spec.MaxSurge = ptr.To[int32](2)
+						obj.Spec.Template.Spec.Image = ptr.To("pingcap/tidb:v8.1.1")
+						return obj
+					}),
+					cluster: fake.FakeObj[v1alpha1.Cluster]("cluster"),
+					dbs: []*v1alpha1.TiDB{
+						fakeAvailableTiDB("aaa-xxx", fake.FakeObj[v1alpha1.TiDBGroup]("aaa"), oldRevision),
+						fakeAvailableTiDB("aaa-yyy", fake.FakeObj[v1alpha1.TiDBGroup]("aaa"), oldRevision),
+					},
+					updateRevision: newRevision,
+				},
+			},
+
+			expectedStatus:  task.SWait,
+			expectedTiDBNum: 4,
+		},
 	}
 
 	for i := range cases {
diff --git a/pkg/controllers/tiproxygroup/tasks/updater.go b/pkg/controllers/tiproxygroup/tasks/updater.go
@@ -86,6 +86,9 @@ func TaskUpdater(state *ReconcileContext, c client.Client, af tracker.AllocateFa
 		noUpdate := false
 		if needRestart {
 			maxSurge, maxUnavailable = 1, 0
+			if proxyg.Spec.MaxSurge != nil {
+				maxSurge = int(*proxyg.Spec.MaxSurge)
+			}
 			noUpdate = true
 		}
 
diff --git a/pkg/controllers/tiproxygroup/tasks/updater_test.go b/pkg/controllers/tiproxygroup/tasks/updater_test.go
@@ -227,6 +227,28 @@ func TestTaskUpdater(t *testing.T) {
 			expectedStatus:     task.SWait,
 			expectedTiProxyNum: 3,
 		},
+		{
+			desc: "rolling restart honors maxSurge",
+			state: &ReconcileContext{
+				State: &state{
+					proxyg: fake.FakeObj("aaa", func(obj *v1alpha1.TiProxyGroup) *v1alpha1.TiProxyGroup {
+						obj.Spec.Replicas = ptr.To[int32](2)
+						obj.Spec.MaxSurge = ptr.To[int32](2)
+						obj.Spec.Template.Spec.Image = ptr.To("pingcap/tiproxy:v8.1.1")
+						return obj
+					}),
+					cluster: fake.FakeObj[v1alpha1.Cluster]("cluster"),
+					proxies: []*v1alpha1.TiProxy{
+						fakeAvailableTiProxy("aaa-xxx", fake.FakeObj[v1alpha1.TiProxyGroup]("aaa"), oldRevision),
+						fakeAvailableTiProxy("aaa-yyy", fake.FakeObj[v1alpha1.TiProxyGroup]("aaa"), oldRevision),
+					},
+					updateRevision: newRevision,
+				},
+			},
+
+			expectedStatus:     task.SWait,
+			expectedTiProxyNum: 4,
+		},
 	}
 
 	for i := range cases {
diff --git a/pkg/utils/hasher/hasher_test.go b/pkg/utils/hasher/hasher_test.go
@@ -42,7 +42,7 @@ func TestHash(t *testing.T) {
 				},
 			},
 			// not important, feel free to change if tidb group is changed
-			expected: "58c864c88",
+			expected: "6c9dc4f5cd",
 		},
 	}
 
diff --git a/tests/e2e/data/tidb.go b/tests/e2e/data/tidb.go
@@ -119,6 +119,12 @@ func WithHotReloadPolicy() GroupPatch[*v1alpha1.TiDBGroup] {
 	})
 }
 
+func WithTiDBMaxSurge(maxSurge int32) GroupPatch[*v1alpha1.TiDBGroup] {
+	return GroupPatchFunc[*v1alpha1.TiDBGroup](func(obj *v1alpha1.TiDBGroup) {
+		obj.Spec.MaxSurge = ptr.To(maxSurge)
+	})
+}
+
 func WithEphemeralVolume() GroupPatch[*v1alpha1.TiDBGroup] {
 	return GroupPatchFunc[*v1alpha1.TiDBGroup](func(obj *v1alpha1.TiDBGroup) {
 		if obj.Spec.Template.Spec.Overlay == nil {
diff --git a/tests/e2e/data/tiproxy.go b/tests/e2e/data/tiproxy.go
@@ -99,6 +99,12 @@ func WithHotReloadPolicyForTiProxy() GroupPatch[*v1alpha1.TiProxyGroup] {
 	})
 }
 
+func WithTiProxyMaxSurge(maxSurge int32) GroupPatch[*v1alpha1.TiProxyGroup] {
+	return GroupPatchFunc[*v1alpha1.TiProxyGroup](func(obj *v1alpha1.TiProxyGroup) {
+		obj.Spec.MaxSurge = ptr.To(maxSurge)
+	})
+}
+
 func WithTiProxyNextGen() GroupPatch[*v1alpha1.TiProxyGroup] {
 	return GroupPatchFunc[*v1alpha1.TiProxyGroup](func(obj *v1alpha1.TiProxyGroup) {
 		obj.Spec.Template.Spec.Version = "v1.4.0"
diff --git a/tests/e2e/framework/framework.go b/tests/e2e/framework/framework.go
@@ -317,12 +317,7 @@ func AsyncWaitPodsRollingUpdateOnce[
 	F client.Object,
 	T runtime.Group,
 ](ctx context.Context, f *Framework, obj F, to int) chan struct{} {
-	maxSurge := 0
-	comp := scope.Component[S]()
-	switch comp {
-	case "tiproxy", "tidb", "ticdc":
-		maxSurge = 1
-	}
+	maxSurge := rollingUpdateMaxSurge(obj)
 	ch := make(chan struct{})
 	nobj := obj.DeepCopyObject().(F)
 	go func() {
@@ -334,6 +329,25 @@ func AsyncWaitPodsRollingUpdateOnce[
 	return ch
 }
 
+func rollingUpdateMaxSurge(obj client.Object) int {
+	switch g := obj.(type) {
+	case *v1alpha1.TiDBGroup:
+		if g.Spec.MaxSurge != nil {
+			return int(*g.Spec.MaxSurge)
+		}
+		return 1
+	case *v1alpha1.TiProxyGroup:
+		if g.Spec.MaxSurge != nil {
+			return int(*g.Spec.MaxSurge)
+		}
+		return 1
+	case *v1alpha1.TiCDCGroup:
+		return 1
+	default:
+		return 0
+	}
+}
+
 func (f *Framework) Describe(fn func(o *desc.Options), optss ...[]desc.Option) {
 	var entries []ginkgo.TableEntry
 	for _, opts := range optss {
diff --git a/tests/e2e/suite/scale/tidb.go b/tests/e2e/suite/scale/tidb.go
@@ -0,0 +1,67 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package scale
+
+import (
+	"context"
+
+	"github.com/onsi/ginkgo/v2"
+
+	"github.com/pingcap/tidb-operator/v2/pkg/runtime"
+	"github.com/pingcap/tidb-operator/v2/pkg/runtime/scope"
+	"github.com/pingcap/tidb-operator/v2/tests/e2e/data"
+	"github.com/pingcap/tidb-operator/v2/tests/e2e/framework"
+	"github.com/pingcap/tidb-operator/v2/tests/e2e/framework/action"
+	"github.com/pingcap/tidb-operator/v2/tests/e2e/framework/desc"
+	"github.com/pingcap/tidb-operator/v2/tests/e2e/label"
+	"github.com/pingcap/tidb-operator/v2/tests/e2e/utils/waiter"
+)
+
+var _ = ginkgo.Describe("Scale TiDB", label.TiDB, label.P0, label.Scale, func() {
+	f := framework.New()
+	f.Setup()
+
+	ginkgo.It("support scale from 3 to 5 and rolling restart at same time with maxSurge 2", label.Update, func(ctx context.Context) {
+		o := desc.DefaultOptions()
+
+		pdg := action.MustCreatePD(ctx, f, o)
+		kvg := action.MustCreateTiKV(ctx, f, o)
+		dbg := action.MustCreateTiDB(ctx, f, o,
+			data.WithReplicas[scope.TiDBGroup](3),
+			data.WithTiDBMaxSurge(2),
+		)
+
+		f.WaitForPDGroupReady(ctx, pdg)
+		f.WaitForTiKVGroupReady(ctx, kvg)
+		f.WaitForTiDBGroupReady(ctx, dbg)
+
+		nctx, cancel := context.WithCancel(ctx)
+		done := framework.AsyncWaitPodsRollingUpdateOnce[scope.TiDBGroup](nctx, f, dbg, 5)
+		defer func() { <-done }()
+		defer cancel()
+
+		// Record the latest old pod creation timestamp so we can verify the rolling restart
+		// actually creates new pods instead of only scaling out.
+		changeTime, err := waiter.MaxPodsCreateTimestamp[scope.TiDBGroup](ctx, f.Client, dbg)
+		f.Must(err)
+
+		// Patch replicas and pod template at the same time to force the updater into the
+		// "scale out + rolling restart" path where spec.maxSurge should be honored.
+		action.MustScaleAndRollingRestart[scope.TiDBGroup](ctx, f, dbg, 5)
+
+		f.Must(waiter.WaitForPodsRecreated(ctx, f.Client, runtime.FromTiDBGroup(dbg), *changeTime, waiter.LongTaskTimeout))
+		f.WaitForTiDBGroupReady(ctx, dbg)
+	})
+})
diff --git a/tests/e2e/suite/scale/tiproxy.go b/tests/e2e/suite/scale/tiproxy.go
@@ -0,0 +1,69 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package scale
+
+import (
+	"context"
+
+	"github.com/onsi/ginkgo/v2"
+
+	"github.com/pingcap/tidb-operator/v2/pkg/runtime"
+	"github.com/pingcap/tidb-operator/v2/pkg/runtime/scope"
+	"github.com/pingcap/tidb-operator/v2/tests/e2e/data"
+	"github.com/pingcap/tidb-operator/v2/tests/e2e/framework"
+	"github.com/pingcap/tidb-operator/v2/tests/e2e/framework/action"
+	"github.com/pingcap/tidb-operator/v2/tests/e2e/framework/desc"
+	"github.com/pingcap/tidb-operator/v2/tests/e2e/label"
+	"github.com/pingcap/tidb-operator/v2/tests/e2e/utils/waiter"
+)
+
+var _ = ginkgo.Describe("Scale TiProxy", label.TiProxy, label.P0, label.Scale, func() {
+	f := framework.New()
+	f.Setup()
+
+	ginkgo.It("support scale from 3 to 5 and rolling restart at same time with maxSurge 2", label.Update, func(ctx context.Context) {
+		o := desc.DefaultOptions()
+
+		pdg := action.MustCreatePD(ctx, f, o)
+		kvg := action.MustCreateTiKV(ctx, f, o)
+		dbg := action.MustCreateTiDB(ctx, f, o)
+		proxyg := action.MustCreateTiProxy(ctx, f, o,
+			data.WithReplicas[scope.TiProxyGroup](3),
+			data.WithTiProxyMaxSurge(2),
+		)
+
+		f.WaitForPDGroupReady(ctx, pdg)
+		f.WaitForTiKVGroupReady(ctx, kvg)
+		f.WaitForTiDBGroupReady(ctx, dbg)
+		f.WaitForTiProxyGroupReady(ctx, proxyg)
+
+		nctx, cancel := context.WithCancel(ctx)
+		done := framework.AsyncWaitPodsRollingUpdateOnce[scope.TiProxyGroup](nctx, f, proxyg, 5)
+		defer func() { <-done }()
+		defer cancel()
+
+		// Record the latest old pod creation timestamp so we can verify the rolling restart
+		// actually creates new pods instead of only scaling out.
+		changeTime, err := waiter.MaxPodsCreateTimestamp[scope.TiProxyGroup](ctx, f.Client, proxyg)
+		f.Must(err)
+
+		// Patch replicas and pod template at the same time to force the updater into the
+		// "scale out + rolling restart" path where spec.maxSurge should be honored.
+		action.MustScaleAndRollingRestart[scope.TiProxyGroup](ctx, f, proxyg, 5)
+
+		f.Must(waiter.WaitForPodsRecreated(ctx, f.Client, runtime.FromTiProxyGroup(proxyg), *changeTime, waiter.LongTaskTimeout))
+		f.WaitForTiProxyGroupReady(ctx, proxyg)
+	})
+})
diff --git a/tests/validation/common_test.go b/tests/validation/common_test.go
diff --git a/tests/validation/tidb_test.go b/tests/validation/tidb_test.go
diff --git a/tests/validation/tiproxy_test.go b/tests/validation/tiproxy_test.go

Original file line number	Diff line number	Diff line change
`@@ -89,6 +89,9 @@ func TaskUpdater(state *ReconcileContext, c client.Client, af tracker.AllocateFa`
`89`	`89`	`noUpdate := false`
`90`	`90`	`if needRestart {`
`91`	`91`	`maxSurge, maxUnavailable = 1, 0`
	`92`	`+ if dbg.Spec.MaxSurge != nil {`
	`93`	`+ maxSurge = int(*dbg.Spec.MaxSurge)`
	`94`	`+ }`
`92`	`95`	`noUpdate = true`
`93`	`96`	`}`
`94`	`97`
Original file line number	Diff line number	Diff line change
`@@ -86,6 +86,9 @@ func TaskUpdater(state *ReconcileContext, c client.Client, af tracker.AllocateFa`
`86`	`86`	`noUpdate := false`
`87`	`87`	`if needRestart {`
`88`	`88`	`maxSurge, maxUnavailable = 1, 0`
	`89`	`+ if proxyg.Spec.MaxSurge != nil {`
	`90`	`+ maxSurge = int(*proxyg.Spec.MaxSurge)`
	`91`	`+ }`
`89`	`92`	`noUpdate = true`
`90`	`93`	`}`
`91`	`94`
Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,7 @@ func TestHash(t *testing.T) {`
`42`	`42`	`},`
`43`	`43`	`},`
`44`	`44`	`// not important, feel free to change if tidb group is changed`
`45`		`- expected: "58c864c88",`
	`45`	`+ expected: "6c9dc4f5cd",`
`46`	`46`	`},`
`47`	`47`	`}`
`48`	`48`