Add readinessProbe config for tidb

docs/api-references/docs.md

@@ -12095,6 +12095,40 @@ string
 </tr>
 </tbody>
 </table>
+<h3 id="tidbprobe">TiDBProbe</h3>
+<p>
+(<em>Appears on:</em>
+<a href="#tidbspec">TiDBSpec</a>)
+</p>
+<p>
+<p>TiDBProbe contains details of probing tidb.
+default probe by TCPPort on 4000.</p>
+</p>
+<table>
+<thead>
+<tr>
+<th>Field</th>
+<th>Description</th>
+</tr>
+</thead>
+<tbody>
+<tr>
+<td>
+<code>type</code></br>
+<em>
+string
+</em>
+</td>
+<td>
+<em>(Optional)</em>
+<p>&ldquo;tcp&rdquo; will use TCP socket to connetct port 4000</p>
+<p>&ldquo;command&rdquo; will probe the status api of tidb.
+This will use curl command to request tidb, before v4.0.9 there is no curl in the image,
+So do not use this before v4.0.9.</p>
+</td>
+</tr>
+</tbody>
+</table>
 <h3 id="tidbservicespec">TiDBServiceSpec</h3>
 <p>
 (<em>Appears on:</em>
@@ -12443,6 +12477,21 @@ events. For the PostStart and PreStop lifecycle handlers, management of the cont
 until the action is complete, unless the container process fails, in which case the handler is aborted.</p>
 </td>
 </tr>
+<tr>
+<td>
+<code>readinessProbe</code></br>
+<em>
+<a href="#tidbprobe">
+TiDBProbe
+</a>
+</em>
+</td>
+<td>
+<em>(Optional)</em>
+<p>ReadinessProbe describes actions that probe the tidb&rsquo;s readiness.
+the default behavior is like setting type as &ldquo;tcp&rdquo;</p>
+</td>
+</tr>
 </tbody>
 </table>
 <h3 id="tidbstatus">TiDBStatus</h3>

manifests/crd.yaml

@@ -6510,6 +6510,11 @@ spec:
                   type: object
                 priorityClassName:
                   type: string
+                readinessProbe:
+                  properties:
+                    type:
+                      type: string
+                  type: object
                 replicas:
                   format: int32
                   type: integer

pkg/apis/pingcap/v1alpha1/types.go

@@ -621,6 +621,29 @@ type TiDBSpec struct {
 	// until the action is complete, unless the container process fails, in which case the handler is aborted.
 	// +optional
 	Lifecycle *corev1.Lifecycle `json:"lifecycle,omitempty"`
+	// ReadinessProbe describes actions that probe the tidb's readiness.
+	// the default behavior is like setting type as "tcp"
+	// +optional
+	ReadinessProbe *TiDBProbe `json:"readinessProbe,omitempty"`
+}
+
+const (
+	TCPProbeType     string = "tcp"
+	CommandProbeType string = "command"
+)
+
+// +k8s:openapi-gen=true
+// TiDBProbe contains details of probing tidb.
+// default probe by TCPPort on 4000.
+type TiDBProbe struct {
+	// "tcp" will use TCP socket to connetct port 4000
+	//
+	// "command" will probe the status api of tidb.
+	// This will use curl command to request tidb, before v4.0.9 there is no curl in the image,
+	// So do not use this before v4.0.9.
+	// +kubebuilder:validation:Enum=tcp,command
+	// +optional
+	Type *string `json:"type,omitempty"` // tcp or command
 }
 
 // +k8s:openapi-gen=true

pkg/manager/member/tidb_member_manager.go

@@ -695,11 +695,7 @@ func getNewTiDBSetForTidbCluster(tc *v1alpha1.TidbCluster, cm *corev1.ConfigMap)
 		Resources:    controller.ContainerResource(tc.Spec.TiDB.ResourceRequirements),
 		Env:          util.AppendEnv(envs, baseTiDBSpec.Env()),
 		ReadinessProbe: &corev1.Probe{
-			Handler: corev1.Handler{
-				TCPSocket: &corev1.TCPSocketAction{
-					Port: intstr.FromInt(4000),
-				},
-			},
+			Handler:             buildTiDBReadinessProbHandler(tc),
 			InitialDelaySeconds: int32(10),
 		},
 	}
@@ -851,6 +847,53 @@ func tidbStatefulSetIsUpgrading(podLister corelisters.PodLister, set *apps.State
 	return false, nil
 }
 
