[deckhouse-cli] remove partial artefacts (#359)

Signed-off-by: Pavel Okhlopkov <pavel.okhlopkov@flant.com>

Author: ldmonster

internal/mirror/chunked/chunk_writer.go

@@ -18,18 +18,31 @@ package chunked
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
 )
 
+// tmpChunkSuffix is appended to in-flight chunk files so an aborted run leaves
+// behind a clearly-named partial file that can be cleaned up later, rather
+// than a half-written .chunk that looks like a finished artifact.
+const tmpChunkSuffix = ".tmp"
+
 type FileWriter struct {
 	chunkSize  int64
 	chunkIndex int
 
 	workingDir   string
 	baseFileName string
 	activeChunk  *os.File
+
+	// writtenChunks tracks paths of all chunk files we have created so far
+	// (relative to workingDir, with the .tmp suffix still attached).  On
+	// Finalize they are renamed to drop the suffix; on Cleanup they are
+	// removed.  Used to guarantee no half-finished artifacts survive a
+	// failed/interrupted run.
+	writtenChunks []string
 }
 
 func NewChunkedFileWriter(chunkSize int64, dirPath, baseFileName string) *FileWriter {
@@ -87,10 +100,46 @@ func (c *FileWriter) Write(p []byte) (int, error) {
 	}
 }
 
+// Close flushes and closes the currently-active chunk. It does NOT promote
+// the .tmp chunk paths to their final names: callers must invoke Finalize for
+// that, or Cleanup to discard everything.  This split lets the caller decide
+// after-the-fact whether the produced bundle is valid (e.g. only after Pack
+// returned without error) and avoids the previous behavior where a Ctrl+C
+// during packing left behind ready-looking .chunk files.
 func (c *FileWriter) Close() error {
 	return c.closeActiveChunk()
 }
 
