feat(tiflash): support tiflash ready api (#6523)

liubog2008 · web-flow · commit 9918f2995bff · 2025-10-28T08:28:55.000Z
diff --git a/api/meta/v1alpha1/feature.go b/api/meta/v1alpha1/feature.go
@@ -18,7 +18,7 @@ package v1alpha1
 // NOTE(liubo02): +enum is not supported now, we have to add all enum into comments
 // NOTE(liubo02): It's supported by https://github.com/kubernetes-sigs/controller-tools/pull/1179
 //
-// +kubebuilder:validation:Enum=FeatureModification;VolumeAttributesClass;DisablePDDefaultReadinessProbe;UsePDReadyAPI;SessionTokenSigning;ClusterSubdomain;TerminableLogTailer;UseTSOReadyAPI;UseSchedulingReadyAPI;UseTiKVReadyAPI;UsePDReadyAPIV2
+// +kubebuilder:validation:Enum=FeatureModification;VolumeAttributesClass;DisablePDDefaultReadinessProbe;UsePDReadyAPI;SessionTokenSigning;ClusterSubdomain;TerminableLogTailer;UseTSOReadyAPI;UseSchedulingReadyAPI;UseTiKVReadyAPI;UsePDReadyAPIV2;UseTiFlashReadyAPI
 // +enum
 type Feature string
 
@@ -98,4 +98,8 @@ const (
 	// UsePDReadyAPIV2 means use PD's /readyz API as the readiness probe.
 	UsePDReadyAPIV2      Feature      = "UsePDReadyAPIV2"
 	UsePDReadyAPIV2Stage FeatureStage = FeatureStageAlpha
+
+	// UseTiFlashReadyAPI means use TiFlash's /readyz API as the readiness probe.
+	UseTiFlashReadyAPI      Feature      = "UseTiFlashReadyAPI"
+	UseTiFlashReadyAPIStage FeatureStage = FeatureStageAlpha
 )
diff --git a/manifests/crd/core.pingcap.com_clusters.yaml b/manifests/crd/core.pingcap.com_clusters.yaml
@@ -97,6 +97,7 @@ spec:
                       - UseSchedulingReadyAPI
                       - UseTiKVReadyAPI
                       - UsePDReadyAPIV2
+                      - UseTiFlashReadyAPI
                       type: string
                   required:
                   - name
@@ -305,6 +306,7 @@ spec:
                       - UseSchedulingReadyAPI
                       - UseTiKVReadyAPI
                       - UsePDReadyAPIV2
+                      - UseTiFlashReadyAPI
                       type: string
                     stage:
                       type: string
diff --git a/manifests/crd/core.pingcap.com_pdgroups.yaml b/manifests/crd/core.pingcap.com_pdgroups.yaml
@@ -115,6 +115,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               replicas:
diff --git a/manifests/crd/core.pingcap.com_pds.yaml b/manifests/crd/core.pingcap.com_pds.yaml
@@ -91,6 +91,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               image:
diff --git a/manifests/crd/core.pingcap.com_replicationworkergroups.yaml b/manifests/crd/core.pingcap.com_replicationworkergroups.yaml
@@ -106,6 +106,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               replicas:
diff --git a/manifests/crd/core.pingcap.com_replicationworkers.yaml b/manifests/crd/core.pingcap.com_replicationworkers.yaml
@@ -91,6 +91,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               image:
diff --git a/manifests/crd/core.pingcap.com_schedulergroups.yaml b/manifests/crd/core.pingcap.com_schedulergroups.yaml
@@ -110,6 +110,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               replicas:
diff --git a/manifests/crd/core.pingcap.com_schedulers.yaml b/manifests/crd/core.pingcap.com_schedulers.yaml
@@ -96,6 +96,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               image:
diff --git a/manifests/crd/core.pingcap.com_schedulinggroups.yaml b/manifests/crd/core.pingcap.com_schedulinggroups.yaml
@@ -105,6 +105,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               replicas:
diff --git a/manifests/crd/core.pingcap.com_schedulings.yaml b/manifests/crd/core.pingcap.com_schedulings.yaml
@@ -91,6 +91,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               image:
diff --git a/manifests/crd/core.pingcap.com_ticdcgroups.yaml b/manifests/crd/core.pingcap.com_ticdcgroups.yaml
@@ -101,6 +101,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               replicas:
diff --git a/manifests/crd/core.pingcap.com_ticdcs.yaml b/manifests/crd/core.pingcap.com_ticdcs.yaml
@@ -88,6 +88,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               image:
diff --git a/manifests/crd/core.pingcap.com_tidbgroups.yaml b/manifests/crd/core.pingcap.com_tidbgroups.yaml
@@ -104,6 +104,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               replicas:
diff --git a/manifests/crd/core.pingcap.com_tidbs.yaml b/manifests/crd/core.pingcap.com_tidbs.yaml
@@ -90,6 +90,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               image:
diff --git a/manifests/crd/core.pingcap.com_tiflashes.yaml b/manifests/crd/core.pingcap.com_tiflashes.yaml
@@ -96,6 +96,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               image:
diff --git a/manifests/crd/core.pingcap.com_tiflashgroups.yaml b/manifests/crd/core.pingcap.com_tiflashgroups.yaml
@@ -103,6 +103,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               replicas:
diff --git a/manifests/crd/core.pingcap.com_tikvgroups.yaml b/manifests/crd/core.pingcap.com_tikvgroups.yaml
@@ -104,6 +104,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               replicas:
diff --git a/manifests/crd/core.pingcap.com_tikvs.yaml b/manifests/crd/core.pingcap.com_tikvs.yaml
@@ -96,6 +96,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               image:
diff --git a/manifests/crd/core.pingcap.com_tiproxies.yaml b/manifests/crd/core.pingcap.com_tiproxies.yaml
@@ -90,6 +90,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               image:
diff --git a/manifests/crd/core.pingcap.com_tiproxygroups.yaml b/manifests/crd/core.pingcap.com_tiproxygroups.yaml
@@ -104,6 +104,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               replicas:
diff --git a/manifests/crd/core.pingcap.com_tsogroups.yaml b/manifests/crd/core.pingcap.com_tsogroups.yaml
@@ -104,6 +104,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               replicas:
diff --git a/manifests/crd/core.pingcap.com_tsos.yaml b/manifests/crd/core.pingcap.com_tsos.yaml
@@ -90,6 +90,7 @@ spec:
                   - UseSchedulingReadyAPI
                   - UseTiKVReadyAPI
                   - UsePDReadyAPIV2
+                  - UseTiFlashReadyAPI
                   type: string
                 type: array
               image:
diff --git a/pkg/controllers/tiflash/tasks/ctx.go b/pkg/controllers/tiflash/tasks/ctx.go
@@ -20,6 +20,7 @@ import (
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"k8s.io/apimachinery/pkg/api/errors"
 
+	metav1alpha1 "github.com/pingcap/tidb-operator/api/v2/meta/v1alpha1"
 	tiflashconfig "github.com/pingcap/tidb-operator/pkg/configs/tiflash"
 	"github.com/pingcap/tidb-operator/pkg/tiflashapi/v1"
 	"github.com/pingcap/tidb-operator/pkg/timanager"
@@ -38,8 +39,6 @@ type ReconcileContext struct {
 	Store       *pdv1.Store
 	StoreLabels []*metapb.StoreLabel
 
-	StoreStatus tiflashapi.Status
-
 	PDSynced bool
 }
 
@@ -105,28 +104,31 @@ func TaskContextInfoFromPD(state *ReconcileContext, cm pdm.PDClientManager, fcm
 			state.StoreLabels = append(state.StoreLabels, &metapb.StoreLabel{Key: k, Value: v})
 		}
 
-		pod := state.Pod()
-
-		var status tiflashapi.Status
-		if pod != nil && statefulset.IsPodReady(pod) {
-			s, err := fc.GetStoreStatus(ctx)
-			if err != nil {
-				return task.Fail().With("failed to get store status: %v", err)
-			}
-			status = s
-			// TODO: check IsBusy ?
-			if status == tiflashapi.Running {
-				state.SetHealthy()
+		if state.FeatureGates().Enabled(metav1alpha1.UseTiFlashReadyAPI) {
+			state.SetHealthy()
+		} else {
+			pod := state.Pod()
+
+			var status tiflashapi.Status
+			if pod != nil && statefulset.IsPodReady(pod) {
+				s, err := fc.GetStoreStatus(ctx)
+				if err != nil {
+					return task.Fail().With("failed to get store status: %v", err)
+				}
+				status = s
+				// TODO: check IsBusy ?
+				if status == tiflashapi.Running {
+					state.SetHealthy()
+				}
 			}
-			state.StoreStatus = status
 		}
 
-		return task.Complete().With("get store info, state: %s, leader count: %v, region count: %v, busy: %v, status: %v",
+		return task.Complete().With("get store info, state: %s, leader count: %v, region count: %v, busy: %v, health: %v",
 			state.GetStoreState(),
 			state.GetLeaderCount(),
 			state.GetRegionCount(),
 			state.IsStoreBusy(),
-			status,
+			state.IsHealthy(),
 		)
 	})
 }
diff --git a/pkg/controllers/tiflash/tasks/pod.go b/pkg/controllers/tiflash/tasks/pod.go
@@ -17,6 +17,7 @@ package tasks
 import (
 	"context"
 	"fmt"
+	"path"
 	"path/filepath"
 
 	"github.com/go-logr/logr"
@@ -216,6 +217,10 @@ func newPod(cluster *v1alpha1.Cluster, tiflash *v1alpha1.TiFlash, store *pdv1.St
 		pod.Labels[v1alpha1.LabelKeyStoreID] = store.ID
 	}
 
+	if fg.Enabled(metav1alpha1.UseTiFlashReadyAPI) {
+		pod.Spec.Containers[0].ReadinessProbe = buildReadinessProbe(cluster, coreutil.TiFlashProxyStatusPort(tiflash))
+	}
+
 	if tiflash.Spec.Overlay != nil {
 		overlay.OverlayPod(pod, tiflash.Spec.Overlay.Pod)
 	}
@@ -265,3 +270,39 @@ func buildLogTailerContainer(
 	}
 	return c
 }
+
+func buildReadinessProbe(cluster *v1alpha1.Cluster, port int32) *corev1.Probe {
+	tlsClusterEnabled := coreutil.IsTLSClusterEnabled(cluster)
+
+	scheme := "http"
+	if tlsClusterEnabled {
+		scheme = "https"
+	}
+
+	readinessURL := fmt.Sprintf("%s://127.0.0.1:%d/tiflash/readyz", scheme, port)
+	var command []string
+	command = append(command, "curl", readinessURL,
+		// Fail silently (no output at all) on server errors
+		// without this if the server return 500, the exist code will be 0
+		// and probe is success.
+		"--fail",
+		// Silent mode, only print error
+		"-sS",
+		// follow 301 or 302 redirect
+		"--location")
+
+	if tlsClusterEnabled {
+		cacert := path.Join(v1alpha1.DirPathClusterTLSTiFlash, corev1.ServiceAccountRootCAKey)
+		cert := path.Join(v1alpha1.DirPathClusterTLSTiFlash, corev1.TLSCertKey)
+		key := path.Join(v1alpha1.DirPathClusterTLSTiFlash, corev1.TLSPrivateKeyKey)
+		command = append(command, "--cacert", cacert, "--cert", cert, "--key", key)
+	}
+
+	return &corev1.Probe{
+		ProbeHandler: corev1.ProbeHandler{
+			Exec: &corev1.ExecAction{
+				Command: command,
+			},
+		},
+	}
+}
diff --git a/pkg/controllers/tiflash/tasks/status.go b/pkg/controllers/tiflash/tasks/status.go
@@ -87,23 +87,6 @@ func TaskStatus(state *ReconcileContext, c client.Client) task.Task {
 	})
 }
 
