From d8e8fda57cbc4fb98c3cac8185bd89ba9e0076b4 Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Mon, 29 Jun 2020 12:08:18 -0700
Subject: [PATCH] push: dedupe push handlers

This allows to work around the containerd issue on
tracking push status on pushing manifest lists.

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 util/push/push.go | 37 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 36 insertions(+), 1 deletion(-)

diff --git a/util/push/push.go b/util/push/push.go
index 428010c39d71..b8d372d7f3f4 100644
--- a/util/push/push.go
+++ b/util/push/push.go
@@ -14,6 +14,7 @@ import (
 	"github.com/containerd/containerd/remotes/docker"
 	"github.com/docker/distribution/reference"
 	"github.com/moby/buildkit/session"
+	"github.com/moby/buildkit/util/flightcontrol"
 	"github.com/moby/buildkit/util/imageutil"
 	"github.com/moby/buildkit/util/progress"
 	"github.com/moby/buildkit/util/resolver"
@@ -73,7 +74,7 @@ func Push(ctx context.Context, sm *session.Manager, cs content.Store, dgst diges
 	handlers := append([]images.Handler{},
 		images.HandlerFunc(annotateDistributionSourceHandler(cs, childrenHandler(cs))),
 		filterHandler,
-		pushUpdateSourceHandler,
+		dedupeHandler(pushUpdateSourceHandler),
 	)
 
 	ra, err := cs.ReaderAt(ctx, desc)
@@ -248,3 +249,37 @@ func updateDistributionSourceHandler(cs content.Store, pushF images.HandlerFunc,
 		return children, nil
 	}), nil
 }
+
+func dedupeHandler(h images.HandlerFunc) images.HandlerFunc {
+	var g flightcontrol.Group
+	res := map[digest.Digest][]ocispec.Descriptor{}
+	var mu sync.Mutex
+
+	return images.HandlerFunc(func(ctx context.Context, desc ocispec.Descriptor) ([]ocispec.Descriptor, error) {
+		res, err := g.Do(ctx, desc.Digest.String(), func(ctx context.Context) (interface{}, error) {
+			mu.Lock()
+			if r, ok := res[desc.Digest]; ok {
+				mu.Unlock()
+				return r, nil
+			}
+			mu.Unlock()
+
+			children, err := h(ctx, desc)
+			if err != nil {
+				return nil, err
+			}
+
+			mu.Lock()
+			res[desc.Digest] = children
+			mu.Unlock()
+			return children, nil
+		})
+		if err != nil {
+			return nil, err
+		}
+		if res == nil {
+			return nil, nil
+		}
+		return res.([]ocispec.Descriptor), nil
+	})
+}