Locking in ProviderGraph to prevent duplicate code generation

src/CoreTests/Bugs/Bug_3083_concurrent_type_generation.cs

@@ -0,0 +1,97 @@
+using System.Collections.Generic;
+using System.Threading.Tasks;
+using Marten;
+using Marten.Internal;
+using Marten.Internal.Storage;
+using Marten.Schema.BulkLoading;
+using Marten.Testing.Harness;
+using Shouldly;
+using Xunit;
+
+namespace CoreTests.Bugs;
+
+public class Bug_3083_concurrent_type_generation: BugIntegrationContext
+{
+    [Fact]
+    public async Task concurrent_type_generation()
+    {
+        var graph = new ProviderGraph(new StoreOptions());
+
+        var tasks = new List<Task<DocumentProvider<SomeDocument>>>();
+
+        for (var i = 0; i < 15; ++i)
+        {
+            var task = Task.Run(() => graph.StorageFor<SomeDocument>());
+
+            tasks.Add(task);
+        }
+
+        var storages = new HashSet<DocumentProvider<SomeDocument>>(ReferenceEqualityComparer.Instance);
+
+        foreach (var task in tasks)
+        {
+            var storage = await task;
+            storages.Add(storage);
+        }
+
+        storages.ShouldHaveSingleItem();
+    }
+
+    [Fact]
+    public async Task concurrent_append_providers()
+    {
+        var graph = new ProviderGraph(new StoreOptions());
+
+        var tasks = new List<Task>();
+
+        var documentProvider1 = new MockDocumentProvider<SomeDocument>();
+        var documentProvider2 = new MockDocumentProvider<OtherDocument>();
+        var documentProvider3 = new MockDocumentProvider<ThirdDocument>();
+        var documentProvider4 = new MockDocumentProvider<ForthDocument>();
+
+        tasks.Add(Task.Run(() => graph.Append(documentProvider1)));
+        tasks.Add(Task.Run(() => graph.Append(documentProvider2)));
+        tasks.Add(Task.Run(() => graph.Append(documentProvider3)));
+        tasks.Add(Task.Run(() => graph.Append(documentProvider4)));
+
+        await Task.WhenAll(tasks);
+
+        graph.StorageFor<SomeDocument>().ShouldBeTheSameAs(documentProvider1);
+        graph.StorageFor<OtherDocument>().ShouldBeTheSameAs(documentProvider2);
+        graph.StorageFor<ThirdDocument>().ShouldBeTheSameAs(documentProvider3);
+        graph.StorageFor<ForthDocument>().ShouldBeTheSameAs(documentProvider4);
+    }
+
+    public class MockDocumentProvider<T>: DocumentProvider<T> where T : notnull
+    {
+        public MockDocumentProvider(): this(null, null, null, null, null)
+        {
+        }
+
+        public MockDocumentProvider(IBulkLoader<T> bulkLoader, IDocumentStorage<T> queryOnly,
+            IDocumentStorage<T> lightweight, IDocumentStorage<T> identityMap, IDocumentStorage<T> dirtyTracking): base(
+            bulkLoader, queryOnly, lightweight, identityMap, dirtyTracking)
+        {
+        }
+    }
+
+    public class SomeDocument
+    {
+        public string Id { get; set; } = string.Empty;
+    }
+
+    public class OtherDocument
+    {
+        public string Id { get; set; } = string.Empty;
+    }
+
+    public class ThirdDocument
+    {
+        public string Id { get; set; } = string.Empty;
+    }
+
+    public class ForthDocument
+    {
+        public string Id { get; set; } = string.Empty;
+    }
+}

src/Marten/Internal/ProviderGraph.cs

@@ -14,6 +14,7 @@ namespace Marten.Internal;
 public class ProviderGraph: IProviderGraph
 {
     private readonly StoreOptions _options;
+    private readonly object _storageLock = new();
     private ImHashMap<Type, object> _storage = ImHashMap<Type, object>.Empty;
 
     public ProviderGraph(StoreOptions options)
@@ -23,7 +24,10 @@ public ProviderGraph(StoreOptions options)
 
     public void Append<T>(DocumentProvider<T> provider) where T : notnull
     {
-        _storage = _storage.AddOrUpdate(typeof(T), provider);
+        lock (_storageLock)
+        {
+            _storage = _storage.AddOrUpdate(typeof(T), provider);
+        }
     }
 
     public DocumentProvider<T> StorageFor<T>() where T : notnull
@@ -35,6 +39,21 @@ public DocumentProvider<T> StorageFor<T>() where T : notnull
             return stored.As<DocumentProvider<T>>();
         }
 
+        lock (_storageLock)
+        {
+            if (_storage.TryFind(documentType, out stored))
+            {
+                return stored.As<DocumentProvider<T>>();
+            }
+
+            return CreateDocumentProvider<T>();
+        }
+    }
+
+    private DocumentProvider<T> CreateDocumentProvider<T>() where T : notnull
+    {
+        var documentType = typeof(T);
+
         if (documentType == typeof(IEvent))
         {
             var rules = _options.CreateGenerationRules();