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

OTEL providers provided by this library to plugins were not properly
shut down. Metrics from the CLI were forcibly flushed in a separate area
and the `buildx` command does its own flushing, but this didn't apply to
traces.

Add shutdown on post run for plugins using `OnFinalize`. This will allow
compose to use the tracer provider from this library rather than roll
its own custom one.

Signed-off-by: Jonathan A. Sternberg <jonathan.sternberg@docker.com>

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(func() { finalizeHook(dockerCli) })
 		})
 		return retErr
 	}
@@ -83,6 +90,59 @@ func RunPlugin(dockerCli *command.DockerCli, plugin *cobra.Command, meta metadat
 	return cmd.Execute()
 }
 
+// finalizeHook is 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 finalizeHook(dockerCli *command.DockerCli) {
+	// 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 +272,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)