From 8bb2c00a932932308ce2c94f3a854e9e460f7e77 Mon Sep 17 00:00:00 2001
From: Yohei Yoshimuta <yoheimuta@gmail.com>
Date: Mon, 6 Nov 2023 16:47:00 +0900
Subject: [PATCH 1/2] Add support for customizing terminationGracePeriodSeconds
 of the VTTablet pods

Signed-off-by: Yohei Yoshimuta <yoheimuta@gmail.com>
---
 deploy/crds/planetscale.com_vitessclusters.yaml      |  6 ++++++
 deploy/crds/planetscale.com_vitesskeyspaces.yaml     |  6 ++++++
 deploy/crds/planetscale.com_vitessshards.yaml        |  3 +++
 docs/api/index.html                                  | 12 ++++++++++++
 pkg/apis/planetscale/v2/vitessshard_types.go         |  4 ++++
 pkg/apis/planetscale/v2/zz_generated.deepcopy.go     |  5 +++++
 pkg/operator/vttablet/pod.go                         |  6 +++++-
 test/integration/vitesscluster/vitesscluster_test.go |  8 ++++++++
 8 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/deploy/crds/planetscale.com_vitessclusters.yaml b/deploy/crds/planetscale.com_vitessclusters.yaml
index 3d6e2ff1..d4daf813 100644
--- a/deploy/crds/planetscale.com_vitessclusters.yaml
+++ b/deploy/crds/planetscale.com_vitessclusters.yaml
@@ -1561,6 +1561,9 @@ spec:
                                                       x-kubernetes-int-or-string: true
                                                     type: object
                                                 type: object
+                                              terminationGracePeriodSeconds:
+                                                format: int64
+                                                type: integer
                                             required:
                                             - resources
                                             type: object
@@ -1960,6 +1963,9 @@ spec:
                                                     x-kubernetes-int-or-string: true
                                                   type: object
                                               type: object
+                                            terminationGracePeriodSeconds:
+                                              format: int64
+                                              type: integer
                                           required:
                                           - resources
                                           type: object
diff --git a/deploy/crds/planetscale.com_vitesskeyspaces.yaml b/deploy/crds/planetscale.com_vitesskeyspaces.yaml
index 9f267325..0512d1a9 100644
--- a/deploy/crds/planetscale.com_vitesskeyspaces.yaml
+++ b/deploy/crds/planetscale.com_vitesskeyspaces.yaml
@@ -609,6 +609,9 @@ spec:
                                                 x-kubernetes-int-or-string: true
                                               type: object
                                           type: object
+                                        terminationGracePeriodSeconds:
+                                          format: int64
+                                          type: integer
                                       required:
                                       - resources
                                       type: object
@@ -1008,6 +1011,9 @@ spec:
                                               x-kubernetes-int-or-string: true
                                             type: object
                                         type: object
+                                      terminationGracePeriodSeconds:
+                                        format: int64
+                                        type: integer
                                     required:
                                     - resources
                                     type: object
diff --git a/deploy/crds/planetscale.com_vitessshards.yaml b/deploy/crds/planetscale.com_vitessshards.yaml
index 4eb7fc4f..cfb4b826 100644
--- a/deploy/crds/planetscale.com_vitessshards.yaml
+++ b/deploy/crds/planetscale.com_vitessshards.yaml
@@ -592,6 +592,9 @@ spec:
                                 x-kubernetes-int-or-string: true
                               type: object
                           type: object
+                        terminationGracePeriodSeconds:
+                          format: int64
+                          type: integer
                       required:
                       - resources
                       type: object
diff --git a/docs/api/index.html b/docs/api/index.html
index 92c2cd3c..ae6f2ac0 100644
--- a/docs/api/index.html
+++ b/docs/api/index.html
@@ -7892,6 +7892,18 @@ <h3 id="planetscale.com/v2.VttabletSpec">VttabletSpec
 to vttablet container</p>
 </td>
 </tr>
