push: update dist repo source for mountFrom

For new built layers, it can be pushed into remote registry. However, it
might not be pulled as base layer. Before this commit, the dist repo
source label for mountFrom call only is updated at pull stage. Without
mountFrom, it will take long time to push existing layers which can be
mountFrom. Based on this case, we should update dist repo source after
pushing layer successfully.

Signed-off-by: Wei Fu <fuweid89@gmail.com>

util/push/push.go

@@ -90,11 +90,15 @@ func Push(ctx context.Context, sm *session.Manager, cs content.Store, dgst diges
 	})
 
 	pushHandler := remotes.PushHandler(pusher, cs)
+	pushUpdateSourceHandler, err := updateDistributionSourceHandler(cs, pushHandler, ref)
+	if err != nil {
+		return err
+	}
 
 	handlers := append([]images.Handler{},
 		images.HandlerFunc(annotateDistributionSourceHandler(cs, childrenHandler(cs))),
 		filterHandler,
-		pushHandler,
+		pushUpdateSourceHandler,
 	)
 
 	ra, err := cs.ReaderAt(ctx, desc)
@@ -234,3 +238,38 @@ func childrenHandler(provider content.Provider) images.HandlerFunc {
 		return descs, nil
 	}
 }
+
+// updateDistributionSourceHandler will update distribution source label after
+// pushing layer successfully.
+//
+// FIXME(fuweid): There is race condition for current design of distribution
+// source label if there are pull/push jobs consuming same layer.
+func updateDistributionSourceHandler(cs content.Store, pushF images.HandlerFunc, ref string) (images.HandlerFunc, error) {
+	updateF, err := docker.AppendDistributionSourceLabel(cs, ref)
+	if err != nil {
+		return nil, err
+	}
+
+	return images.HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		var islayer bool
+
+		switch desc.MediaType {
+		case images.MediaTypeDockerSchema2Layer, images.MediaTypeDockerSchema2LayerGzip,
+			ocispec.MediaTypeImageLayer, ocispec.MediaTypeImageLayerGzip:
+			islayer = true
+		}
+
+		children, err := pushF(ctx, desc)
+		if err != nil {
+			return nil, err
+		}
+
+		// update distribution source to layer
+		if islayer {
+			if _, err := updateF(ctx, desc); err != nil {
+				logrus.Warnf("failed to update distribution source for layer %v: %v", desc.Digest, err)
+			}
+		}
+		return children, nil
+	}), nil
+}