diff --git a/cmd/svcinit/BUILD.bazel b/cmd/svcinit/BUILD.bazel index 7b4a22a..ee5c5a7 100644 --- a/cmd/svcinit/BUILD.bazel +++ b/cmd/svcinit/BUILD.bazel @@ -60,7 +60,10 @@ go_library( go_test( name = "svcinit_test", - srcs = ["reserve_reusable_port_test.go"], + srcs = [ + "reserve_reusable_port_test.go", + "target_args_test.go", + ], embed = [":svcinit_lib"], ) diff --git a/cmd/svcinit/main.go b/cmd/svcinit/main.go index 6114f6c..7d131e3 100644 --- a/cmd/svcinit/main.go +++ b/cmd/svcinit/main.go @@ -16,6 +16,7 @@ import ( "os/exec" "os/signal" "runtime" + "sort" "strconv" "strings" "sync" @@ -44,6 +45,9 @@ var ( shouldKeepServicesUp = os.Getenv("SVCINIT_KEEP_SERVICES_UP") == "True" ) +const delegatedTargetFlag = "--target_arg" +const delegatedTargetEnvFlag = "--target_env" + // Assigned by x_def var getAssignedPortRlocationPath string @@ -117,8 +121,12 @@ func main() { isOneShot := !shouldHotReload && testLabel != "" && !shouldKeepServicesUp - unversionedSpecs, err := readServiceSpecs(serviceSpecsPath) + unversionedSpecs, aliases, err := readServiceSpecs(serviceSpecsPath) must(err) + if testLabel == "" { + err = appendDelegatedTargetConfig(unversionedSpecs, aliases, os.Args[1:]) + must(err) + } // Make sure we grab the svcctl port before we assign test ports, // otherwise we might steal an assigned port by accident. @@ -284,19 +292,23 @@ func main() { log.Println(ibazelCmd) // Restart any services as needed. - unversionedSpecs, err := readServiceSpecs(serviceSpecsPath) + unversionedSpecs, aliases, err := readServiceSpecs(serviceSpecsPath) must(err) + if testLabel == "" { + err = appendDelegatedTargetConfig(unversionedSpecs, aliases, os.Args[1:]) + must(err) + } serviceSpecs, err := augmentServiceSpecs(unversionedSpecs, ports, svcctlPortStr) must(err) testCancel() - // This is a brittle way of draining a channel in a nonblocking way, - // consider instead signalling cancellation of the services with a - // context, letting them close the channel, and using a waitgroup to - // wait for them to exit. - // See: https://github.com/hermeticbuild/rules_itest/issues/72 + // This is a brittle way of draining a channel in a nonblocking way, + // consider instead signalling cancellation of the services with a + // context, letting them close the channel, and using a waitgroup to + // wait for them to exit. + // See: https://github.com/hermeticbuild/rules_itest/issues/72 Drain: for { select { @@ -363,14 +375,26 @@ func main() { func readServiceSpecs( path string, ) ( - map[string]svclib.ServiceSpec, error, + map[string]svclib.ServiceSpec, map[string][]string, error, ) { data, err := os.ReadFile(path) must(err) - var serviceSpecs map[string]svclib.ServiceSpec - err = json.Unmarshal(data, &serviceSpecs) - return serviceSpecs, err + var graph struct { + Services map[string]svclib.ServiceSpec `json:"services"` + Aliases map[string][]string `json:"aliases"` + } + err = json.Unmarshal(data, &graph) + if err != nil { + return nil, nil, err + } + if graph.Services == nil { + graph.Services = map[string]svclib.ServiceSpec{} + } + if graph.Aliases == nil { + graph.Aliases = map[string][]string{} + } + return graph.Services, graph.Aliases, nil } func assignPorts( @@ -652,6 +676,202 @@ func replaceAll(s string, replacements []Replacement) string { return s } +type delegatedTargetConfig struct { + Args map[string][]string + Env map[string]map[string]string +} + +func appendDelegatedTargetConfig(serviceSpecs map[string]svclib.ServiceSpec, aliases map[string][]string, rawArgs []string) error { + config, err := parseDelegatedTargetConfig(rawArgs, serviceSpecs, aliases) + if err != nil { + return err + } + + for label, args := range config.Args { + spec := serviceSpecs[label] + spec.Args = append(spec.Args, args...) + serviceSpecs[label] = spec + } + for label, env := range config.Env { + spec := serviceSpecs[label] + if spec.Env == nil { + spec.Env = map[string]string{} + } + for key, value := range env { + spec.Env[key] = value + } + serviceSpecs[label] = spec + } + + return nil +} + +func parseDelegatedTargetConfig(rawArgs []string, serviceSpecs map[string]svclib.ServiceSpec, aliases map[string][]string) (delegatedTargetConfig, error) { + config := delegatedTargetConfig{ + Args: map[string][]string{}, + Env: map[string]map[string]string{}, + } + currentTarget := "" + + for i := 0; i < len(rawArgs); i++ { + arg := rawArgs[i] + switch { + case arg == delegatedTargetFlag: + if i+1 >= len(rawArgs) { + return delegatedTargetConfig{}, fmt.Errorf("missing target after %s", delegatedTargetFlag) + } + label, err := resolveDelegatedTarget(rawArgs[i+1], serviceSpecs, aliases) + if err != nil { + return delegatedTargetConfig{}, err + } + currentTarget = label + if _, ok := config.Args[label]; !ok { + config.Args[label] = nil + } + i++ + case arg == delegatedTargetEnvFlag: + if i+2 >= len(rawArgs) { + return delegatedTargetConfig{}, fmt.Errorf("expected %s ", delegatedTargetEnvFlag) + } + label, err := resolveDelegatedTarget(rawArgs[i+1], serviceSpecs, aliases) + if err != nil { + return delegatedTargetConfig{}, err + } + key, value, err := parseDelegatedEnvAssignment(rawArgs[i+2]) + if err != nil { + return delegatedTargetConfig{}, err + } + if _, ok := config.Env[label]; !ok { + config.Env[label] = map[string]string{} + } + config.Env[label][key] = value + i += 2 + default: + if currentTarget == "" { + return delegatedTargetConfig{}, fmt.Errorf("unexpected argument %q: expected %s before delegated args", arg, delegatedTargetFlag) + } + config.Args[currentTarget] = append(config.Args[currentTarget], arg) + } + } + + return config, nil +} + +func parseDelegatedTargetArgs(rawArgs []string, serviceSpecs map[string]svclib.ServiceSpec, aliases map[string][]string) (map[string][]string, error) { + config, err := parseDelegatedTargetConfig(rawArgs, serviceSpecs, aliases) + if err != nil { + return nil, err + } + + return config.Args, nil +} + +func parseDelegatedEnvAssignment(raw string) (string, string, error) { + key, value, ok := strings.Cut(raw, "=") + if !ok || key == "" { + return "", "", fmt.Errorf("invalid delegated env assignment %q: expected KEY=VALUE", raw) + } + return key, value, nil +} + +func resolveDelegatedTarget(target string, serviceSpecs map[string]svclib.ServiceSpec, aliases map[string][]string) (string, error) { + matches := []string{} + groupMatches := []string{} + for label, spec := range serviceSpecs { + if !delegatedTargetMatches(target, label) { + continue + } + if spec.Type == "group" { + groupMatches = append(groupMatches, label) + continue + } + matches = append(matches, label) + } + for alias, labels := range aliases { + if !delegatedTargetMatches(target, alias) { + continue + } + for _, label := range labels { + spec, ok := serviceSpecs[label] + if !ok { + return "", fmt.Errorf("delegated target alias %q points at unknown itest target %q", alias, label) + } + if spec.Type == "group" { + groupMatches = append(groupMatches, alias+" -> "+label) + continue + } + matches = append(matches, label) + } + } + + matches = uniqueSorted(matches) + groupMatches = uniqueSorted(groupMatches) + sort.Strings(matches) + sort.Strings(groupMatches) + + switch len(matches) { + case 1: + return matches[0], nil + case 0: + if len(groupMatches) > 0 { + return "", fmt.Errorf("delegated target %q refers to a non-executable itest_service_group: %s", target, strings.Join(groupMatches, ", ")) + } + return "", fmt.Errorf("delegated target %q not found. Available executable itest targets: %s", target, strings.Join(executableItestTargets(serviceSpecs), ", ")) + default: + return "", fmt.Errorf("delegated target %q is ambiguous. Matches: %s", target, strings.Join(matches, ", ")) + } +} + +func delegatedTargetMatches(target string, label string) bool { + if target == label { + return true + } + + withoutModule := strings.TrimPrefix(label, "@@") + if target == withoutModule { + return true + } + + colon := strings.LastIndex(label, ":") + if colon >= 0 { + targetName := label[colon+1:] + if target == targetName || target == ":"+targetName { + return true + } + } + + return false +} + +func uniqueSorted(values []string) []string { + if len(values) == 0 { + return nil + } + + seen := map[string]struct{}{} + unique := make([]string, 0, len(values)) + for _, value := range values { + if _, ok := seen[value]; ok { + continue + } + seen[value] = struct{}{} + unique = append(unique, value) + } + sort.Strings(unique) + return unique +} + +func executableItestTargets(serviceSpecs map[string]svclib.ServiceSpec) []string { + targets := make([]string, 0, len(serviceSpecs)) + for label, spec := range serviceSpecs { + if spec.Type == "group" { + continue + } + targets = append(targets, label) + } + return uniqueSorted(targets) +} + func buildTestEnv(ports svclib.Ports) ([]string, error) { testEnvPath, err := runfiles.Rlocation(os.Getenv("SVCINIT_TEST_ENV_RLOCATION_PATH")) if err != nil { diff --git a/cmd/svcinit/target_args/BUILD.bazel b/cmd/svcinit/target_args/BUILD.bazel new file mode 100644 index 0000000..6aa1af8 --- /dev/null +++ b/cmd/svcinit/target_args/BUILD.bazel @@ -0,0 +1,62 @@ +load("@rules_go//go:def.bzl", "go_binary") +load("@rules_itest//:itest.bzl", "itest_task") +load(":target_args_test.bzl", "target_args_test") + +package(default_visibility = ["//visibility:public"]) + +go_binary( + name = "capture_args", + srcs = ["capture_args.go"], +) + +itest_task( + name = "dep_task", + args = ["dep-base"], + env = { + "OUTPUT_FILE": "dep_args.txt", + "TASK_NAME": "@@//cmd/svcinit/target_args:dep_task", + }, + exe = ":capture_args", +) + +alias( + name = "dep_alias", + actual = ":dep_task", +) + +itest_task( + name = "top_task", + args = ["top-base"], + deps = [":dep_alias"], + env = { + "OUTPUT_FILE": "top_args.txt", + "TASK_NAME": "@@//cmd/svcinit/target_args:top_task", + }, + exe = ":capture_args", +) + +target_args_test( + name = "target_args_test", + argv = [ + "--target_arg", + "top_task", + "--top-flag", + "--top-pair", + "with value", + "--target_env", + "//cmd/svcinit/target_args:top_task", + "INJECTED_ENV=top env value", + "--target_arg", + "dep_alias", + "--dep-flag", + "--dep-other", + "--target_env", + "//cmd/svcinit/target_args:dep_alias", + "INJECTED_ENV=dep env value", + ], + expected_files = { + "dep_args.txt": "TASK_NAME=@@//cmd/svcinit/target_args:dep_task\nINJECTED_ENV=dep env value\nARGS=dep-base\n--dep-flag\n--dep-other", + "top_args.txt": "TASK_NAME=@@//cmd/svcinit/target_args:top_task\nINJECTED_ENV=top env value\nARGS=top-base\n--top-flag\n--top-pair\nwith value", + }, + target_under_test = ":top_task", +) diff --git a/cmd/svcinit/target_args/capture_args.go b/cmd/svcinit/target_args/capture_args.go new file mode 100644 index 0000000..a7ae2ea --- /dev/null +++ b/cmd/svcinit/target_args/capture_args.go @@ -0,0 +1,25 @@ +package main + +import ( + "fmt" + "os" + "path/filepath" + "strings" +) + +func main() { + outputFile := os.Getenv("OUTPUT_FILE") + if outputFile == "" { + panic("OUTPUT_FILE must be set") + } + + outputPath := filepath.Join(os.Getenv("TEST_TMPDIR"), outputFile) + content := fmt.Sprintf("TASK_NAME=%s\nINJECTED_ENV=%s\nARGS=%s\n", + os.Getenv("TASK_NAME"), + os.Getenv("INJECTED_ENV"), + strings.Join(os.Args[1:], "\n"), + ) + if err := os.WriteFile(outputPath, []byte(content), 0o600); err != nil { + panic(err) + } +} diff --git a/cmd/svcinit/target_args/target_args_test.bzl b/cmd/svcinit/target_args/target_args_test.bzl new file mode 100644 index 0000000..c5ca60f --- /dev/null +++ b/cmd/svcinit/target_args/target_args_test.bzl @@ -0,0 +1,113 @@ +def _shell_quote(value): + return "'" + value.replace("'", "'\"'\"'") + "'" + +def _target_args_test_impl(ctx): + target_env = ctx.attr.target_under_test[RunEnvironmentInfo].environment + target_path = ctx.executable.target_under_test.short_path + + env_lines = [ + "export {}={}".format(key, _shell_quote(value)) + for key, value in sorted(target_env.items()) + ] + cleanup_lines = [ + "rm -f \"${{TEST_TMPDIR}}/{}\"".format(path) + for path in sorted(ctx.attr.expected_files.keys()) + ] + expected_paths = [ + "\"${{TEST_TMPDIR}}/{}\"".format(path) + for path in sorted(ctx.attr.expected_files.keys()) + ] + check_blocks = [] + for path, content in sorted(ctx.attr.expected_files.items()): + check_blocks.append("""expected_path="${{TEST_TMPDIR}}/{path}.expected" +actual_path="${{TEST_TMPDIR}}/{path}" +cat <<'EOF' > "${{expected_path}}" +{content} +EOF +if [[ ! -f "${{actual_path}}" ]]; then + echo "missing output file: ${{actual_path}}" >&2 + exit 1 +fi +diff -u "${{expected_path}}" "${{actual_path}}" +""".format(path = path, content = content)) + + 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 + +target_path="${{TEST_SRCDIR}}/${{TEST_WORKSPACE}}/{target_path}" + +{cleanup} +{env} +unset TEST_TARGET TEST_SIZE TEST_TIMEOUT XML_OUTPUT_FILE + +"${{target_path}}" {argv} >"${{TEST_TMPDIR}}/target.log" 2>&1 & +target_pid=$! +trap 'kill "${{target_pid}}" 2>/dev/null || true; wait "${{target_pid}}" 2>/dev/null || true' EXIT + +for _ in $(seq 1 100); do + missing=0 + for expected_file in {expected_paths}; do + if [[ ! -f "${{expected_file}}" ]]; then + missing=1 + break + fi + done + if [[ "${{missing}}" -eq 0 ]]; then + break + fi + if ! kill -0 "${{target_pid}}" 2>/dev/null; then + cat "${{TEST_TMPDIR}}/target.log" + echo "target exited before producing expected files" >&2 + exit 1 + fi + sleep 0.1 +done + +if [[ "${{missing}}" -ne 0 ]]; then + cat "${{TEST_TMPDIR}}/target.log" + echo "timed out waiting for expected files" >&2 + exit 1 +fi + +kill "${{target_pid}}" +wait "${{target_pid}}" || true +trap - EXIT + +{checks} +""".format( + target_path = target_path, + cleanup = "\n".join(cleanup_lines), + env = "\n".join(env_lines), + expected_paths = " ".join(expected_paths), + argv = " ".join([_shell_quote(arg) for arg in ctx.attr.argv]), + checks = "\n".join(check_blocks), + ), + ) + + runfiles = ctx.runfiles(files = [ctx.executable.target_under_test]) + runfiles = runfiles.merge(ctx.attr.target_under_test.default_runfiles) + + return [ + DefaultInfo( + executable = script, + runfiles = runfiles, + ), + ] + +target_args_test = rule( + implementation = _target_args_test_impl, + attrs = { + "argv": attr.string_list(), + "expected_files": attr.string_dict(), + "target_under_test": attr.label( + executable = True, + cfg = "target", + providers = [RunEnvironmentInfo], + ), + }, + test = True, +) diff --git a/cmd/svcinit/target_args_test.go b/cmd/svcinit/target_args_test.go new file mode 100644 index 0000000..de10734 --- /dev/null +++ b/cmd/svcinit/target_args_test.go @@ -0,0 +1,98 @@ +package main + +import ( + "reflect" + "strings" + "testing" + + "rules_itest/svclib" +) + +func TestParseDelegatedTargetArgs(t *testing.T) { + serviceSpecs := map[string]svclib.ServiceSpec{ + "@@//cmd/svcinit/target_args:dep_task": {Type: "task"}, + "@@//cmd/svcinit/target_args:top_task": {Type: "task"}, + } + aliases := map[string][]string{ + "@@//cmd/svcinit/target_args:dep_alias": {"@@//cmd/svcinit/target_args:dep_task"}, + } + + got, err := parseDelegatedTargetArgs([]string{ + "--target_arg", "top_task", + "--top-flag", + "--top-pair", "with value", + "--target_arg", "dep_alias", + "--dep-flag", + }, serviceSpecs, aliases) + if err != nil { + t.Fatalf("parseDelegatedTargetArgs() error = %v", err) + } + + want := map[string][]string{ + "@@//cmd/svcinit/target_args:top_task": {"--top-flag", "--top-pair", "with value"}, + "@@//cmd/svcinit/target_args:dep_task": {"--dep-flag"}, + } + if !reflect.DeepEqual(got, want) { + t.Fatalf("parseDelegatedTargetArgs() mismatch (-want +got):\nwant: %#v\ngot: %#v", want, got) + } +} + +func TestParseDelegatedTargetConfigIncludesEnv(t *testing.T) { + serviceSpecs := map[string]svclib.ServiceSpec{ + "@@//cmd/svcinit/target_args:dep_task": {Type: "task"}, + "@@//cmd/svcinit/target_args:top_task": {Type: "task"}, + } + aliases := map[string][]string{ + "@@//cmd/svcinit/target_args:dep_alias": {"@@//cmd/svcinit/target_args:dep_task"}, + } + + got, err := parseDelegatedTargetConfig([]string{ + "--target_arg", "//cmd/svcinit/target_args:top_task", + "--top-flag", + "--target_env", "//cmd/svcinit/target_args:top_task", "TOP_ENV=hello world", + "--target_env", "dep_alias", "DEP_ENV=1", + }, serviceSpecs, aliases) + if err != nil { + t.Fatalf("parseDelegatedTargetConfig() error = %v", err) + } + + want := delegatedTargetConfig{ + Args: map[string][]string{ + "@@//cmd/svcinit/target_args:top_task": {"--top-flag"}, + }, + Env: map[string]map[string]string{ + "@@//cmd/svcinit/target_args:top_task": {"TOP_ENV": "hello world"}, + "@@//cmd/svcinit/target_args:dep_task": {"DEP_ENV": "1"}, + }, + } + if !reflect.DeepEqual(got, want) { + t.Fatalf("parseDelegatedTargetConfig() mismatch (-want +got):\nwant: %#v\ngot: %#v", want, got) + } +} + +func TestParseDelegatedEnvAssignment(t *testing.T) { + key, value, err := parseDelegatedEnvAssignment("FOO=") + if err != nil { + t.Fatalf("parseDelegatedEnvAssignment() error = %v", err) + } + if key != "FOO" || value != "" { + t.Fatalf("parseDelegatedEnvAssignment() = (%q, %q), want (%q, %q)", key, value, "FOO", "") + } +} + +func TestResolveDelegatedTargetRejectsGroupAlias(t *testing.T) { + serviceSpecs := map[string]svclib.ServiceSpec{ + "@@//cmd/svcinit/target_args:group": {Type: "group"}, + } + aliases := map[string][]string{ + "@@//cmd/svcinit/target_args:group_alias": {"@@//cmd/svcinit/target_args:group"}, + } + + _, err := resolveDelegatedTarget("group_alias", serviceSpecs, aliases) + if err == nil { + t.Fatal("resolveDelegatedTarget() error = nil, want group rejection") + } + if !strings.Contains(err.Error(), "non-executable itest_service_group") { + t.Fatalf("resolveDelegatedTarget() error = %q, want group rejection", err) + } +} diff --git a/docs/itest.md b/docs/itest.md index 28f1b59..cac7cc7 100644 --- a/docs/itest.md +++ b/docs/itest.md @@ -19,6 +19,22 @@ query:enable-reload --@rules_itest//:enable_per_service_reload In addition, if the `hot_reloadable` attribute is set on an `itest_service`, the service manager will forward the ibazel hot-reload notification over stdin instead of restarting the service. +When using `bazel run` on an executable `itest_service`, `itest_task`, or `itest_service_group`, you can +delegate additional arguments and environment variables to executable `itest_*` targets in the dependency graph: + +```text +bazel run //path/to:target -- \ + --target_env //path/to:some_task EXTRA_ENV=value \ + --target_arg //path/to:some_task --flag value --other-flag "two words" \ + --target_env //path/to:service_alias EXTRA_ENV=value "OTHER_ENV=value with space" \ + --target_arg //path/to:service_alias --other-flag +``` + +Each `--target_arg` section appends its remaining arguments to that target's configured `args` until the next +`--target_arg`. `--target_env KEY=VALUE` injects or overrides an environment variable for that executable +target. Targets are typically passed as their full label. Package-relative labels, target names, and Bazel alias +that appears in the graph are also accepted. + # Reusable port reservations For each service with `so_reuseport_aware = True`, the service manager adds @@ -186,4 +202,3 @@ All [common binary attributes](https://bazel.build/reference/be/common-definitio | port_aliases | Port aliases allow you to 're-export' another service's port as belonging to this service group. This can be used to create abstractions (such as an itest_service combined with an itest_task) but not leak their implementation through how client code accesses port names. | Dictionary: String -> String | optional | `{}` | | services | Services/tasks that comprise this group. Can be `itest_service`, `itest_task`, or `itest_service_group`. | List of labels | optional | `[]` | | test | The underlying test target to execute once the services have been brought up and healthchecked. | Label | optional | `None` | - diff --git a/examples/target_args/BUILD.bazel b/examples/target_args/BUILD.bazel new file mode 100644 index 0000000..20869f0 --- /dev/null +++ b/examples/target_args/BUILD.bazel @@ -0,0 +1,40 @@ +load("@rules_go//go:def.bzl", "go_binary") +load("@rules_itest//:itest.bzl", "itest_service_group", "itest_task") + +go_binary( + name = "print_args", + srcs = ["main.go"], +) + +itest_task( + name = "migrate", + args = ["--base-flag=migrate-default"], + env = {"TARGET_ARGS_TASK_NAME": "//target_args:migrate"}, + exe = ":print_args", +) + +alias( + name = "migrate_alias", + actual = ":migrate", +) + +itest_task( + name = "seed", + args = ["--base-flag=seed-default"], + env = {"TARGET_ARGS_TASK_NAME": "//target_args:seed"}, + exe = ":print_args", +) + +alias( + name = "seed_alias", + actual = ":seed", +) + +itest_service_group( + name = "workflow", + hygienic = False, + services = [ + ":migrate_alias", + ":seed_alias", + ], +) diff --git a/examples/target_args/README.md b/examples/target_args/README.md new file mode 100644 index 0000000..6ac72c7 --- /dev/null +++ b/examples/target_args/README.md @@ -0,0 +1,26 @@ +# Delegating `bazel run` arguments to itest targets + +This example shows how `bazel run` can append extra arguments and inject extra environment variables into specific `itest_*` targets in the graph, including Bazel aliases. + +Run it from the `examples/` workspace: + +```text +bazel run //target_args:workflow -- \ + --target_env //target_args:migrate_alias TARGET_ARGS_INJECTED_ENV=migrate-env \ + --target_arg //target_args:migrate_alias --migration=users --dry-run \ + --target_env //target_args:seed_alias "TARGET_ARGS_INJECTED_ENV=seed env" \ + --target_arg //target_args:seed_alias "--users=alice bob" +``` + +Expected output includes the `TARGET_ARGS_` environment variables for each task, followed by its delegated args. For example: + +```text +@@//target_args:migrate> TARGET_ARGS_INJECTED_ENV=migrate-env +@@//target_args:migrate> TARGET_ARGS_TASK_NAME=//target_args:migrate +@@//target_args:migrate> ARGS=["--base-flag=migrate-default" "--migration=users" "--dry-run"] +@@//target_args:seed> TARGET_ARGS_INJECTED_ENV=seed env +@@//target_args:seed> TARGET_ARGS_TASK_NAME=//target_args:seed +@@//target_args:seed> ARGS=["--base-flag=seed-default" "--users=alice bob"] +``` + +`bazel run` mode keeps the service manager alive after the tasks finish, so press `Ctrl-C` once the output is printed. diff --git a/examples/target_args/main.go b/examples/target_args/main.go new file mode 100644 index 0000000..d86e759 --- /dev/null +++ b/examples/target_args/main.go @@ -0,0 +1,21 @@ +package main + +import ( + "fmt" + "os" + "sort" + "strings" +) + +func main() { + env := os.Environ() + sort.Strings(env) + + for _, entry := range env { + if !strings.HasPrefix(entry, "TARGET_ARGS_") { + continue + } + fmt.Println(entry) + } + fmt.Printf("ARGS=%q\n", os.Args[1:]) +} diff --git a/itest.bzl b/itest.bzl index b24ab3d..a78c646 100644 --- a/itest.bzl +++ b/itest.bzl @@ -43,6 +43,12 @@ def named_port_alias(label, name): """ return _to_relative_named_port(label, name) +def _normalize_labels(labels): + return [ + str(native.package_relative_label(label)) + for label in labels + ] + def itest_service(name, tags = [], hygienic = True, named_ports = [], **kwargs): if "port" in kwargs: fail("Do not specify `port`, instead set it via the `%s` flag" % (name + ".port")) @@ -64,6 +70,7 @@ def itest_service(name, tags = [], hygienic = True, named_ports = [], **kwargs): _itest_service( name = name, + dep_labels_internal = _normalize_labels(kwargs.get("deps", [])), tags = tags + ["ibazel_notify_changes"], port = name + ".port", named_ports = named_ports_attr, @@ -78,6 +85,7 @@ def itest_service(name, tags = [], hygienic = True, named_ports = [], **kwargs): def itest_service_group(name, tags = [], hygienic = True, **kwargs): _itest_service_group( + service_labels_internal = _normalize_labels(kwargs.get("services", [])), name = name, tags = tags + ["ibazel_notify_changes"], **kwargs @@ -91,6 +99,7 @@ def itest_service_group(name, tags = [], hygienic = True, **kwargs): def itest_task(name, tags = [], hygienic = True, **kwargs): _itest_task( + dep_labels_internal = _normalize_labels(kwargs.get("deps", [])), name = name, tags = tags + ["ibazel_notify_changes"], **kwargs @@ -110,4 +119,8 @@ def _hygiene_test(name, **kwargs): **kwargs ) -service_test = _service_test +def service_test(**kwargs): + _service_test( + service_labels_internal = _normalize_labels(kwargs.get("services", [])), + **kwargs + ) diff --git a/private/itest.bzl b/private/itest.bzl index 921aded..834b80f 100644 --- a/private/itest.bzl +++ b/private/itest.bzl @@ -18,6 +18,21 @@ query:enable-reload --@rules_itest//:enable_per_service_reload In addition, if the `hot_reloadable` attribute is set on an `itest_service`, the service manager will forward the ibazel hot-reload notification over stdin instead of restarting the service. +When using `bazel run` on an executable `itest_service`, `itest_task`, or `itest_service_group`, you can +delegate additional arguments and environment variables to executable `itest_*` targets in the dependency graph: + +``` +bazel run //path/to:target -- \ + --target_env //path/to:some_task EXTRA_ENV=value \ + --target_arg //path/to:some_task --flag value --other-flag "two words" \ + --target_arg //path/to:service_alias --other-flag +``` + +Each `--target_arg` section appends its remaining arguments to that target's configured `args` until the next +`--target_arg`. `--target_env KEY=VALUE` injects or overrides an environment variable for that executable +target. Targets are typically passed as their full label. Package-relative labels, target names, and Bazel alias +that appears in the graph are also accepted. + # Reusable port reservations For each service with `so_reuseport_aware = True`, the service manager adds @@ -49,15 +64,27 @@ _ServiceGroupInfo = provider( doc = "Info about a service group", fields = { "deferred": "Flag if this service/task/group should be deferred or not", + "aliases": "Dict of alias labels to underlying target labels", + "roots": "Labels represented by this target", "services": "Dict of services/tasks", }, ) -def _collect_services(deps): +def _collect_graph(deps, dep_labels): + if len(deps) != len(dep_labels): + fail("Internal error: dep label metadata length mismatch for %s" % dep_labels) + services = {} - for dep in deps: - services |= dep[_ServiceGroupInfo].services - return services + aliases = {} + for dep, dep_label in zip(deps, dep_labels): + info = dep[_ServiceGroupInfo] + services |= info.services + aliases |= info.aliases + + if dep_label not in info.roots: + aliases[dep_label] = info.roots + + return services, aliases def _run_environment(ctx, service_specs_file): return { @@ -117,6 +144,7 @@ _itest_binary_attrs = { providers = [_ServiceGroupInfo], doc = "Services/tasks that must be started before this service/task can be started. Can be `itest_service`, `itest_task`, or `itest_service_group`.", ), + "dep_labels_internal": attr.string_list(doc = "Internal"), } | _svcinit_attrs def _compute_env(ctx, underlying_target): @@ -169,10 +197,10 @@ def _itest_binary_impl(ctx, extra_service_spec_kwargs, extra_exe_runfiles = []): **extra_service_spec_kwargs ) - services = _collect_services(ctx.attr.deps) + services, aliases = _collect_graph(ctx.attr.deps, ctx.attr.dep_labels_internal) services[service.label] = service - service_specs_file = _create_svcinit_actions(ctx, services) + service_specs_file = _create_svcinit_actions(ctx, services, aliases) direct_runfiles = ctx.files.data + [service_specs_file] if version_file: @@ -184,7 +212,12 @@ def _itest_binary_impl(ctx, extra_service_spec_kwargs, extra_exe_runfiles = []): return [ RunEnvironmentInfo(environment = _run_environment(ctx, service_specs_file)), DefaultInfo(runfiles = runfiles), - _ServiceGroupInfo(services = services, deferred = ctx.attr.deferred), + _ServiceGroupInfo( + services = services, + aliases = aliases, + roots = [service.label], + deferred = ctx.attr.deferred, + ), ] def _validate_duration(name, s): @@ -355,7 +388,7 @@ All [common binary attributes](https://bazel.build/reference/be/common-definitio def _itest_service_group_impl(ctx): _validate_deferred(ctx, ctx.attr.services) - services = _collect_services(ctx.attr.services) + services, aliases = _collect_graph(ctx.attr.services, ctx.attr.service_labels_internal) service = struct( type = "group", @@ -365,7 +398,7 @@ def _itest_service_group_impl(ctx): ) services[service.label] = service - service_specs_file = _create_svcinit_actions(ctx, services) + service_specs_file = _create_svcinit_actions(ctx, services, aliases) runfiles = ctx.runfiles([service_specs_file]) runfiles = runfiles.merge_all(_services_runfiles(ctx)) @@ -373,7 +406,12 @@ def _itest_service_group_impl(ctx): return [ RunEnvironmentInfo(environment = _run_environment(ctx, service_specs_file)), DefaultInfo(runfiles = runfiles), - _ServiceGroupInfo(services = services, deferred = ctx.attr.deferred), + _ServiceGroupInfo( + services = services, + aliases = aliases, + roots = [service.label], + deferred = ctx.attr.deferred, + ), ] _itest_service_group_attrs = _svcinit_attrs | { @@ -391,6 +429,7 @@ Use the functions `port_alias` and `named_port_alias` to reference ports from th providers = [_ServiceGroupInfo], doc = "Services/tasks that comprise this group. Can be `itest_service`, `itest_task`, or `itest_service_group`.", ), + "service_labels_internal": attr.string_list(doc = "Internal"), } itest_service_group = rule( @@ -405,7 +444,7 @@ forcing the services within the group to define a specific startup ordering with It can bring up multiple services with a single `bazel run` command, which is useful for creating dev environments.""", ) -def _create_svcinit_actions(ctx, services): +def _create_svcinit_actions(ctx, services, aliases): ctx.actions.symlink( output = ctx.outputs.executable, target_file = ctx.executable._svcinit, @@ -414,7 +453,10 @@ def _create_svcinit_actions(ctx, services): # Avoid expanding during analysis phase. service_content = ctx.actions.args() service_content.set_param_file_format("multiline") - service_content.add_all([services], map_each = json.encode) + service_content.add_all([{ + "services": services, + "aliases": aliases, + }], map_each = json.encode) service_specs_file = ctx.actions.declare_file(ctx.label.name + ".service_specs.json") ctx.actions.write( @@ -425,9 +467,11 @@ def _create_svcinit_actions(ctx, services): return service_specs_file def _service_test_impl(ctx): + services, aliases = _collect_graph(ctx.attr.services, ctx.attr.service_labels_internal) service_specs_file = _create_svcinit_actions( ctx, - _collect_services(ctx.attr.services), + services, + aliases, ) env_file = ctx.actions.declare_file(ctx.label.name + ".env.json") @@ -460,6 +504,7 @@ _service_test_attrs = { doc = "The service manager will merge these variables into the environment when spawning the underlying binary.", ), "data": attr.label_list(allow_files = True), + "service_labels_internal": attr.string_list(doc = "Internal"), ## This is taken directly from rules_go: https://github.com/bazel-contrib/rules_go/blob/85eef05357c9421eaa568d101e62355384bc49bb/go/private/rules/test.bzl#L442-L457 # Required for Bazel to merge coverage reports for Go and other # languages into a single report per test. diff --git a/runner/runner.go b/runner/runner.go index daa85f9..e96ca82 100644 --- a/runner/runner.go +++ b/runner/runner.go @@ -8,6 +8,8 @@ import ( "os/exec" "reflect" "runtime" + "strconv" + "strings" "sync" "syscall" "time" @@ -51,6 +53,18 @@ func colorize(s svclib.VersionedServiceSpec) string { return s.Colorize(s.Label) } +func formatCommandArgs(args []string) string { + if len(args) == 0 { + return "[]" + } + + quoted := make([]string, len(args)) + for i, arg := range args { + quoted[i] = strconv.Quote(arg) + } + return "[" + strings.Join(quoted, " ") + "]" +} + func (r *Runner) StartAll(serviceErrCh chan error) ([]topological.Task, error) { tasks := allTasks(r.serviceInstances, func(ctx context.Context, service *ServiceInstance) error { if service.Type == "group" { @@ -65,7 +79,7 @@ func (r *Runner) StartAll(serviceErrCh chan error) ([]topological.Task, error) { if terseOutput { log.Printf("Starting %s\n", colorize(service.VersionedServiceSpec)) } else { - log.Printf("Starting %s %v\n", colorize(service.VersionedServiceSpec), service.cmd.Args[1:]) + log.Printf("Starting %s %s\n", colorize(service.VersionedServiceSpec), formatCommandArgs(service.cmd.Args[1:])) } startErr := service.Start(ctx) diff --git a/runner/service_instance.go b/runner/service_instance.go index e645d96..405a404 100644 --- a/runner/service_instance.go +++ b/runner/service_instance.go @@ -9,7 +9,6 @@ import ( "net/http" "os" "os/exec" - "strings" "sync" "syscall" "time" @@ -160,7 +159,7 @@ func (s *ServiceInstance) HealthCheck(ctx context.Context, expectedStartDuration if terseOutput { log.Printf("CMD Healthchecking %s\n", coloredLabel) } else { - log.Printf("CMD Healthchecking %s (pid %d) : %s %v\n", coloredLabel, s.Pid(), s.Colorize(s.HealthCheckLabel), strings.Join(s.HealthCheckArgs, " ")) + log.Printf("CMD Healthchecking %s (pid %d) : %s %s\n", coloredLabel, s.Pid(), s.Colorize(s.HealthCheckLabel), formatCommandArgs(s.HealthCheckArgs)) } }