Add rusage wrapper to correctly report usage for services that spawn children (like js_binary)

cmd/rusage_wrapper/BUILD.bazel

@@ -0,0 +1,14 @@
+load("@rules_go//go:def.bzl", "go_binary", "go_library")
+
+go_library(
+    name = "rusage_wrapper_lib",
+    srcs = ["main.go"],
+    importpath = "rules_itest/cmd/rusage_wrapper",
+    visibility = ["//visibility:private"],
+)
+
+go_binary(
+    name = "rusage_wrapper",
+    embed = [":rusage_wrapper_lib"],
+    visibility = ["//visibility:public"],
+)

cmd/rusage_wrapper/main.go

@@ -0,0 +1,68 @@
+package main
+
+import (
+	"encoding/json"
+	"os"
+	"os/exec"
+	"os/signal"
+	"syscall"
+)
+
+func must(err error) {
+	if err != nil {
+		panic(err)
+	}
+}
+
+func main() {
+	signals := make(chan os.Signal, 1)
+	signal.Notify(signals,
+		syscall.SIGINT,
+		syscall.SIGTERM,
+		syscall.SIGQUIT,
+		syscall.SIGHUP,
+		syscall.SIGUSR1,
+		syscall.SIGUSR2,
+		syscall.SIGPIPE,
+		syscall.SIGALRM,
+		syscall.SIGTSTP,
+		syscall.SIGCONT,
+	)
+
+	rusageFile := os.Args[1]
+	bin := os.Args[2]
+	args := os.Args[3:]
+
+	cmd := exec.Command(bin, args...)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+
+	err := cmd.Start()
+	must(err)
+
+	go func() {
+		for sig := range signals {
+			if sig == syscall.SIGUSR2 {
+				sig = syscall.SIGKILL
+			}
+			err := cmd.Process.Signal(sig)
+			must(err)
+		}
+	}()
+
+	err = cmd.Wait()
+	if _, ok := err.(*exec.ExitError); !ok {
+		panic(err)
+	}
+
+	var usage syscall.Rusage
+	syscall.Getrusage(syscall.RUSAGE_CHILDREN, &usage)
+
+	data, err := json.Marshal(usage)
+	must(err)
+
+	err = os.WriteFile(rusageFile, data, 0600)
+	must(err)
+
+	os.Exit(cmd.ProcessState.ExitCode())
+}

cmd/svcinit/main.go

@@ -33,6 +33,10 @@ func must(err error) {
 	}
 }
 
