Merge pull request #247 from puppetlabs/gh-225/main/refactor-telemetry

(GH-225) Refactor telemetry and add error tracing

cmd/build/build.go

@@ -8,7 +8,6 @@ import (
 	"github.com/puppetlabs/pdkgo/internal/pkg/gzip"
 	"github.com/puppetlabs/pdkgo/internal/pkg/pct"
 	"github.com/puppetlabs/pdkgo/internal/pkg/tar"
-	"github.com/puppetlabs/pdkgo/pkg/telemetry"
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/afero"
 	"github.com/spf13/cobra"
@@ -63,10 +62,6 @@ func preExecute(cmd *cobra.Command, args []string) error {
 }
 
 func execute(cmd *cobra.Command, args []string) error {
-	_, span := telemetry.NewSpan(cmd.Context(), "build")
-	defer telemetry.EndSpan(span)
-	telemetry.AddStringSpanAttribute(span, "name", "build")
-
 	gzipArchiveFilePath, err := builder.Build(sourceDir, targetDir)
 
 	if err != nil {

cmd/install/install.go

@@ -5,7 +5,6 @@ import (
 
 	"github.com/puppetlabs/pdkgo/internal/pkg/pct"
 	"github.com/puppetlabs/pdkgo/internal/pkg/utils"
-	"github.com/puppetlabs/pdkgo/pkg/telemetry"
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
@@ -40,10 +39,6 @@ func (ic *InstallCommand) CreateCommand() *cobra.Command {
 }
 
 func (ic *InstallCommand) executeInstall(cmd *cobra.Command, args []string) error {
-	_, span := telemetry.NewSpan(cmd.Context(), "install")
-	defer telemetry.EndSpan(span)
-	telemetry.AddStringSpanAttribute(span, "name", "install")
-
 	templateInstallationPath, err := ic.PctInstaller.Install(ic.TemplatePkgPath, ic.InstallPath, ic.Force)
 	if err != nil {
 		return err

cmd/new/new.go

@@ -156,10 +156,8 @@ func getApplicationInfo(appVersionString string) pct.PDKInfo {
 }
 
 func execute(cmd *cobra.Command, args []string) error {
-	_, span := telemetry.NewSpan(cmd.Context(), "new")
-	defer telemetry.EndSpan(span)
-	telemetry.AddStringSpanAttribute(span, "name", "new")
-
+	span := telemetry.GetSpanFromContext(cmd.Context())
+	// Add template to span if needed
 	if len(args) == 1 {
 		telemetry.AddStringSpanAttribute(span, "template", args[0])
 	}

cmd/root/root.go

@@ -3,6 +3,7 @@ package root
 import (
 	"os"
 	"path/filepath"
+	"strings"
 
 	homedir "github.com/mitchellh/go-homedir"
 	"github.com/puppetlabs/pdkgo/internal/pkg/utils"
@@ -85,3 +86,17 @@ func InitConfig() {
 		log.Trace().Msgf("Using config file: %s", viper.ConfigFileUsed())
 	}
 }
+
+// Returns the cobra command called, e.g. new or install
+// and also the fully formatted command as passed with arguments/flags.
+// Idea borrowed from carolynvs/porter:
+// https://github.com/carolynvs/porter/blob/ccca10a63627e328616c1006600153da8411a438/cmd/porter/main.go
+func GetCalledCommand(cmd *cobra.Command) (string, string) {
+	calledCommandStr := os.Args[1]
+
+	// Also figure out the full called command from the CLI
+	// Is there anything sensitive you could leak here? 🤔
+	calledCommandArgs := strings.Join(os.Args[1:], " ")
+
+	return calledCommandStr, calledCommandArgs
+}

main.go

@@ -30,11 +30,17 @@ var (
 func main() {
 	// Telemetry must be initialized before anything else;
 	// If the telemetry build tag was not passed, this is all null ops
-	ctx, traceProvider, span := telemetry.Start(context.Background(), honeycomb_api_key, honeycomb_dataset)
-	defer telemetry.ShutDown(ctx, traceProvider, span)
+	ctx, traceProvider, parentSpan := telemetry.Start(context.Background(), honeycomb_api_key, honeycomb_dataset, "pct")
 
 	var rootCmd = root.CreateRootCommand()
 
+	// Get the command called and its arguments;
+	// The arguments are only necessary if we want to
+	// hand them off as an attribute to the parent span:
+	// do we? Otherwise we just need the calledCommand
+	calledCommand, calledCommandArguments := root.GetCalledCommand(rootCmd)
+	telemetry.AddStringSpanAttribute(parentSpan, "arguments", calledCommandArguments)
+
 	var verCmd = appver.CreateVersionCommand(version, date, commit)
 	v := appver.Format(version, date, commit)
 	rootCmd.Version = v
@@ -66,6 +72,18 @@ func main() {
 	// new
 	rootCmd.AddCommand(new.CreateCommand())
 
+	// initialize
 	cobra.OnInitialize(root.InitLogger, root.InitConfig)
-	cobra.CheckErr(rootCmd.ExecuteContext(ctx))
+
+	// instrument & execute called command
+	ctx, childSpan := telemetry.NewSpan(ctx, calledCommand)
+	err := rootCmd.ExecuteContext(ctx)
+	telemetry.RecordSpanError(childSpan, err)
+	telemetry.EndSpan(childSpan)
+
+	// Send all events
+	telemetry.ShutDown(ctx, traceProvider, parentSpan)
+
+	// Handle exiting with/out errors.
+	cobra.CheckErr(err)
 }

pkg/telemetry/telemetry.go

@@ -22,7 +22,7 @@ import (
 	"go.opentelemetry.io/otel/sdk/resource"
 )
 
-func Start(ctx context.Context, honeycomb_api_key string, honeycomb_dataset string) (context.Context, *sdktrace.TracerProvider, trace.Span) {
+func Start(ctx context.Context, honeycomb_api_key string, honeycomb_dataset string, rootSpanName string) (context.Context, *sdktrace.TracerProvider, trace.Span) {
 	var tp *sdktrace.TracerProvider
 	// if telemetry is turned on and honeycomb is configured, hook it up
 	api_key_set := honeycomb_api_key != "not_set" && honeycomb_api_key != ""
@@ -43,7 +43,7 @@ func Start(ctx context.Context, honeycomb_api_key string, honeycomb_dataset stri
 		res, err := resource.New(ctx,
 			resource.WithAttributes(
 				// the service name used to display traces in backends
-				semconv.ServiceNameKey.String("PCT"),
+				semconv.ServiceNameKey.String("PDK"),
 			),
 		)
 		if err != nil {
@@ -73,10 +73,10 @@ func Start(ctx context.Context, honeycomb_api_key string, honeycomb_dataset stri
 		// No spans will be reported, maybe it's best to return nils or error?
 	}
 
-	tracer := otel.Tracer("pct")
+	tracer := otel.GetTracerProvider().Tracer("")
 
 	var span trace.Span
-	ctx, span = tracer.Start(ctx, "execution")
+	ctx, span = tracer.Start(ctx, rootSpanName)
 
 	// The Protected ID is hashed base on application name to prevent any
 	// accidental leakage of a reversable ID.
@@ -94,13 +94,24 @@ func EndSpan(span trace.Span) {
 	span.End()
 }
 
+// Returns the current span from context
+func GetSpanFromContext(ctx context.Context) trace.Span {
+	return trace.SpanFromContext(ctx)
+}
+
 // Create a new span;
 // the span will need to be closed somewhere up the call stack
 func NewSpan(ctx context.Context, name string) (context.Context, trace.Span) {
 	tracer := otel.GetTracerProvider().Tracer("")
 	return tracer.Start(ctx, name)
 }
 
+// Records an error to the span;
+// If err is nil, this function does nothing
+func RecordSpanError(span trace.Span, err error) {
+	span.RecordError(err)
+}
+
 // Create a new attribute and attach it to the specified span
 func AddStringSpanAttribute(span trace.Span, key string, value string) {
 	attr := attribute.Key(key)

pkg/telemetry/telemetry_noop.go

@@ -7,7 +7,7 @@ import (
 	"context"
 )
 
-func Start(ctx context.Context, honeycomb_api_key string, honeycomb_dataset string) (context.Context, string, string) {
+func Start(ctx context.Context, honeycomb_api_key string, honeycomb_dataset string, rootSpanName string) (context.Context, string, string) {
 	// deliberately does nothing
 	return ctx, "", ""
 }
@@ -16,11 +16,20 @@ func EndSpan(span string) {
 	// deliberately does nothing
 }
 
+func GetSpanFromContext(ctx context.Context) string {
+	// deliberately does nothing
+	return ""
+}
+
 func NewSpan(ctx context.Context, name string) (context.Context, string) {
 	// deliberately does nothing
 	return ctx, ""
 }
 
+func RecordSpanError(span string, err error) {
+	// deliberately does nothing
+}
+
 func AddStringSpanAttribute(span string, key string, value string) {
 	// deliberately does nothing
 }