fix: Always cleanup temporary directory (#41)

Previously, if the command was interrupted, the
temporary directory would be left behind.

cleanup/cleaner.go

@@ -0,0 +1,63 @@
+// Copyright 2021 D2iQ, Inc. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package cleanup
+
+import (
+	"context"
+	"os"
+	"os/signal"
+	"sync"
+)
+
+type Cleaner interface {
+	Cleanup()
+	AddCleanupFn(f func())
+}
+
+func NewCleaner() Cleaner {
+	return &cleaner{}
+}
+
+type cleaner struct {
+	sigNotifier sync.Once
+	mu          sync.RWMutex
+	cleanups    []func()
+}
+
+func (c *cleaner) setupSignalHandling() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
+	go func() {
+		<-ctx.Done()
+		stop()
+		c.doCleanup()
+		p, _ := os.FindProcess(os.Getpid())
+		if err := p.Signal(os.Interrupt); err != nil {
+			panic(err)
+		}
+	}()
+}
+
+func (c *cleaner) Cleanup() {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	c.doCleanup()
+}
+
+func (c *cleaner) doCleanup() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	for _, f := range c.cleanups {
+		f()
+	}
+}
+
+func (c *cleaner) AddCleanupFn(f func()) {
+	c.sigNotifier.Do(c.setupSignalHandling)
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.cleanups = append(c.cleanups, f)
+}

cmd/create/imagebundle/image_bundle.go

@@ -15,6 +15,7 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/mesosphere/mindthegap/archive"
+	"github.com/mesosphere/mindthegap/cleanup"
 	"github.com/mesosphere/mindthegap/config"
 	"github.com/mesosphere/mindthegap/docker/registry"
 	"github.com/mesosphere/mindthegap/skopeo"
@@ -62,12 +63,16 @@ func NewCommand(out output.Output) *cobra.Command {
 				return fmt.Errorf("failed to determine where to create temporary directory: %w", err)
 			}
 
+			cleaner := cleanup.NewCleaner()
+			defer cleaner.Cleanup()
+
 			tempDir, err := os.MkdirTemp(filepath.Dir(outputFileAbs), ".image-bundle-*")
 			if err != nil {
 				out.EndOperation(false)
 				return fmt.Errorf("failed to create temporary directory: %w", err)
 			}
-			defer os.RemoveAll(tempDir)
+			cleaner.AddCleanupFn(func() { _ = os.RemoveAll(tempDir) })
+
 			out.EndOperation(true)
 
 			out.StartOperation("Starting temporary Docker registry")
@@ -85,7 +90,7 @@ func NewCommand(out output.Output) *cobra.Command {
 			out.EndOperation(true)
 
 			skopeoRunner, skopeoCleanup := skopeo.NewRunner()
-			defer func() { _ = skopeoCleanup() }()
+			cleaner.AddCleanupFn(func() { _ = skopeoCleanup() })
 
 			for registryName, registryConfig := range cfg {
 				var skopeoOpts []skopeo.SkopeoOption

cmd/importcmd/imagebundle/image_bundle.go

@@ -13,6 +13,7 @@ import (
 	"github.com/mholt/archiver/v3"
 	"github.com/spf13/cobra"
 
+	"github.com/mesosphere/mindthegap/cleanup"
 	"github.com/mesosphere/mindthegap/config"
 	"github.com/mesosphere/mindthegap/containerd"
 	"github.com/mesosphere/mindthegap/docker/registry"
@@ -27,13 +28,16 @@ func NewCommand(out output.Output) *cobra.Command {
 	cmd := &cobra.Command{
 		Use: "image-bundle",
 		RunE: func(cmd *cobra.Command, args []string) error {
+			cleaner := cleanup.NewCleaner()
+			defer cleaner.Cleanup()
+
 			out.StartOperation("Creating temporary directory")
 			tempDir, err := os.MkdirTemp("", ".image-bundle-*")
 			if err != nil {
 				out.EndOperation(false)
 				return fmt.Errorf("failed to create temporary directory: %w", err)
 			}
-			defer os.RemoveAll(tempDir)
+			cleaner.AddCleanupFn(func() { _ = os.RemoveAll(tempDir) })
 			out.EndOperation(true)
 
 			out.StartOperation("Unarchiving image bundle")

cmd/push/imagebundle/image_bundle.go

@@ -13,6 +13,7 @@ import (
 	"github.com/mholt/archiver/v3"
 	"github.com/spf13/cobra"
 
+	"github.com/mesosphere/mindthegap/cleanup"
 	"github.com/mesosphere/mindthegap/config"
 	"github.com/mesosphere/mindthegap/docker/registry"
 	"github.com/mesosphere/mindthegap/skopeo"
@@ -28,13 +29,16 @@ func NewCommand(out output.Output) *cobra.Command {
 	cmd := &cobra.Command{
 		Use: "image-bundle",
 		RunE: func(cmd *cobra.Command, args []string) error {
+			cleaner := cleanup.NewCleaner()
+			defer cleaner.Cleanup()
+
 			out.StartOperation("Creating temporary directory")
 			tempDir, err := os.MkdirTemp("", ".image-bundle-*")
 			if err != nil {
 				out.EndOperation(false)
 				return fmt.Errorf("failed to create temporary directory: %w", err)
 			}
-			defer os.RemoveAll(tempDir)
+			cleaner.AddCleanupFn(func() { _ = os.RemoveAll(tempDir) })
 			out.EndOperation(true)
 
 			out.StartOperation("Unarchiving image bundle")
@@ -69,7 +73,7 @@ func NewCommand(out output.Output) *cobra.Command {
 			out.EndOperation(true)
 
 			skopeoRunner, skopeoCleanup := skopeo.NewRunner()
-			defer func() { _ = skopeoCleanup() }()
+			cleaner.AddCleanupFn(func() { _ = skopeoCleanup() })
 
 			skopeoStdout, skopeoStderr, err := skopeoRunner.AttemptToLoginToRegistry(context.TODO(), destRegistry)
 			if err != nil {

cmd/serve/imagebundle/image_bundle.go

@@ -11,6 +11,7 @@ import (
 	"github.com/mholt/archiver/v3"
 	"github.com/spf13/cobra"
 
+	"github.com/mesosphere/mindthegap/cleanup"
 	"github.com/mesosphere/mindthegap/docker/registry"
 )
 
@@ -26,13 +27,16 @@ func NewCommand(out output.Output) *cobra.Command {
 	cmd := &cobra.Command{
 		Use: "image-bundle",
 		RunE: func(cmd *cobra.Command, args []string) error {
+			cleaner := cleanup.NewCleaner()
+			defer cleaner.Cleanup()
 			out.StartOperation("Creating temporary directory")
 			tempDir, err := os.MkdirTemp("", ".image-bundle-*")
 			if err != nil {
 				out.EndOperation(false)
 				return fmt.Errorf("failed to create temporary directory: %w", err)
 			}
-			defer os.RemoveAll(tempDir)
+			cleaner.AddCleanupFn(func() { _ = os.RemoveAll(tempDir) })
+
 			out.EndOperation(true)
 
 			out.StartOperation("Unarchiving image bundle")