From f07834d1854233a300ec0b97d6cd52b8338402f4 Mon Sep 17 00:00:00 2001
From: Laura Brehm <laurabrehm@hey.com>
Date: Thu, 2 May 2024 03:21:59 +0100
Subject: [PATCH] OTel: add `command.time` metric to plugin commands

Signed-off-by: Laura Brehm <laurabrehm@hey.com>
---
 cli-plugins/plugin/plugin.go   | 18 ++++++++++++++++++
 cli/command/telemetry_utils.go | 29 ++++++++++++++++++++++-------
 cmd/docker/docker.go           |  2 +-
 3 files changed, 41 insertions(+), 8 deletions(-)

diff --git a/cli-plugins/plugin/plugin.go b/cli-plugins/plugin/plugin.go
index 5b1204f3cb..d4f38d74cd 100644
--- a/cli-plugins/plugin/plugin.go
+++ b/cli-plugins/plugin/plugin.go
@@ -46,6 +46,24 @@ func RunPlugin(dockerCli *command.DockerCli, plugin *cobra.Command, meta manager
 				opts = append(opts, withPluginClientConn(plugin.Name()))
 			}
 			err = tcmd.Initialize(opts...)
+			ogRunE := cmd.RunE
+			if ogRunE == nil {
+				ogRun := cmd.Run
+				// necessary because error will always be nil here
+				// see: https://github.com/golangci/golangci-lint/issues/1379
+				//nolint:unparam
+				ogRunE = func(cmd *cobra.Command, args []string) error {
+					ogRun(cmd, args)
+					return nil
+				}
+				cmd.Run = nil
+			}
+			cmd.RunE = func(cmd *cobra.Command, args []string) error {
+				stopInstrumentation := dockerCli.StartInstrumentation(cmd)
+				err := ogRunE(cmd, args)
+				stopInstrumentation(err)
+				return err
+			}
 		})
 		return err
 	}
diff --git a/cli/command/telemetry_utils.go b/cli/command/telemetry_utils.go
index 5749b4f18c..87b976ce69 100644
--- a/cli/command/telemetry_utils.go
+++ b/cli/command/telemetry_utils.go
@@ -26,8 +26,7 @@ func BaseCommandAttributes(cmd *cobra.Command, streams Streams) []attribute.KeyV
 // Note: this should be the last func to wrap/modify the PersistentRunE/RunE funcs before command execution.
 //
 // can also be used for spans!
-func (cli *DockerCli) InstrumentCobraCommands(cmd *cobra.Command, mp metric.MeterProvider) {
-	meter := getDefaultMeter(mp)
+func (cli *DockerCli) InstrumentCobraCommands(ctx context.Context, cmd *cobra.Command) {
 	// If PersistentPreRunE is nil, make it execute PersistentPreRun and return nil by default
 	ogPersistentPreRunE := cmd.PersistentPreRunE
 	if ogPersistentPreRunE == nil {
@@ -55,10 +54,9 @@ func (cli *DockerCli) InstrumentCobraCommands(cmd *cobra.Command, mp metric.Mete
 		}
 		cmd.RunE = func(cmd *cobra.Command, args []string) error {
 			// start the timer as the first step of every cobra command
-			baseAttrs := BaseCommandAttributes(cmd, cli)
-			stopCobraCmdTimer := startCobraCommandTimer(cmd, meter, baseAttrs)
+			stopInstrumentation := cli.StartInstrumentation(cmd)
 			cmdErr := ogRunE(cmd, args)
-			stopCobraCmdTimer(cmdErr)
+			stopInstrumentation(cmdErr)
 			return cmdErr
 		}
 
@@ -66,8 +64,17 @@ func (cli *DockerCli) InstrumentCobraCommands(cmd *cobra.Command, mp metric.Mete
 	}
 }
 
-func startCobraCommandTimer(cmd *cobra.Command, meter metric.Meter, attrs []attribute.KeyValue) func(err error) {
-	ctx := cmd.Context()
+// StartInstrumentation instruments CLI commands with the individual metrics and spans configured.
+// It's the main command OTel utility, and new command-related metrics should be added to it.
+// It should be called immediately before command execution, and returns a stopInstrumentation function
+// that must be called with the error resulting from the command execution.
+func (cli *DockerCli) StartInstrumentation(cmd *cobra.Command) (stopInstrumentation func(error)) {
+	baseAttrs := BaseCommandAttributes(cmd, cli)
+	return startCobraCommandTimer(cli.MeterProvider(), baseAttrs)
+}
+
+func startCobraCommandTimer(mp metric.MeterProvider, attrs []attribute.KeyValue) func(err error) {
+	meter := getDefaultMeter(mp)
 	durationCounter, _ := meter.Float64Counter(
 		"command.time",
 		metric.WithDescription("Measures the duration of the cobra command"),
@@ -76,12 +83,20 @@ func startCobraCommandTimer(cmd *cobra.Command, meter metric.Meter, attrs []attr
 	start := time.Now()
 
 	return func(err error) {
+		// Use a new context for the export so that the command being cancelled
+		// doesn't affect the metrics, and we get metrics for cancelled commands.
+		ctx, cancel := context.WithTimeout(context.Background(), exportTimeout)
+		defer cancel()
+
 		duration := float64(time.Since(start)) / float64(time.Millisecond)
 		cmdStatusAttrs := attributesFromError(err)
 		durationCounter.Add(ctx, duration,
 			metric.WithAttributes(attrs...),
 			metric.WithAttributes(cmdStatusAttrs...),
 		)
+		if mp, ok := mp.(MeterProvider); ok {
+			mp.ForceFlush(ctx)
+		}
 	}
 }
 
diff --git a/cmd/docker/docker.go b/cmd/docker/docker.go
index 7103997f2d..27eb0263bd 100644
--- a/cmd/docker/docker.go
+++ b/cmd/docker/docker.go
@@ -315,7 +315,7 @@ func runDocker(ctx context.Context, dockerCli *command.DockerCli) error {
 		fmt.Fprint(dockerCli.Err(), "Warning: Unexpected OTEL error, metrics may not be flushed")
 	}
 
-	dockerCli.InstrumentCobraCommands(cmd, mp)
+	dockerCli.InstrumentCobraCommands(ctx, cmd)
 
 	var envs []string
 	args, os.Args, envs, err = processAliases(dockerCli, cmd, args, os.Args)