Merge pull request #125 from mksergiy/master

fix potential race conditions in ModelObjectPool
rebus-org · Dec 18, 2024 · 6d44441 · 6d44441
2 parents 2b07403 + 12c15c3
commit 6d44441
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/Rebus.RabbitMq.Tests/Examples/CheckMaximumDegreeOfParallelism.cs b/Rebus.RabbitMq.Tests/Examples/CheckMaximumDegreeOfParallelism.cs
@@ -50,7 +50,8 @@ public async Task TryParallelism(int maxParallelism)
             .Transport(t =>
             {
                 t.UseRabbitMq(RabbitMqTransportFactory.ConnectionString, queueName)
-                    .Prefetch(maxNumberOfMessagesToPrefetch: 3 * maxParallelism);
+                    .Prefetch(maxNumberOfMessagesToPrefetch: 3 * maxParallelism)
+                    .SetPublisherConfirms(false);
             })
             .Options(o =>
             {

diff --git a/Rebus.RabbitMq/Internals/ModelObjectPool.cs b/Rebus.RabbitMq/Internals/ModelObjectPool.cs
@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Concurrent;
+using System.Threading;
 using RabbitMQ.Client;
 
 namespace Rebus.Internals;
@@ -14,6 +15,7 @@ class ModelObjectPool : IDisposable
     readonly WriterModelPoolPolicy _policy;
 
     int _maxEntries;
+    int _currentCount = 0; 
 
     public ModelObjectPool(WriterModelPoolPolicy policy, int maxEntries)
     {
@@ -26,12 +28,22 @@ public void SetMaxEntries(int maxEntries)
         _maxEntries = maxEntries;
     }
 
-    public IModel Get() => _availableObjects.TryTake(out var model) ? model : _policy.Create();
+    public IModel Get()
+    {
+        if (_availableObjects.TryTake(out var model))
+        {
+            Interlocked.Decrement(ref _currentCount);
+            return model;
+        }
+
+        return _policy.Create();
+    }
 
     public void Return(IModel model)
     {
-        if (_availableObjects.Count >= _maxEntries)
+        if (Interlocked.Increment(ref _currentCount) > _maxEntries)
         {
+            Interlocked.Decrement(ref _currentCount);
             model.SafeDrop();
         }
         else