+func duration(t syscall.Timeval) time.Duration {
+	return time.Duration(t.Nano()) * time.Nanosecond
+}
+
 func main() {
 	serviceSpecsPath, err := runfiles.Rlocation(os.Getenv("SVCINIT_SERVICE_SPECS_RLOCATION_PATH"))
 	must(err)
@@ -217,11 +221,11 @@ func main() {
 
 		if isOneShot {
 			buf.WriteString("Target\tUser Time\tSystem Time\n")
-			states, err := r.StopAll()
+			rusages, err := r.StopAll()
 			must(err)
-			for label, state := range states {
-				buf.WriteString(fmt.Sprintf("%s\t%s\t%s\n",
-					label, state.UserTime(), state.SystemTime()))
+			for label, rusage := range rusages {
+				buf.WriteString(fmt.Sprintf("%s\t%v\t%v\n",
+					label, duration(rusage.Utime), duration(rusage.Stime)))
 			}
 		} else {
 			buf.WriteString("Target\tStartup Time\n")
@@ -383,13 +387,25 @@ func augmentServiceSpecs(
 ) (
 	map[string]svclib.VersionedServiceSpec, error,
 ) {
+	rusageWrapperBin, err := runfiles.Rlocation(os.Getenv("RUSAGE_WRAPPER_BIN_RLOCATION_PATH"))
+	if err != nil {
+		return nil, err
+	}
+
 	tmpDir := os.Getenv("TMPDIR")
 	socketDir := os.Getenv("SOCKET_DIR")
 
 	versionedServiceSpecs := make(map[string]svclib.VersionedServiceSpec, len(serviceSpecs))
 	for label, serviceSpec := range serviceSpecs {
+		f, err := os.CreateTemp("", "rusage")
+		must(err)
+
+		err = f.Close()
+		must(err)
+
 		s := svclib.VersionedServiceSpec{
 			ServiceSpec: serviceSpec,
+			RusageFile:  f.Name(),
 		}
 
 		if s.Type == "group" {
@@ -401,7 +417,8 @@ func augmentServiceSpecs(
 		if err != nil {
 			return nil, err
 		}
-		s.Exe = exePath
+		s.Exe = rusageWrapperBin
+		s.Args = append([]string{s.RusageFile, exePath}, s.Args...)
 
 		if s.HealthCheck != "" {
 			healthCheckPath, err := runfiles.Rlocation(serviceSpec.HealthCheck)

private/itest.bzl

@@ -56,6 +56,7 @@ def _run_environment(ctx, service_specs_file):
         "SVCINIT_SERVICE_SPECS_RLOCATION_PATH": to_rlocation_path(ctx, service_specs_file),
         "SVCINIT_ENABLE_PER_SERVICE_RELOAD": str(ctx.attr._enable_per_service_reload[BuildSettingInfo].value),
         "SVCINIT_GET_ASSIGNED_PORT_BIN_RLOCATION_PATH": to_rlocation_path(ctx, ctx.executable._get_assigned_port),
+        "RUSAGE_WRAPPER_BIN_RLOCATION_PATH": to_rlocation_path(ctx, ctx.executable._rusage_wrapper),
     }
 
 def _services_runfiles(ctx, services_attr_name = "services"):
@@ -65,6 +66,7 @@ def _services_runfiles(ctx, services_attr_name = "services"):
     ] + [
         ctx.attr._svcinit.default_runfiles,
         ctx.attr._get_assigned_port.default_runfiles,
+        ctx.attr._rusage_wrapper.default_runfiles,
     ]
 
 _svcinit_attrs = {
@@ -78,6 +80,11 @@ _svcinit_attrs = {
         executable = True,
         cfg = "target",
     ),
+    "_rusage_wrapper": attr.label(
+        default = "//cmd/rusage_wrapper",
+        executable = True,
+        cfg = "target",
+    ),
     "_enable_per_service_reload": attr.label(
         default = "//:enable_per_service_reload",
     ),

runner/runner.go

@@ -2,6 +2,7 @@ package runner
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"log"
 	"os"
@@ -91,28 +92,40 @@ func (r *Runner) StartAll(serviceErrCh chan error) ([]topological.Task, error) {
 	return starter.CriticalPath(), err
 }
 
-func (r *Runner) StopAll() (map[string]*os.ProcessState, error) {
+func (r *Runner) StopAll() (map[string]syscall.Rusage, error) {
 	tasks := allTasks(r.serviceInstances, func(ctx context.Context, service *ServiceInstance) error {
 		if service.Type == "group" {
 			return nil
 		}
 		log.Printf("Stopping %s\n", colorize(service.VersionedServiceSpec))
-		service.Stop(syscall.SIGKILL)
+		// Rusage wrapper reinterprets SIGUSR2 as SIGKILL for the underlying binary.
+		service.Stop(syscall.SIGUSR2)
 		return nil
 	})
 	stopper := topological.NewReversedRunner(tasks)
 	err := stopper.Run(r.ctx)
 
-	states := make(map[string]*os.ProcessState)
+	rusages := make(map[string]syscall.Rusage)
 
 	for _, serviceInstance := range r.serviceInstances {
 		if serviceInstance.Type == "group" {
 			continue
 		}
-		states[serviceInstance.Label] = serviceInstance.ProcessState()
+		data, err := os.ReadFile(serviceInstance.RusageFile)
+		if err != nil {
+			log.Print("Could not read rusage for", serviceInstance.Label)
+			continue
+		}
+
+		var rusage syscall.Rusage
+		err = json.Unmarshal(data, &rusage)
+		if err != nil {
+			return nil, err
+		}
+		rusages[serviceInstance.Label] = rusage
 	}
 
-	return states, err
+	return rusages, err
 }
 
 func (r *Runner) GetStartDurations() map[string]time.Duration {

svclib/types.go

@@ -27,8 +27,9 @@ type ServiceSpec struct {
 // Our internal representation.
 type VersionedServiceSpec struct {
 	ServiceSpec
-	Version string
-	Color   string
+	Version    string
+	Color      string
+	RusageFile string
 }
 
 func (v VersionedServiceSpec) Colorize(label string) string {

tests/svcctl/svcctl_test.go

@@ -75,7 +75,7 @@ func TestSvcctl(t *testing.T) {
 		t.Errorf("Got status code %d, want %d", resp.StatusCode, http.StatusOK)
 	}
 
-	// Wait for speedy service to stop with exit code -1
+	// Wait for speedy service to stop with exit code 255
 	params = url.Values{}
 	params.Add("service", "@@//:_speedy_service")
 	resp, err = http.Get(svcctlHost + "/v0/wait?" + params.Encode())
@@ -86,13 +86,13 @@ func TestSvcctl(t *testing.T) {
 		t.Errorf("Got status code %d, want %d", resp.StatusCode, http.StatusOK)
 	}
 
-	// Terminated by signal, so the exit code should be -1 (except on Windows)
+	// Terminated by signal, so the exit code should be 255 (except on Windows)
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
 		t.Errorf("Failed to read response body: %v", err)
 	}
 
-	wantCode := "-1"
+	wantCode := "255"
 	if runtime.GOOS == "windows" {
 		wantCode = "1"
 	}