Move SystemManager & span handling to root command (#2026)

Moving the creation & clean up of the SystemManager and root span avoids
having the same code copied between all the commands. This also gives
the opportunity to give the root spans a more coherent name rather than
the code location.

This also fixes an issue with the wasm command which didn't apply the
root commands PersistentPreRun.

Fixes #2023

Author: wjam

cmd/bacalhau/cancel.go

@@ -4,9 +4,6 @@ import (
 	"fmt"
 
 	"github.com/filecoin-project/bacalhau/pkg/bacerrors"
-	"github.com/filecoin-project/bacalhau/pkg/telemetry"
-
-	"github.com/filecoin-project/bacalhau/pkg/system"
 	"github.com/filecoin-project/bacalhau/pkg/util/templates"
 	"github.com/spf13/cobra"
 	"k8s.io/kubectl/pkg/util/i18n"
@@ -52,14 +49,8 @@ func newCancelCmd() *cobra.Command {
 }
 
 func cancel(cmd *cobra.Command, cmdArgs []string, options *CancelOptions) error {
-	cm := system.NewCleanupManager()
-	defer cm.Cleanup()
 	ctx := cmd.Context()
 
-	ctx, span := system.NewRootSpan(ctx, system.GetTracer(), "cmd/bacalhau.cancel")
-	defer span.End()
-	cm.RegisterCallback(telemetry.Cleanup)
-
 	var err error
 
 	requestedJobID := cmdArgs[0]

cmd/bacalhau/create.go

@@ -8,10 +8,8 @@ import (
 	"strings"
 	"time"
 
-	"github.com/filecoin-project/bacalhau/pkg/downloader/util"
-	"github.com/filecoin-project/bacalhau/pkg/telemetry"
-
 	"github.com/filecoin-project/bacalhau/pkg/bacerrors"
+	"github.com/filecoin-project/bacalhau/pkg/downloader/util"
 	jobutils "github.com/filecoin-project/bacalhau/pkg/job"
 	"github.com/filecoin-project/bacalhau/pkg/model"
 	"github.com/filecoin-project/bacalhau/pkg/system"
@@ -83,13 +81,9 @@ func newCreateCmd() *cobra.Command {
 }
 
 func create(cmd *cobra.Command, cmdArgs []string, OC *CreateOptions) error { //nolint:funlen,gocyclo
-	cm := system.NewCleanupManager()
-	defer cm.Cleanup()
 	ctx := cmd.Context()
 
-	ctx, rootSpan := system.NewRootSpan(ctx, system.GetTracer(), "cmd/bacalhau.create")
-	defer rootSpan.End()
-	cm.RegisterCallback(telemetry.Cleanup)
+	cm := cmd.Context().Value(systemManagerKey).(*system.CleanupManager)
 
 	// Custom unmarshaller
 	// https://stackoverflow.com/questions/70635636/unmarshaling-yaml-into-different-struct-based-off-yaml-field?rq=1

cmd/bacalhau/describe.go

@@ -5,8 +5,6 @@ import (
 	"fmt"
 
 	"github.com/filecoin-project/bacalhau/pkg/bacerrors"
-	"github.com/filecoin-project/bacalhau/pkg/system"
-	"github.com/filecoin-project/bacalhau/pkg/telemetry"
 	"github.com/filecoin-project/bacalhau/pkg/util/templates"
 	"github.com/spf13/cobra"
 	"k8s.io/kubectl/pkg/util/i18n"
@@ -78,18 +76,12 @@ func newDescribeCmd() *cobra.Command {
 }
 
 func describe(cmd *cobra.Command, cmdArgs []string, OD *DescribeOptions) error {
-	cm := system.NewCleanupManager()
-	defer cm.Cleanup()
 	ctx := cmd.Context()
 
 	if err := cmd.ParseFlags(cmdArgs[1:]); err != nil {
 		Fatal(cmd, fmt.Sprintf("Failed to parse flags: %v\n", err), 1)
 	}
 
-	ctx, rootSpan := system.NewRootSpan(ctx, system.GetTracer(), "cmd/bacalhau.describe")
-	defer rootSpan.End()
-	cm.RegisterCallback(telemetry.Cleanup)
-
 	var err error
 	inputJobID := cmdArgs[0]
 	if inputJobID == "" {

cmd/bacalhau/devstack.go

@@ -127,12 +127,9 @@ func newDevStackCmd() *cobra.Command {
 }
 
 func runDevstack(cmd *cobra.Command, ODs *devstack.DevStackOptions, OS *ServeOptions, IsNoop bool) error {
-	cm := system.NewCleanupManager()
-	defer cm.Cleanup()
 	ctx := cmd.Context()
 
-	ctx, rootSpan := system.NewRootSpan(ctx, system.GetTracer(), "cmd/bacalhau.runDevstack")
-	defer rootSpan.End()
+	cm := ctx.Value(systemManagerKey).(*system.CleanupManager)
 
 	if config.DevstackShouldWriteEnvFile() {
 		cm.RegisterCallback(cleanupDevstackDotEnv)

cmd/bacalhau/docker_run.go

@@ -9,9 +9,7 @@ import (
 	jobutils "github.com/filecoin-project/bacalhau/pkg/job"
 	"github.com/filecoin-project/bacalhau/pkg/model"
 	"github.com/filecoin-project/bacalhau/pkg/system"
-	"github.com/filecoin-project/bacalhau/pkg/telemetry"
 	"github.com/filecoin-project/bacalhau/pkg/util/templates"
-	"github.com/filecoin-project/bacalhau/pkg/version"
 	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
 	"k8s.io/kubectl/pkg/util/i18n"
@@ -123,17 +121,9 @@ func NewDockerRunOptions() *DockerRunOptions {
 
 func newDockerCmd() *cobra.Command {
 	dockerCmd := &cobra.Command{
-		Use:   "docker",
-		Short: "Run a docker job on the network (see run subcommand)",
-		PersistentPreRunE: func(cmd *cobra.Command, _ []string) error {
-			// Check that the server version is compatible with the client version
-			serverVersion, _ := GetAPIClient().Version(cmd.Context()) // Ok if this fails, version validation will skip
-			if err := ensureValidVersion(cmd.Context(), version.Get(), serverVersion); err != nil {
-				cmd.Println(err.Error())
-				return err
-			}
-			return nil
-		},
+		Use:               "docker",
+		Short:             "Run a docker job on the network (see run subcommand)",
+		PersistentPreRunE: checkVersion,
 	}
 
 	dockerCmd.AddCommand(newDockerRunCmd())
@@ -280,13 +270,9 @@ func newDockerRunCmd() *cobra.Command { //nolint:funlen
 }
 
 func dockerRun(cmd *cobra.Command, cmdArgs []string, ODR *DockerRunOptions) error {
-	cm := system.NewCleanupManager()
-	defer cm.Cleanup()
 	ctx := cmd.Context()
 
-	ctx, rootSpan := system.NewRootSpan(ctx, system.GetTracer(), "cmd/bacalhau.dockerRun")
-	defer rootSpan.End()
-	cm.RegisterCallback(telemetry.Cleanup)
+	cm := cmd.Context().Value(systemManagerKey).(*system.CleanupManager)
 
 	j, err := CreateJob(cmdArgs, ODR)
 	if err != nil {

cmd/bacalhau/get.go

@@ -5,8 +5,6 @@ import (
 
 	"github.com/filecoin-project/bacalhau/pkg/downloader/util"
 	"github.com/filecoin-project/bacalhau/pkg/model"
-	"github.com/filecoin-project/bacalhau/pkg/telemetry"
-
 	"github.com/filecoin-project/bacalhau/pkg/system"
 	"github.com/filecoin-project/bacalhau/pkg/util/templates"
 	"github.com/pkg/errors"
@@ -60,13 +58,9 @@ func newGetCmd() *cobra.Command {
 }
 
 func get(cmd *cobra.Command, cmdArgs []string, OG *GetOptions) error {
-	cm := system.NewCleanupManager()
-	defer cm.Cleanup()
 	ctx := cmd.Context()
 
-	ctx, span := system.NewRootSpan(ctx, system.GetTracer(), "cmd/bacalhau.get")
-	defer span.End()
-	cm.RegisterCallback(telemetry.Cleanup)
+	cm := cmd.Context().Value(systemManagerKey).(*system.CleanupManager)
 
 	var err error
 

cmd/bacalhau/id.go

@@ -5,8 +5,6 @@ import (
 	"os"
 
 	"github.com/filecoin-project/bacalhau/pkg/libp2p"
-	"github.com/filecoin-project/bacalhau/pkg/system"
-	"github.com/filecoin-project/bacalhau/pkg/telemetry"
 	"github.com/spf13/cobra"
 )
 
@@ -31,11 +29,6 @@ func newIDCmd() *cobra.Command {
 }
 
 func id(_ *cobra.Command, OS *ServeOptions) error {
-	// Cleanup manager ensures that resources are freed before exiting:
-	cm := system.NewCleanupManager()
-	cm.RegisterCallback(telemetry.Cleanup)
-	defer cm.Cleanup()
-
 	libp2pHost, err := libp2p.NewHost(OS.SwarmPort)
 	if err != nil {
 		return err

cmd/bacalhau/list.go

@@ -8,7 +8,6 @@ import (
 	"github.com/filecoin-project/bacalhau/pkg/job"
 	"github.com/filecoin-project/bacalhau/pkg/model"
 	"github.com/filecoin-project/bacalhau/pkg/system"
-	"github.com/filecoin-project/bacalhau/pkg/telemetry"
 	"github.com/filecoin-project/bacalhau/pkg/util/templates"
 	"github.com/jedib0t/go-pretty/v6/table"
 	"github.com/rs/zerolog/log"
@@ -148,14 +147,8 @@ func (c *ColumnEnum) Set(v string) error {
 }
 
 func list(cmd *cobra.Command, OL *ListOptions) error {
-	cm := system.NewCleanupManager()
-	defer cm.Cleanup()
 	ctx := cmd.Context()
 
-	ctx, rootSpan := system.NewRootSpan(ctx, system.GetTracer(), "cmd/bacalhau.list")
-	defer rootSpan.End()
-	cm.RegisterCallback(telemetry.Cleanup)
-
 	log.Ctx(ctx).Debug().Msgf("Table filter flag set to: %s", OL.IDFilter)
 	log.Ctx(ctx).Debug().Msgf("Table limit flag set to: %d", OL.MaxJobs)
 	log.Ctx(ctx).Debug().Msgf("Table output format flag set to: %s", OL.OutputFormat)

cmd/bacalhau/root.go

@@ -2,16 +2,20 @@ package bacalhau
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"strconv"
 	"strings"
 
 	"github.com/filecoin-project/bacalhau/pkg/config"
 	"github.com/filecoin-project/bacalhau/pkg/logger"
 	"github.com/filecoin-project/bacalhau/pkg/system"
+	"github.com/filecoin-project/bacalhau/pkg/telemetry"
+	"github.com/filecoin-project/bacalhau/pkg/version"
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
+	"go.opentelemetry.io/otel/trace"
 )
 
 var apiHost string
@@ -54,7 +58,29 @@ func NewRootCmd() *cobra.Command {
 		Short: "Compute over data",
 		Long:  `Compute over data`,
 		PersistentPreRun: func(cmd *cobra.Command, args []string) {
+			ctx := cmd.Context()
+
 			logger.ConfigureLogging(loggingMode)
+
+			cm := system.NewCleanupManager()
+			cm.RegisterCallback(telemetry.Cleanup)
+			ctx = context.WithValue(ctx, systemManagerKey, cm)
+
+			var names []string
+			root := cmd
+			for ; root.HasParent(); root = root.Parent() {
+				names = append([]string{root.Name()}, names...)
+			}
+			name := fmt.Sprintf("bacalhau.%s", strings.Join(names, "."))
+			ctx, span := system.NewRootSpan(ctx, system.GetTracer(), name)
+			ctx = context.WithValue(ctx, spanKey, span)
+
+			cmd.SetContext(ctx)
+		},
+		PersistentPostRun: func(cmd *cobra.Command, args []string) {
+			ctx := cmd.Context()
+			ctx.Value(spanKey).(trace.Span).End()
+			ctx.Value(systemManagerKey).(*system.CleanupManager).Cleanup()
 		},
 	}
 	// ====== Start a job
@@ -162,3 +188,28 @@ func Execute() {
 		Fatal(RootCmd, err.Error(), 1)
 	}
 }
+
+type contextKey struct {
+	name string
+}
+
+var systemManagerKey = contextKey{name: "context key for storing the system manager"}
+var spanKey = contextKey{name: "context key for storing the root span"}
+
+func checkVersion(cmd *cobra.Command, args []string) error {
+	// corba doesn't do PersistentPreRun{,E} chaining yet
+	// https://github.com/spf13/cobra/issues/252
+	root := cmd
+	for ; root.HasParent(); root = root.Parent() {
+	}
+	root.PersistentPreRun(cmd, args)
+
+	// Check that the server version is compatible with the client version
+	serverVersion, _ := GetAPIClient().Version(cmd.Context()) // Ok if this fails, version validation will skip
+	if err := ensureValidVersion(cmd.Context(), version.Get(), serverVersion); err != nil {
+		Fatal(cmd, fmt.Sprintf("version validation failed: %s", err), 1)
+		return err
+	}
+
+	return nil
+}

cmd/bacalhau/run.go

@@ -1,27 +1,15 @@
 package bacalhau
 
 import (
-	"fmt"
-
-	"github.com/filecoin-project/bacalhau/pkg/version"
 	"github.com/spf13/cobra"
 )
 
 func newRunCmd() *cobra.Command {
 	runCmd := &cobra.Command{
-		Use:    "run",
-		Short:  "Run a job on the network (see subcommands for supported flavors)",
-		PreRun: applyPorcelainLogLevel,
-		PersistentPreRunE: func(cmd *cobra.Command, _ []string) error {
-			// Check that the server version is compatible with the client version
-			serverVersion, _ := GetAPIClient().Version(cmd.Context()) // Ok if this fails, version validation will skip
-			if err := ensureValidVersion(cmd.Context(), version.Get(), serverVersion); err != nil {
-				Fatal(cmd, fmt.Sprintf("version validation failed: %s", err), 1)
-				return err
-			}
-
-			return nil
-		},
+		Use:               "run",
+		Short:             "Run a job on the network (see subcommands for supported flavors)",
+		PreRun:            applyPorcelainLogLevel,
+		PersistentPreRunE: checkVersion,
 	}
 	runCmd.AddCommand(newRunPythonCmd())
 	return runCmd

cmd/bacalhau/run_python.go

@@ -5,10 +5,8 @@ import (
 	"time"
 
 	"github.com/filecoin-project/bacalhau/pkg/downloader/util"
-	"github.com/filecoin-project/bacalhau/pkg/model"
-	"github.com/filecoin-project/bacalhau/pkg/telemetry"
-
 	"github.com/filecoin-project/bacalhau/pkg/job"
+	"github.com/filecoin-project/bacalhau/pkg/model"
 	"github.com/filecoin-project/bacalhau/pkg/storage/inline"
 	"github.com/filecoin-project/bacalhau/pkg/system"
 	"github.com/filecoin-project/bacalhau/pkg/util/templates"
@@ -165,13 +163,9 @@ func newRunPythonCmd() *cobra.Command {
 }
 
 func runPython(cmd *cobra.Command, cmdArgs []string, OLR *LanguageRunOptions) error {
-	cm := system.NewCleanupManager()
-	defer cm.Cleanup()
 	ctx := cmd.Context()
 
-	ctx, rootSpan := system.NewRootSpan(ctx, system.GetTracer(), "cmd/bacalhau.runPython")
-	defer rootSpan.End()
-	cm.RegisterCallback(telemetry.Cleanup)
+	cm := cmd.Context().Value(systemManagerKey).(*system.CleanupManager)
 
 	// error if determinism is false
 	if !OLR.Deterministic {

cmd/bacalhau/serve.go

@@ -10,16 +10,14 @@ import (
 	"time"
 
 	"github.com/filecoin-project/bacalhau/pkg/compute/capacity"
+	"github.com/filecoin-project/bacalhau/pkg/ipfs"
 	"github.com/filecoin-project/bacalhau/pkg/jobstore/inmemory"
 	"github.com/filecoin-project/bacalhau/pkg/libp2p"
 	"github.com/filecoin-project/bacalhau/pkg/libp2p/rcmgr"
 	"github.com/filecoin-project/bacalhau/pkg/logger"
-	filecoinlotus "github.com/filecoin-project/bacalhau/pkg/publisher/filecoin_lotus"
-	"github.com/filecoin-project/bacalhau/pkg/telemetry"
-
-	"github.com/filecoin-project/bacalhau/pkg/ipfs"
 	"github.com/filecoin-project/bacalhau/pkg/model"
 	"github.com/filecoin-project/bacalhau/pkg/node"
+	filecoinlotus "github.com/filecoin-project/bacalhau/pkg/publisher/filecoin_lotus"
 	"github.com/filecoin-project/bacalhau/pkg/system"
 	"github.com/filecoin-project/bacalhau/pkg/util/templates"
 	"github.com/multiformats/go-multiaddr"
@@ -304,18 +302,13 @@ func newServeCmd() *cobra.Command {
 
 //nolint:funlen,gocyclo
 func serve(cmd *cobra.Command, OS *ServeOptions) error {
-	// Cleanup manager ensures that resources are freed before exiting:
-	cm := system.NewCleanupManager()
-	cm.RegisterCallback(telemetry.Cleanup)
-	defer cm.Cleanup()
+	cm := cmd.Context().Value(systemManagerKey).(*system.CleanupManager)
 
+	// TODO this should be for all commands
 	// Context ensures main goroutine waits until killed with ctrl+c:
 	ctx, cancel := signal.NotifyContext(cmd.Context(), ShutdownSignals...)
 	defer cancel()
 
-	ctx, rootSpan := system.NewRootSpan(ctx, system.GetTracer(), "cmd/bacalhau.serve")
-	defer rootSpan.End()
-
 	isComputeNode, isRequesterNode := false, false
 	for _, nodeType := range OS.NodeType {
 		if nodeType == "compute" {