From f97946c4e84e2a03105a3b4e8db643dea3319859 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?T=C3=B5nis=20Tiigi?= <tonistiigi@gmail.com>
Date: Thu, 7 Mar 2024 06:45:57 +0000
Subject: [PATCH] builder-next: fix missing lock in ensurelayer
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Tõnis Tiigi <tonistiigi@gmail.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 73120ea70b2ee..fc83058339c7b 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 a0d28ad984ba4..510ffefb49406 100644
--- a/builder/builder-next/adapters/snapshot/snapshot.go
+++ b/builder/builder-next/adapters/snapshot/snapshot.go
@@ -17,6 +17,7 @@ import (
 	"github.com/moby/buildkit/identity"
 	"github.com/moby/buildkit/snapshot"
 	"github.com/moby/buildkit/util/leaseutil"
+	"github.com/moby/locker"
 	"github.com/opencontainers/go-digest"
 	"github.com/pkg/errors"
 	bolt "go.etcd.io/bbolt"
@@ -51,10 +52,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
@@ -71,10 +73,11 @@ func NewSnapshotter(opt Opt, prevLM leases.Manager, ns string) (snapshot.Snapsho
 	}
 
 	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(),
 	}
 
 	slm := newLeaseManager(s, prevLM)