+// Finalize promotes every successfully-written chunk from its temporary path
+// (<base>.NNNN.chunk.tmp) to the final path (<base>.NNNN.chunk).  Must be
+// called only after Close and only on a successful pack.
+func (c *FileWriter) Finalize() error {
+	var errs []error
+	for _, tmpRelPath := range c.writtenChunks {
+		tmp := filepath.Join(c.workingDir, tmpRelPath)
+		final := filepath.Join(c.workingDir, finalChunkName(tmpRelPath))
+		if err := os.Rename(tmp, final); err != nil {
+			errs = append(errs, fmt.Errorf("rename %s -> %s: %w", tmp, final, err))
+		}
+	}
+	c.writtenChunks = nil
+	return errors.Join(errs...)
+}
+
+// Cleanup removes any temporary chunk files this writer has created. Safe to
+// call multiple times. Use this when the operation failed or was cancelled so
+// no partial artifacts leak into the user's bundle directory.
+func (c *FileWriter) Cleanup() {
+	if c.activeChunk != nil {
+		_ = c.activeChunk.Close()
+		c.activeChunk = nil
+	}
+	for _, tmpRelPath := range c.writtenChunks {
+		_ = os.Remove(filepath.Join(c.workingDir, tmpRelPath))
+	}
+	c.writtenChunks = nil
+}
+
 func (c *FileWriter) swapActiveChunk() error {
 	if c.activeChunk != nil {
 		if err := c.closeActiveChunk(); err != nil {
@@ -99,12 +148,14 @@ func (c *FileWriter) swapActiveChunk() error {
 		c.chunkIndex++
 	}
 
-	newChunk, err := os.Create(filepath.Join(c.workingDir, fmt.Sprintf("%s.%04d.chunk", c.baseFileName, c.chunkIndex)))
+	tmpName := fmt.Sprintf("%s.%04d.chunk%s", c.baseFileName, c.chunkIndex, tmpChunkSuffix)
+	newChunk, err := os.Create(filepath.Join(c.workingDir, tmpName))
 	if err != nil {
 		return fmt.Errorf("Create new chunk file: %w", err)
 	}
 
 	c.activeChunk = newChunk
+	c.writtenChunks = append(c.writtenChunks, tmpName)
 	return nil
 }
 
@@ -116,6 +167,15 @@ func (c *FileWriter) closeActiveChunk() error {
 		if err := c.activeChunk.Close(); err != nil {
 			return fmt.Errorf("Close chunk: %w", err)
 		}
+		c.activeChunk = nil
 	}
 	return nil
 }
+
+// finalChunkName strips the temporary suffix from a chunk file name.
+func finalChunkName(tmpName string) string {
+	if len(tmpName) > len(tmpChunkSuffix) && tmpName[len(tmpName)-len(tmpChunkSuffix):] == tmpChunkSuffix {
+		return tmpName[:len(tmpName)-len(tmpChunkSuffix)]
+	}
+	return tmpName
+}

internal/mirror/chunked/chunk_writer_test.go

@@ -29,30 +29,78 @@ import (
 )
 
 func TestChunkedFileWriterHappyPath(t *testing.T) {
-	workingDir := filepath.Join(os.TempDir(), "chunk_test")
-	require.NoError(t, os.MkdirAll(workingDir, 0o777))
-	t.Cleanup(func() {
-		_ = os.RemoveAll(workingDir)
-	})
+	workingDir := t.TempDir()
 
 	const testDatasetSize, chunkSize = 10 * 1024 * 1024, 3 * 1024 * 1024
 	sourceFile := make([]byte, testDatasetSize)
 	bytesGenerated, err := rand.Reader.Read(sourceFile)
 	require.NoError(t, err)
 	require.Equal(t, testDatasetSize, bytesGenerated)
 
+	w := NewChunkedFileWriter(chunkSize, workingDir, "d8.tar")
 	bytesWritten, err := io.CopyBuffer(
-		NewChunkedFileWriter(chunkSize, workingDir, "d8.tar"),
+		w,
 		bytes.NewReader(sourceFile),
 		make([]byte, 512*1024),
 	)
 	require.NoError(t, err)
 	require.Equal(t, int64(bytesGenerated), bytesWritten)
+	require.NoError(t, w.Close())
+	require.NoError(t, w.Finalize())
 
 	validateSizes(t, workingDir, testDatasetSize, chunkSize)
 	compareHashes(t, sourceFile, testDatasetSize, workingDir)
 }
 
+// TestChunkedFileWriterCleanupRemovesPartialChunks documents the contract
+// that callers rely on to keep the bundle directory clean after an
+// interrupted pack: every chunk file written so far must disappear from disk
+// when Cleanup is invoked, and crucially no file with the final .chunk name
+// must ever appear unless Finalize was called.
+func TestChunkedFileWriterCleanupRemovesPartialChunks(t *testing.T) {
+	workingDir := t.TempDir()
+
+	w := NewChunkedFileWriter(1024, workingDir, "module-foo.tar")
+	_, err := w.Write(bytes.Repeat([]byte("a"), 4096))
+	require.NoError(t, err)
+
+	entries, err := os.ReadDir(workingDir)
+	require.NoError(t, err)
+	for _, e := range entries {
+		require.True(t, hasSuffix(e.Name(), tmpChunkSuffix),
+			"chunk %q should still be staged with %q suffix before Finalize", e.Name(), tmpChunkSuffix)
+	}
+
+	w.Cleanup()
+
+	entries, err = os.ReadDir(workingDir)
+	require.NoError(t, err)
+	require.Empty(t, entries, "Cleanup must remove every staged chunk")
+}
+
+// TestChunkedFileWriterFinalizePromotesAllChunks verifies the full lifecycle:
+// staged .tmp chunks must turn into their final .chunk names after Finalize.
+func TestChunkedFileWriterFinalizePromotesAllChunks(t *testing.T) {
+	workingDir := t.TempDir()
+
+	w := NewChunkedFileWriter(1024, workingDir, "module-foo.tar")
+	_, err := w.Write(bytes.Repeat([]byte("a"), 4096))
+	require.NoError(t, err)
+	require.NoError(t, w.Close())
+	require.NoError(t, w.Finalize())
+
+	entries, err := os.ReadDir(workingDir)
+	require.NoError(t, err)
+	for _, e := range entries {
+		require.Equal(t, ".chunk", filepath.Ext(e.Name()),
+			"chunk %q must lose the staging suffix after Finalize", e.Name())
+	}
+}
+
+func hasSuffix(name, suffix string) bool {
+	return len(name) >= len(suffix) && name[len(name)-len(suffix):] == suffix
+}
+
 func compareHashes(t *testing.T, sourceFile []byte, testDatasetSize int, workingDir string) {
 	t.Helper()
 

internal/mirror/installer/installer.go

@@ -21,13 +21,12 @@ import (
 	"fmt"
 	"io"
 	"log/slog"
-	"os"
 	"path/filepath"
 	"time"
 
 	dkplog "github.com/deckhouse/deckhouse/pkg/log"
 
-	"github.com/deckhouse/deckhouse-cli/internal/mirror/chunked"
+	"github.com/deckhouse/deckhouse-cli/internal/mirror/pack"
 	"github.com/deckhouse/deckhouse-cli/internal/mirror/puller"
 	"github.com/deckhouse/deckhouse-cli/pkg/libmirror/bundle"
 	"github.com/deckhouse/deckhouse-cli/pkg/libmirror/util/log"
@@ -178,28 +177,10 @@ func (svc *Service) pullInstaller(ctx context.Context) error {
 		return fmt.Errorf("pull installer: %w", err)
 	}
 
-	if err := logger.Process("Pack Deckhouse images into platform.tar", func() error {
-		bundleChunkSize := svc.options.BundleChunkSize
-		bundleDir := svc.options.BundleDir
-
-		var installer io.Writer = chunked.NewChunkedFileWriter(
-			bundleChunkSize,
-			bundleDir,
-			"installer.tar",
-		)
-
-		if bundleChunkSize == 0 {
-			installer, err = os.Create(filepath.Join(bundleDir, "installer.tar"))
-			if err != nil {
-				return fmt.Errorf("create installer.tar: %w", err)
-			}
-		}
-
-		if err := bundle.Pack(context.Background(), svc.layout.workingDir, installer); err != nil {
-			return fmt.Errorf("pack installer.tar: %w", err)
-		}
-
-		return nil
+	if err := logger.Process("Pack Deckhouse images into installer.tar", func() error {
+		return pack.Bundle(ctx, svc.options.BundleDir, "installer.tar", svc.options.BundleChunkSize, func(w io.Writer) error {
+			return bundle.Pack(ctx, svc.layout.workingDir, w)
+		})
 	}); err != nil {
 		return err
 	}

internal/mirror/modules/modules.go

@@ -23,7 +23,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"os"
 	"path/filepath"
 	"regexp"
 	"strings"
@@ -35,7 +34,7 @@ import (
 	"github.com/deckhouse/deckhouse/pkg/registry/client"
 
 	"github.com/deckhouse/deckhouse-cli/internal"
-	"github.com/deckhouse/deckhouse-cli/internal/mirror/chunked"
+	"github.com/deckhouse/deckhouse-cli/internal/mirror/pack"
 	"github.com/deckhouse/deckhouse-cli/internal/mirror/puller"
 	"github.com/deckhouse/deckhouse-cli/pkg/libmirror/bundle"
 	"github.com/deckhouse/deckhouse-cli/pkg/libmirror/layouts"
@@ -237,11 +236,27 @@ func (svc *Service) pullModules(ctx context.Context) error {
 		processName = "Pull Extra Images"
 	}
 
+	// pullCancelled is set when the user interrupts mid-pull. In that case we
+	// still want the post-processing + packing phase to finalize whatever was
+	// successfully downloaded so far, rather than throw it all away and leave
+	// the user with nothing for a multi-hour pull.
+	pullCancelled := false
 	err = logger.Process(processName, func() error {
 		for i, module := range filteredModules {
+			if err := ctx.Err(); err != nil {
+				logger.Warnf("Pull cancelled; %d/%d modules attempted, will pack already-downloaded modules", i, len(filteredModules))
+				pullCancelled = true
+				return nil
+			}
+
 			logger.Infof("[%d/%d] Processing module: %s", i+1, len(filteredModules), module.name)
 
 			if err := svc.pullSingleModule(ctx, module); err != nil {
+				if isContextErr(err) {
+					logger.Warnf("Pull of module %s cancelled, will pack already-downloaded modules", module.name)
+					pullCancelled = true
+					return nil
+				}
 				return fmt.Errorf("pull module %s: %w", module.name, err)
 			}
 		}
@@ -256,6 +271,16 @@ func (svc *Service) pullModules(ctx context.Context) error {
 		return nil
 	}
 
+	// Decouple the post-pull phase from the cancellation context so that the
+	// last bit of packing finalizes for the modules that *did* download.
+	// Without this, a Ctrl+C during pull would still surface as ctx.Canceled
+	// inside bundle.PackWithPrefix and we would discard a multi-GB download
+	// that was already on disk.
+	postCtx := ctx
+	if pullCancelled {
+		postCtx = context.WithoutCancel(ctx)
+	}
+
 	err = logger.Process("Processing modules image indexes", func() error {
 		for _, l := range svc.layout.AsList() {
 			err = layouts.SortIndexManifests(l)
@@ -279,13 +304,19 @@ func (svc *Service) pullModules(ctx context.Context) error {
 	}
 
 	// Pack each module into separate tar
-	if err := svc.packModules(filteredModules); err != nil {
+	if err := svc.packModules(postCtx, filteredModules); err != nil {
 		return err
 	}
 
 	return nil
 }
 
+// isContextErr reports whether err is one of the context cancellation errors,
+// either directly or wrapped.
+func isContextErr(err error) bool {
+	return errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded)
+}
+
 func (svc *Service) pullSingleModule(ctx context.Context, module moduleData) error {
 	downloadList := NewImageDownloadList(filepath.Join(svc.rootURL, "modules", module.name))
 	svc.modulesDownloadList.list[module.name] = downloadList
@@ -879,13 +910,19 @@ func (svc *Service) applyChannelAliases(moduleName string) error {
 	return nil
 }
 
-func (svc *Service) packModules(modules []moduleData) error {
+func (svc *Service) packModules(ctx context.Context, modules []moduleData) error {
 	logger := svc.userLogger
 
 	bundleDir := svc.options.BundleDir
 	bundleChunkSize := svc.options.BundleChunkSize
 
 	for _, module := range modules {
+		// Honor cancellation between modules so a Ctrl+C during the pack
+		// phase doesn't keep producing more (potentially useless) tars.
+		if err := ctx.Err(); err != nil {
+			return err
+		}
+
 		pkgName := "module-" + module.name + ".tar"
 
 		if err := logger.Process(fmt.Sprintf("Pack %s", pkgName), func() error {
@@ -899,24 +936,14 @@ func (svc *Service) packModules(modules []moduleData) error {
 				return nil
 			}
 
-			var pkg io.Writer = chunked.NewChunkedFileWriter(bundleChunkSize, bundleDir, pkgName)
-			if bundleChunkSize == 0 {
-				f, err := os.Create(filepath.Join(bundleDir, pkgName))
-				if err != nil {
-					return fmt.Errorf("create %s: %w", pkgName, err)
-				}
-				pkg = f
-			}
-
 			// Pack from the module's working directory with prefix to create correct registry structure.
 			// This ensures the tar contains paths like "modules/<name>/index.json" instead of just "index.json".
 			moduleDir := filepath.Join(svc.layout.workingDir, module.name)
 			tarPrefix := filepath.Join("modules", module.name)
-			if err := bundle.PackWithPrefix(context.Background(), moduleDir, tarPrefix, pkg); err != nil {
-				return fmt.Errorf("pack module %s: %w", pkgName, err)
-			}
 
-			return nil
+			return pack.Bundle(ctx, bundleDir, pkgName, bundleChunkSize, func(w io.Writer) error {
+				return bundle.PackWithPrefix(ctx, moduleDir, tarPrefix, w)
+			})
 		}); err != nil {
 			return err
 		}