WIP: Staged layer creation

storage/drivers/overlay/overlay.go

@@ -2371,16 +2371,6 @@ func (d *Driver) DifferTarget(id string) (string, error) {
 
 // ApplyDiff applies the new layer into a root
 func (d *Driver) ApplyDiff(id, parent string, options graphdriver.ApplyDiffOpts) (size int64, err error) {
-	if !d.isParent(id, parent) {
-		if d.options.ignoreChownErrors {
-			options.IgnoreChownErrors = d.options.ignoreChownErrors
-		}
-		if d.options.forceMask != nil {
-			options.ForceMask = d.options.forceMask
-		}
-		return d.naiveDiff.ApplyDiff(id, parent, options)
-	}
-
 	idMappings := options.Mappings
 	if idMappings == nil {
 		idMappings = &idtools.IDMappings{}

storage/internal/tempdir/tempdir.go

@@ -91,6 +91,19 @@ type TempDir struct {
 	counter uint64
 }
 
+type stageAddition struct {
+	source string
+}
+
+// Commit the staged content into its final destination by using os.Rename().
+// That means the dest must be on the same on the same fs as the root directory
+// that was given to NewTempDir() and the dest must not exist yet.
+// Commit must only be called once per instance returned from the
+// StageAddition() call.
+func (s *stageAddition) Commit(destination string) error {
+	return os.Rename(s.source, destination)
+}
+
 // CleanupTempDirFunc is a function type that can be returned by operations
 // which need to perform cleanup actions later.
 type CleanupTempDirFunc func() error
@@ -190,6 +203,30 @@ func NewTempDir(rootDir string) (*TempDir, error) {
 	return td, nil
 }
 
+// StageAddition creates a new temp directory which is then passed as argument to the
+// given callback function. The function should be used to populate the directory with
+// content.
+// On success StageAddition returns a type with the Commit() function, that function then
+// must be used to move the content from the temp directory into its final location.
+//
+// The caller MUST ensure .Cleanup() is called after Commit().
+// If the TempDir has been cleaned up, this method will return an error.
+func (td *TempDir) StageAddition(callback func(path string) error) (*stageAddition, error) {
+	if td.tempDirLock == nil {
+		return nil, fmt.Errorf("temp dir instance not initialized or already cleaned up")
+	}
+	fileName := fmt.Sprintf("%d-", td.counter) + "addition"
+	tmpAddPath := filepath.Join(td.tempDirPath, fileName)
+	if err := os.Mkdir(tmpAddPath, 0o700); err != nil {
+		return nil, fmt.Errorf("creating temp directory for addition failed: %w", err)
+	}
+	td.counter++
+	if err := callback(tmpAddPath); err != nil {
+		return nil, err
+	}
+	return &stageAddition{source: tmpAddPath}, nil
+}
+
 // StageDeletion moves the specified file into the instance's temporary directory.
 // The temporary directory must already exist (created during NewTempDir).
 // Files are renamed with a counter-based prefix (e.g., "0-filename", "1-filename") to ensure uniqueness.

storage/internal/tempdir/tempdir_test.go

@@ -1,6 +1,7 @@
 package tempdir
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -271,3 +272,44 @@ func TestTempDirFileNaming(t *testing.T) {
 		assert.True(t, found, "Expected file %s not found", expectedName)
 	}
 }
+
+func TestStageAddition(t *testing.T) {
+	rootDir := t.TempDir()
+	td, err := NewTempDir(rootDir)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		assert.NoError(t, td.Cleanup())
+	})
+
+	sa, err := td.StageAddition(func(path string) error {
+		f, err := os.Create(filepath.Join(path, "file1"))
+		require.NoError(t, err)
+		require.NoError(t, f.Close())
+		return nil
+	})
+	require.NoError(t, err)
+
+	// need to use a dest which does not exist yet
+	dest := filepath.Join(t.TempDir(), "dest")
+
+	err = sa.Commit(dest)
+	require.NoError(t, err)
+	assert.FileExists(t, filepath.Join(dest, "file1"))
+	assert.NoDirExists(t, sa.source)
+}
+
+func TestStageAdditionCallbackError(t *testing.T) {
+	rootDir := t.TempDir()
+	td, err := NewTempDir(rootDir)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		assert.NoError(t, td.Cleanup())
+	})
+
+	customErr := errors.New("custom error")
+
+	_, err = td.StageAddition(func(path string) error {
+		return customErr
+	})
+	require.ErrorIs(t, err, customErr)
+}

storage/layers.go