+<tr>
+<td>
+<code>terminationGracePeriodSeconds</code></br>
+<em>
+int64
+</em>
+</td>
+<td>
+<p>TerminationGracePeriodSeconds can optionally be used to customize
+terminationGracePeriodSeconds of the vttablet pod.</p>
+</td>
+</tr>
 </tbody>
 </table>
 <h3 id="planetscale.com/v2.WorkflowState">WorkflowState
diff --git a/pkg/apis/planetscale/v2/vitessshard_types.go b/pkg/apis/planetscale/v2/vitessshard_types.go
index e3f0dca8..1a799681 100644
--- a/pkg/apis/planetscale/v2/vitessshard_types.go
+++ b/pkg/apis/planetscale/v2/vitessshard_types.go
@@ -292,6 +292,10 @@ type VttabletSpec struct {
 	// +kubebuilder:validation:Schemaless
 	// +kubebuilder:pruning:PreserveUnknownFields
 	Lifecycle corev1.Lifecycle `json:"lifecycle,omitempty"`
+
+	// TerminationGracePeriodSeconds can optionally be used to customize
+	// terminationGracePeriodSeconds of the vttablet pod.
+	TerminationGracePeriodSeconds *int64 `json:"terminationGracePeriodSeconds,omitempty"`
 }
 
 // MysqldSpec configures the local MySQL server within a tablet.
diff --git a/pkg/apis/planetscale/v2/zz_generated.deepcopy.go b/pkg/apis/planetscale/v2/zz_generated.deepcopy.go
index ff017d72..941c22f1 100644
--- a/pkg/apis/planetscale/v2/zz_generated.deepcopy.go
+++ b/pkg/apis/planetscale/v2/zz_generated.deepcopy.go
@@ -2623,6 +2623,11 @@ func (in *VttabletSpec) DeepCopyInto(out *VttabletSpec) {
 		}
 	}
 	in.Lifecycle.DeepCopyInto(&out.Lifecycle)
+	if in.TerminationGracePeriodSeconds != nil {
+		in, out := &in.TerminationGracePeriodSeconds, &out.TerminationGracePeriodSeconds
+		*out = new(int64)
+		**out = **in
+	}
 }
 
 // DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new VttabletSpec.
diff --git a/pkg/operator/vttablet/pod.go b/pkg/operator/vttablet/pod.go
index da2341c1..992ae0f6 100644
--- a/pkg/operator/vttablet/pod.go
+++ b/pkg/operator/vttablet/pod.go
@@ -313,7 +313,11 @@ func UpdatePod(obj *corev1.Pod, spec *Spec) {
 		obj.Spec.SecurityContext.FSGroup = pointer.Int64Ptr(planetscalev2.DefaultVitessFSGroup)
 	}
 
-	obj.Spec.TerminationGracePeriodSeconds = pointer.Int64Ptr(terminationGracePeriodSeconds)
+	if spec.Vttablet.TerminationGracePeriodSeconds != nil {
+		obj.Spec.TerminationGracePeriodSeconds = spec.Vttablet.TerminationGracePeriodSeconds
+	} else {
+		obj.Spec.TerminationGracePeriodSeconds = pointer.Int64Ptr(terminationGracePeriodSeconds)
+	}
 
 	// In both the case of the user injecting their own affinity and the default, we
 	// simply override the pod's existing affinity configuration.
diff --git a/test/integration/vitesscluster/vitesscluster_test.go b/test/integration/vitesscluster/vitesscluster_test.go
index 0483efb8..b7f2440d 100644
--- a/test/integration/vitesscluster/vitesscluster_test.go
+++ b/test/integration/vitesscluster/vitesscluster_test.go
@@ -51,6 +51,8 @@ spec:
           - cell: cell1
             type: replica
             replicas: 3
