Stop services when startup fails

Author: bmcdonnel

.bazelrc

@@ -12,9 +12,9 @@ common --nolegacy_external_runfiles
 common --incompatible_autoload_externally=
 common --repo_env=BAZEL_DO_NOT_DETECT_CPP_TOOLCHAIN=1
 common --repo_env=BAZEL_NO_APPLE_CPP_TOOLCHAIN=1
+common --@rules_go//go/config:pure
 
 common:enable-reload --@rules_itest//:enable_per_service_reload
 
 # TODO: Remove once transitives are fixed
 common --incompatible_autoload_externally=+@rules_cc
-

cmd/svcinit/main.go

@@ -16,6 +16,7 @@ import (
 	"os/signal"
 	"strconv"
 	"strings"
+	"sync"
 	"syscall"
 	"text/tabwriter"
 	"time"
@@ -140,6 +141,17 @@ func main() {
 	r, err := runner.New(ctx, serviceSpecs)
 	must(err)
 
+	stopAll := sync.OnceValues(func() (map[string]*os.ProcessState, error) {
+		states, err := r.StopAll()
+		cancelFunc()
+		return states, err
+	})
+	mustStopAll := func() map[string]*os.ProcessState {
+		states, err := stopAll()
+		must(err)
+		return states
+	}
+
 	servicesErrCh := make(chan error, len(unversionedSpecs))
 
 	go func() {
@@ -168,10 +180,12 @@ func main() {
 
 	criticalPath, err := r.StartAll(servicesErrCh)
 	if errors.Is(err, context.Canceled) {
-		_, err := r.StopAll()
-		must(err)
+		mustStopAll()
 		return
 	}
+	if err != nil {
+		mustStopAll()
+	}
 	must(err)
 
 	// API is                 NewWriter(output io.Writer, minwidth, tabwidth, padding int, padchar byte, flags uint) *Writer
@@ -261,8 +275,7 @@ func main() {
 		select {
 		case <-ctx.Done():
 			log.Println("Shutting down services.")
-			_, err := r.StopAll()
-			must(err)
+			mustStopAll()
 			log.Println("Cleaning up.")
 			return
 		case ibazelCmd := <-interactiveCh:
@@ -302,20 +315,21 @@ func main() {
 			if testErr != nil {
 				log.Printf("Encountered error during test run: %s\n", testErr)
 				if isOneShot {
+					mustStopAll()
 					os.Exit(1)
 				}
 			}
 		case serviceErr := <-servicesErrCh:
 			log.Print(serviceErr)
 			if isOneShot {
+				mustStopAll()
 				log.Fatal("Service exited uncleanly, marking test as failed.\n\n")
 			}
 		}
 
 		if isOneShot {
 			buf.WriteString("Target\tUser Time\tSystem Time\n")
-			states, err := r.StopAll()
-			must(err)
+			states := mustStopAll()
 			for label, state := range states {
 				buf.WriteString(fmt.Sprintf("%s\t%s\t%s\n",
 					label, state.UserTime(), state.SystemTime()))
@@ -637,4 +651,3 @@ func buildTestEnv(ports svclib.Ports) ([]string, error) {
 
 	return baseEnv, nil
 }
-

tests/startup_failure/BUILD.bazel

@@ -1,6 +1,7 @@
 load("@rules_itest//:itest.bzl", "itest_service", "service_test")
 load("@rules_shell//shell:sh_binary.bzl", "sh_binary")
 load("//:must_fail.bzl", "must_fail")
+load(":cleanup_failure_test.bzl", "cleanup_failure_test")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -9,6 +10,17 @@ sh_binary(
     srcs = ["crash_immediately.sh"],
 )
 
+sh_binary(
+    name = "cleanup_on_sigterm",
+    srcs = ["cleanup_on_sigterm.sh"],
+)
+
+cleanup_failure_test(
+    name = "cleanup_on_startup_failure_test",
+    test = ":dependent_startup_failure_test",
+    timeout = "short",
+)
+
 itest_service(
     name = "immediate_exit_service",
     autoassign_port = True,
@@ -26,3 +38,31 @@ must_fail(
     test = "immediate_exit_service_hygiene_test",
     timeout = "short",
 )
+
+itest_service(
+    name = "cleanup_service",
+    exe = ":cleanup_on_sigterm",
+    shutdown_timeout = "5s",
+    tags = ["manual"],
+)
+
+itest_service(
+    name = "dependent_immediate_exit_service",
+    autoassign_port = True,
+    deps = [":cleanup_service"],
+    env = {
+        "PORT": "$${PORT}",
+    },
+    exe = ":crash_immediately",
+    health_check_interval = "100ms",
+    http_health_check_address = "http://127.0.0.1:$${PORT}/health",
+    tags = ["manual"],
+)
+
+service_test(
+    name = "dependent_startup_failure_test",
+    services = [":dependent_immediate_exit_service"],
+    tags = ["manual"],
+    test = "@rules_itest//:exit0",
+    timeout = "short",
+)

tests/startup_failure/cleanup_failure_test.bzl

@@ -0,0 +1,72 @@
+def _shell_quote(value):
+    return "'" + value.replace("'", "'\"'\"'") + "'"
+
+def _cleanup_failure_test_impl(ctx):
+    test_env = ctx.attr.test[RunEnvironmentInfo].environment
+    test_path = ctx.executable.test.short_path
+
+    env_lines = [
+        "export {}={}".format(key, _shell_quote(value))
+        for key, value in sorted(test_env.items())
+    ]
+
+    script = ctx.actions.declare_file(ctx.label.name + ".sh")
+    ctx.actions.write(
+        output = script,
+        is_executable = True,
+        content = """#!/usr/bin/env bash
+set -euo pipefail
+
+cleanup_marker="${{TEST_TMPDIR}}/startup_failure_cleanup_marker"
+log_file="${{TEST_TMPDIR}}/dependent_startup_failure.log"
+test_path="${{TEST_SRCDIR}}/${{TEST_WORKSPACE}}/{test_path}"
+
+rm -f "${{cleanup_marker}}"
+
+{env}
+
+set +e
+"${{test_path}}" >"${{log_file}}" 2>&1
+status=$?
+set -e
+
+cat "${{log_file}}"
+
+if [[ "${{status}}" -eq 0 ]]; then
+    echo "expected dependent startup failure test to fail" >&2
+    exit 1
+fi
+
+if [[ ! -f "${{cleanup_marker}}" ]]; then
+    echo "expected already-started service to receive shutdown before svcinit exited" >&2
+    exit 1
+fi
+
+echo "cleanup marker found: $(cat "${{cleanup_marker}}")"
+""".format(
+            env = "\n".join(env_lines),
+            test_path = test_path,
+        ),
+    )
+
+    runfiles = ctx.runfiles(files = [ctx.executable.test])
+    runfiles = runfiles.merge(ctx.attr.test.default_runfiles)
+
+    return [
+        DefaultInfo(
+            executable = script,
+            runfiles = runfiles,
+        ),
+    ]
+
+cleanup_failure_test = rule(
+    implementation = _cleanup_failure_test_impl,
+    attrs = {
+        "test": attr.label(
+            executable = True,
+            cfg = "target",
+            providers = [RunEnvironmentInfo],
+        ),
+    },
+    test = True,
+)

tests/startup_failure/cleanup_on_sigterm.sh

@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cleanup_marker="${TEST_TMPDIR}/startup_failure_cleanup_marker"
+
+cleanup() {
+    echo "cleanup service received shutdown"
+    echo "cleanup ran" >"${cleanup_marker}"
+    exit 0
+}
+
+trap cleanup TERM INT
+
+while true; do
+    sleep 1
+done