RSDK-12623: Viam cli can work with trace files from viam-server

cli/app.go

@@ -8,6 +8,7 @@ import (
 	"regexp"
 	"strings"
 
+	"github.com/samber/lo"
 	"github.com/urfave/cli/v2"
 
 	"go.viam.com/rdk/logging"
@@ -189,6 +190,14 @@ var commonPartFlags = []cli.Flag{
 	},
 }
 
+var commonOtlpFlags = []cli.Flag{
+	&cli.StringFlag{
+		Name:        "endpoint",
+		DefaultText: "localhost:4317",
+		Usage:       "OTLP endpoint in host:port format",
+	},
+}
+
 // matches all uppercase characters that follow lowercase chars and aren't at the [0] index of a string.
 // This is useful for converting camel case into kabob case when getting values out of a CLI Context
 // based on a flag name, and putting them into a struct with a camel cased field name.
@@ -456,6 +465,63 @@ var app = &cli.App{
 		},
 	},
 	Commands: []*cli.Command{
+		{
+			Name:  "traces",
+			Usage: "Work with viam-server traces",
+			Subcommands: []*cli.Command{
+				{
+					Name:      "import-local",
+					Usage:     "Import traces from a local viam server trace file to an OTLP endpoint.",
+					ArgsUsage: "<traces file>",
+					Flags:     commonOtlpFlags,
+					Action:    createCommandWithT(traceImportLocalAction),
+				},
+				{
+					Name: "import-remote",
+					Description: `
+In order to use the import-remote command, the machine must have a valid shell type service.
+Organization and location are required flags if using name (rather than ID) for the part.
+Note: There is no progress meter while copying is in progress.
+`,
+					Usage: "Import traces from a remote viam machine to an OTLP endpoint.",
+					Flags: lo.Flatten([][]cli.Flag{
+						commonOtlpFlags,
+						commonPartFlags,
+					}),
+					Action: createCommandWithT(traceImportRemoteAction),
+				},
+				{
+					Name:      "print-local",
+					Usage:     "Print traces in a local file to the console",
+					ArgsUsage: "<traces file>",
+					Action:    createCommandWithT(tracePrintLocalAction),
+				},
+				{
+					Name:  "print-remote",
+					Usage: "Print traces from a remote viam machine to the console",
+					Description: `
+In order to use the print-remote command, the machine must have a valid shell type service.
+Organization and location are required flags if using name (rather than ID) for the part.
+Note: There is no progress meter while copying is in progress.
+`,
+					Flags:  commonPartFlags,
+					Action: createCommandWithT(tracePrintRemoteAction),
+				},
+				{
+					Usage:     "Download traces from a viam machine and save them to disk",
+					Name:      "get-remote",
+					ArgsUsage: "[target]",
+					Description: `
+In order to use the get-remote command, the machine must have a valid shell type service.
+Organization and location are required flags if using name (rather than ID) for the part.
+If [target] is not specified then the traces file will be saved to the current working directory.
+Note: There is no progress meter while copying is in progress.
+`,
+					Flags:  commonPartFlags,
+					Action: createCommandWithT(traceGetRemoteAction),
+				},
+			},
+		},
 		{
 			Name: "login",
 			// NOTE(benjirewis): maintain `auth` as an alias for backward compatibility.

cli/traces.go

@@ -0,0 +1,311 @@
+package cli
+
+import (
+	stderrors "errors"
+	"os"
+	"path"
+	"path/filepath"
+	"time"
+
+	"github.com/pkg/errors"
+	"github.com/urfave/cli/v2"
+	"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc"
+	tracepb "go.opentelemetry.io/proto/otlp/trace/v1"
+	"go.viam.com/utils/perf"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	"go.viam.com/rdk/logging"
+	"go.viam.com/rdk/protoutils"
+	"go.viam.com/rdk/services/shell"
+)
+
+var tracesPath = path.Join("~", ".viam", "trace")
+
+type traceGetRemoteArgs struct {
+	Organization string
+	Location     string
+	Machine      string
+	Part         string
+}
+
+type traceImportRemoteArgs struct {
+	Organization string
+	Location     string
+	Machine      string
+	Part         string
+	Endpoint     string
+}
+
+type traceImportLocalArgs struct {
+	Endpoint string
+}
+
+func traceImportRemoteAction(ctx *cli.Context, args traceImportRemoteArgs) error {
+	client, err := newViamClient(ctx)
+	if err != nil {
+		return err
+	}
+
+	globalArgs, err := getGlobalArgs(ctx)
+	if err != nil {
+		return err
+	}
+	logger := globalArgs.createLogger()
+	tmp, err := os.MkdirTemp("", "viamtraceimport")
+	if err != nil {
+		return err
+	}
+	//nolint: errcheck
+	defer os.RemoveAll(tmp)
+	if err := client.tracesGetRemoteAction(
+		ctx,
+		traceGetRemoteArgs{
+			Organization: args.Organization,
+			Location:     args.Location,
+			Machine:      args.Machine,
+			Part:         args.Part,
+		},
+		tmp,
+		false,
+		globalArgs.Debug,
+		logger,
+	); err != nil {
+		return err
+	}
+
+	return traceImportLocal(ctx, traceImportLocalArgs{
+		Endpoint: args.Endpoint,
+	},
+		filepath.Join(tmp, "traces"),
+	)
+}
+
+func (c *viamClient) tracesGetRemoteAction(
+	ctx *cli.Context,
+	flagArgs traceGetRemoteArgs,
+	target string,
+	getAll bool,
+	debug bool,
+	logger logging.Logger,
+) error {
+	part, err := c.robotPart(flagArgs.Organization, flagArgs.Location, flagArgs.Machine, flagArgs.Part)
+	if err != nil {
+		return err
+	}
+	// Intentional use of path instead of filepath: Windows understands both / and
+	// \ as path separators, and we don't want a cli running on Windows to send
+	// a path using \ to a *NIX machine.
+	src := path.Join(tracesPath, part.Id)
+	// if getAll is set then download the entire directory, including rotated
+	// files. Otherwise just get the current file.
+	if !getAll {
+		src = path.Join(src, "traces")
+	}
+	gArgs, err := getGlobalArgs(ctx)
+	quiet := err == nil && gArgs != nil && gArgs.Quiet
+	var startTime time.Time
+	if !quiet {
+		startTime = time.Now()
+		printf(ctx.App.Writer, "Saving to %s ...", path.Join(target))
+	}
+	if err := c.copyFilesFromMachine(
+		flagArgs.Organization,
+		flagArgs.Location,
+		flagArgs.Machine,
+		flagArgs.Part,
+		debug,
+		true,
+		false,
+		[]string{src},
+		target,
+		logger,
+	); err != nil {
+		if statusErr := status.Convert(err); statusErr != nil &&
+			statusErr.Code() == codes.InvalidArgument &&
+			statusErr.Message() == shell.ErrMsgDirectoryCopyRequestNoRecursion {
+			return errDirectoryCopyRequestNoRecursion
+		}
+		return err
+	}
+	if !quiet {
+		printf(ctx.App.Writer, "Download finished in %s.", time.Since(startTime))
+	}
+	return nil
+}
+
+func tracePrintRemoteAction(
+	ctx *cli.Context,
+	args traceGetRemoteArgs,
+) error {
+	client, err := newViamClient(ctx)
+	if err != nil {
+		return err
+	}
+
+	globalArgs, err := getGlobalArgs(ctx)
+	if err != nil {
+		return err
+	}
+	logger := globalArgs.createLogger()
+	tmp, err := os.MkdirTemp("", "viamtraceimport")
+	if err != nil {
+		return err
+	}
+	//nolint: errcheck
+	defer os.RemoveAll(tmp)
+	if err := client.tracesGetRemoteAction(
+		ctx,
+		args,
+		tmp,
+		false,
+		globalArgs.Debug,
+		logger,
+	); err != nil {
+		return err
+	}
+	return tracePrintLocal(ctx, filepath.Join(tmp, "traces"))
+}
+
+func getSingularArg(ctx *cli.Context) (string, error) {
+	cliArgs := ctx.Args().Slice()
+	var result string
+	switch numArgs := len(cliArgs); numArgs {
+	case 1:
+		result = cliArgs[0]
+	default:
+		return "", wrongNumArgsError{have: numArgs, min: 1}
+	}
+	return result, nil
+}
+
+func traceGetRemoteAction(ctx *cli.Context, args traceGetRemoteArgs) error {
+	cliArgs := ctx.Args().Slice()
+	var targetPath string
+	switch numArgs := len(cliArgs); numArgs {
+	case 0:
+		var err error
+		targetPath, err = os.Getwd()
+		if err != nil {
+			return err
+		}
+	case 1:
+		targetPath = cliArgs[0]
+	default:
+		return wrongNumArgsError{numArgs, 0, 1}
+	}
+
+	client, err := newViamClient(ctx)
+	if err != nil {
+		return err
+	}
+
+	globalArgs, err := getGlobalArgs(ctx)
+	if err != nil {
+		return err
+	}
+	logger := globalArgs.createLogger()
+
+	return client.tracesGetRemoteAction(ctx, args, targetPath, true, globalArgs.Debug, logger)
+}
+
+func tracePrintLocalAction(
+	ctx *cli.Context,
+	_ struct{},
+) error {
+	target, err := getSingularArg(ctx)
+	if err != nil {
+		return err
+	}
+	return tracePrintLocal(ctx, target)
+}
+
+func tracePrintLocal(
+	ctx *cli.Context,
+	source string,
+) error {
+	//nolint: gosec
+	traceFile, err := os.Open(source)
+	if err != nil {
+		if os.IsNotExist(err) {
+			printf(ctx.App.Writer, "No traces found")
+			return nil
+		}
+		return errors.Wrap(err, "failed to open trace file")
+	}
+	traceReader := protoutils.NewDelimitedProtoReader[tracepb.ResourceSpans](traceFile)
+	//nolint: errcheck
+	defer traceReader.Close()
+
+	devExporter := perf.NewOtelDevelopmentExporter()
+	if err := devExporter.Start(); err != nil {
+		return err
+	}
+	defer devExporter.Stop()
+	var msg tracepb.ResourceSpans
+	err = nil
+	for resource := range traceReader.AllWithMemory(&msg) {
+		for _, scope := range resource.ScopeSpans {
+			err = stderrors.Join(err, devExporter.ExportOTLPSpans(ctx.Context, scope.Spans))
+		}
+	}
+	return err
+}
+
+func traceImportLocalAction(
+	ctx *cli.Context,
+	args traceImportLocalArgs,
+) error {
+	target, err := getSingularArg(ctx)
+	if err != nil {
+		return err
+	}
+	return traceImportLocal(ctx, args, target)
+}
+
+func traceImportLocal(
+	ctx *cli.Context,
+	args traceImportLocalArgs,
+	source string,
+) error {
+	//nolint: gosec
+	traceFile, err := os.Open(source)
+	if err != nil {
+		if os.IsNotExist(err) {
+			printf(ctx.App.Writer, "No traces found")
+			return nil
+		}
+		return errors.Wrap(err, "failed to open trace file")
+	}
+	traceReader := protoutils.NewDelimitedProtoReader[tracepb.ResourceSpans](traceFile)
+	//nolint: errcheck
+	defer traceReader.Close()
+	endpoint := args.Endpoint
+	if endpoint == "" {
+		endpoint = "localhost:4317"
+	}
+	otlpClient := otlptracegrpc.NewClient(
+		otlptracegrpc.WithEndpoint(endpoint),
+		otlptracegrpc.WithInsecure(),
+	)
+	if err := otlpClient.Start(ctx.Context); err != nil {
+		return err
+	}
+	//nolint: errcheck
+	defer otlpClient.Stop(ctx.Context)
+	var msg tracepb.ResourceSpans
+	msgSuccess := 0
+	msgTotal := 0
+	printf(ctx.App.Writer, "Importing spans to %v...", endpoint)
+	for span := range traceReader.AllWithMemory(&msg) {
+		msgTotal++
+		err := otlpClient.UploadTraces(ctx.Context, []*tracepb.ResourceSpans{span})
+		if err != nil {
+			printf(ctx.App.Writer, "Error uploading trace: %v", err)
+		} else {
+			msgSuccess++
+		}
+	}
+	printf(ctx.App.Writer, "Imported %d/%d messages", msgSuccess, msgTotal)
+	return nil
+}