From dfa530a0332a0111248e43023cf97c0c4534bd02 Mon Sep 17 00:00:00 2001
From: Sergiy Matviychuk <sergiy.matviychuk@renome.ua>
Date: Sun, 1 Dec 2024 23:58:08 +0200
Subject: [PATCH] fix potential race conditions in ModelObjectPool

---
 Rebus.RabbitMq/Internals/ModelObjectPool.cs | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/Rebus.RabbitMq/Internals/ModelObjectPool.cs b/Rebus.RabbitMq/Internals/ModelObjectPool.cs
index 55271a1..4996037 100644
--- 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