From 37545cc644344dcb576cba67eb7b6f51a463d31e 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

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>
---
 .../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 73120ea70b..fc83058339 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 a0d28ad984..510ffefb49 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)