-func TaskStoreStatus(state *ReconcileContext) task.Task {
-	return task.NameTaskFunc("StoreStatus", func(ctx context.Context) task.Result {
-		needUpdate := state.IsStatusChanged()
-		tiflash := state.TiFlash()
-		if state.Store != nil {
-			needUpdate = compare.SetIfChanged(&tiflash.Status.ID, state.Store.ID) || needUpdate
-		}
-		needUpdate = compare.SetIfChanged(&tiflash.Status.State, string(state.GetStoreState())) || needUpdate
-		if needUpdate {
-			state.SetStatusChanged()
-			return task.Complete().With("store state is changed")
-		}
-
-		return task.Complete().With("store state is not changed")
-	})
-}
-
 func syncSuspendCond(tiflash *v1alpha1.TiFlash) bool {
 	// always set it as unsuspended
 	return meta.SetStatusCondition(&tiflash.Status.Conditions, metav1.Condition{
diff --git a/pkg/controllers/tiflash/tasks/status_test.go b/pkg/controllers/tiflash/tasks/status_test.go
@@ -27,7 +27,6 @@ import (
 
 	"github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
 	"github.com/pingcap/tidb-operator/pkg/client"
-	"github.com/pingcap/tidb-operator/pkg/tiflashapi/v1"
 	pdv1 "github.com/pingcap/tidb-operator/pkg/timanager/apis/pd/v1"
 	"github.com/pingcap/tidb-operator/pkg/utils/fake"
 	"github.com/pingcap/tidb-operator/pkg/utils/task/v3"
@@ -129,7 +128,6 @@ func TestTaskStatus(t *testing.T) {
 				Store: &pdv1.Store{
 					ID: fakeTiFlashName,
 				},
-				StoreStatus: tiflashapi.Running,
 			},
 
 			expectedStatus: task.SComplete,
@@ -278,7 +276,6 @@ func TestTaskStatus(t *testing.T) {
 				Store: &pdv1.Store{
 					ID: fakeTiFlashName,
 				},
-				StoreStatus: tiflashapi.Ready,
 			},
 			expectedStatus: task.SRetry,
 			expectedObj: fake.FakeObj(fakeTiFlashName, func(obj *v1alpha1.TiFlash) *v1alpha1.TiFlash {
diff --git a/pkg/controllers/tikv/tasks/pod.go b/pkg/controllers/tikv/tasks/pod.go
@@ -271,14 +271,14 @@ func newPod(cluster *v1alpha1.Cluster, tikv *v1alpha1.TiKV, store *pdv1.Store, g
 		pod.Labels[v1alpha1.LabelKeyStoreID] = store.ID
 	}
 
-	if tikv.Spec.Overlay != nil {
-		overlay.OverlayPod(pod, tikv.Spec.Overlay.Pod)
-	}
-
 	if g.Enabled(metav1alpha1.UseTiKVReadyAPI) {
 		pod.Spec.Containers[0].ReadinessProbe = buildReadinessProbeWithReadyAPI(cluster, coreutil.TiKVStatusPort(tikv))
 	}
 
+	if tikv.Spec.Overlay != nil {
+		overlay.OverlayPod(pod, tikv.Spec.Overlay.Pod)
+	}
+
 	reloadable.MustEncodeLastTiKVTemplate(tikv, pod)
 	return pod
 }
diff --git a/pkg/features/reload.go b/pkg/features/reload.go
@@ -68,6 +68,9 @@ var unreloadable = map[meta.Feature][]meta.Component{
 	meta.UsePDReadyAPIV2: {
 		meta.ComponentPD,
 	},
+	meta.UseTiFlashReadyAPI: {
+		meta.ComponentTiFlash,
+	},
 }
 
 func Reloadable(c meta.Component, update, current []meta.Feature) bool {
diff --git a/tests/e2e/suite/availability/tiflash.go b/tests/e2e/suite/availability/tiflash.go
diff --git a/tests/e2e/suite/cluster/tls.go b/tests/e2e/suite/cluster/tls.go
