go_test: Allow not registering a SIGTERM handler

docs/go/core/rules.md

@@ -334,7 +334,7 @@ This declares a set of source files and related dependencies that can be embedde
 <pre>
 go_test(<a href="#go_test-name">name</a>, <a href="#go_test-cdeps">cdeps</a>, <a href="#go_test-cgo">cgo</a>, <a href="#go_test-clinkopts">clinkopts</a>, <a href="#go_test-copts">copts</a>, <a href="#go_test-cppopts">cppopts</a>, <a href="#go_test-cxxopts">cxxopts</a>, <a href="#go_test-data">data</a>, <a href="#go_test-deps">deps</a>, <a href="#go_test-embed">embed</a>, <a href="#go_test-embedsrcs">embedsrcs</a>, <a href="#go_test-env">env</a>,
         <a href="#go_test-env_inherit">env_inherit</a>, <a href="#go_test-gc_goopts">gc_goopts</a>, <a href="#go_test-gc_linkopts">gc_linkopts</a>, <a href="#go_test-goarch">goarch</a>, <a href="#go_test-goos">goos</a>, <a href="#go_test-gotags">gotags</a>, <a href="#go_test-importpath">importpath</a>, <a href="#go_test-linkmode">linkmode</a>, <a href="#go_test-msan">msan</a>, <a href="#go_test-pure">pure</a>,
-        <a href="#go_test-race">race</a>, <a href="#go_test-rundir">rundir</a>, <a href="#go_test-srcs">srcs</a>, <a href="#go_test-static">static</a>, <a href="#go_test-x_defs">x_defs</a>)
+        <a href="#go_test-race">race</a>, <a href="#go_test-register_timeout_handler">register_timeout_handler</a>, <a href="#go_test-rundir">rundir</a>, <a href="#go_test-srcs">srcs</a>, <a href="#go_test-static">static</a>, <a href="#go_test-x_defs">x_defs</a>)
 </pre>
 
 This builds a set of tests that can be run with `bazel test`.<br><br>
@@ -396,6 +396,7 @@ This builds a set of tests that can be run with `bazel test`.<br><br>
 | <a id="go_test-msan"></a>msan |  Controls whether code is instrumented for memory sanitization. May be one of             <code>on</code>, <code>off</code>, or <code>auto</code>. Not available when cgo is             disabled. In most cases, it's better to control this on the command line with             <code>--@io_bazel_rules_go//go/config:msan</code>. See [mode attributes], specifically             [msan].   | String | optional | "auto" |
 | <a id="go_test-pure"></a>pure |  Controls whether cgo source code and dependencies are compiled and linked,             similar to setting <code>CGO_ENABLED</code>. May be one of <code>on</code>, <code>off</code>,             or <code>auto</code>. If <code>auto</code>, pure mode is enabled when no C/C++             toolchain is configured or when cross-compiling. It's usually better to             control this on the command line with             <code>--@io_bazel_rules_go//go/config:pure</code>. See [mode attributes], specifically             [pure].   | String | optional | "auto" |
 | <a id="go_test-race"></a>race |  Controls whether code is instrumented for race detection. May be one of             <code>on</code>, <code>off</code>, or <code>auto</code>. Not available when cgo is             disabled. In most cases, it's better to control this on the command line with             <code>--@io_bazel_rules_go//go/config:race</code>. See [mode attributes], specifically             [race].   | String | optional | "auto" |
+| <a id="go_test-register_timeout_handler"></a>register_timeout_handler |  Whether to register a SIGTERM handler to give improved output on test timeouts.<br><br>            Setting this to False may be useful if your test itself also handles signals.   | Boolean | optional | True |
 | <a id="go_test-rundir"></a>rundir |  A directory to cd to before the test is run.             This should be a path relative to the root directory of the             repository in which the test is defined, which can be the main or an             external repository.<br><br>            The default behaviour is to change to the relative path             corresponding to the test's package, which replicates the normal             behaviour of <code>go test</code> so it is easy to write compatible tests.<br><br>            Setting it to <code>.</code> makes the test behave the normal way for a bazel             test, except that the working directory is always that of the test's             repository, which is not necessarily the main repository.<br><br>            Note: If runfile symlinks are disabled (such as on Windows by             default), the test will run in the working directory set by Bazel,             which is the subdirectory of the runfiles directory corresponding to             the main repository.   | String | optional | "" |
 | <a id="go_test-srcs"></a>srcs |  The list of Go source files that are compiled to create the package.             Only <code>.go</code> and <code>.s</code> files are permitted, unless the <code>cgo</code>             attribute is set, in which case,             <code>.c .cc .cpp .cxx .h .hh .hpp .hxx .inc .m .mm</code>             files are also permitted. Files may be filtered at build time             using Go [build constraints].   | <a href="https://bazel.build/concepts/labels">List of labels</a> | optional | [] |
 | <a id="go_test-static"></a>static |  Controls whether a binary is statically linked. May be one of <code>on</code>,             <code>off</code>, or <code>auto</code>. Not available on all platforms or in all             modes. It's usually better to control this on the command line with             <code>--@io_bazel_rules_go//go/config:static</code>. See [mode attributes],             specifically [static].   | String | optional | "auto" |

go/private/rules/test.bzl

