diff --git a/cleanup/cleaner.go b/cleanup/cleaner.go new file mode 100644 index 00000000..a11cc401 --- /dev/null +++ b/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) +} diff --git a/cmd/create/imagebundle/image_bundle.go b/cmd/create/imagebundle/image_bundle.go index 77c4a76c..72f30c19 100644 --- a/cmd/create/imagebundle/image_bundle.go +++ b/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 diff --git a/cmd/importcmd/imagebundle/image_bundle.go b/cmd/importcmd/imagebundle/image_bundle.go index 651cf62b..cadf7654 100644 --- a/cmd/importcmd/imagebundle/image_bundle.go +++ b/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") diff --git a/cmd/push/imagebundle/image_bundle.go b/cmd/push/imagebundle/image_bundle.go index b3a4798a..364863b7 100644 --- a/cmd/push/imagebundle/image_bundle.go +++ b/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 { diff --git a/cmd/serve/imagebundle/image_bundle.go b/cmd/serve/imagebundle/image_bundle.go index 12251804..bc66402f 100644 --- a/cmd/serve/imagebundle/image_bundle.go +++ b/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")