@@ -339,7 +339,7 @@ type rwLayerStore interface {
 	CleanupStagingDirectory(stagingDirectory string) error
 
 	// applyDiffFromStagingDirectory uses diffOutput.Target to create the diff.
-	applyDiffFromStagingDirectory(id string, diffOutput *drivers.DriverWithDifferOutput, options *drivers.ApplyDiffWithDifferOpts) error
+	applyDiffFromStagingDirectory(layer *Layer, diffOutput *drivers.DriverWithDifferOutput, options *drivers.ApplyDiffWithDifferOpts) error
 
 	// DifferTarget gets the location where files are stored for the layer.
 	DifferTarget(id string) (string, error)
@@ -1569,12 +1569,12 @@ func (r *layerStore) create(id string, parentLayer *Layer, names []string, mount
 
 	size = -1
 	if diff != nil {
-		if size, err = r.applyDiffWithOptions(layer.ID, moreOptions, diff); err != nil {
+		if size, err = r.applyDiffWithOptions(layer, moreOptions, diff); err != nil {
 			cleanupFailureContext = "applying layer diff"
 			return nil, -1, err
 		}
 	} else if slo != nil {
-		if err := r.applyDiffFromStagingDirectory(layer.ID, slo.DiffOutput, slo.DiffOptions); err != nil {
+		if err := r.applyDiffFromStagingDirectory(layer, slo.DiffOutput, slo.DiffOptions); err != nil {
 			cleanupFailureContext = "applying staged directory diff"
 			return nil, -1, err
 		}
@@ -2395,20 +2395,19 @@ func updateDigestMap(m *map[digest.Digest][]string, oldvalue, newvalue digest.Di
 
 // Requires startWriting.
 func (r *layerStore) ApplyDiff(to string, diff io.Reader) (size int64, err error) {
-	return r.applyDiffWithOptions(to, nil, diff)
+	layer, ok := r.lookup(to)
+	if !ok {
+		return -1, ErrLayerUnknown
+	}
+	return r.applyDiffWithOptions(layer, nil, diff)
 }
 
 // Requires startWriting.
-func (r *layerStore) applyDiffWithOptions(to string, layerOptions *LayerOptions, diff io.Reader) (size int64, err error) {
+func (r *layerStore) applyDiffWithOptions(layer *Layer, layerOptions *LayerOptions, diff io.Reader) (size int64, err error) {
 	if !r.lockfile.IsReadWrite() {
 		return -1, fmt.Errorf("not allowed to modify layer contents at %q: %w", r.layerdir, ErrStoreIsReadOnly)
 	}
 
-	layer, ok := r.lookup(to)
-	if !ok {
-		return -1, ErrLayerUnknown
-	}
-
 	header := make([]byte, 10240)
 	n, err := diff.Read(header)
 	if err != nil && err != io.EOF {
@@ -2553,15 +2552,11 @@ func (r *layerStore) DifferTarget(id string) (string, error) {
 }
 
 // Requires startWriting.
-func (r *layerStore) applyDiffFromStagingDirectory(id string, diffOutput *drivers.DriverWithDifferOutput, options *drivers.ApplyDiffWithDifferOpts) error {
+func (r *layerStore) applyDiffFromStagingDirectory(layer *Layer, diffOutput *drivers.DriverWithDifferOutput, options *drivers.ApplyDiffWithDifferOpts) error {
 	ddriver, ok := r.driver.(drivers.DriverWithDiffer)
 	if !ok {
 		return ErrNotSupported
 	}
-	layer, ok := r.lookup(id)
-	if !ok {
-		return ErrLayerUnknown
-	}
 	if options == nil {
 		options = &drivers.ApplyDiffWithDifferOpts{
 			ApplyDiffOpts: drivers.ApplyDiffOpts{
@@ -2622,9 +2617,9 @@ func (r *layerStore) applyDiffFromStagingDirectory(id string, diffOutput *driver
 		}
 	}
 	for k, v := range diffOutput.BigData {
-		if err := r.SetBigData(id, k, bytes.NewReader(v)); err != nil {
-			if err2 := r.deleteWhileHoldingLock(id); err2 != nil {
-				logrus.Errorf("While recovering from a failure to set big data, error deleting layer %#v: %v", id, err2)
+		if err := r.SetBigData(layer.ID, k, bytes.NewReader(v)); err != nil {
+			if err2 := r.deleteWhileHoldingLock(layer.ID); err2 != nil {
+				logrus.Errorf("While recovering from a failure to set big data, error deleting layer %#v: %v", layer.ID, err2)
 			}
 			return err
 		}

storage/store.go

@@ -3177,7 +3177,7 @@ func (s *store) ApplyStagedLayer(args ApplyStagedLayerOptions) (*Layer, error) {
 		// This code path exists only for cmd/containers/storage.applyDiffUsingStagingDirectory; we have tests that
 		// assume layer creation and applying a staged layer are separate steps. Production pull code always uses the
 		// other path, where layer creation is atomic.
-		return layer, rlstore.applyDiffFromStagingDirectory(args.ID, args.DiffOutput, args.DiffOptions)
+		return layer, rlstore.applyDiffFromStagingDirectory(layer, args.DiffOutput, args.DiffOptions)
 	}
 
 	// if the layer doesn't exist yet, try to create it.