ctr-remote: ensure cancel cleanly when recieves signals during conver…

…sion

Signed-off-by: Kohei Tokunaga <ktokunaga.mail@gmail.com>

cmd/ctr-remote/commands/convert.go

@@ -22,6 +22,7 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"os/signal"
 
 	"github.com/containerd/containerd/cmd/ctr/commands"
 	"github.com/containerd/containerd/images/converter"
@@ -174,6 +175,17 @@ When '--all-platforms' is given all images in a manifest list must be available.
 		}
 		defer cancel()
 
+		sigCh := make(chan os.Signal, 1)
+		signal.Notify(sigCh, os.Interrupt)
+		go func() {
+			// Cleanly cancel conversion
+			select {
+			case s := <-sigCh:
+				logrus.Infof("Got %v", s)
+				cancel()
+			case <-ctx.Done():
+			}
+		}()
 		newImg, err := converter.Convert(ctx, client, targetRef, srcRef, convertOpts...)
 		if err != nil {
 			return err

cmd/ctr-remote/commands/optimize.go

@@ -24,6 +24,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"os/signal"
 	"time"
 
 	"github.com/containerd/containerd"
@@ -157,6 +158,18 @@ var OptimizeCommand = cli.Command{
 			f = wrapper(f)
 		}
 		layerConvertFunc := logWrapper(f)
+
+		sigCh := make(chan os.Signal, 1)
+		signal.Notify(sigCh, os.Interrupt)
+		go func() {
+			// Cleanly cancel conversion
+			select {
+			case s := <-sigCh:
+				logrus.Infof("Got %v", s)
+				cancel()
+			case <-ctx.Done():
+			}
+		}()
 		convertOpts = append(convertOpts, converter.WithLayerConvertFunc(layerConvertFunc))
 		newImg, err := converter.Convert(ctx, client, targetRef, srcRef, convertOpts...)
 		if err != nil {

estargz/build.go

@@ -26,6 +26,7 @@ import (
 	"archive/tar"
 	"bytes"
 	"compress/gzip"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -47,6 +48,7 @@ type options struct {
 	prioritizedFiles       []string
 	missedPrioritizedFiles *[]string
 	compression            Compression
+	ctx                    context.Context
 }
 
 type Option func(o *options) error
@@ -103,6 +105,14 @@ func WithCompression(compression Compression) Option {
 	}
 }
 
+// WithContext specifies a context that can be used for clean canceleration.
+func WithContext(ctx context.Context) Option {
+	return func(o *options) error {
+		o.ctx = ctx
+		return nil
+	}
+}
+
 // Blob is an eStargz blob.
 type Blob struct {
 	io.ReadCloser
@@ -138,12 +148,29 @@ func Build(tarBlob *io.SectionReader, opt ...Option) (_ *Blob, rErr error) {
 		opts.compression = newGzipCompressionWithLevel(opts.compressionLevel)
 	}
 	layerFiles := newTempFiles()
+	ctx := opts.ctx
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	done := make(chan struct{})
+	defer close(done)
+	go func() {
+		select {
+		case <-done:
+			// nop
+		case <-ctx.Done():
+			layerFiles.CleanupAll()
+		}
+	}()
 	defer func() {
 		if rErr != nil {
 			if err := layerFiles.CleanupAll(); err != nil {
 				rErr = fmt.Errorf("failed to cleanup tmp files: %v: %w", err, rErr)
 			}
 		}
+		if cErr := ctx.Err(); cErr != nil {
+			rErr = fmt.Errorf("error from context %q: %w", cErr, rErr)
+		}
 	}()
 	tarBlob, err := decompressBlob(tarBlob, layerFiles)
 	if err != nil {
@@ -505,8 +532,9 @@ func newTempFiles() *tempFiles {
 }
 
 type tempFiles struct {
-	files   []*os.File
-	filesMu sync.Mutex
+	files       []*os.File
+	filesMu     sync.Mutex
+	cleanupOnce sync.Once
 }
 
 func (tf *tempFiles) TempFile(dir, pattern string) (*os.File, error) {
@@ -520,7 +548,14 @@ func (tf *tempFiles) TempFile(dir, pattern string) (*os.File, error) {
 	return f, nil
 }
 
-func (tf *tempFiles) CleanupAll() error {
+func (tf *tempFiles) CleanupAll() (err error) {
+	tf.cleanupOnce.Do(func() {
+		err = tf.cleanupAll()
+	})
+	return
+}
+
+func (tf *tempFiles) cleanupAll() error {
 	tf.filesMu.Lock()
 	defer tf.filesMu.Unlock()
 	var allErr []error

nativeconverter/estargz/estargz.go

@@ -77,7 +77,7 @@ func LayerConvertFunc(opts ...estargz.Option) converter.ConvertFunc {
 		}
 		defer ra.Close()
 		sr := io.NewSectionReader(ra, 0, desc.Size)
-		blob, err := estargz.Build(sr, opts...)
+		blob, err := estargz.Build(sr, append(opts, estargz.WithContext(ctx))...)
 		if err != nil {
 			return nil, err
 		}

nativeconverter/zstdchunked/zstdchunked.go

@@ -114,7 +114,7 @@ func LayerConvertFunc(opts ...estargz.Option) converter.ConvertFunc {
 				Metadata:         metadata,
 			},
 		}))
-		blob, err := estargz.Build(uncompressedSR, opts...)
+		blob, err := estargz.Build(uncompressedSR, append(opts, estargz.WithContext(ctx))...)
 		if err != nil {
 			return nil, err
 		}