From 80cae1b039c4ebf0cb704ba81c7fc5b09aca4782 Mon Sep 17 00:00:00 2001 From: jeffluoo Date: Mon, 9 Nov 2026 19:10:16 +0000 Subject: [PATCH] cli: add --trace to trace a command end to end ax-server cannot join a slow or failed command to the client that issued it, because the CLI sends no trace context. With --trace, a command now starts a sampled root span, prints its trace ID to stderr, and sends a W3C traceparent on every AX RPC, so ax-server and Substrate, which sample ParentBased, record the whole request. - Without --trace the CLI creates no tracer and sends no trace context, so ax-server keeps the sampling decision for CLI traffic. - Each RPC is a client span under the command span, and resolving the ax-server address (which may start a port-forward) is a child span. Only traceparent is propagated, never baggage. - CLI spans are exported over OTLP only when an OTLP endpoint is set. OTEL_TRACES_EXPORTER follows ax-server's rules, and the resource reads OTEL_SERVICE_NAME and OTEL_RESOURCE_ATTRIBUTES. Export waits at most 2 s and warns on failure; OpenTelemetry errors print as warnings. - ax ssh returns the remote exit code instead of exiting the process, so the command span ends and its port-forward is cleaned up. --- README.md | 1 + cmd/ax/main.go | 25 ++- cmd/ax/main_test.go | 10 ++ cmd/ax/trace.go | 264 +++++++++++++++++++++++++++++ cmd/ax/trace_test.go | 390 +++++++++++++++++++++++++++++++++++++++++++ go.mod | 17 +- go.sum | 18 +- 7 files changed, 712 insertions(+), 13 deletions(-) create mode 100644 cmd/ax/trace.go create mode 100644 cmd/ax/trace_test.go diff --git a/README.md b/README.md index bd6ac673..7d6ed7a0 100644 --- a/README.md +++ b/README.md @@ -213,6 +213,7 @@ ax --context=dev-cluster get tasks | `-n`, `--namespace` | Kubernetes namespace where AX is installed | `ax-system` | | `--context` | Kubernetes context to target | active `kubectx` / `current-context` | | `--server` | Control plane address, bypassing auto-detection | derived from kube context, or `$AX_SERVER` | +| `--trace` | Trace the command end to end and print its trace ID to stderr. Its RPCs send a sampled W3C `traceparent`, so ax-server records them; CLI spans are exported over OTLP when `OTEL_EXPORTER_OTLP_ENDPOINT` is set. Without it, the CLI sends no trace context and ax-server samples on its own | off | ## License diff --git a/cmd/ax/main.go b/cmd/ax/main.go index d07cc4ac..ecb0f89e 100644 --- a/cmd/ax/main.go +++ b/cmd/ax/main.go @@ -94,15 +94,18 @@ func main() { } } + activeTrace = startCommandTrace(cmd, opts.trace, os.Stderr) + // Resolve the AX server URL (auto-tunneling to active kube context if not explicitly set) + endResolve := activeTrace.startStep("resolve ax-server address") serverURL, err := tunnel.EnsureServerURL(tunnel.Options{ ServerURL: explicitServer, Context: kubeContext, Namespace: axNamespace, }) + endResolve(err) if err != nil { - fmt.Fprintf(os.Stderr, "Error: %v\n", err) - os.Exit(1) + os.Exit(finish(activeTrace, err, os.Stderr)) } switch cmd { @@ -125,14 +128,14 @@ func main() { case "ssh": err = runSSH(serverURL, atespace, kubeContext, axNamespace, cleanArgs) default: + activeTrace.end(fmt.Errorf("unknown command %q", cmd)) fmt.Fprintf(os.Stderr, "unknown command: %s\n", cmd) printUsage() os.Exit(1) } - if err != nil { - fmt.Fprintf(os.Stderr, "Error: %v\n", err) - os.Exit(1) + if code := finish(activeTrace, err, os.Stderr); code != 0 { + os.Exit(code) } } @@ -146,6 +149,7 @@ type globalArgs struct { server string kubeContext string namespace string + trace bool } // parseGlobalArgs extracts the command name and the global flags from args. It @@ -192,6 +196,8 @@ func parseGlobalArgs(args []string) globalArgs { } } else if strings.HasPrefix(arg, "--namespace=") { g.namespace = strings.TrimPrefix(arg, "--namespace=") + } else if arg == "--trace" { + g.trace = true } else if g.cmd == "" && !strings.HasPrefix(arg, "-") { g.cmd = arg } else { @@ -233,13 +239,15 @@ Flags: -a, --atespace string Task atespace scope (default: "default") -n, --namespace string Kubernetes namespace where AX is installed (default: "ax-system") --context string Kubernetes context (defaults to active kubectx / current-context) - --server string AX API server address (default: auto-detected from kube context or $AX_SERVER)`) + --server string AX API server address (default: auto-detected from kube context or $AX_SERVER) + --trace Trace this command end to end and print its trace ID to stderr`) } func getAXClient(serverURL string) (v1alpha1.AXClient, *grpc.ClientConn, error) { target := strings.TrimPrefix(serverURL, "http://") target = strings.TrimPrefix(target, "https://") - conn, err := grpc.NewClient(target, grpc.WithTransportCredentials(insecure.NewCredentials())) + opts := append([]grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())}, activeTrace.dialOptions()...) + conn, err := grpc.NewClient(target, opts...) if err != nil { return nil, nil, fmt.Errorf("connecting to %s: %w", target, err) } @@ -1090,6 +1098,7 @@ func runSSH(serverURL, atespace, kubeContext, axNamespace string, args []string) defer signal.Stop(rawSigs) sigs := escalateInterrupts(rawSigs, interruptExitWindow, func() { fmt.Fprintln(os.Stderr, "\nax ssh: interrupted twice, exiting") + activeTrace.end(errors.New("ax ssh: interrupted twice")) os.Exit(130) }) @@ -1105,7 +1114,7 @@ func runSSH(serverURL, atespace, kubeContext, axNamespace string, args []string) } if exitCode != 0 { - os.Exit(exitCode) + return exitCodeError{code: exitCode} } return nil diff --git a/cmd/ax/main_test.go b/cmd/ax/main_test.go index 8641ae22..e63ee69b 100644 --- a/cmd/ax/main_test.go +++ b/cmd/ax/main_test.go @@ -439,6 +439,16 @@ func TestParseGlobalArgs(t *testing.T) { namespace: "ax-system", }, }, + { + name: "trace flag", + args: []string{"--trace", "get", "tasks"}, + want: globalArgs{cmd: "get", args: []string{"tasks"}, atespace: "default", namespace: "ax-system", trace: true}, + }, + { + name: "trace after -- belongs to the remote command", + args: []string{"ssh", "task123", "--", "app", "--trace"}, + want: globalArgs{cmd: "ssh", args: []string{"task123", "--", "app", "--trace"}, atespace: "default", namespace: "ax-system"}, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { diff --git a/cmd/ax/trace.go b/cmd/ax/trace.go new file mode 100644 index 00000000..abbc81d3 --- /dev/null +++ b/cmd/ax/trace.go @@ -0,0 +1,264 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "context" + "errors" + "fmt" + "io" + "os" + "slices" + "strings" + "time" + + "go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/codes" + "go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc" + metricnoop "go.opentelemetry.io/otel/metric/noop" + "go.opentelemetry.io/otel/propagation" + "go.opentelemetry.io/otel/sdk/resource" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + semconv "go.opentelemetry.io/otel/semconv/v1.43.0" + "go.opentelemetry.io/otel/trace" + "google.golang.org/grpc" +) + +const ( + // traceShutdownTimeout bounds how long the CLI waits to export its spans. + traceShutdownTimeout = 2 * time.Second + + tracesExporterEnv = "OTEL_TRACES_EXPORTER" + exporterOTLP = "otlp" + exporterNone = "none" +) + +// commandTrace is the root span of one `ax --trace` command. Every RPC the +// command makes joins its trace and sends a sampled W3C traceparent, so +// ax-server and Substrate, which sample ParentBased, record the whole request. +type commandTrace struct { + tp *sdktrace.TracerProvider + span trace.Span + stderr io.Writer +} + +// activeTrace is the trace of the running command. It is nil without --trace: +// the CLI then creates no tracer, sends no trace context, and leaves the +// sampling decision to ax-server. +var activeTrace *commandTrace + +// startCommandTrace starts the root span for cmd when force (--trace) is set, +// prints its trace ID to stderr, and reports OpenTelemetry errors there as +// warnings. Spans are exported over OTLP only when an OTLP endpoint is set in +// the environment. It returns nil without force. +func startCommandTrace(cmd string, force bool, stderr io.Writer) *commandTrace { + if !force { + return nil + } + otel.SetErrorHandler(otel.ErrorHandlerFunc(func(err error) { + fmt.Fprintf(stderr, "Warning: OpenTelemetry: %v\n", err) + })) + t := newCommandTrace(cmd, envSpanExporter(stderr), stderr) + fmt.Fprintf(stderr, "Trace ID: %s\n", t.span.SpanContext().TraceID()) + return t +} + +// newCommandTrace starts a sampled root span named "ax ". A nil exporter +// keeps spans local; the trace context still reaches ax-server. +func newCommandTrace(cmd string, exporter sdktrace.SpanExporter, stderr io.Writer) *commandTrace { + opts := []sdktrace.TracerProviderOption{ + sdktrace.WithSampler(sdktrace.AlwaysSample()), + sdktrace.WithResource(commandResource(stderr)), + } + if exporter != nil { + opts = append(opts, sdktrace.WithBatcher(exporter)) + } + tp := sdktrace.NewTracerProvider(opts...) + _, span := tp.Tracer("github.com/google/ax/cmd/ax").Start(context.Background(), "ax "+cmd) + return &commandTrace{tp: tp, span: span, stderr: stderr} +} + +// commandResource describes the CLI the way ax-server describes itself: +// service.name "ax", the SDK attributes, and OTEL_SERVICE_NAME and +// OTEL_RESOURCE_ATTRIBUTES on top. +func commandResource(stderr io.Writer) *resource.Resource { + res, err := resource.New(context.Background(), + resource.WithTelemetrySDK(), + resource.WithSchemaURL(semconv.SchemaURL), + resource.WithAttributes(semconv.ServiceName("ax")), + resource.WithFromEnv(), + ) + if err != nil { + fmt.Fprintf(stderr, "Warning: OpenTelemetry resource: %v\n", err) + } + if res == nil { + res = resource.Empty() + } + return res +} + +// envSpanExporter returns an OTLP gRPC exporter configured from the standard +// OTEL_EXPORTER_OTLP_* variables, or nil when no endpoint is set or +// OTEL_TRACES_EXPORTER turns OTLP off. +func envSpanExporter(stderr io.Writer) sdktrace.SpanExporter { + value := os.Getenv(tracesExporterEnv) + names, err := resolveExporters(tracesExporterEnv, value, []string{exporterOTLP}, []string{exporterOTLP}) + if err != nil { + fmt.Fprintf(stderr, "Warning: %v; using %q\n", err, strings.Join(names, ",")) + } + if !slices.Contains(names, exporterOTLP) { + return nil + } + if os.Getenv("OTEL_EXPORTER_OTLP_ENDPOINT") == "" && os.Getenv("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT") == "" { + return nil + } + exporter, err := otlptracegrpc.New(context.Background()) + if err != nil { + fmt.Fprintf(stderr, "Warning: not exporting the command trace: %v\n", err) + return nil + } + return exporter +} + +// resolveExporters reads value as a comma-separated exporter list, with the +// rules ax-server uses for the same variables. Names are trimmed and +// lowercased; empty and unknown names are skipped with an error, and the rest +// apply. A value that is empty, names no known exporter, or lists none with an +// exporter keeps def. The returned names are the ones to use even when err is +// not nil. +func resolveExporters(env, value string, known, def []string) ([]string, error) { + if strings.TrimSpace(value) == "" { + return slices.Clone(def), nil + } + var got, skipped []string + none := false + for item := range strings.SplitSeq(value, ",") { + switch name := strings.ToLower(strings.TrimSpace(item)); { + case name == exporterNone: + none = true + case slices.Contains(known, name): + if !slices.Contains(got, name) { + got = append(got, name) + } + default: + skipped = append(skipped, name) + } + } + var skipErr error + if len(skipped) > 0 { + skipErr = fmt.Errorf("%s: skipped empty or unknown exporters %q", env, skipped) + } + switch { + case none && len(got) > 0: + return slices.Clone(def), fmt.Errorf("%s %q lists none with an exporter", env, value) + case none: + return nil, skipErr + case len(got) == 0: + return slices.Clone(def), fmt.Errorf("%s %q names no known exporter", env, value) + } + return got, skipErr +} + +// dialOptions instruments a client connection: each RPC becomes a client span +// under the command span and carries only the traceparent header (no baggage). +func (t *commandTrace) dialOptions() []grpc.DialOption { + if t == nil { + return nil + } + return []grpc.DialOption{ + grpc.WithStatsHandler(otelgrpc.NewClientHandler( + otelgrpc.WithTracerProvider(t.tp), + otelgrpc.WithMeterProvider(metricnoop.NewMeterProvider()), + otelgrpc.WithPropagators(propagation.TraceContext{}), + )), + grpc.WithChainUnaryInterceptor(func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error { + return invoker(t.parent(ctx), method, req, reply, cc, opts...) + }), + grpc.WithChainStreamInterceptor(func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) { + return streamer(t.parent(ctx), desc, cc, method, opts...) + }), + } +} + +// parent puts the command span in ctx unless ctx already carries a span. The +// commands build RPC contexts from context.Background, so without this each +// RPC would start a trace of its own. +func (t *commandTrace) parent(ctx context.Context) context.Context { + if trace.SpanContextFromContext(ctx).IsValid() { + return ctx + } + return trace.ContextWithSpan(ctx, t.span) +} + +// startStep starts a child span of the command for one local step, such as +// resolving the ax-server address, and returns the function that ends it. +func (t *commandTrace) startStep(name string) func(error) { + if t == nil { + return func(error) {} + } + _, span := t.tp.Tracer("github.com/google/ax/cmd/ax").Start(trace.ContextWithSpan(context.Background(), t.span), name) + return func(err error) { + recordError(span, err) + span.End() + } +} + +// end records err on the command span, ends it, and flushes any exporter. +func (t *commandTrace) end(err error) { + if t == nil { + return + } + recordError(t.span, err) + t.span.End() + ctx, cancel := context.WithTimeout(context.Background(), traceShutdownTimeout) + defer cancel() + if err := t.tp.Shutdown(ctx); err != nil { + fmt.Fprintf(t.stderr, "Warning: could not export the command trace: %v\n", err) + } +} + +func recordError(span trace.Span, err error) { + if err == nil { + return + } + span.RecordError(err) + span.SetStatus(codes.Error, err.Error()) +} + +// exitCodeError makes the CLI exit with code without printing an error, as +// when the command run by `ax ssh` exits non-zero. +type exitCodeError struct { + code int +} + +func (e exitCodeError) Error() string { + return fmt.Sprintf("remote command exited with code %d", e.code) +} + +// finish ends the command trace, reports err, and returns the exit code. The +// trace ends first, so a command that fails still exports its spans. +func finish(t *commandTrace, err error, stderr io.Writer) int { + t.end(err) + if err == nil { + return 0 + } + var exitErr exitCodeError + if errors.As(err, &exitErr) { + return exitErr.code + } + fmt.Fprintf(stderr, "Error: %v\n", err) + return 1 +} diff --git a/cmd/ax/trace_test.go b/cmd/ax/trace_test.go new file mode 100644 index 00000000..013a3c88 --- /dev/null +++ b/cmd/ax/trace_test.go @@ -0,0 +1,390 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "bytes" + "context" + "errors" + "io" + "log" + "net" + "os" + "slices" + "strings" + "testing" + + "github.com/google/ax/pkg/apis/v1alpha1" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/codes" + sdktrace "go.opentelemetry.io/otel/sdk/trace" + "go.opentelemetry.io/otel/sdk/trace/tracetest" + "go.opentelemetry.io/otel/trace" + "google.golang.org/grpc" + "google.golang.org/grpc/metadata" + "google.golang.org/protobuf/types/known/emptypb" +) + +// traceCaptureServer records the request metadata of every RPC it serves. +type traceCaptureServer struct { + v1alpha1.UnimplementedAXServer + headers chan metadata.MD +} + +func (s *traceCaptureServer) ListTasks(ctx context.Context, _ *v1alpha1.ListTasksRequest) (*v1alpha1.ListTasksResponse, error) { + md, _ := metadata.FromIncomingContext(ctx) + s.headers <- md + return &v1alpha1.ListTasksResponse{}, nil +} + +// startTraceCaptureServer serves the AX API, and records the metadata of +// streams to unregistered methods too. +func startTraceCaptureServer(t *testing.T) (string, <-chan metadata.MD) { + t.Helper() + headers := make(chan metadata.MD, 4) + srv := grpc.NewServer(grpc.UnknownServiceHandler(func(_ any, stream grpc.ServerStream) error { + md, _ := metadata.FromIncomingContext(stream.Context()) + headers <- md + return nil + })) + v1alpha1.RegisterAXServer(srv, &traceCaptureServer{headers: headers}) + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + t.Cleanup(srv.Stop) + go func() { _ = srv.Serve(listener) }() + return listener.Addr().String(), headers +} + +// useTrace makes ct the active command trace for the rest of the test. +func useTrace(t *testing.T, ct *commandTrace) { + t.Helper() + activeTrace = ct + t.Cleanup(func() { activeTrace = nil }) +} + +// discardStdout drops what the command prints. +func discardStdout(t *testing.T) { + t.Helper() + devNull, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0) + if err != nil { + t.Fatal(err) + } + stdout := os.Stdout + os.Stdout = devNull + t.Cleanup(func() { + os.Stdout = stdout + _ = devNull.Close() + }) +} + +// clearOTelEnv keeps the developer's OTel settings out of the test, and +// restores the global error handler that startCommandTrace replaces. +func clearOTelEnv(t *testing.T) { + t.Helper() + for _, name := range []string{ + "OTEL_TRACES_EXPORTER", "OTEL_TRACES_SAMPLER", "OTEL_TRACES_SAMPLER_ARG", + "OTEL_EXPORTER_OTLP_ENDPOINT", "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", + "OTEL_SERVICE_NAME", "OTEL_RESOURCE_ATTRIBUTES", + } { + // t.Setenv records the value to restore; the SDK treats a set but + // empty OTEL_TRACES_SAMPLER as an unknown sampler, so unset it. + t.Setenv(name, "") + if err := os.Unsetenv(name); err != nil { + t.Fatal(err) + } + } + t.Cleanup(func() { + otel.SetErrorHandler(otel.ErrorHandlerFunc(func(err error) { log.Print(err) })) + }) +} + +// keptSpans is an in-memory exporter that keeps its spans after the tracer +// provider shuts down; tracetest.InMemoryExporter clears them on Shutdown. +type keptSpans struct{ *tracetest.InMemoryExporter } + +func (keptSpans) Shutdown(context.Context) error { return nil } + +func newKeptSpans() keptSpans { return keptSpans{tracetest.NewInMemoryExporter()} } + +// failingExporter fails to shut down, as an exporter does when its collector +// is unreachable. +type failingExporter struct{} + +func (failingExporter) ExportSpans(context.Context, []sdktrace.ReadOnlySpan) error { return nil } +func (failingExporter) Shutdown(context.Context) error { + return errors.New("context deadline exceeded") +} + +// traceparent returns the fields of the single traceparent header in md. +func traceparent(t *testing.T, md metadata.MD) (traceID, spanID, flags string) { + t.Helper() + values := md.Get("traceparent") + if len(values) != 1 { + t.Fatalf("expected one traceparent header, got %q", values) + } + parts := strings.Split(values[0], "-") + if len(parts) != 4 || parts[0] != "00" { + t.Fatalf("malformed traceparent %q", values[0]) + } + return parts[1], parts[2], parts[3] +} + +func TestTraceFlagSamplesCommandAndPrintsTraceID(t *testing.T) { + clearOTelEnv(t) + discardStdout(t) + addr, headers := startTraceCaptureServer(t) + + var stderr bytes.Buffer + ct := startCommandTrace("get", true, &stderr) + useTrace(t, ct) + defer ct.end(nil) + + if err := runGet(addr, "default", []string{"tasks"}); err != nil { + t.Fatal(err) + } + md := <-headers + traceID, _, flags := traceparent(t, md) + if want := ct.span.SpanContext().TraceID().String(); traceID != want { + t.Errorf("traceparent trace ID = %s, want the command trace %s", traceID, want) + } + if flags != "01" { + t.Errorf("traceparent flags = %s, want 01 (sampled) with --trace", flags) + } + if got := md.Get("baggage"); len(got) != 0 { + t.Errorf("expected no baggage header, got %q", got) + } + if want := "Trace ID: " + traceID + "\n"; stderr.String() != want { + t.Errorf("stderr = %q, want %q", stderr.String(), want) + } +} + +func TestNoTraceWithoutFlag(t *testing.T) { + clearOTelEnv(t) + var stderr bytes.Buffer + if ct := startCommandTrace("get", false, &stderr); ct != nil { + ct.end(nil) + t.Fatal("expected no command trace without --trace") + } + if stderr.Len() != 0 { + t.Errorf("expected no output without --trace, got %q", stderr.String()) + } +} + +func TestNoTraceparentWithoutCommandTrace(t *testing.T) { + discardStdout(t) + addr, headers := startTraceCaptureServer(t) + useTrace(t, nil) + + if err := runGet(addr, "default", []string{"tasks"}); err != nil { + t.Fatal(err) + } + if got := (<-headers).Get("traceparent"); len(got) != 0 { + t.Errorf("expected no traceparent without --trace, got %q", got) + } +} + +func TestCommandTraceExportsSpans(t *testing.T) { + clearOTelEnv(t) + discardStdout(t) + addr, headers := startTraceCaptureServer(t) + + exporter := newKeptSpans() + ct := newCommandTrace("get", exporter, io.Discard) + useTrace(t, ct) + + ct.startStep("resolve ax-server address")(nil) + if err := runGet(addr, "default", []string{"tasks"}); err != nil { + t.Fatal(err) + } + _, sentSpanID, _ := traceparent(t, <-headers) + ct.end(nil) + + byName := map[string]tracetest.SpanStub{} + var client *tracetest.SpanStub + spans := exporter.GetSpans() + for i, s := range spans { + byName[s.Name] = s + if s.SpanKind == trace.SpanKindClient && strings.HasSuffix(s.Name, "/ListTasks") { + client = &spans[i] + } + } + root, ok := byName["ax get"] + step, stepOK := byName["resolve ax-server address"] + if !ok || !stepOK || client == nil { + t.Fatalf("expected the command, resolve, and ListTasks client spans, got %d spans: %+v", len(spans), spans) + } + if root.Parent.IsValid() { + t.Errorf("command span has parent %v, want a root span", root.Parent) + } + for _, child := range []tracetest.SpanStub{step, *client} { + if child.Parent.SpanID() != root.SpanContext.SpanID() { + t.Errorf("%s parent = %s, want the command span %s", child.Name, child.Parent.SpanID(), root.SpanContext.SpanID()) + } + } + if got := client.SpanContext.SpanID().String(); got != sentSpanID { + t.Errorf("traceparent parent ID = %s, want the client span %s", sentSpanID, got) + } + attrs := root.Resource.Set() + if v, ok := attrs.Value("service.name"); !ok || v.AsString() != "ax" { + t.Errorf("service.name = %v, want ax", v) + } + if v, ok := attrs.Value("telemetry.sdk.language"); !ok || v.AsString() != "go" { + t.Errorf("telemetry.sdk.language = %v, want go", v) + } +} + +func TestCommandResourceReadsEnv(t *testing.T) { + clearOTelEnv(t) + t.Setenv("OTEL_SERVICE_NAME", "ax-dev") + t.Setenv("OTEL_RESOURCE_ATTRIBUTES", "deployment.environment.name=laptop") + + attrs := commandResource(io.Discard).Set() + if v, _ := attrs.Value("service.name"); v.AsString() != "ax-dev" { + t.Errorf("service.name = %q, want OTEL_SERVICE_NAME ax-dev", v.AsString()) + } + if v, _ := attrs.Value("deployment.environment.name"); v.AsString() != "laptop" { + t.Errorf("deployment.environment.name = %q, want laptop", v.AsString()) + } +} + +func TestCommandTraceRecordsError(t *testing.T) { + exporter := newKeptSpans() + newCommandTrace("delete", exporter, io.Discard).end(errors.New("task not found")) + + spans := exporter.GetSpans() + if len(spans) != 1 { + t.Fatalf("expected one span, got %d", len(spans)) + } + if spans[0].Status.Code != codes.Error || spans[0].Status.Description != "task not found" { + t.Errorf("status = %+v, want Error with the command error", spans[0].Status) + } +} + +func TestEndWarnsWhenExportFails(t *testing.T) { + var stderr bytes.Buffer + newCommandTrace("get", failingExporter{}, &stderr).end(nil) + + if want := "Warning: could not export the command trace: context deadline exceeded\n"; stderr.String() != want { + t.Errorf("stderr = %q, want %q", stderr.String(), want) + } +} + +func TestSDKErrorsPrintAsWarnings(t *testing.T) { + clearOTelEnv(t) + t.Setenv("OTEL_TRACES_SAMPLER", "bogus") + + var stderr bytes.Buffer + startCommandTrace("get", true, &stderr).end(nil) + + for line := range strings.Lines(stderr.String()) { + if !strings.HasPrefix(line, "Trace ID: ") && !strings.HasPrefix(line, "Warning: OpenTelemetry: ") { + t.Errorf("unexpected stderr line %q", line) + } + } + if !strings.Contains(stderr.String(), "Warning: OpenTelemetry: ") { + t.Errorf("expected the SDK's sampler error as a warning, got %q", stderr.String()) + } +} + +func TestStreamRPCJoinsCommandTrace(t *testing.T) { + addr, headers := startTraceCaptureServer(t) + + ct := newCommandTrace("logs", nil, io.Discard) + useTrace(t, ct) + defer ct.end(nil) + + _, conn, err := getAXClient(addr) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + stream, err := conn.NewStream(context.Background(), &grpc.StreamDesc{ServerStreams: true}, "/ax.v1alpha1.AX/StreamTaskLogs") + if err != nil { + t.Fatal(err) + } + if err := stream.CloseSend(); err != nil { + t.Fatal(err) + } + if err := stream.RecvMsg(&emptypb.Empty{}); !errors.Is(err, io.EOF) { + t.Fatalf("RecvMsg = %v, want io.EOF", err) + } + traceID, _, _ := traceparent(t, <-headers) + if want := ct.span.SpanContext().TraceID().String(); traceID != want { + t.Errorf("stream traceparent trace ID = %s, want %s", traceID, want) + } +} + +func TestResolveExporters(t *testing.T) { + known := []string{"otlp"} + def := []string{"otlp"} + tests := []struct { + value string + want []string + wantErr bool + }{ + {value: "", want: def}, + {value: " ", want: def}, + {value: "otlp", want: []string{"otlp"}}, + {value: " OTLP ", want: []string{"otlp"}}, + {value: "none", want: nil}, + {value: " None ", want: nil}, + {value: "otlp,console", want: []string{"otlp"}, wantErr: true}, + {value: "otlp,", want: []string{"otlp"}, wantErr: true}, + {value: "console", want: def, wantErr: true}, + {value: "none,otlp", want: def, wantErr: true}, + {value: "none,console", want: nil, wantErr: true}, + } + for _, tt := range tests { + got, err := resolveExporters("OTEL_TRACES_EXPORTER", tt.value, known, def) + if !slices.Equal(got, tt.want) || (err != nil) != tt.wantErr { + t.Errorf("resolveExporters(%q) = %q, %v; want %q, error %v", tt.value, got, err, tt.want, tt.wantErr) + } + } +} + +func TestFinish(t *testing.T) { + t.Run("remote exit code", func(t *testing.T) { + exporter := newKeptSpans() + ct := newCommandTrace("ssh", exporter, io.Discard) + var stderr bytes.Buffer + + if code := finish(ct, exitCodeError{code: 3}, &stderr); code != 3 { + t.Errorf("finish = %d, want the remote exit code 3", code) + } + if stderr.Len() != 0 { + t.Errorf("expected no error message for a remote exit code, got %q", stderr.String()) + } + spans := exporter.GetSpans() + if len(spans) != 1 || spans[0].Status.Code != codes.Error { + t.Errorf("expected the command span to be ended and exported with an error before finish returns, got %+v", spans) + } + }) + t.Run("error", func(t *testing.T) { + var stderr bytes.Buffer + if code := finish(nil, errors.New("boom"), &stderr); code != 1 { + t.Errorf("finish = %d, want 1", code) + } + if stderr.String() != "Error: boom\n" { + t.Errorf("stderr = %q, want %q", stderr.String(), "Error: boom\n") + } + }) + t.Run("success", func(t *testing.T) { + if code := finish(nil, nil, io.Discard); code != 0 { + t.Errorf("finish = %d, want 0", code) + } + }) +} diff --git a/go.mod b/go.mod index 968855e5..7ab38a39 100644 --- a/go.mod +++ b/go.mod @@ -7,20 +7,33 @@ require ( github.com/cenkalti/backoff/v5 v5.0.3 github.com/jackc/pgx/v5 v5.11.0 go.ate.dev/substrate v1.0.0 + go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.70.0 + go.opentelemetry.io/otel v1.45.0 + go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.45.0 + go.opentelemetry.io/otel/metric v1.45.0 + go.opentelemetry.io/otel/sdk v1.45.0 + go.opentelemetry.io/otel/trace v1.45.0 google.golang.org/grpc v1.83.2 google.golang.org/protobuf v1.36.12 gopkg.in/yaml.v3 v3.0.1 ) require ( + github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/go-logr/logr v1.4.4 // indirect + github.com/go-logr/stdr v1.2.2 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/grpc-ecosystem/grpc-gateway/v2 v2.30.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect github.com/jackc/puddle/v2 v2.2.2 // indirect - github.com/kr/pretty v0.3.1 // indirect - github.com/rogpeppe/go-internal v1.14.1 // indirect + go.opentelemetry.io/auto/sdk v1.2.1 // indirect + go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.45.0 // indirect + go.opentelemetry.io/proto/otlp v1.11.0 // indirect golang.org/x/net v0.60.0 // indirect golang.org/x/sync v0.23.0 // indirect golang.org/x/sys v0.48.0 // indirect golang.org/x/text v0.42.0 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20260803160001-6ac0973c030d // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260803160001-6ac0973c030d // indirect ) diff --git a/go.sum b/go.sum index 9c33c634..6300ed97 100644 --- a/go.sum +++ b/go.sum @@ -4,8 +4,8 @@ github.com/cenkalti/backoff/v5 v5.0.3 h1:ZN+IMa753KfX5hd8vVaMixjnqRZ3y8CuJKRKj1x github.com/cenkalti/backoff/v5 v5.0.3/go.mod h1:rkhZdG3JZukswDf7f0cwqPNk4K0sa+F97BxZthm/crw= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= -github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= github.com/go-logr/logr v1.4.4 h1:tG4xh9yMsRCAiodLVTxyrkzSZ9+o0L1Kg/+cPVcbP/8= github.com/go-logr/logr v1.4.4/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= @@ -16,6 +16,8 @@ github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.30.0 h1:/Tnpcb2E0Pz/tN9s3bfEY2Q8ePCEX9iuS+cneUwncnw= +github.com/grpc-ecosystem/grpc-gateway/v2 v2.30.0/go.mod h1:zOBXOsUaBSjKgmH4OGzV1esUpR3oUSCPYVd2cUBjKYY= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo= @@ -28,9 +30,7 @@ github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= -github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= github.com/rogpeppe/go-internal v1.14.1 h1:UQB4HGPB6osV0SQTLymcB4TgvyWu6ZyliaW0tI/otEQ= github.com/rogpeppe/go-internal v1.14.1/go.mod h1:MaRKkUm5W0goXpeCfT7UZI6fk/L7L7so1lCWt35ZSgc= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= @@ -42,8 +42,14 @@ go.ate.dev/substrate v1.0.0 h1:7zFdRCb4e4OHoJUUWugBCQ7m1GLHHb1SxJo05HqEDjw= go.ate.dev/substrate v1.0.0/go.mod h1:iBSRX82TClE5CSsD80HoYsxVP81NSpf0zKKTW+2s/+0= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= +go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.70.0 h1:oECp5f+hN7nkwjU/8BxQ/q23bGPb8FIrD839owX222E= +go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.70.0/go.mod h1:DqEFwLumhzMBDQv9PcWbyoDxHI/4lAk6CM4nJBH39sc= go.opentelemetry.io/otel v1.45.0 h1:pdrWmLHofpubmArBv1LgFSv1Z0Ie/ppdZzu+kUN5EeU= go.opentelemetry.io/otel v1.45.0/go.mod h1:XZxIqPapzEYnhNSScF5DIqXhm/rYi0FzCe2XddAwZfQ= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.45.0 h1:QRefszxJmfPdjXUUm3j6iDzY03mTPXMjqErFqQ67vUg= +go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.45.0/go.mod h1:Tiz03lTBVBrm7eWZBOidzEaYaJa8tjwGUGv6d8mlTyk= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.45.0 h1:fG5MCxGz8+2VtrN/WgqSpJFctVz24gpxj8CxkKmc8Ww= +go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc v1.45.0/go.mod h1:BmAYTn+3ysbRe+IU2msxmf5Rx3g6DHvex+tWI3LdhYI= go.opentelemetry.io/otel/metric v1.45.0 h1:7Eg1uH7CJ5cXv9is6tnBe1FI6rj1nwUdbFypRm3br/M= go.opentelemetry.io/otel/metric v1.45.0/go.mod h1:HAPbm1nd3p1PmFH7v2dR+6BjXxw+Lq4a2+pndMAm08s= go.opentelemetry.io/otel/sdk v1.45.0 h1:4VVSMgQ83dUgW2aoX5f6JgLvHwIvzcuLnF9lUdCSpCw= @@ -52,6 +58,10 @@ go.opentelemetry.io/otel/sdk/metric v1.45.0 h1:oVFszMfyj1Am6s24Vtc7wBb8BKLcwepJj go.opentelemetry.io/otel/sdk/metric v1.45.0/go.mod h1:vUWUxDZvu1WVRj8JA8S0AdhsPrZoDpA2DdZauIh4mDA= go.opentelemetry.io/otel/trace v1.45.0 h1:l/mP6Uv7oNO7/TblbhpbgMidxhq1uO/rPsikOyVhxag= go.opentelemetry.io/otel/trace v1.45.0/go.mod h1:qoJJA2xNMnxRrdISU/kLtfUH2wNeQbiv+jhs/CxI8bc= +go.opentelemetry.io/proto/otlp v1.11.0 h1:5rrYs0Ykyj50sdU/JU0x8etU+LubXWb+gED6TbEdMIk= +go.opentelemetry.io/proto/otlp v1.11.0/go.mod h1:SmVizdCOAm3XBtG1g1NnOdhW6jtddT72hLMhv8VwA8E= +go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg= golang.org/x/net v0.60.0 h1:79p50tfZlm0J9YfoDsSi639qSXNGVwEzOPLCxM2FsYU= @@ -64,6 +74,8 @@ golang.org/x/text v0.42.0 h1:JbOZXgfeCPU9gacVtYliJqOhD+zhrEqK4LfdpmlUZqI= golang.org/x/text v0.42.0/go.mod h1:ojzP1Z+2QtioaF8DTtO8K5q7JWVVYwZKenzujK0Zd0E= gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= +google.golang.org/genproto/googleapis/api v0.0.0-20260803160001-6ac0973c030d h1:FarXi840EJWSHYTN3ERkADbPWjl307+FGrA22KAVjjc= +google.golang.org/genproto/googleapis/api v0.0.0-20260803160001-6ac0973c030d/go.mod h1:K/+WGbmBY7aNW1HDw1fJnKYo10i0DkAX6pows00dLig= google.golang.org/genproto/googleapis/rpc v0.0.0-20260803160001-6ac0973c030d h1:IL4hdHzcUv2l/gcg98/Rj3FbtE6axwqslOW8SW0C+S0= google.golang.org/genproto/googleapis/rpc v0.0.0-20260803160001-6ac0973c030d/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= google.golang.org/grpc v1.83.2 h1:EManeRomTObA0BU7I8vXgg/78uE5MJ9M8B39EX2WscU=