Attempt to gracefully shutdown services first, then hard kill

Author: darkrift

BUILD.bazel

@@ -1,4 +1,4 @@
-load("@bazel_skylib//rules:common_settings.bzl", "bool_flag")
+load("@bazel_skylib//rules:common_settings.bzl", "bool_flag", "string_flag")
 
 alias(
     name = "gazelle",
@@ -11,6 +11,12 @@ cc_binary(
     visibility = ["//visibility:public"],
 )
 
+cc_test(
+    name = "exit0_test",
+    srcs = [":exit0.c"],
+    visibility = ["//visibility:public"],
+)
+
 bool_flag(
     name = "allow_configuring_tmpdir",
     build_setting_default = False,
@@ -36,3 +42,16 @@ bool_flag(
     build_setting_default = False,
     visibility = ["//visibility:public"],
 )
+
+# Default value for shutdown_timeout in itest_service rules.
+string_flag(
+    name = "shutdown_timeout",
+    build_setting_default = "1s",
+    visibility = ["//visibility:public"],
+)
+
+bool_flag(
+    name = "error_on_forceful_shutdown",
+    build_setting_default = False,
+    visibility = ["//visibility:public"],
+)

docs/itest.md

@@ -42,9 +42,10 @@ This can be used in conjunction with the `/v0/port` API to let other tools inter
 <pre>
 load("@rules_itest//private:itest.bzl", "itest_service")
 
-itest_service(<a href="#itest_service-name">name</a>, <a href="#itest_service-deps">deps</a>, <a href="#itest_service-data">data</a>, <a href="#itest_service-autoassign_port">autoassign_port</a>, <a href="#itest_service-env">env</a>, <a href="#itest_service-exe">exe</a>, <a href="#itest_service-expected_start_duration">expected_start_duration</a>, <a href="#itest_service-health_check">health_check</a>,
-              <a href="#itest_service-health_check_args">health_check_args</a>, <a href="#itest_service-health_check_interval">health_check_interval</a>, <a href="#itest_service-health_check_timeout">health_check_timeout</a>, <a href="#itest_service-hot_reloadable">hot_reloadable</a>,
-              <a href="#itest_service-http_health_check_address">http_health_check_address</a>, <a href="#itest_service-named_ports">named_ports</a>, <a href="#itest_service-so_reuseport_aware">so_reuseport_aware</a>)
+itest_service(<a href="#itest_service-name">name</a>, <a href="#itest_service-deps">deps</a>, <a href="#itest_service-data">data</a>, <a href="#itest_service-autoassign_port">autoassign_port</a>, <a href="#itest_service-env">env</a>, <a href="#itest_service-error_on_forceful_shutdown">error_on_forceful_shutdown</a>, <a href="#itest_service-exe">exe</a>,
+              <a href="#itest_service-expected_start_duration">expected_start_duration</a>, <a href="#itest_service-health_check">health_check</a>, <a href="#itest_service-health_check_args">health_check_args</a>, <a href="#itest_service-health_check_interval">health_check_interval</a>,
+              <a href="#itest_service-health_check_timeout">health_check_timeout</a>, <a href="#itest_service-hot_reloadable">hot_reloadable</a>, <a href="#itest_service-http_health_check_address">http_health_check_address</a>, <a href="#itest_service-named_ports">named_ports</a>,
+              <a href="#itest_service-shutdown_signal">shutdown_signal</a>, <a href="#itest_service-shutdown_timeout">shutdown_timeout</a>, <a href="#itest_service-so_reuseport_aware">so_reuseport_aware</a>)
 </pre>
 
 An itest_service is a binary that is intended to run for the duration of the integration test. Examples include databases, HTTP/RPC servers, queue consumers, external service mocks, etc.
@@ -61,6 +62,7 @@ All [common binary attributes](https://bazel.build/reference/be/common-definitio
 | <a id="itest_service-data"></a>data |  -   | <a href="https://bazel.build/concepts/labels">List of labels</a> | optional |  `[]`  |
 | <a id="itest_service-autoassign_port"></a>autoassign_port |  If true, the service manager will pick a free port and assign it to the service. The port will be interpolated into `$${PORT}` in the service's `http_health_check_address` and `args`. It will also be exported under the target's fully qualified label in the service-port mapping.<br><br>The assigned ports for all services are available for substitution in `http_health_check_address` and `args` (in case one service needs the address for another one.) For example, the following substitution: `args = ["-client-addr", "127.0.0.1:$${@@//label/for:service}"]`<br><br>The service-port mapping is a JSON string -> int map propagated through the `ASSIGNED_PORTS` env var. For example, a port can be retrieved with the following JS code: `JSON.parse(process.env["ASSIGNED_PORTS"])["@@//label/for:service"]`.<br><br>Alternately, the env will also contain the location of a binary that can return the port, for contexts without a readily-accessible JSON parser. For example, the following Bash command: `PORT=$($GET_ASSIGNED_PORT_BIN @@//label/for:service)`   | Boolean | optional |  `False`  |
 | <a id="itest_service-env"></a>env |  The service manager will merge these variables into the environment when spawning the underlying binary.   | <a href="https://bazel.build/rules/lib/dict">Dictionary: String -> String</a> | optional |  `{}`  |
+| <a id="itest_service-error_on_forceful_shutdown"></a>error_on_forceful_shutdown |  If set to True, the service manager will fail the service_test if the service had to be forcefully killed if the signal was not SIGKILL and after the shutdown timeout elapsed.<br><br>This needs to be False to have coverage of your services but don't want a them to be graceful at shutdown   | Boolean | optional |  `True`  |
 | <a id="itest_service-exe"></a>exe |  The binary target to run.   | <a href="https://bazel.build/concepts/labels">Label</a> | required |  |
 | <a id="itest_service-expected_start_duration"></a>expected_start_duration |  How long the service expected to take before passing a healthcheck. Any failing health checks before this duration elapses will not be logged.   | String | optional |  `"0s"`  |
 | <a id="itest_service-health_check"></a>health_check |  If set, the service manager will execute this binary to check if the service came up in a healthy state. This check will be retried until it exits with a 0 exit code. When used in conjunction with autoassigned ports, use one of the methods described in `autoassign_port` to locate the service.   | <a href="https://bazel.build/concepts/labels">Label</a> | optional |  `None`  |
@@ -70,6 +72,8 @@ All [common binary attributes](https://bazel.build/reference/be/common-definitio
 | <a id="itest_service-hot_reloadable"></a>hot_reloadable |  If set to True, the service manager will propagate ibazel's reload notification over stdin instead of restarting the service. See the ruleset docstring for more info on using ibazel   | Boolean | optional |  `False`  |
 | <a id="itest_service-http_health_check_address"></a>http_health_check_address |  If set, the service manager will send an HTTP request to this address to check if the service came up in a healthy state. This check will be retried until it returns a 200 HTTP code. When used in conjunction with autoassigned ports, `$${@@//label/for:service:port_name}` can be used in the address. Example: `http_health_check_address = "http://127.0.0.1:$${@@//label/for:service:port_name}",`   | String | optional |  `""`  |
 | <a id="itest_service-named_ports"></a>named_ports |  For each element of the list, the service manager will pick a free port and assign it to the service. The port's fully-qualified name is the service's fully-qualified label and the port name, separated by a colon. For example, a port assigned with `named_ports = ["http_port"]` will be assigned a fully-qualified name of `@@//label/for:service:http_port`.<br><br>Named ports are accessible through the service-port mapping. For more details, see `autoassign_port`.   | List of strings | optional |  `[]`  |
+| <a id="itest_service-shutdown_signal"></a>shutdown_signal |  The signal to send to the service when it needs to be shut down. Valid values are: SIGTERM and SIGKILL. SIGTERM is necessary to have proper coverage of services which needs to be gracefully terminated   | String | optional |  `"SIGTERM"`  |
+| <a id="itest_service-shutdown_timeout"></a>shutdown_timeout |  The duration to wait after sending the shutdown signal before forcefully killing the service. The syntax is based on common time duration with a number, followed by the time unit. For example, `200ms`, `1s`, `2m`, `3h`, `4d`.   | String | optional |  `"1s"`  |
 | <a id="itest_service-so_reuseport_aware"></a>so_reuseport_aware |  If set, the service manager will not release the autoassigned port. The service binary must use SO_REUSEPORT when binding it. This reduces the possibility of port collisions when running many service_tests in parallel, or when code binds port 0 without being aware of the port assignment mechanism.<br><br>Must only be set when `autoassign_port` is enabled or `named_ports` are used.   | Boolean | optional |  `False`  |
 
 

private/itest.bzl

@@ -138,7 +138,6 @@ def _itest_binary_impl(ctx, extra_service_spec_kwargs, extra_exe_runfiles = []):
         for arg in ctx.attr.args
     ]
 
-
     if version_file:
         extra_service_spec_kwargs["version_file"] = to_rlocation_path(ctx, version_file)
 
@@ -198,6 +197,9 @@ def _itest_service_impl(ctx):
         "expected_start_duration": ctx.attr.expected_start_duration,
         "health_check_interval": ctx.attr.health_check_interval,
         "health_check_timeout": ctx.attr.health_check_timeout,
+        "shutdown_signal": ctx.attr.shutdown_signal,
+        "shutdown_timeout": str(ctx.attr.shutdown_timeout[BuildSettingInfo].value),
+        "error_on_forceful_shutdown": bool(ctx.attr.error_on_forceful_shutdown[BuildSettingInfo].value),
     }
     extra_exe_runfiles = []
 
@@ -278,6 +280,21 @@ _itest_service_attrs = _itest_binary_attrs | {
         This check will be retried until it returns a 200 HTTP code. When used in conjunction with autoassigned ports, `$${@@//label/for:service:port_name}` can be used in the address.
         Example: `http_health_check_address = "http://127.0.0.1:$${@@//label/for:service:port_name}",`""",
     ),
+    "shutdown_signal": attr.string(
+        default = "SIGTERM",
+        doc = "The signal to send to the service when it needs to be shut down. Valid values are: SIGTERM and SIGKILL. SIGTERM is necessary to have proper coverage of services which needs to be gracefully terminated",
+        values = ["SIGTERM", "SIGKILL"],
+    ),
+    "shutdown_timeout": attr.label(
+        default = "//:shutdown_timeout",
+        doc = "The duration to wait after sending the shutdown signal before forcefully killing the service. The syntax is based on common time duration with a number, followed by the time unit. For example, `200ms`, `1s`, `2m`, `3h`, `4d`.",
+    ),
+    "error_on_forceful_shutdown": attr.label(
+        default = "//:error_on_forceful_shutdown",
+        doc = """If set to True, the service manager will fail the service_test if the service had to be forcefully killed if the signal was not SIGKILL and after the shutdown timeout elapsed.
+
+        This needs to be False to have coverage of your services but don't want a them to be graceful at shutdown""",
+    ),
 }
 
 itest_service = rule(

runner/runner.go

@@ -9,7 +9,6 @@ import (
 	"reflect"
 	"runtime"
 	"sync"
-	"syscall"
 	"time"
 
 	"rules_itest/logger"
@@ -100,12 +99,11 @@ func (r *Runner) StartAll(serviceErrCh chan error) ([]topological.Task, error) {
 
 func (r *Runner) StopAll() (map[string]*os.ProcessState, error) {
 	tasks := allTasks(r.serviceInstances, func(ctx context.Context, service *ServiceInstance) error {
-		if service.Type == "group" {
+		if service.Type == "group" || service.Type == "task" {
 			return nil
 		}
 		log.Printf("Stopping %s\n", colorize(service.VersionedServiceSpec))
-		service.Stop(syscall.SIGKILL)
-		return nil
+		return service.Stop(nil)
 	})
 	stopper := topological.NewReversedRunner(tasks)
 	err := stopper.Run(r.ctx)
@@ -185,10 +183,10 @@ func (r *Runner) UpdateSpecs(serviceSpecs ServiceSpecs, ibazelCmd []byte) error
 
 	for _, label := range updateActions.toStopLabels {
 		serviceInstance := r.serviceInstances[label]
-		if serviceInstance.Type == "group" {
+		if serviceInstance.Type == "group" || serviceInstance.Type == "task" {
 			continue
 		}
-		serviceInstance.Stop(syscall.SIGKILL)
+		serviceInstance.Stop(nil)
 		delete(r.serviceInstances, label)
 	}
 

runner/service_instance.go

@@ -2,6 +2,7 @@ package runner
 
 import (
 	"context"
+	"errors"
 	"io"
 	"log"
 	"net/http"
@@ -34,6 +35,8 @@ type ServiceInstance struct {
 	done                 bool
 }
 
+var forcedKillError = errors.New("process forcefully killed with SIGKILL, to circumvent this error, set shutdown_signal to SIGKILL")
+
 func (s *ServiceInstance) Start(ctx context.Context) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -197,7 +200,7 @@ func (s *ServiceInstance) Error() error {
 	return s.runErr
 }
 
-func (s *ServiceInstance) Stop(sig syscall.Signal) error {
+func (s *ServiceInstance) Stop(sig *syscall.Signal) error {
 	func() {
 		s.mu.Lock()
 		defer s.mu.Unlock()
@@ -208,13 +211,74 @@ func (s *ServiceInstance) Stop(sig syscall.Signal) error {
 		return nil
 	}
 
-	err := killGroup(s.cmd, sig)
+	if s.cmd.ProcessState.ExitCode() == 0 {
+		return nil
+	}
+
+	var signal syscall.Signal
+	if sig == nil {
+		if s.ShutdownSignal == "SIGKILL" {
+			signal = syscall.SIGKILL
+		} else if s.ShutdownSignal == "SIGTERM" {
+			signal = syscall.SIGTERM
+		} else {
+			// Default to SIGKILL if unspecified or unrecognized. In case we add new values to itest.bzl but forget to add it here
+			signal = syscall.SIGKILL
+		}
+	} else {
+		signal = *sig
+	}
+
+	err := killGroup(s.cmd, signal)
 	if err != nil {
 		return err
 	}
 
-	for !s.isDone() {
-		time.Sleep(5 * time.Millisecond)
+	if signal == syscall.SIGKILL {
+		log.Printf("Sent SIGKILL to %s\n", s.Colorize(s.Label))
+		for !s.isDone() {
+			time.Sleep(5 * time.Millisecond)
+		}
+	} else {
+		log.Printf("Sent SIGTERM to %s\n", s.Colorize(s.Label))
+
+		shutdownTimeout, err := time.ParseDuration(s.VersionedServiceSpec.ShutdownTimeout)
+		if err != nil {
+			log.Printf("failed to parse health check timeout, falling back to no timeout: %v", err)
+			return err
+		}
+
+		log.Printf("Sent SIGTERM to %s, waiting for %s to stop gracefully\n", s.Colorize(s.Label), s.VersionedServiceSpec.ShutdownTimeout)
+
+		waitFor := time.After(shutdownTimeout)
+		for {
+			// Check if the process has exited
+			if s.isDone() {
+				break
+			}
+
+			select {
+			case <-waitFor:
+				log.Printf("%s did not exit within %s, sending SIGKILL. If you are trying to collect coverage, you will most likely miss stats, try increasing the shutdown timeout.\n", s.Colorize(s.Label), shutdownTimeout)
+
+				err := killGroup(s.cmd, syscall.SIGKILL)
+				if err != nil {
+					return err
+				}
+
+				for !s.isDone() {
+					time.Sleep(5 * time.Millisecond)
+				}
+
+				if s.ErrorOnForcefulShutdown {
+					return forcedKillError
+				}
+
+				return nil
+			default:
+				time.Sleep(5 * time.Millisecond)
+			}
+		}
 	}
 
 	return nil

svcctl/svcctl.go

@@ -83,17 +83,21 @@ func handleStart(ctx context.Context, r *runner.Runner, serviceErrCh chan error,
 	w.WriteHeader(http.StatusOK)
 }
 
+func Ptr[T any](v T) *T {
+	return &v
+}
+
 func handleKill(ctx context.Context, r *runner.Runner, _ chan error, w http.ResponseWriter, req *http.Request) {
-	sig := syscall.SIGKILL
+	var sig *syscall.Signal // Use the default defined signal configured on the service.
 	params := req.URL.Query()
 	signal := params.Get("signal")
 	if signal != "" {
 		// Currently only SIGTERM and SIGKILL are supported.
 		switch signal {
 		case "SIGTERM":
-			sig = syscall.SIGTERM
+			sig = Ptr(syscall.SIGTERM)
 		case "SIGKILL":
-			sig = syscall.SIGKILL
+			sig = Ptr(syscall.SIGKILL)
 		default:
 			http.Error(w, "unsupported signal", http.StatusBadRequest)
 		}

svclib/types.go

@@ -1,29 +1,34 @@
 package svclib
 
-import "rules_itest/logger"
+import (
+	"rules_itest/logger"
+)
 
 // Created by Starlark
 type ServiceSpec struct {
 	// Type can be "service", "task", or "group".
-	Type                   string            `json:"type"`
-	Label                  string            `json:"label"`
-	Args                   []string          `json:"args"`
-	Env                    map[string]string `json:"env"`
-	Exe                    string            `json:"exe"`
-	HttpHealthCheckAddress string            `json:"http_health_check_address"`
-	ExpectedStartDuration  string            `json:"expected_start_duration"`
-	HealthCheck            string            `json:"health_check"`
-	HealthCheckLabel       string            `json:"health_check_label"`
-	HealthCheckArgs        []string          `json:"health_check_args"`
-	HealthCheckInterval    string            `json:"health_check_interval"`
-	HealthCheckTimeout     string            `json:"health_check_timeout"`
-	VersionFile            string            `json:"version_file"`
-	Deps                   []string          `json:"deps"`
-	AutoassignPort         bool              `json:"autoassign_port"`
-	SoReuseportAware       bool              `json:"so_reuseport_aware"`
-	NamedPorts             []string          `json:"named_ports"`
-	HotReloadable          bool              `json:"hot_reloadable"`
-	PortAliases            map[string]string `json:"port_aliases"`
+	Type                    string            `json:"type"`
+	Label                   string            `json:"label"`
+	Args                    []string          `json:"args"`
+	Env                     map[string]string `json:"env"`
+	Exe                     string            `json:"exe"`
+	HttpHealthCheckAddress  string            `json:"http_health_check_address"`
+	ExpectedStartDuration   string            `json:"expected_start_duration"`
+	HealthCheck             string            `json:"health_check"`
+	HealthCheckLabel        string            `json:"health_check_label"`
+	HealthCheckArgs         []string          `json:"health_check_args"`
+	HealthCheckInterval     string            `json:"health_check_interval"`
+	HealthCheckTimeout      string            `json:"health_check_timeout"`
+	VersionFile             string            `json:"version_file"`
+	Deps                    []string          `json:"deps"`
+	AutoassignPort          bool              `json:"autoassign_port"`
+	SoReuseportAware        bool              `json:"so_reuseport_aware"`
+	NamedPorts              []string          `json:"named_ports"`
+	HotReloadable           bool              `json:"hot_reloadable"`
+	PortAliases             map[string]string `json:"port_aliases"`
+	ShutdownSignal          string            `json:"shutdown_signal"`
+	ShutdownTimeout         string            `json:"shutdown_timeout"`
+	ErrorOnForcefulShutdown bool              `json:"error_on_forceful_shutdown"`
 }
 
 // Our internal representation.