From 06000fe66b6123c37686384d73557356bd33dd7b Mon Sep 17 00:00:00 2001 From: David Zbarsky Date: Wed, 28 Aug 2024 10:07:50 -0400 Subject: [PATCH] Add rusage wrapper to correctly report usage for services that spawn children (like js_binary) --- cmd/rusage_wrapper/BUILD.bazel | 14 +++++++ cmd/rusage_wrapper/main.go | 68 ++++++++++++++++++++++++++++++++++ cmd/svcinit/main.go | 27 +++++++++++--- private/itest.bzl | 7 ++++ runner/runner.go | 23 +++++++++--- svclib/types.go | 5 ++- tests/svcctl/svcctl_test.go | 6 +-- 7 files changed, 135 insertions(+), 15 deletions(-) create mode 100644 cmd/rusage_wrapper/BUILD.bazel create mode 100644 cmd/rusage_wrapper/main.go diff --git a/cmd/rusage_wrapper/BUILD.bazel b/cmd/rusage_wrapper/BUILD.bazel new file mode 100644 index 0000000..a428e1b --- /dev/null +++ b/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"], +) diff --git a/cmd/rusage_wrapper/main.go b/cmd/rusage_wrapper/main.go new file mode 100644 index 0000000..2a78228 --- /dev/null +++ b/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()) +} diff --git a/cmd/svcinit/main.go b/cmd/svcinit/main.go index 9b873e5..9ddf488 100644 --- a/cmd/svcinit/main.go +++ b/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) diff --git a/private/itest.bzl b/private/itest.bzl index e5cbd0c..a219d67 100644 --- a/private/itest.bzl +++ b/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", ), diff --git a/runner/runner.go b/runner/runner.go index 7b5f3c6..121c145 100644 --- a/runner/runner.go +++ b/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 { diff --git a/svclib/types.go b/svclib/types.go index c80032b..ba0ee2e 100644 --- a/svclib/types.go +++ b/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 { diff --git a/tests/svcctl/svcctl_test.go b/tests/svcctl/svcctl_test.go index 7a5dbe7..2491b1f 100644 --- a/tests/svcctl/svcctl_test.go +++ b/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" }