Protect AutoMappingStore from Threading issues.

samus · Aug 30, 2010 · f662081 · f662081
1 parent 20fa57e
commit f662081
Showing 1 changed file with 33 additions and 12 deletions.
diff --git a/source/MongoDB/Configuration/Mapping/AutoMappingStore.cs b/source/MongoDB/Configuration/Mapping/AutoMappingStore.cs
@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Threading;
 using MongoDB.Configuration.Mapping.Auto;
 using MongoDB.Configuration.Mapping.Model;
 
@@ -11,6 +12,7 @@ public class AutoMappingStore : IMappingStore
     {
         private readonly IAutoMapper _autoMapper;
         private readonly Dictionary<Type, IClassMap> _autoMaps;
+        private readonly ReaderWriterLockSlim _lock = new ReaderWriterLockSlim(LockRecursionPolicy.SupportsRecursion);
         private readonly IMappingStore _wrappedMappingStore;
 
         /// <summary>
@@ -71,23 +73,42 @@ public AutoMappingStore(IAutoMapper autoMapper, IMappingStore mappingStore)
         /// <returns></returns>
         public IClassMap GetClassMap(Type classType)
         {
-            IClassMap classMap;
+            try
-
-            if(_autoMaps.TryGetValue(classType, out classMap))
-                return classMap;
-
-            if(_wrappedMappingStore != null)
             {
-                classMap = _wrappedMappingStore.GetClassMap(classType);
+                _lock.EnterUpgradeableReadLock();
-                if(classMap != null)
+
+                IClassMap classMap;
+                if(_autoMaps.TryGetValue(classType, out classMap))
                     return classMap;
-            }
 
-            classMap = _autoMapper.CreateClassMap(classType, GetClassMap);
+                if(_wrappedMappingStore != null)
+                {
+                    classMap = _wrappedMappingStore.GetClassMap(classType);
+                    if(classMap != null)
+                        return classMap;
+                }
+
+                classMap = _autoMapper.CreateClassMap(classType, GetClassMap);
 
-            _autoMaps.Add(classType, classMap);
+                try
+                {
+                    _lock.EnterWriteLock();
 
-            return classMap;
+                    _autoMaps.Add(classType, classMap);
+
+                    return classMap;
+                }
+                finally
+                {
+                    if(_lock.IsWriteLockHeld)
+                        _lock.ExitWriteLock();
+                }
+            }
+            finally
+            {
+                if(_lock.IsUpgradeableReadLockHeld)
+                    _lock.ExitUpgradeableReadLock();
+            }
         }
     }
 }