From cc911bfec4f52d93d1c47cc92a3bc03ec8f9cb33 Mon Sep 17 00:00:00 2001
From: Harsh Modi <hmodi@aviatrix.com>
Date: Wed, 8 May 2024 18:35:12 +0000
Subject: [PATCH] generate_test_main: Move timeout handling back to bzltestutil

This gives us a determinstic location to ignore for goleak, which we
previously had, but was removed as part of #3920.
---
 go/tools/builders/generate_test_main.go | 24 +------------
 go/tools/bzltestutil/BUILD.bazel        |  1 +
 go/tools/bzltestutil/timeout.go         | 45 +++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 23 deletions(-)
 create mode 100644 go/tools/bzltestutil/timeout.go

diff --git a/go/tools/builders/generate_test_main.go b/go/tools/builders/generate_test_main.go
index 83001fbd40..0e1bdbf8cd 100644
--- a/go/tools/builders/generate_test_main.go
+++ b/go/tools/builders/generate_test_main.go
@@ -98,13 +98,11 @@ import (
 	"log"
 	"os"
 	"os/exec"
-	"os/signal"
 {{if .TestMain}}
 	"reflect"
 {{end}}
 	"strconv"
 	"strings"
-	"syscall"
 	"testing"
 	"testing/internal/testdeps"
 
@@ -241,27 +239,7 @@ func main() {
 	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. (1)
-		// We do not panic (like native go test does) because users may legitimately want to use SIGTERM
-		// in tests.
-		//
-		// signal.Notify is used to ensure that there is a no-op signal handler registered.
-		// Avoid using signal.Ignore here: despite the name, it's only used to unregister handlers that
-		// were previously registered by signal.Notify. See (2) for more information.
-		//
-		// (1): https://github.com/golang/go/blob/e816eb50140841c524fd07ecb4eaa078954eb47c/src/testing/testing.go#L2351
-		// (2): https://github.com/bazelbuild/rules_go/pull/3929
-		c := make(chan os.Signal, 1)
-		signal.Notify(c, syscall.SIGTERM)
-		go func() {
-			<-c
-		}()
+		bzltestutil.RegisterTimeoutHandler()
 	}
 
 	{{if not .TestMain}}
diff --git a/go/tools/bzltestutil/BUILD.bazel b/go/tools/bzltestutil/BUILD.bazel
index a97a7516b9..e6b3637d3a 100644
--- a/go/tools/bzltestutil/BUILD.bazel
+++ b/go/tools/bzltestutil/BUILD.bazel
@@ -5,6 +5,7 @@ go_tool_library(
     srcs = [
         "lcov.go",
         "test2json.go",
+        "timeout.go",
         "wrap.go",
         "xml.go",
     ],
diff --git a/go/tools/bzltestutil/timeout.go b/go/tools/bzltestutil/timeout.go
new file mode 100644
index 0000000000..e94f67046b
--- /dev/null
+++ b/go/tools/bzltestutil/timeout.go
@@ -0,0 +1,45 @@
+// Copyright 2024 The Bazel Authors. All rights reserved.
+//
+// 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 bzltestutil
+
+import (
+	"os"
+	"os/signal"
+	"syscall"
+)
+
+func RegisterTimeoutHandler() {
+	// 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. (1)
+	// We do not panic (like native go test does) because users may legitimately want to use SIGTERM
+	// in tests.
+	//
+	// signal.Notify is used to ensure that there is a no-op signal handler registered.
+	// Avoid using signal.Ignore here: despite the name, it's only used to unregister handlers that
+	// were previously registered by signal.Notify. See (2) for more information.
+	//
+	// (1): https://github.com/golang/go/blob/e816eb50140841c524fd07ecb4eaa078954eb47c/src/testing/testing.go#L2351
+	// (2): https://github.com/bazelbuild/rules_go/pull/3929
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, syscall.SIGTERM)
+	go func() {
+		<-c
+	}()
+}