From 97da4a3c6a0ce1c90dad8d1cab40cedc35a86c42 Mon Sep 17 00:00:00 2001
From: David Piegza <697113+davidpiegza@users.noreply.github.com>
Date: Wed, 9 Aug 2023 14:07:05 +0000
Subject: [PATCH] Add terminating state

Signed-off-by: David Piegza <697113+davidpiegza@users.noreply.github.com>
---
 go/vt/vtgate/plugin_mysql_server.go | 5 +++++
 go/vt/vtgate/vtgate.go              | 8 ++++++++
 2 files changed, 13 insertions(+)

diff --git a/go/vt/vtgate/plugin_mysql_server.go b/go/vt/vtgate/plugin_mysql_server.go
index f880b6c0412..072d56afdc1 100644
--- a/go/vt/vtgate/plugin_mysql_server.go
+++ b/go/vt/vtgate/plugin_mysql_server.go
@@ -183,6 +183,11 @@ func startSpan(ctx context.Context, query, label string) (trace.Span, context.Co
 
 func (vh *vtgateHandler) ComQuery(c *mysql.Conn, query string, callback func(*sqltypes.Result) error) error {
 	ctx := context.Background()
+
+	if terminating {
+		return vterrors.New(vtrpcpb.Code_ABORTED, "Terminating vtgate")
+	}
+
 	var cancel context.CancelFunc
 	if *mysqlQueryTimeout != 0 {
 		ctx, cancel = context.WithTimeout(ctx, *mysqlQueryTimeout)
diff --git a/go/vt/vtgate/vtgate.go b/go/vt/vtgate/vtgate.go
index de3fba2fbe0..20146bea46b 100644
--- a/go/vt/vtgate/vtgate.go
+++ b/go/vt/vtgate/vtgate.go
@@ -83,6 +83,8 @@ var (
 	// healthCheckTimeout is the timeout on the RPC call to tablets
 	healthCheckTimeout = time.Minute
 
+	terminating bool
+
 	// System settings related flags
 	sysVarSetEnabled = true
 	setVarEnabled    = true
@@ -328,6 +330,8 @@ func Init(
 		}
 	})
 	servenv.OnTerm(func() {
+		terminating = true
+
 		if st != nil && enableSchemaChangeSignal {
 			st.Stop()
 		}
@@ -405,6 +409,10 @@ func (vtg *VTGate) registerDebugHealthHandler() {
 // IsHealthy returns nil if server is healthy.
 // Otherwise, it returns an error indicating the reason.
 func (vtg *VTGate) IsHealthy() error {
+	if terminating {
+		return errors.New("Terminating")
+	}
+
 	return nil
 }