From 9aef7dec264f23a9157f329cab30376f041cabc7 Mon Sep 17 00:00:00 2001
From: Manan Gupta <manan@planetscale.com>
Date: Tue, 23 Jul 2024 12:07:29 +0530
Subject: [PATCH] feat: add ability to write debug code without affecting
 production status

Signed-off-by: Manan Gupta <manan@planetscale.com>
---
 Makefile                |  1 +
 go/testing/debug2PC.go  | 23 ++++++++++++++++++++++
 go/testing/modes.go     | 42 +++++++++++++++++++++++++++++++++++++++++
 go/vt/vtgate/tx_conn.go |  6 ++++++
 4 files changed, 72 insertions(+)
 create mode 100644 go/testing/debug2PC.go
 create mode 100644 go/testing/modes.go

diff --git a/Makefile b/Makefile
index e46f5d1ca5a..939b43571f8 100644
--- a/Makefile
+++ b/Makefile
@@ -94,6 +94,7 @@ endif
 	CGO_ENABLED=0 go build \
 		    -trimpath $(EXTRA_BUILD_FLAGS) $(VT_GO_PARALLEL) \
 		    -ldflags "$(EXTRA_BUILD_LDFLAGS) $(shell tools/build_version_flags.sh)" \
+		    -tags "$(EXTRA_BUILD_TAGS)" \
 		    -o ${VTROOTBIN} ./go/...
 ifndef NOVTADMINBUILD
 	echo "Building VTAdmin Web, disable VTAdmin build by setting 'NOVTADMINBUILD'"
diff --git a/go/testing/debug2PC.go b/go/testing/debug2PC.go
new file mode 100644
index 00000000000..6888ef17759
--- /dev/null
+++ b/go/testing/debug2PC.go
@@ -0,0 +1,23 @@
+//go:build debug2PC
+
+/*
+Copyright 2024 The Vitess Authors.
+
+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 testing
+
+func init() {
+	Mode = Debug2PC
+}
diff --git a/go/testing/modes.go b/go/testing/modes.go
new file mode 100644
index 00000000000..314e59d00c2
--- /dev/null
+++ b/go/testing/modes.go
@@ -0,0 +1,42 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+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 testing
+
+type TestingMode int
+
+const (
+	Production TestingMode = iota
+	Debug2PC
+)
+
+// Mode is used to define whether we are in Production mode or testing mode.
+// We use go build directives to include files that change the testing mode when certain tags are provided
+// while building binaries. This allows to have debugging code written in normal code flow without affecting
+// production performance.
+// Usage - To add debug code, put it behind testing.Mode != Production etc checks.
+var Mode TestingMode = Production
+
+// String is an unused function, but can be used in debugging to see what the testing mode is.
+func (tm TestingMode) String() string {
+	switch tm {
+	case Production:
+		return "Production"
+	case Debug2PC:
+		return "Debug2PC"
+	}
+	return "Unknown"
+}
diff --git a/go/vt/vtgate/tx_conn.go b/go/vt/vtgate/tx_conn.go
index 2eccdc54992..eda0c4717db 100644
--- a/go/vt/vtgate/tx_conn.go
+++ b/go/vt/vtgate/tx_conn.go
@@ -21,8 +21,10 @@ import (
 	"fmt"
 	"strings"
 	"sync"
+	"time"
 
 	"vitess.io/vitess/go/mysql/sqlerror"
+	"vitess.io/vitess/go/testing"
 	"vitess.io/vitess/go/vt/concurrency"
 	"vitess.io/vitess/go/vt/dtids"
 	"vitess.io/vitess/go/vt/log"
@@ -99,6 +101,10 @@ func (txc *TxConn) Commit(ctx context.Context, session *SafeSession) error {
 	}
 
 	if twopc {
+		if testing.Mode == testing.Debug2PC {
+			// Custom debug code
+			time.Sleep(3 * time.Second)
+		}
 		return txc.commit2PC(ctx, session)
 	}
 	return txc.commitNormal(ctx, session)