@@ -112,6 +112,10 @@ def _go_test_impl(ctx):
         "l_test=" + external_source.library.importpath,
     )
     arguments.add("-pkgname", internal_source.library.importpath)
+
+    if not ctx.attr.register_timeout_handler:
+        # Boolean Go flags require =, so we can't just use the attr as an arg value, we need to convert it ourselves.
+        arguments.add("-register_timeout_handler=false")
     arguments.add_all(go_srcs, before_each = "-src", format_each = "l=%s")
     ctx.actions.run(
         inputs = go_srcs,
@@ -415,6 +419,12 @@ _go_test_kwargs = {
             See [Cross compilation] for more information.
             """,
         ),
+        "register_timeout_handler": attr.bool(
+            default = True,
+            doc = """Whether to register a SIGTERM handler to give improved output on test timeouts.
+
+            Setting this to False may be useful if your test itself also handles signals.""",
+        ),
         "_go_context_data": attr.label(default = "//:go_context_data", cfg = go_transition),
         "_testmain_additional_deps": attr.label_list(
             providers = [GoLibrary],

go/tools/builders/generate_test_main.go

@@ -63,15 +63,16 @@ type Example struct {
 
 // Cases holds template data.
 type Cases struct {
-	Imports     []*Import
-	Tests       []TestCase
-	Benchmarks  []TestCase
-	FuzzTargets []TestCase
-	Examples    []Example
-	TestMain    string
-	CoverMode   string
-	CoverFormat string
-	Pkgname     string
+	Imports                []*Import
+	Tests                  []TestCase
+	Benchmarks             []TestCase
+	FuzzTargets            []TestCase
+	Examples               []Example
+	TestMain               string
+	CoverMode              string
+	CoverFormat            string
+	Pkgname                string
+	RegisterTimeoutHandler bool
 }
 
 // Version returns whether v is a supported Go version (like "go1.18").
@@ -235,7 +236,9 @@ func main() {
 		}
 	}
 	{{end}}
+	{{if .RegisterTimeoutHandler}}
 	bzltestutil.RegisterTimeoutHandler()
+	{{end}}
 	{{if not .TestMain}}
 	res := m.Run()
 	{{else}}
@@ -261,6 +264,7 @@ func genTestMain(args []string) error {
 	coverMode := flags.String("cover_mode", "", "the coverage mode to use")
 	coverFormat := flags.String("cover_format", "", "the coverage report type to generate (go_cover or lcov)")
 	pkgname := flags.String("pkgname", "", "package name of test")
+	registerTimeoutHandler := flags.Bool("register_timeout_handler", true, "whether to register a SIGTERM handler for improved timeout reporting")
 	flags.Var(&imports, "import", "Packages to import")
 	flags.Var(&sources, "src", "Sources to process for tests")
 	if err := flags.Parse(args); err != nil {
@@ -269,6 +273,12 @@ func genTestMain(args []string) error {
 	if err := goenv.checkFlags(); err != nil {
 		return err
 	}
+	// Boolean flags must be specified using `--foo=bar`, not `--foo bar`.
+	// Verify that there are no unexpected positional arguments, which may be the case if e.g. someone naively called
+	// args.add("-register_timeout_handler", ctx.attr.register_timeout_handler)
+	if len(flag.Args()) > 0 {
+		return fmt.Errorf("unexpected positional args: %v", flag.Args())
+	}
 	// Process import args
 	importMap := map[string]*Import{}
 	for _, imp := range imports {
@@ -309,9 +319,10 @@ func genTestMain(args []string) error {
 	}
 
 	cases := Cases{
-		CoverFormat: *coverFormat,
-		CoverMode:   *coverMode,
-		Pkgname:     *pkgname,
+		CoverFormat:            *coverFormat,
+		CoverMode:              *coverMode,
+		Pkgname:                *pkgname,
+		RegisterTimeoutHandler: *registerTimeoutHandler,
 	}
 
 	testFileSet := token.NewFileSet()

tests/core/go_test/BUILD.bazel

@@ -250,6 +250,16 @@ go_test(
     shard_count = 2,
 )
 
+go_test(
+    name = "sigterm_handler_test",
+    srcs = ["sigterm_handler_test.go"],
+    register_timeout_handler = False,
+    target_compatible_with = select({
+        "@platforms//os:windows": ["@platforms//:incompatible"],
+        "//conditions:default": [],
+    }),
+)
+
 go_bazel_test(
     name = "env_inherit_test",
     srcs = ["env_inherit_test.go"],

tests/core/go_test/sigterm_handler_test.go

@@ -0,0 +1,36 @@
+package sigterm_handler_test
+
+import (
+	"os"
+	"os/signal"
+	"sync"
+	"syscall"
+	"testing"
+)
+
+func TestRegisterSignalHandler(t *testing.T) {
+	called := false
+	var wg sync.WaitGroup
+
+	wg.Add(1)
+
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, syscall.SIGTERM)
+
+	go func() {
+		switch <-c {
+		case syscall.SIGTERM:
+			called = true
+			wg.Done()
+		}
+	}()
+
+	if err := syscall.Kill(os.Getpid(), syscall.SIGTERM); err != nil {
+		t.Fatalf("Failed to send SIGTERM: %v", err)
+	}
+	wg.Wait()
+
+	if !called {
+		t.Fatal("Our handler has not run")
+	}
+}