From 33d7a2c9840e4f07a7caa09618572744211e23e1 Mon Sep 17 00:00:00 2001
From: tiancaiamao <tiancaiamao@gmail.com>
Date: Wed, 28 Feb 2024 10:42:30 +0800
Subject: [PATCH] tools/check: adpot staged build to make ut a little bit
 faster (#51374)

ref pingcap/tidb#31880
---
 tools/check/go-compile-without-link.sh | 33 ++++++++++++++++++++++++++
 tools/check/ut.go                      | 18 ++++++++++++++
 2 files changed, 51 insertions(+)
 create mode 100644 tools/check/go-compile-without-link.sh

diff --git a/tools/check/go-compile-without-link.sh b/tools/check/go-compile-without-link.sh
new file mode 100644
index 0000000000000..9ad1059747ec6
--- /dev/null
+++ b/tools/check/go-compile-without-link.sh
@@ -0,0 +1,33 @@
+#!/bin/bash
+# Copyright 2024 PingCAP, Inc.
+#
+# 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.
+
+# See https://gist.github.com/howardjohn/c0f5d0bc293ef7d7fada533a2c9ffaf4
+# Usage: go test -exec=true -toolexec=go-compile-without-link -vet=off ./...
+# Preferably as an alias like `alias go-test-compile='go test -exec=true -toolexec=go-compile-without-link -vet=off'`
+# This will compile all tests, but not link them (which is the least cacheable part)
+
+if [[ "${2}" == "-V=full" ]]; then
+  "$@"
+  exit 0
+fi
+case "$(basename ${1})" in
+  link)
+    # Output a dummy file
+    touch "${3}"
+    ;;
+  # We could skip vet as well, but it can be done with -vet=off if desired
+  *)
+    "$@"
+esac
diff --git a/tools/check/ut.go b/tools/check/ut.go
index c28c535377c40..34f826f957f6f 100644
--- a/tools/check/ut.go
+++ b/tools/check/ut.go
@@ -869,8 +869,26 @@ func buildTestBinary(pkg string) error {
 	return nil
 }
 
+func generateBuildCache() error {
+	// cd cmd/tidb-server && go test -tags intest -exec true -vet off -toolexec=go-compile-without-link
+	cmd := exec.Command("go", "test", "-tags=intest", "-exec=true", "-vet=off")
+	goCompileWithoutLink := fmt.Sprintf("-toolexec=%s/tools/check/go-compile-without-link.sh", workDir)
+	cmd.Args = append(cmd.Args, goCompileWithoutLink)
+	cmd.Dir = path.Join(workDir, "cmd/tidb-server")
+	if err := cmd.Run(); err != nil {
+		return withTrace(err)
+	}
+	return nil
+}
+
 // buildTestBinaryMulti is much faster than build the test packages one by one.
 func buildTestBinaryMulti(pkgs []string) error {
+	// staged build, generate the build cache for all the tests first, then generate the test binary.
+	// This way is faster than generating test binaries directly, because the cache can be used.
+	if err := generateBuildCache(); err != nil {
+		return withTrace(err)
+	}
+
 	// go test --exec=xprog -cover -vet=off --count=0 $(pkgs)
 	xprogPath := path.Join(workDir, "tools/bin/xprog")
 	packages := make([]string, 0, len(pkgs))