From 20ce84796e4b7bdea532d328aeac7f4a3b104ab5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Norbert=20=C5=BDof=C3=A1k?= <Norbert.Zofak@seznam.cz>
Date: Tue, 13 Feb 2024 17:10:48 +0100
Subject: [PATCH] Non-deterministic collection initialization fixed

---
 .../Utils/CollectionBase.cs                   | 76 ++++++++++---------
 1 file changed, 39 insertions(+), 37 deletions(-)

diff --git a/Orleans.Providers.MongoDB/Utils/CollectionBase.cs b/Orleans.Providers.MongoDB/Utils/CollectionBase.cs
index 9dcbd22..77d9d91 100644
--- a/Orleans.Providers.MongoDB/Utils/CollectionBase.cs
+++ b/Orleans.Providers.MongoDB/Utils/CollectionBase.cs
@@ -22,12 +22,25 @@ public class CollectionBase<TEntity>
 
         private readonly IMongoDatabase mongoDatabase;
         private readonly IMongoClient mongoClient;
-        private readonly Lazy<IMongoCollection<TEntity>> mongoCollection;
+        private readonly Action<MongoCollectionSettings> collectionConfigurator;
+        private IMongoCollection<TEntity> mongoCollection;
+        private readonly object mongoCollectionInitializerLock = new();
         private readonly bool createShardKey;
 
         protected IMongoCollection<TEntity> Collection
         {
-            get { return mongoCollection.Value; }
+            get
+            {
+                if (mongoCollection == null)
+                {
+                    lock (mongoCollectionInitializerLock)
+                    {
+                        mongoCollection ??= CreateCollection(collectionConfigurator);
+                    }
+                }
+
+                return mongoCollection;
+            }
         }
 
         protected IMongoDatabase Database
@@ -44,9 +57,9 @@ protected CollectionBase(IMongoClient mongoClient, string databaseName,
             Action<MongoCollectionSettings> collectionConfigurator, bool createShardKey)
         {
             this.mongoClient = mongoClient;
+            this.collectionConfigurator = collectionConfigurator;
 
             mongoDatabase = mongoClient.GetDatabase(databaseName);
-            mongoCollection = CreateCollection(collectionConfigurator);
 
             this.createShardKey = createShardKey;
         }
@@ -65,51 +78,40 @@ protected virtual void SetupCollection(IMongoCollection<TEntity> collection)
         {
         }
 
-        private Lazy<IMongoCollection<TEntity>> CreateCollection(Action<MongoCollectionSettings> collectionConfigurator)
+        private IMongoCollection<TEntity> CreateCollection(Action<MongoCollectionSettings> collectionConfigurator)
         {
-            return new Lazy<IMongoCollection<TEntity>>(() =>
-            {
-                var collectionFilter = new ListCollectionNamesOptions
-                {
-                    Filter = Builders<BsonDocument>.Filter.Eq("name", CollectionName())
-                };
+            var collectionName = CollectionName();
 
-                if (!mongoDatabase.ListCollectionNames(collectionFilter).Any())
-                {
-                    mongoDatabase.CreateCollection(CollectionName());
-                }
-
-                var collectionSettings = CollectionSettings() ?? new MongoCollectionSettings();
+            var collectionSettings = CollectionSettings() ?? new MongoCollectionSettings();
 
-                collectionConfigurator?.Invoke(collectionSettings);
+            collectionConfigurator?.Invoke(collectionSettings);
 
-                var databaseCollection = mongoDatabase.GetCollection<TEntity>(
-                    CollectionName(),
-                    collectionSettings);
+            var databaseCollection = mongoDatabase.GetCollection<TEntity>(
+                collectionName,
+                collectionSettings);
 
-                if (this.createShardKey)
+            if (createShardKey)
+            {
+                try
                 {
-                    try
+                    mongoClient.GetDatabase("admin").RunCommand<BsonDocument>(new BsonDocument
                     {
-                        Database.RunCommand<BsonDocument>(new BsonDocument
+                        ["shardCollection"] = $"{mongoDatabase.DatabaseNamespace.DatabaseName}.{collectionName}",
+                        ["key"] = new BsonDocument
                         {
-                            ["key"] = new BsonDocument
-                            {
-                                ["_id"] = "hashed"
-                            },
-                            ["shardCollection"] = $"{mongoDatabase.DatabaseNamespace.DatabaseName}.{CollectionName()}"
-                        });
-                    }
-                    catch (MongoException)
-                    {
-                        // Shared key probably created already.
-                    }
+                            ["_id"] = "hashed"
+                        }
+                    });
+                }
+                catch (MongoException)
+                {
+                    // Shared key probably created already.
                 }
+            }
 
-                SetupCollection(databaseCollection);
+            SetupCollection(databaseCollection);
 
-                return databaseCollection;
-            });
+            return databaseCollection;
         }
     }
 }
\ No newline at end of file