Add support for passing a context when cleaning up

This adds a second type of clean up function - one that takes a
`context.Context`. This makes it easier to avoid problems of using
a cancelled context when cleaning up, as the logic doesn't need to be
spread about, and makes it easy to migrate a portion of log statements
that aren't using `log.Ctx` to do so.

Part of #2001

Author: wjam

cmd/bacalhau/base_test.go

@@ -52,6 +52,6 @@ func (s *BaseSuite) SetupTest() {
 func (s *BaseSuite) TearDownTest() {
 	Fatal = FatalErrorHandler
 	if s.node != nil {
-		s.node.CleanupManager.Cleanup()
+		s.node.CleanupManager.Cleanup(context.Background())
 	}
 }

cmd/bacalhau/root.go

@@ -80,7 +80,7 @@ func NewRootCmd() *cobra.Command {
 		PersistentPostRun: func(cmd *cobra.Command, args []string) {
 			ctx := cmd.Context()
 			ctx.Value(spanKey).(trace.Span).End()
-			ctx.Value(systemManagerKey).(*system.CleanupManager).Cleanup()
+			ctx.Value(systemManagerKey).(*system.CleanupManager).Cleanup(cmd.Context())
 		},
 	}
 	// ====== Start a job
@@ -154,12 +154,12 @@ func Execute() {
 
 	err := viper.BindEnv("API_HOST")
 	if err != nil {
-		log.Fatal().Msgf("API_HOST was set, but could not bind.")
+		log.Ctx(RootCmd.Context()).Fatal().Msgf("API_HOST was set, but could not bind.")
 	}
 
 	err = viper.BindEnv("API_PORT")
 	if err != nil {
-		log.Fatal().Msgf("API_PORT was set, but could not bind.")
+		log.Ctx(RootCmd.Context()).Fatal().Msgf("API_PORT was set, but could not bind.")
 	}
 
 	viper.AutomaticEnv()
@@ -174,7 +174,7 @@ func Execute() {
 		var parseErr error
 		apiPort, parseErr = strconv.Atoi(envAPIPort.(string))
 		if parseErr != nil {
-			log.Fatal().Msgf("could not parse API_PORT into an int. %s", envAPIPort)
+			log.Ctx(RootCmd.Context()).Fatal().Msgf("could not parse API_PORT into an int. %s", envAPIPort)
 		}
 	}
 

cmd/bacalhau/serve.go

@@ -481,7 +481,7 @@ func ipfsClient(ctx context.Context, OS *ServeOptions, cm *system.CleanupManager
 		if err != nil {
 			return ipfs.Client{}, fmt.Errorf("error creating IPFS node: %s", err)
 		}
-		cm.RegisterCallback(ipfsNode.Close)
+		cm.RegisterCallbackWithContext(ipfsNode.Close)
 		client := ipfsNode.Client()
 
 		swarmAddresses, err := client.SwarmAddresses(ctx)

cmd/bacalhau/serve_test.go

@@ -44,7 +44,9 @@ func (s *ServeSuite) SetupTest() {
 	s.Require().NoError(system.InitConfigForTesting(s.T()))
 
 	cm := system.NewCleanupManager()
-	s.T().Cleanup(cm.Cleanup)
+	s.T().Cleanup(func() {
+		cm.Cleanup(context.Background())
+	})
 
 	node, err := ipfs.NewLocalNode(context.Background(), cm, []string{})
 	s.NoError(err)
@@ -69,7 +71,7 @@ func (s *ServeSuite) writeToServeChannel(rootCmd *cobra.Command, port int) {
 
 	rootCmd.SetArgs(args)
 
-	log.Trace().Msgf("Command to execute: %v", rootCmd.CalledAs())
+	log.Ctx(rootCmd.Context()).Trace().Msgf("Command to execute: %v", rootCmd.CalledAs())
 
 	_, err := rootCmd.ExecuteC()
 	s.Require().NoError(err)

dashboard/api/cmd/dashboard/serve.go

@@ -87,7 +87,7 @@ func serve(cmd *cobra.Command, options *ServeOptions) error {
 	// Cleanup manager ensures that resources are freed before exiting:
 	cm := system.NewCleanupManager()
 	cm.RegisterCallback(telemetry.Cleanup)
-	defer cm.Cleanup()
+	defer cm.Cleanup(cmd.Context())
 	ctx := cmd.Context()
 
 	if options.ServerOptions.JWTSecret == "" {
@@ -121,7 +121,7 @@ func serve(cmd *cobra.Command, options *ServeOptions) error {
 	if err != nil {
 		return err
 	}
-	cm.RegisterCallback(model.Stop)
+	cm.RegisterCallbackWithContext(model.Stop)
 
 	// Start transport layer
 	err = libp2p.ConnectToPeersContinuously(ctx, cm, libp2pHost, peers)

dashboard/api/pkg/model/api.go

@@ -169,9 +169,9 @@ func (api *ModelAPI) Start(ctx context.Context) error {
 	return nil
 }
 
-func (api *ModelAPI) Stop() error {
+func (api *ModelAPI) Stop(ctx context.Context) error {
 	if api.cleanupFunc != nil {
-		api.cleanupFunc(context.Background())
+		api.cleanupFunc(ctx)
 	}
 	return nil
 }

dashboard/api/pkg/server/utils.go

@@ -34,21 +34,6 @@ func GetRequestBody[T any](w http.ResponseWriter, r *http.Request) (*T, error) {
 	return &requestBody, nil
 }
 
-func HTTPGet[T any](url string) (T, error) {
-	var data T
-	//nolint:gosec
-	resp, err := http.Get(url)
-	if err != nil {
-		return data, err
-	}
-	err = json.NewDecoder(resp.Body).Decode(&data)
-	if err != nil {
-		return data, err
-	}
-	resp.Body.Close()
-	return data, nil
-}
-
 func generateJWT(
 	secret string,
 	username string,

go.mod

@@ -94,7 +94,6 @@ require (
 	golang.org/x/crypto v0.5.0
 	golang.org/x/exp v0.0.0-20221106115401-f9659909a136
 	golang.org/x/mod v0.7.0
-	golang.org/x/net v0.7.0
 	k8s.io/apimachinery v0.26.1
 	k8s.io/kubectl v0.26.1
 	modernc.org/sqlite v1.20.2
@@ -333,6 +332,7 @@ require (
 	go.uber.org/dig v1.14.1 // indirect
 	go.uber.org/fx v1.17.1 // indirect
 	go4.org v0.0.0-20201209231011-d4a079459e60 // indirect
+	golang.org/x/net v0.7.0 // indirect
 	golang.org/x/oauth2 v0.1.0 // indirect
 	golang.org/x/sync v0.1.0 // indirect
 	golang.org/x/sys v0.5.0 // indirect

pkg/config/config.go

@@ -55,17 +55,12 @@ type contextKey int
 
 const (
 	getVolumeSizeRequestTimeoutKey contextKey = iota
-	downloadCidRequestTimeoutKey
 )
 
 const (
 	// by default we wait 2 minutes for the IPFS network to resolve a CID
 	// tests will override this using config.SetVolumeSizeRequestTimeout(2)
-	getVolumeSizeRequestTimeout time.Duration = 2 * time.Minute
-
-	// by default we wait 5 minutes for the IPFS network to download a CID
-	// tests will override this using config.SetVolumeSizeRequestTimeout(2)
-	downloadCidRequestTimeout time.Duration = 5 * time.Minute
+	getVolumeSizeRequestTimeout = 2 * time.Minute
 )
 
 // how long do we wait for a volume size request to timeout
@@ -87,19 +82,6 @@ func SetVolumeSizeRequestTimeout(ctx context.Context, value time.Duration) conte
 	return context.WithValue(ctx, getVolumeSizeRequestTimeoutKey, value)
 }
 
-// how long do we wait for a cid to download
-func GetDownloadCidRequestTimeout(ctx context.Context) time.Duration {
-	value := ctx.Value(downloadCidRequestTimeoutKey)
-	if value == nil {
-		value = downloadCidRequestTimeout
-	}
-	return value.(time.Duration)
-}
-
-func SetDownloadCidRequestTimeout(ctx context.Context, value time.Duration) context.Context {
-	return context.WithValue(ctx, downloadCidRequestTimeoutKey, value)
-}
-
 // by default we wait 5 minutes for a URL to download
 // tests will override this using config.SetDownloadURLRequestTimeoutSeconds(2)
 var downloadURLRequestTimeoutSeconds int64 = 300
@@ -135,7 +117,7 @@ func GetConfigPath() string {
 		// e.g. /home/francesca/.bacalhau
 		dirname, err := os.UserHomeDir()
 		if err != nil {
-			log.Fatal().Err(err)
+			log.Fatal().Err(err).Send()
 		}
 		d = filepath.Join(dirname, suffix)
 	} else {
@@ -144,7 +126,7 @@ func GetConfigPath() string {
 	}
 	// create dir if not exists
 	if err := os.MkdirAll(d, util.OS_USER_RWX); err != nil {
-		log.Fatal().Err(err)
+		log.Fatal().Err(err).Send()
 	}
 	return d
 }

pkg/devstack/devstack.go

@@ -131,7 +131,7 @@ func NewDevStack(
 			return nil, err
 		}
 
-		cm.RegisterCallback(lotus.Close)
+		cm.RegisterCallbackWithContext(lotus.Close)
 
 		if err := lotus.start(ctx); err != nil { //nolint:govet
 			return nil, err
@@ -323,9 +323,9 @@ func NewDevStack(
 	}
 
 	// only start profiling after we've set everything up!
-	profiler := StartProfiling(options.CPUProfilingFile, options.MemoryProfilingFile)
+	profiler := StartProfiling(ctx, options.CPUProfilingFile, options.MemoryProfilingFile)
 	if profiler != nil {
-		cm.RegisterCallback(profiler.Close)
+		cm.RegisterCallbackWithContext(profiler.Close)
 	}
 
 	return &DevStack{

pkg/devstack/devstack_ipfs.go

@@ -47,20 +47,3 @@ func NewDevStackIPFS(ctx context.Context, cm *system.CleanupManager, count int)
 
 	return stack, nil
 }
-
-func (stack *DevStackIPFS) PrintNodeInfo() {
-	logString := `
--------------------------------
-ipfs
--------------------------------
-
-command="add -q testdata/grep_file.txt"
-	`
-	for _, node := range stack.IPFSClients {
-		logString += fmt.Sprintf(`
-cid=$(ipfs --api %s ipfs $command)
-curl -XPOST %s`, node.APIAddress(), node.APIAddress())
-	}
-
-	log.Trace().Msg(logString + "\n")
-}

pkg/devstack/lotus.go

@@ -2,6 +2,7 @@ package devstack
 
 import (
 	"archive/tar"
+	"context"
 	"fmt"
 	"io"
 	"os"
@@ -18,7 +19,6 @@ import (
 	"github.com/filecoin-project/bacalhau/pkg/util/closer"
 	"github.com/hashicorp/go-multierror"
 	"github.com/rs/zerolog/log"
-	"golang.org/x/net/context"
 )
 
 const defaultImage = "ghcr.io/bacalhau-project/lotus-filecoin-image:v0.0.2"
@@ -111,7 +111,7 @@ func (l *LotusNode) start(ctx context.Context) error {
 	}
 
 	if err := l.waitForLotusToBeHealthy(ctx); err != nil {
-		if err := l.Close(); err != nil { //nolint:govet
+		if err := l.Close(ctx); err != nil { //nolint:govet
 			log.Ctx(ctx).Err(err).Msgf(`Problem occurred when giving up waiting for Lotus to become healthy`)
 		}
 		return err
@@ -196,12 +196,12 @@ ListenAddress = "/ip4/0.0.0.0/tcp/%s/http"
 	return nil
 }
 
-func (l *LotusNode) Close() error {
+func (l *LotusNode) Close(ctx context.Context) error {
 	var errs error
 
 	defer closer.CloseWithLogOnError("Docker client", l.client)
 	if l.container != "" {
-		if err := l.client.RemoveContainer(context.Background(), l.container); err != nil {
+		if err := l.client.RemoveContainer(ctx, l.container); err != nil {
 			errs = multierror.Append(errs, err)
 		}
 	}

pkg/devstack/profiler.go

@@ -1,7 +1,7 @@
 package devstack
 
 import (
-	"io"
+	"context"
 	"os"
 	"runtime"
 	"runtime/pprof"
@@ -15,31 +15,31 @@ type profiler struct {
 	memoryFile string
 }
 
-func StartProfiling(cpuFile, memoryFile string) io.Closer {
+func StartProfiling(ctx context.Context, cpuFile, memoryFile string) CloserWithContext {
 	// do a GC before we start profiling
 	runtime.GC()
 
-	log.Trace().Msg("============= STARTING PROFILING ============")
+	log.Ctx(ctx).Trace().Msg("============= STARTING PROFILING ============")
 
 	var f *os.File
 	if cpuFile != "" {
 		var err error
 		f, err = os.Create(cpuFile)
 		if err != nil {
-			log.Debug().Err(err).Str("Path", cpuFile).Msg("could not create CPU profile")
+			log.Ctx(ctx).Debug().Err(err).Str("Path", cpuFile).Msg("could not create CPU profile")
 			return nil
 		}
 		if err := pprof.StartCPUProfile(f); err != nil {
 			closer.CloseWithLogOnError(cpuFile, f)
-			log.Debug().Err(err).Msg("could not start CPU profile")
+			log.Ctx(ctx).Debug().Err(err).Msg("could not start CPU profile")
 			return nil
 		}
 	}
 
 	return &profiler{cpuFile: f, memoryFile: memoryFile}
 }
 
-func (p *profiler) Close() error {
+func (p *profiler) Close(ctx context.Context) error {
 	// stop profiling now, just before we clean up, if we're profiling.
 	log.Trace().Msg("============= STOPPING PROFILING ============")
 	if p.cpuFile != nil {
@@ -50,18 +50,23 @@ func (p *profiler) Close() error {
 	if p.memoryFile != "" {
 		f, err := os.Create(p.memoryFile)
 		if err != nil {
-			log.Debug().Err(err).Str("Path", p.memoryFile).Msg("could not create memory profile")
+			log.Ctx(ctx).Debug().Err(err).Str("Path", p.memoryFile).Msg("could not create memory profile")
 			return nil
 		}
 		defer closer.CloseWithLogOnError(p.memoryFile, f) // error handling omitted for example
 
 		runtime.GC() // get up-to-date statistics
 		if err := pprof.WriteHeapProfile(f); err != nil {
-			log.Debug().Err(err).Msg("could not write memory profile")
+			log.Ctx(ctx).Debug().Err(err).Msg("could not write memory profile")
 		}
 	}
 
 	return nil
 }
 
-var _ io.Closer = (*profiler)(nil)
+var _ CloserWithContext = (*profiler)(nil)
+
+type CloserWithContext interface {
+	// Close closes the resource.
+	Close(context.Context) error
+}

pkg/downloader/download_test.go

@@ -41,7 +41,9 @@ func (ds *DownloaderSuite) SetupSuite() {
 // Before each test
 func (ds *DownloaderSuite) SetupTest() {
 	ds.cm = system.NewCleanupManager()
-	ds.T().Cleanup(ds.cm.Cleanup)
+	ds.T().Cleanup(func() {
+		ds.cm.Cleanup(context.Background())
+	})
 
 	ctx, cancel := context.WithCancel(context.Background())
 	ds.T().Cleanup(cancel)

pkg/downloader/ipfs/downloader.go

@@ -40,7 +40,7 @@ func (ipfsDownloader *Downloader) FetchResult(ctx context.Context, result model.
 		return err
 	}
 	defer func() {
-		if closeErr := n.Close(); closeErr != nil {
+		if closeErr := n.Close(ctx); closeErr != nil {
 			log.Ctx(ctx).Error().Err(closeErr).Msg("Failed to close IPFS node")
 		}
 	}()