From 5aa1ff9afad56ef0cf4acd983ff441c8048c0ba3 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/components/engine/builder/builder-next/adapters/snapshot/layer.go b/components/engine/builder/builder-next/adapters/snapshot/layer.go
index ffde5eec..13847d5a 100644
--- a/components/engine/builder/builder-next/adapters/snapshot/layer.go
+++ b/components/engine/builder/builder-next/adapters/snapshot/layer.go
@@ -13,6 +13,9 @@ import (
 )
 
 func (s *snapshotter) EnsureLayer(ctx context.Context, key string) ([]layer.DiffID, error) {
+	s.layerCreateLocker.Lock(key)
+	defer s.layerCreateLocker.Unlock(key)
+
 	if l, err := s.getLayer(key, true); err != nil {
 		return nil, err
 	} else if l != nil {
diff --git a/components/engine/builder/builder-next/adapters/snapshot/snapshot.go b/components/engine/builder/builder-next/adapters/snapshot/snapshot.go
index c1388da7..2b1d33d7 100644
--- a/components/engine/builder/builder-next/adapters/snapshot/snapshot.go
+++ b/components/engine/builder/builder-next/adapters/snapshot/snapshot.go
@@ -11,6 +11,7 @@ import (
 	"github.com/containerd/containerd/snapshots"
 	"github.com/docker/docker/daemon/graphdriver"
 	"github.com/docker/docker/layer"
+	"github.com/docker/docker/pkg/locker"
 	"github.com/moby/buildkit/identity"
 	"github.com/moby/buildkit/snapshot"
 	digest "github.com/opencontainers/go-digest"
@@ -43,10 +44,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
 }
 
 var _ snapshot.SnapshotterBase = &snapshotter{}
@@ -65,10 +67,11 @@ func NewSnapshotter(opt Opt) (snapshot.SnapshotterBase, error) {
 	}
 
 	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(),
 	}
 	return s, nil
 }
-- 
2.33.0