Merge pull request #1862 from zregvart/issue/EC-733

Moar tracing

cmd/root.go

@@ -18,7 +18,8 @@ package cmd
 
 import (
 	"context"
-	"os"
+
+	log "github.com/sirupsen/logrus"
 
 	"github.com/enterprise-contract/ec-cli/cmd/fetch"
 	"github.com/enterprise-contract/ec-cli/cmd/initialize"
@@ -42,8 +43,11 @@ var RootCmd = root.NewRootCmd()
 // This is called by main.main(). It only needs to happen once to the rootCmd.
 func Execute() {
 	if err := RootCmd.ExecuteContext(context.Background()); err != nil {
-		os.Exit(1)
+		root.OnExit()
+		log.Fatalf("error executing command: %v", err)
 	}
+
+	root.OnExit()
 }
 
 func init() {

cmd/root/root_cmd.go

@@ -18,28 +18,40 @@ package root
 
 import (
 	"context"
+	"fmt"
 	"io"
+	"os"
+	"runtime"
+	"runtime/pprof"
+	"sync"
 	"time"
 
 	hd "github.com/MakeNowJust/heredoc"
-	"github.com/sirupsen/logrus"
+	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 
 	"github.com/enterprise-contract/ec-cli/internal/kubernetes"
 	"github.com/enterprise-contract/ec-cli/internal/logging"
 )
 
-var cancel context.CancelFunc
-
 var (
 	quiet         bool = false
 	verbose       bool = false
 	debug         bool = false
 	trace         bool = false
 	globalTimeout      = 5 * time.Minute
 	logfile       string
+	OnExit        func() = func() {}
 )
 
+type customDeadlineExceededError struct{}
+
+func (customDeadlineExceededError) Error() string {
+	return fmt.Sprintf("exceeded allowed execution time of %s, the timeout can be adjusted using the --timeout command line argument", globalTimeout)
+}
+func (customDeadlineExceededError) Timeout() bool   { return true }
+func (customDeadlineExceededError) Temporary() bool { return true }
+
 func NewRootCmd() *cobra.Command {
 	rootCmd := &cobra.Command{
 		Use:   "ec",
@@ -56,19 +68,56 @@ func NewRootCmd() *cobra.Command {
 		PersistentPreRun: func(cmd *cobra.Command, _ []string) {
 			logging.InitLogging(verbose, quiet, debug, trace, logfile)
 
+			// set a custom message for context.DeadlineExceeded error
+			context.DeadlineExceeded = customDeadlineExceededError{}
+
 			// Create a new context now that flags have been parsed so a custom timeout can be used.
-			ctx := cmd.Context()
-			ctx, cancel = context.WithTimeout(ctx, globalTimeout)
+			ctx, cancel := context.WithTimeout(cmd.Context(), globalTimeout)
 			cmd.SetContext(ctx)
-		},
 
-		PersistentPostRun: func(cmd *cobra.Command, _ []string) {
-			if f, ok := logrus.StandardLogger().Out.(io.Closer); ok {
-				f.Close()
-			}
-			if cancel != nil {
-				cancel()
+			// if trace is enabled setup CPU profiling
+			var cpuprofile *os.File
+			if trace {
+				var err error
+				if cpuprofile, err = os.CreateTemp("", "cpuprofile.*"); err != nil {
+					log.Fatal("could not create CPU profile: ", err)
+				}
+				if err := pprof.StartCPUProfile(cpuprofile); err != nil {
+					log.Fatal("could not start CPU profile: ", err)
+				}
 			}
+
+			OnExit = sync.OnceFunc(func() {
+				if trace {
+					// dump memory profile
+					if memprofile, err := os.CreateTemp("", "memprofile.*"); err != nil {
+						log.Fatal("could not create memory profile: ", err)
+					} else {
+						defer memprofile.Close()
+						runtime.GC()
+						if err := pprof.WriteHeapProfile(memprofile); err != nil {
+							log.Fatal("could not start CPU profile: ", err)
+						}
+
+						log.Tracef("wrote memory profile to: %s", memprofile.Name())
+					}
+
+					// dump the CPU profile
+					pprof.StopCPUProfile()
+					if cpuprofile != nil {
+						_ = cpuprofile.Close() // ignore errors
+						log.Tracef("wrote CPU profile to: %s", cpuprofile.Name())
+					}
+				}
+
+				// perform resource cleanup
+				if f, ok := log.StandardLogger().Out.(io.Closer); ok {
+					f.Close()
+				}
+				if cancel != nil {
+					cancel()
+				}
+			})
 		},
 	}
 

internal/utils/oci/client.go

@@ -19,6 +19,7 @@ package oci
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"os"
 	"path"
 	"strconv"
@@ -47,6 +48,23 @@ const clientContextKey contextKey = "ec.oci.client"
 
 var imgCache = sync.OnceValue(initCache)
 
+type tracingRoundTripper struct{}
+
+func (t *tracingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	log.Tracef("START: %s %s", req.Method, req.URL)
+	resp, err := remote.DefaultTransport.RoundTrip(req)
+
+	log.Tracef("DONE: %s %s (%d)", req.Method, req.URL, resp.ContentLength)
+
+	return resp, err
+}
+
+func init() {
+	if log.IsLevelEnabled(log.TraceLevel) {
+		imageRefTransport = remote.WithTransport(&tracingRoundTripper{})
+	}
+}
+
 func initCache() cache.Cache {
 	// if a value was set and it is parsed as false, turn the cache off
 	if v, err := strconv.ParseBool(os.Getenv("EC_CACHE")); err == nil && !v {