+            vttablet:
+              terminationGracePeriodSeconds: 60
             mysqld: {}
             dataVolumeClaimTemplate:
               accessModes: [ReadWriteOnce]
@@ -219,9 +221,15 @@ func verifyBasicVitessShard(f *framework.Fixture, ns, cluster, keyspace, shard s
 	// Each vttablet Pod should have a PVC.
 	for i := range cell1Pods.Items {
 		f.MustGet(ns, cell1Pods.Items[i].Name, &corev1.PersistentVolumeClaim{})
+		if *cell1Pods.Items[i].Spec.TerminationGracePeriodSeconds != 60 {
+			f.Fatalf("TerminationGracePeriodSeconds should be 60, but got %d", *cell1Pods.Items[i].Spec.TerminationGracePeriodSeconds)
+		}
 	}
 	for i := range cell2Pods.Items {
 		f.MustGet(ns, cell2Pods.Items[i].Name, &corev1.PersistentVolumeClaim{})
+		if *cell2Pods.Items[i].Spec.TerminationGracePeriodSeconds != 1800 {
+			f.Fatalf("TerminationGracePeriodSeconds should be 1800, but got %d", *cell2Pods.Items[i].Spec.TerminationGracePeriodSeconds)
+		}
 	}
 
 	// VitessShard creates vtbackup-init Pod/PVC.

From befb94e903b83c2eb1992ab02e875b1d407c9e4f Mon Sep 17 00:00:00 2001
From: Yohei Yoshimuta <yoheimuta@gmail.com>
Date: Thu, 30 Nov 2023 11:54:50 +0900
Subject: [PATCH 2/2] Rename terminationGracePeriodSeconds constant

Signed-off-by: Yohei Yoshimuta <yoheimuta@gmail.com>
---
 pkg/operator/vttablet/constants.go | 4 ++--
 pkg/operator/vttablet/pod.go       | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/pkg/operator/vttablet/constants.go b/pkg/operator/vttablet/constants.go
index 94a388f7..98f5fc48 100644
--- a/pkg/operator/vttablet/constants.go
+++ b/pkg/operator/vttablet/constants.go
@@ -43,12 +43,12 @@ const (
 	serviceMap          = "grpc-queryservice,grpc-tabletmanager,grpc-updatestream"
 	healthCheckInterval = 5 * time.Second
 
-	// terminationGracePeriodSeconds is how long Kubernetes will wait for the
+	// defaultTerminationGracePeriodSeconds is how long Kubernetes will wait for the
 	// tablet processes (vttablet, mysqlctld) to terminate gracefully after
 	// sending SIGTERM, before resorting to SIGKILL. We set this fairly high
 	// because mysqld needs time to flush buffers to disk in order to shut down
 	// cleanly.
-	terminationGracePeriodSeconds = 30 * 60 // 30min
+	defaultTerminationGracePeriodSeconds = 30 * 60 // 30min
 
 	grpcMaxMessageSize = 64 * 1024 * 1024 // 64 MiB
 
diff --git a/pkg/operator/vttablet/pod.go b/pkg/operator/vttablet/pod.go
index 992ae0f6..ba4b154c 100644
--- a/pkg/operator/vttablet/pod.go
+++ b/pkg/operator/vttablet/pod.go
@@ -316,7 +316,7 @@ func UpdatePod(obj *corev1.Pod, spec *Spec) {
 	if spec.Vttablet.TerminationGracePeriodSeconds != nil {
 		obj.Spec.TerminationGracePeriodSeconds = spec.Vttablet.TerminationGracePeriodSeconds
 	} else {
-		obj.Spec.TerminationGracePeriodSeconds = pointer.Int64Ptr(terminationGracePeriodSeconds)
+		obj.Spec.TerminationGracePeriodSeconds = pointer.Int64Ptr(defaultTerminationGracePeriodSeconds)
 	}
 
 	// In both the case of the user injecting their own affinity and the default, we