Gracefully shut down services on test failure and context cancellation

BUILD.bazel

@@ -18,6 +18,12 @@ cc_test(
     visibility = ["//visibility:public"],
 )
 
+cc_binary(
+    name = "exit1",
+    srcs = ["exit1.c"],
+    visibility = ["//visibility:public"],
+)
+
 bool_flag(
     name = "allow_configuring_tmpdir",
     build_setting_default = False,

cmd/svcinit/main.go

@@ -139,6 +139,7 @@ func main() {
 
 	r, err := runner.New(ctx, serviceSpecs)
 	must(err)
+	defer r.StopAll()
 
 	servicesErrCh := make(chan error, len(unversionedSpecs))
 
@@ -238,6 +239,8 @@ func main() {
 
 		fmt.Println()
 
+		exitCode := 0
+
 		if shouldHotReload && !enablePerServiceReload {
 			fmt.Println()
 			fmt.Println("###########################################################################################")
@@ -295,13 +298,18 @@ func main() {
 			if testErr != nil {
 				log.Printf("Encountered error during test run: %s\n", testErr)
 				if isOneShot {
-					os.Exit(1)
+					exitCode = 1
 				}
 			}
 		case serviceErr := <-servicesErrCh:
 			log.Print(serviceErr)
 			if isOneShot {
-				log.Fatal("Service exited uncleanly, marking test as failed.\n\n")
+				log.Println("Service exited uncleanly, marking test as failed.")
+				exitCode = 1
+				testCancel()
+				if testLabel != "" {
+					<-testErrCh // Wait for test process to exit
+				}
 			}
 		}
 
@@ -321,7 +329,7 @@ func main() {
 			}
 		}
 
-		if testLabel != "" {
+		if testLabel != "" && testCmd.ProcessState != nil {
 			buf.WriteString(fmt.Sprintf("%s\t%s\t%s\n",
 				testLabel, testCmd.ProcessState.UserTime(), testCmd.ProcessState.SystemTime()))
 		}
@@ -333,6 +341,9 @@ func main() {
 		must(err)
 
 		if isOneShot {
+			if exitCode != 0 {
+				os.Exit(exitCode)
+			}
 			break
 		}
 	}

exit1.c

@@ -0,0 +1 @@
+int main() { return 1; }

runner/runner.go

@@ -263,8 +263,18 @@ func initializeServiceCmd(ctx context.Context, instance *ServiceInstance) error
 	// Even if a child process exits, Wait will block until the I/O pipes are closed.
 	// They may have been forwarded to an orphaned child, so we disable that behavior to unblock exit.
 	if s.Type == "service" && !s.Deferred {
-		// We need a bit of grace period to allow I/O pipes to close on our end.
-		cmd.WaitDelay = 50 * time.Millisecond
+		// Don't kill the process on context cancellation; StopAll() handles
+		// orderly shutdown in reverse dependency order.
+		cmd.Cancel = func() error { return nil }
+
+		// Use the configured shutdown timeout as WaitDelay so the process has
+		// time to exit after StopAll() sends the shutdown signal. Fall back to
+		// a small grace period for I/O pipe closing.
+		waitDelay := 50 * time.Millisecond
+		if shutdownTimeout, err := time.ParseDuration(s.ShutdownTimeout); err == nil && shutdownTimeout > waitDelay {
+			waitDelay = shutdownTimeout
+		}
+		cmd.WaitDelay = waitDelay
 	}
 
 	instance.cmd = cmd

tests/graceful_shutdown_on_failure/BUILD.bazel

@@ -0,0 +1,32 @@
+load("@rules_go//go:def.bzl", "go_binary")
+load("@rules_itest//:itest.bzl", "itest_service", "service_test")
+load(":check_graceful_shutdown.bzl", "check_graceful_shutdown_test")
+
+go_binary(
+    name = "service",
+    srcs = ["main.go"],
+    tags = ["manual"],
+)
+
+itest_service(
+    name = "graceful_service",
+    args = ["$${PORT}"],
+    autoassign_port = True,
+    exe = ":service",
+    http_health_check_address = "http://127.0.0.1:$${PORT}",
+    tags = ["manual"],
+)
+
+service_test(
+    name = "_failing_test_with_service",
+    services = [
+        ":graceful_service",
+    ],
+    tags = ["manual"],
+    test = "@rules_itest//:exit1",
+)
+
+check_graceful_shutdown_test(
+    name = "graceful_shutdown_on_failure_test",
+    service_test = ":_failing_test_with_service",
+)

tests/graceful_shutdown_on_failure/check_graceful_shutdown.bzl

@@ -0,0 +1,68 @@
+"""Rule to verify that services are gracefully shut down when a test fails."""
+
+load("@rules_shell//shell:sh_test.bzl", "sh_test")
+
+def _gen_check_script_impl(ctx):
+    """Generates a shell script that exports RunEnvironmentInfo env vars and checks for graceful shutdown."""
+    inner_env = ctx.attr.service_test[RunEnvironmentInfo].environment
+
+    script = ctx.actions.declare_file(ctx.label.name + ".sh")
+    env_exports = "\n".join([
+        'export %s="%s"' % (k, v)
+        for k, v in inner_env.items()
+    ])
+
+    ctx.actions.write(
+        output = script,
+        content = """#!/bin/bash
+set -uo pipefail
+
+# Set the env vars that the service_test needs (from RunEnvironmentInfo).
+{env_exports}
+
+# Run the service_test binary ($1 is resolved via $(location) by Bazel).
+# It is expected to fail because the inner test exits 1.
+"$1" || true
+
+# Poll for the shutdown marker with a timeout.
+for i in $(seq 1 20); do
+    if [ -f "$TEST_TMPDIR/shutdown_marker" ]; then
+        echo "PASS: Service received SIGTERM and shut down gracefully"
+        exit 0
+    fi
+    sleep 0.25
+done
+
+echo "FAIL: Service did NOT receive SIGTERM — graceful shutdown was skipped"
+exit 1
+""".format(env_exports = env_exports),
+        is_executable = True,
+    )
+
+    return [DefaultInfo(files = depset([script]))]
+
+_gen_check_script = rule(
+    implementation = _gen_check_script_impl,
+    attrs = {
+        "service_test": attr.label(
+            mandatory = True,
+        ),
+    },
+)
+
+def check_graceful_shutdown_test(name, service_test, **kwargs):
+    """Verifies that services receive SIGTERM even when the inner test fails."""
+    _gen_check_script(
+        name = name + "_script",
+        service_test = service_test,
+        testonly = True,
+        tags = ["manual"],
+    )
+
+    sh_test(
+        name = name,
+        srcs = [":" + name + "_script"],
+        args = ["$(location %s)" % service_test],
+        data = [service_test],
+        **kwargs
+    )

tests/graceful_shutdown_on_failure/main.go

@@ -0,0 +1,30 @@
+package main
+
+import (
+	"fmt"
+	"net/http"
+	"os"
+	"os/signal"
+	"syscall"
+)
+
+func main() {
+	port := os.Args[1]
+
+	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	})
+
+	go http.ListenAndServe("127.0.0.1:"+port, nil)
+
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, syscall.SIGTERM)
+	<-sigCh
+
+	markerPath := os.Getenv("TEST_TMPDIR") + "/shutdown_marker"
+	if err := os.WriteFile(markerPath, []byte("shutdown"), 0644); err != nil {
+		fmt.Printf("Failed to write shutdown marker: %v\n", err)
+		os.Exit(1)
+	}
+	fmt.Println("Graceful shutdown completed")
+}