+func buildTiDBReadinessProbHandler(tc *v1alpha1.TidbCluster) corev1.Handler {
+	if tc.Spec.TiDB.ReadinessProbe != nil {
+		if tp := tc.Spec.TiDB.ReadinessProbe.Type; tp != nil {
+			if *tp == v1alpha1.CommandProbeType {
+				command := buildTiDBProbeCommand(tc)
+				return corev1.Handler{
+					Exec: &corev1.ExecAction{
+						Command: command,
+					},
+				}
+			}
+		}
+	}
+
+	// fall to default case v1alpha1.TCPProbeType
+	return corev1.Handler{
+		TCPSocket: &corev1.TCPSocketAction{
+			Port: intstr.FromInt(4000),
+		},
+	}
+}
+
+func buildTiDBProbeCommand(tc *v1alpha1.TidbCluster) (command []string) {
+	host := "127.0.0.1"
+
+	readinessURL := fmt.Sprintf("%s://%s:10080/status", tc.Scheme(), host)
+	command = append(command, "curl")
+	command = append(command, 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.
+	command = append(command, "--fail")
+	// follow 301 or 302 redirect
+	command = append(command, "--location")
+
+	if tc.IsTLSClusterEnabled() {
+		cacert := path.Join(clusterCertPath, tlsSecretRootCAKey)
+		cert := path.Join(clusterCertPath, corev1.TLSCertKey)
+		key := path.Join(clusterCertPath, corev1.TLSPrivateKeyKey)
+		command = append(command, "--cacert", cacert)
+		command = append(command, "--cert", cert)
+		command = append(command, "--key", key)
+	}
+	return
+}
+
 func tlsClientSecretName(tc *v1alpha1.TidbCluster) string {
 	return fmt.Sprintf("%s-server-secret", controller.TiDBMemberName(tc.Name))
 }

pkg/manager/member/tidb_member_manager_test.go

@@ -16,6 +16,7 @@ package member
 import (
 	"context"
 	"fmt"
+	"path"
 	"strings"
 	"testing"
 	"time"
@@ -2157,6 +2158,72 @@ func TestTiDBShouldRecover(t *testing.T) {
 	}
 }
 
+func TestBuildTiDBProbeHandler(t *testing.T) {
+	g := NewGomegaWithT(t)
+
+	defaultHandler := corev1.Handler{
+		TCPSocket: &corev1.TCPSocketAction{
+			Port: intstr.FromInt(4000),
+		},
+	}
+
+	execHandler := corev1.Handler{
+		Exec: &corev1.ExecAction{
+			Command: []string{
+				"curl",
+				"http://127.0.0.1:10080/status",
+				"--fail",
+				"--location",
+			},
+		},
+	}
+
+	sslExecHandler := corev1.Handler{
+		Exec: &corev1.ExecAction{
+			Command: []string{
+				"curl",
+				"https://127.0.0.1:10080/status",
+				"--fail",
+				"--location",
+				"--cacert", path.Join(clusterCertPath, tlsSecretRootCAKey),
+				"--cert", path.Join(clusterCertPath, corev1.TLSCertKey),
+				"--key", path.Join(clusterCertPath, corev1.TLSPrivateKeyKey),
+			},
+		},
+	}
+
+	tc := &v1alpha1.TidbCluster{
+		Spec: v1alpha1.TidbClusterSpec{
+			TiDB: &v1alpha1.TiDBSpec{},
+		},
+	}
+
+	// test default
+	get := buildTiDBReadinessProbHandler(tc)
+	g.Expect(get).Should(Equal(defaultHandler))
+
+	// test set command type & not tls
+	tc.Spec.TiDB.ReadinessProbe = &v1alpha1.TiDBProbe{
+		Type: pointer.StringPtr(v1alpha1.CommandProbeType),
+	}
+	get = buildTiDBReadinessProbHandler(tc)
+	g.Expect(get).Should(Equal(execHandler))
+
+	// test command type and tls
+	tc.Spec.TLSCluster = &v1alpha1.TLSCluster{
+		Enabled: true,
+	}
+	get = buildTiDBReadinessProbHandler(tc)
+	g.Expect(get).Should(Equal(sslExecHandler))
+
+	// test tcp type
+	tc.Spec.TiDB.ReadinessProbe = &v1alpha1.TiDBProbe{
+		Type: pointer.StringPtr(v1alpha1.TCPProbeType),
+	}
+	get = buildTiDBReadinessProbHandler(tc)
+	g.Expect(get).Should(Equal(defaultHandler))
+}
+
 func mustConfig(x interface{}) *v1alpha1.TiDBConfigWraper {
 	data, err := MarshalTOML(x)
 	if err != nil {