From 8ffaef6d61835b2867794eb18c5890e8d9f48d38 Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Wed, 6 Mar 2024 23:11:32 -0800
Subject: [PATCH] builder-next: fix missing lock in ensurelayer
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When this was called concurrently from the moby image
exporter there could be a data race where a layer was
written to the refs map when it was already there.

In that case the reference count got mixed up and on
release only one of these layers was actually released.

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
(cherry picked from commit 37545cc644344dcb576cba67eb7b6f51a463d31e)
Signed-off-by: Paweł Gronowski <pawel.gronowski@docker.com>
---
 .../builder-next/adapters/snapshot/layer.go   |  3 +++
 .../adapters/snapshot/snapshot.go             | 19 +++++++++++--------
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/builder/builder-next/adapters/snapshot/layer.go b/builder/builder-next/adapters/snapshot/layer.go
index e606a8b472..f0a3a62ea1 100644
--- a/builder/builder-next/adapters/snapshot/layer.go
+++ b/builder/builder-next/adapters/snapshot/layer.go
@@ -22,6 +22,9 @@ func (s *snapshotter) GetDiffIDs(ctx context.Context, key string) ([]layer.DiffI
 }
 
 func (s *snapshotter) EnsureLayer(ctx context.Context, key string) ([]layer.DiffID, error) {
+	s.layerCreateLocker.Lock(key)
+	defer s.layerCreateLocker.Unlock(key)
+
 	diffIDs, err := s.GetDiffIDs(ctx, key)
 	if err != nil {
 		return nil, err
diff --git a/builder/builder-next/adapters/snapshot/snapshot.go b/builder/builder-next/adapters/snapshot/snapshot.go
index 0827af8210..f297e9663e 100644
--- a/builder/builder-next/adapters/snapshot/snapshot.go
+++ b/builder/builder-next/adapters/snapshot/snapshot.go
@@ -16,6 +16,7 @@ import (
 	"github.com/docker/docker/pkg/idtools"
 	"github.com/moby/buildkit/identity"
 	"github.com/moby/buildkit/snapshot"
+	"github.com/moby/locker"
 	"github.com/opencontainers/go-digest"
 	"github.com/pkg/errors"
 	bolt "go.etcd.io/bbolt"
@@ -48,10 +49,11 @@ type checksumCalculator interface {
 type snapshotter struct {
 	opt Opt
 
-	refs map[string]layer.Layer
-	db   *bolt.DB
-	mu   sync.Mutex
-	reg  graphIDRegistrar
+	refs              map[string]layer.Layer
+	db                *bolt.DB
+	mu                sync.Mutex
+	reg               graphIDRegistrar
+	layerCreateLocker *locker.Locker
 }
 
 // NewSnapshotter creates a new snapshotter
@@ -68,10 +70,11 @@ func NewSnapshotter(opt Opt, prevLM leases.Manager) (snapshot.Snapshotter, lease
 	}
 
 	s := &snapshotter{
-		opt:  opt,
-		db:   db,
-		refs: map[string]layer.Layer{},
-		reg:  reg,
+		opt:               opt,
+		db:                db,
+		refs:              map[string]layer.Layer{},
+		reg:               reg,
+		layerCreateLocker: locker.New(),
 	}
 
 	lm := newLeaseManager(s, prevLM)