cli-plugins/plugin: add OnFinalize hook for shutting down otel providers

cli-plugins/plugin/plugin.go

@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"os"
 	"sync"
+	"time"
 
 	"github.com/docker/cli/cli"
 	"github.com/docker/cli/cli-plugins/metadata"
@@ -32,6 +33,11 @@ import (
 // called.
 var PersistentPreRunE func(*cobra.Command, []string) error
 
+// finalizeHookTimeout is the timeout given for running the finalize hook.
+// It contains the timeout for exporting metrics and a little bit extra in case
+// shutdown does not respond to the export timeout promptly.
+const finalizeHookTimeout = command.ExportTimeout + 10*time.Millisecond
+
 // RunPlugin executes the specified plugin command
 func RunPlugin(dockerCli *command.DockerCli, plugin *cobra.Command, meta metadata.Metadata) error {
 	tcmd := newPluginCommand(dockerCli, plugin, meta)
@@ -69,6 +75,7 @@ func RunPlugin(dockerCli *command.DockerCli, plugin *cobra.Command, meta metadat
 				stopInstrumentation(err)
 				return err
 			}
+			cobra.OnFinalize(makeFinalizeHook(dockerCli))
 		})
 		return retErr
 	}
@@ -83,6 +90,61 @@ func RunPlugin(dockerCli *command.DockerCli, plugin *cobra.Command, meta metadat
 	return cmd.Execute()
 }
 
+// makeFinalizeHook creates a hook to be run when the CLI has finished executing the run command.
+// This happens after any post run methods but also always runs even when an error happens.
+func makeFinalizeHook(dockerCli *command.DockerCli) func() {
+	return func() {
+		// 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(), command.ExportTimeout)
+		defer cancel()
+
+		var (
+			wg   sync.WaitGroup
+			wait bool
+		)
+		if tp, ok := dockerCli.TracerProvider().(command.TracerProvider); ok {
+			wg.Go(func() {
+				if err := tp.Shutdown(ctx); err != nil {
+					otel.Handle(err)
+				}
+			})
+			wait = true
+		}
+
+		if mp, ok := dockerCli.MeterProvider().(command.MeterProvider); ok {
+			wg.Go(func() {
+				if err := mp.Shutdown(ctx); err != nil {
+					otel.Handle(err)
+				}
+			})
+			wait = true
+		}
+
+		// Wait for the wait group in another goroutine if any were created.
+		if !wait {
+			return
+		}
+
+		// Spawn a goroutine to wait for the wait group.
+		// If the shutdown does not promptly hang up with the
+		// cancellation of the context, we leave this goroutine
+		// orphaned so the process can exit.
+		done := make(chan struct{})
+		go func() {
+			defer close(done)
+			wg.Wait()
+		}()
+
+		select {
+		case <-done:
+		case <-time.After(finalizeHookTimeout):
+			// Ensure we still exit even if the shutdown doesn't work properly.
+			// We give 10 extra milliseconds for the cleanup to happen.
+		}
+	}
+}
+
 // Run is the top-level entry point to the CLI plugin framework. It should
 // be called from the plugin's "main()" function. It initializes a new
 // [command.DockerCli] instance with the given options before calling
@@ -212,8 +274,8 @@ func newMetadataSubcommand(plugin *cobra.Command, meta metadata.Metadata) *cobra
 		Use:    metadata.MetadataSubcommandName,
 		Hidden: true,
 		// Suppress the global/parent PersistentPreRunE, which
-		// needlessly initializes the client and tries to
-		// connect to the daemon.
+		// needlessly initializes the client and tries to connect to the
+		// daemon.
 		PersistentPreRun: func(cmd *cobra.Command, args []string) {},
 		RunE: func(cmd *cobra.Command, args []string) error {
 			enc := json.NewEncoder(os.Stdout)

cli/command/telemetry.go

@@ -20,7 +20,7 @@ import (
 	"go.opentelemetry.io/otel/trace"
 )
 
-const exportTimeout = 50 * time.Millisecond
+const ExportTimeout = 50 * time.Millisecond
 
 // TracerProvider is an extension of the trace.TracerProvider interface for CLI programs.
 type TracerProvider interface {
@@ -205,7 +205,7 @@ func (r *cliReader) Shutdown(ctx context.Context) error {
 	// We don't want CLI metrics to prevent the CLI from exiting
 	// so if there's some kind of issue we need to abort pretty
 	// quickly.
-	ctx, cancel := context.WithTimeout(ctx, exportTimeout)
+	ctx, cancel := context.WithTimeout(ctx, ExportTimeout)
 	defer cancel()
 
 	return r.ForceFlush(ctx)

cli/command/telemetry_docker.go

@@ -109,7 +109,7 @@ func dockerSpanExporter(ctx context.Context, cli Cli) []sdktrace.TracerProviderO
 		otel.Handle(err)
 		return nil
 	}
-	return []sdktrace.TracerProviderOption{sdktrace.WithBatcher(exp, sdktrace.WithExportTimeout(exportTimeout))}
+	return []sdktrace.TracerProviderOption{sdktrace.WithBatcher(exp, sdktrace.WithExportTimeout(ExportTimeout))}
 }
 
 func dockerMetricExporter(ctx context.Context, cli Cli) []sdkmetric.Option {

cli/command/telemetry_utils.go

@@ -85,7 +85,7 @@ func startCobraCommandTimer(mp metric.MeterProvider, attrs []attribute.KeyValue)
 	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)
+		ctx, cancel := context.WithTimeout(context.Background(), ExportTimeout)
 		defer cancel()
 
 		duration := float64(time.Since(start)) / float64(time.Millisecond)