Setting test.timeout according to TEST_TIMEOUT

go/tools/builders/generate_test_main.go

@@ -98,11 +98,13 @@ import (
 	"log"
 	"os"
 	"os/exec"
+	"os/signal"
 {{if .TestMain}}
 	"reflect"
 {{end}}
 	"strconv"
 	"strings"
+	"syscall"
 	"testing"
 	"testing/internal/testdeps"
 
@@ -235,7 +237,19 @@ func main() {
 		}
 	}
 	{{end}}
-	bzltestutil.RegisterTimeoutHandler()
+
+	testTimeout := os.Getenv("TEST_TIMEOUT")
+	if testTimeout != "" {
+		flag.Lookup("test.timeout").Value.Set(testTimeout+"s")
+		// If Bazel sends a SIGTERM because the test timed out, it sends it to all child processes. Because
+		// we set -test.timeout according to the TEST_TIMEOUT, we need to ignore the signal so the test has
+		// time to properly produce the output (e.g. stack trace). It will be killed by Bazel after the grace
+		// period (15s) expires.
+		// If TEST_TIMEOUT is not set (e.g., when the test binary is run by Delve for debugging), we don't
+		// ignore SIGTERM so it can be properly terminated.
+		signal.Ignore(syscall.SIGTERM)
+	}
+
 	{{if not .TestMain}}
 	res := m.Run()
 	{{else}}

go/tools/bzltestutil/BUILD.bazel

@@ -5,7 +5,6 @@ go_tool_library(
     srcs = [
         "lcov.go",
         "test2json.go",
-        "timeout.go",
         "wrap.go",
         "xml.go",
     ],

go/tools/bzltestutil/wrap.go

@@ -125,10 +125,11 @@ func Wrap(pkg string) error {
 		exePath = filepath.Join(chdir.TestExecDir, exePath)
 	}
 
-	// If Bazel sends a SIGTERM because the test timed out, it sends it to all child processes. As
-	// a result, the child process will print stack traces of all Go routines and we want the
-	// wrapper to be around to capute and forward this output. Thus, we need to ignore the signal
-	// and will be killed by Bazel after the grace period instead.
+	// If Bazel sends a SIGTERM because the test timed out, it sends it to all child processes. However,
+	// we want the wrapper to be around to capute and forward the test output when this happens. Thus,
+	// we need to ignore the signal. This wrapper will natually ends after the Go test ends, either by
+	// SIGTERM or the time set by -test.timeout expires. If that doesn't happen, the test and this warpper
+	// will be killed by Bazel after the grace period (15s) expires.
 	signal.Ignore(syscall.SIGTERM)
 
 	cmd := exec.Command(exePath, args...)

tests/core/go_test/timeout_test.go

@@ -41,11 +41,19 @@ func TestTimeout(t *testing.T) {
 		t.Skip("stack traces on timeouts are not yet supported on Windows")
 	}
 
-	if err := bazel_testing.RunBazel("test", "//:timeout_test", "--test_timeout=3"); err == nil {
+	var stderr string
+	if err := bazel_testing.RunBazel("test", "//:timeout_test", "--test_timeout=3", "--test_arg=-test.v"); err == nil {
 		t.Fatal("expected bazel test to fail")
 	} else if exitErr, ok := err.(*bazel_testing.StderrExitError); !ok || exitErr.Err.ExitCode() != 3 {
 		t.Fatalf("expected bazel test to fail with exit code 3", err)
+	} else {
+		stderr = string(exitErr.Err.Stderr)
 	}
+
+	if !strings.Contains(stderr, "TIMEOUT: //:timeout_test") {
+		t.Errorf("expect Bazel to report the test timed out: \n%s", stderr)
+	}
+
 	p, err := bazel_testing.BazelOutput("info", "bazel-testlogs")
 	if err != nil {
 		t.Fatalf("could not find testlogs root: %s", err)
@@ -57,10 +65,21 @@ func TestTimeout(t *testing.T) {
 	}
 
 	testLog := string(b)
-	if !strings.Contains(testLog, "Received SIGTERM, printing stack traces of all goroutines:") {
-		t.Fatalf("test log does not contain expected header:\n%s", testLog)
+	if !strings.Contains(testLog, "panic: test timed out after 3s") {
+		t.Errorf("test log does not contain expected header:\n%s", testLog)
 	}
 	if !strings.Contains(testLog, "timeout_test.neverTerminates(") {
-		t.Fatalf("test log does not contain expected stack trace:\n%s", testLog)
+		t.Errorf("test log does not contain expected stack trace:\n%s", testLog)
+	}
+
+	path = filepath.Join(strings.TrimSpace(string(p)), "timeout_test/test.xml")
+	b, err = os.ReadFile(path)
+	if err != nil {
+		t.Fatalf("could not read test XML: %s", err)
+	}
+
+	testXML := string(b)
+	if !strings.Contains(testXML, `<testcase classname="timeout_test" name="TestFoo"`) {
+		t.Errorf("test XML does not contain expected element:\n%s", testXML)
 	}
 }