storage: support staged addition of layers

The extracting of the tar under the store lock is a bottleneck as many
concurrent processes might hold the locks for a long time on big layers.

To address this move the layer extraction before we take the locks if
possible. Currently this only work when using the overlay driver as the
implementation requires driver specifc details in order for a rename()
to work.

Signed-off-by: Paul Holzinger <[email protected]>

Author: Luap99

storage/layers.go

@@ -199,8 +199,26 @@ type DiffOptions struct {
 // stagedLayerOptions are the options passed to .create to populate a staged
 // layer
 type stagedLayerOptions struct {
+	// These are used via the zstd:chunked pull paths
 	DiffOutput  *drivers.DriverWithDifferOutput
 	DiffOptions *drivers.ApplyDiffWithDifferOpts
+
+	// stagedLayerExtraction is used by the normal tar layer extraction.
+	stagedLayerExtraction *maybeStagedLayerExtraction
+}
+
+// maybeStagedLayerExtraction is a helper to encapsulate details around extracting
+// a layer potentially before we even take a look if the driver implements the
+// ApplyDiffStaging interface.
+type maybeStagedLayerExtraction struct {
+	// diff contains the tar archive, can be compressed
+	diff    io.Reader
+	staging drivers.ApplyDiffStaging
+	result  *applyDiffResult
+
+	cleanupFuncs   []tempdir.CleanupTempDirFunc
+	stagedTarSplit *tempdir.StageAddition
+	stagedLayer    *tempdir.StageAddition
 }
 
 type applyDiffResult struct {
@@ -300,7 +318,7 @@ type rwLayerStore interface {
 	// underlying drivers do not themselves distinguish between writeable
 	// and read-only layers.  Returns the new layer structure and the size of the
 	// diff which was applied to its parent to initialize its contents.
-	create(id string, parent *Layer, names []string, mountLabel string, options map[string]string, moreOptions *LayerOptions, writeable bool, diff io.Reader, slo *stagedLayerOptions) (*Layer, int64, error)
+	create(id string, parent *Layer, names []string, mountLabel string, options map[string]string, moreOptions *LayerOptions, writeable bool, slo *stagedLayerOptions) (*Layer, int64, error)
 
 	// updateNames modifies names associated with a layer based on (op, names).
 	updateNames(id string, names []string, op updateNameOperation) error
@@ -1391,7 +1409,7 @@ func (r *layerStore) pickStoreLocation(volatile, writeable bool) layerLocations
 }
 
 // Requires startWriting.
-func (r *layerStore) create(id string, parentLayer *Layer, names []string, mountLabel string, options map[string]string, moreOptions *LayerOptions, writeable bool, diff io.Reader, slo *stagedLayerOptions) (layer *Layer, size int64, err error) {
+func (r *layerStore) create(id string, parentLayer *Layer, names []string, mountLabel string, options map[string]string, moreOptions *LayerOptions, writeable bool, slo *stagedLayerOptions) (layer *Layer, size int64, err error) {
 	if moreOptions == nil {
 		moreOptions = &LayerOptions{}
 	}
@@ -1580,15 +1598,28 @@ 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 {
-			cleanupFailureContext = "applying layer diff"
-			return nil, -1, err
-		}
-	} else if slo != nil {
-		if err := r.applyDiffFromStagingDirectory(layer.ID, slo.DiffOutput, slo.DiffOptions); err != nil {
-			cleanupFailureContext = "applying staged directory diff"
-			return nil, -1, err
+	if slo != nil {
+		if slo.stagedLayerExtraction != nil {
+			if slo.stagedLayerExtraction.result != nil {
+				// The layer is staged, just commit it and update the metadata.
+				if err := slo.stagedLayerExtraction.commitLayer(r, layer.ID); err != nil {
+					cleanupFailureContext = "commiting staged layer diff"
+					return nil, -1, err
+				}
+				applyDiffResultToLayer(r, layer, moreOptions, slo.stagedLayerExtraction.result)
+			} else {
+				// The diff was not staged, apply it now here instead.
+				if size, err = r.applyDiffWithOptions(layer.ID, moreOptions, slo.stagedLayerExtraction.diff); err != nil {
+					cleanupFailureContext = "applying layer diff"
+					return nil, -1, err
+				}
+			}
+		} else {
+			// staging logic for the chunked pull path
+			if err := r.applyDiffFromStagingDirectory(layer.ID, slo.DiffOutput, slo.DiffOptions); err != nil {
+				cleanupFailureContext = "applying staged directory diff"
+				return nil, -1, err
+			}
 		}
 	} else {
 		// applyDiffWithOptions() would have updated r.bycompressedsum
@@ -2417,6 +2448,78 @@ func createTarSplitFile(r *layerStore, layerID string) (*os.File, error) {
 	return os.OpenFile(r.tspath(layerID), os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
 }
 
+// newMaybeStagedLayerExtraction initlaizes a new maybeStagedLayerExtraction. The caller
+// must call .cleanup() to remove any temporary files.
+func newMaybeStagedLayerExtraction(diff io.Reader, driver drivers.Driver) *maybeStagedLayerExtraction {
+	m := &maybeStagedLayerExtraction{
+		diff: diff,
+	}
+	if d, ok := driver.(drivers.ApplyDiffStaging); ok {
+		m.staging = d
+	}
+	return m
+}
+
+func (sl *maybeStagedLayerExtraction) cleanup() error {
+	return tempdir.CleanupTemporaryDirectories(sl.cleanupFuncs...)
+}
+
+// stageWithUnlockedStore stages the layer content without needing the store locked.
+// If the driver does not support stage addition then this is a NOP and does nothing.
+func (sl *maybeStagedLayerExtraction) stageWithUnlockedStore(r *layerStore, layerOptions *LayerOptions) error {
+	if sl.staging == nil {
+		// driver does not implement stage addition
+		return nil
+	}
+	td, err := tempdir.NewTempDir(filepath.Join(r.layerdir, tempDirPath))
+	if err != nil {
+		return err
+	}
+	sl.cleanupFuncs = append(sl.cleanupFuncs, td.Cleanup)
+
+	stageTarSplit, err := td.StageAddition()
+	if err != nil {
+		return err
+	}
+	sl.stagedTarSplit = stageTarSplit
+
+	f, err := os.OpenFile(stageTarSplit.Path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	result, err := applyDiff(layerOptions, sl.diff, f, func(payload io.Reader) (int64, error) {
+		cleanup, stageLayer, size, err := sl.staging.StartStagingDiffToApply(drivers.ApplyDiffOpts{
+			Diff:     payload,
+			Mappings: idtools.NewIDMappingsFromMaps(layerOptions.UIDMap, layerOptions.GIDMap),
+			// FIXME: What to do here? We have no lock and assigned label yet.
+			// Overlayfs should not need it anyway so this seems fine for now.
+			MountLabel: "",
+		})
+		sl.cleanupFuncs = append(sl.cleanupFuncs, cleanup)
+		sl.stagedLayer = stageLayer
+		return size, err
+	})
+	if err != nil {
+		return err
+	}
+
+	sl.result = result
+	return nil
+}
+
+// commitLayer() commits the content that was staged in stageWithUnlockedStore()
+//
+// Requires startWriting.
+func (sl *maybeStagedLayerExtraction) commitLayer(r *layerStore, layerID string) error {
+	err := sl.stagedTarSplit.Commit(r.tspath(layerID))
+	if err != nil {
+		return err
+	}
+	return sl.staging.CommitStagedLayer(layerID, sl.stagedLayer)
+}
+
 func applyDiff(layerOptions *LayerOptions, diff io.Reader, tarSplitFile *os.File, applyDriverFunc func(io.Reader) (int64, error)) (*applyDiffResult, error) {
 	header := make([]byte, 10240)
 	n, err := diff.Read(header)

storage/store.go

@@ -1452,7 +1452,7 @@ func (s *store) canUseShifting(uidmap, gidmap []idtools.IDMap) bool {
 // On entry:
 // - rlstore must be locked for writing
 // - rlstores MUST NOT be locked
-func (s *store) putLayer(rlstore rwLayerStore, rlstores []roLayerStore, id, parent string, names []string, mountLabel string, writeable bool, lOptions *LayerOptions, diff io.Reader, slo *stagedLayerOptions) (*Layer, int64, error) {
+func (s *store) putLayer(rlstore rwLayerStore, rlstores []roLayerStore, id, parent string, names []string, mountLabel string, writeable bool, lOptions *LayerOptions, slo *stagedLayerOptions) (*Layer, int64, error) {
 	var parentLayer *Layer
 	var options LayerOptions
 	if lOptions != nil {
@@ -1533,19 +1533,39 @@ func (s *store) putLayer(rlstore rwLayerStore, rlstores []roLayerStore, id, pare
 			GIDMap:         copySlicePreferringNil(gidMap),
 		}
 	}
-	return rlstore.create(id, parentLayer, names, mountLabel, nil, &options, writeable, diff, slo)
+	return rlstore.create(id, parentLayer, names, mountLabel, nil, &options, writeable, slo)
 }
 
 func (s *store) PutLayer(id, parent string, names []string, mountLabel string, writeable bool, lOptions *LayerOptions, diff io.Reader) (*Layer, int64, error) {
 	rlstore, rlstores, err := s.bothLayerStoreKinds()
 	if err != nil {
 		return nil, -1, err
 	}
+
+	var slo *stagedLayerOptions
+
+	if diff != nil {
+		m := newMaybeStagedLayerExtraction(diff, s.graphDriver)
+		defer func() {
+			if err := m.cleanup(); err != nil {
+				logrus.Errorf("Error cleaning up temporary directories: %v", err)
+			}
+		}()
+		// FIXME: type case should be safe for now but really there should be a better way to do this
+		err = m.stageWithUnlockedStore(rlstore.(*layerStore), lOptions)
+		if err != nil {
+			return nil, -1, err
+		}
+		slo = &stagedLayerOptions{
+			stagedLayerExtraction: m,
+		}
+	}
+
 	if err := rlstore.startWriting(); err != nil {
 		return nil, -1, err
 	}
 	defer rlstore.stopWriting()
-	return s.putLayer(rlstore, rlstores, id, parent, names, mountLabel, writeable, lOptions, diff, nil)
+	return s.putLayer(rlstore, rlstores, id, parent, names, mountLabel, writeable, lOptions, slo)
 }
 
 func (s *store) CreateLayer(id, parent string, names []string, mountLabel string, writeable bool, options *LayerOptions) (*Layer, error) {
@@ -1753,7 +1773,7 @@ func (s *store) imageTopLayerForMapping(image *Image, ristore roImageStore, rlst
 		}
 	}
 	layerOptions.TemplateLayer = layer.ID
-	mappedLayer, _, err := rlstore.create("", parentLayer, nil, layer.MountLabel, nil, &layerOptions, false, nil, nil)
+	mappedLayer, _, err := rlstore.create("", parentLayer, nil, layer.MountLabel, nil, &layerOptions, false, nil)
 	if err != nil {
 		return nil, fmt.Errorf("creating an ID-mapped copy of layer %q: %w", layer.ID, err)
 	}
@@ -1924,7 +1944,7 @@ func (s *store) CreateContainer(id string, names []string, image, layer, metadat
 		options.Flags[mountLabelFlag] = mountLabel
 	}
 
-	clayer, _, err := rlstore.create(layer, imageTopLayer, nil, mlabel, options.StorageOpt, layerOptions, true, nil, nil)
+	clayer, _, err := rlstore.create(layer, imageTopLayer, nil, mlabel, options.StorageOpt, layerOptions, true, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -3186,7 +3206,7 @@ func (s *store) ApplyStagedLayer(args ApplyStagedLayerOptions) (*Layer, error) {
 		DiffOutput:  args.DiffOutput,
 		DiffOptions: args.DiffOptions,
 	}
-	layer, _, err = s.putLayer(rlstore, rlstores, args.ID, args.ParentLayer, args.Names, args.MountLabel, args.Writeable, args.LayerOptions, nil, &slo)
+	layer, _, err = s.putLayer(rlstore, rlstores, args.ID, args.ParentLayer, args.Names, args.MountLabel, args.Writeable, args.LayerOptions, &slo)
 	return layer, err
 }
 

storage/userns.go

@@ -197,7 +197,7 @@ outer:
 
 	// We need to create a temporary layer so we can mount it and lookup the
 	// maximum IDs used.
-	clayer, _, err := rlstore.create("", topLayer, nil, "", nil, layerOptions, false, nil, nil)
+	clayer, _, err := rlstore.create("", topLayer, nil, "", nil, layerOptions, false, nil)
 	if err != nil {
 		return 0, err
 	}