From 2d48b2c42d925bc9395205bdf3c199d418b96fdb Mon Sep 17 00:00:00 2001
From: yaroslavtykhonchuk <yaroslav.tykhonchuk@youscan.io>
Date: Fri, 15 Dec 2023 16:50:57 +0200
Subject: [PATCH] Do not fail on disposed timer

When we have a race condition and timer becomes disposed, the other thread may try to access timer and received ObjectDisposedException. That's why we set timer reference to null before disposing it, and make timer nullable. This is not 100% thread safe code, but want to test how it will behave.
---
 AzureBatchQueue/TimerBatch.cs | 27 +++++++++++++++++++--------
 1 file changed, 19 insertions(+), 8 deletions(-)

diff --git a/AzureBatchQueue/TimerBatch.cs b/AzureBatchQueue/TimerBatch.cs
index 76d2a54..00998da 100644
--- a/AzureBatchQueue/TimerBatch.cs
+++ b/AzureBatchQueue/TimerBatch.cs
@@ -12,7 +12,7 @@ internal class TimerBatch<T>
     readonly ILogger logger;
     readonly ConcurrentDictionary<string, BatchItem<T>> items;
 
-    readonly Timer timer;
+    Timer? timer;
     BatchCompletedResult? completedResult;
 
     public TimerBatch(BatchQueue<T> batchQueue, QueueMessage<T[]> msg, int maxDequeueCount, ILogger logger)
@@ -36,7 +36,7 @@ async Task Flush()
     {
         try
         {
-            await timer.DisposeAsync();
+            DisposeTimer();
 
             await DoFlush();
         }
@@ -57,6 +57,9 @@ async Task Flush()
 
         async Task DoFlush()
         {
+            if (completedResult != null)
+                return;
+
             if (items.IsEmpty)
             {
                 completedResult = BatchCompletedResult.FullyProcessed;
@@ -77,6 +80,16 @@ async Task DoFlush()
         }
     }
 
+    /// <summary>
+    /// Set timer reference to null, so that call to timer in Complete() will not throw ObjectDisposedException
+    /// </summary>
+    void DisposeTimer()
+    {
+        var timerCopy = timer;
+        timer = null;
+        timerCopy.Dispose();
+    }
+
     QueueMessage<T[]> Message()
     {
         var notCompletedItems = items.Values.Select(x => x.Item).ToArray();
@@ -93,13 +106,11 @@ public BatchItemCompleteResult Complete(string itemId)
         if (!res)
             throw new ItemNotFoundException(itemId);
 
-        if (items.IsEmpty)
-        {
-            timer.Change(TimeSpan.Zero, Timeout.InfiniteTimeSpan);
-            return BatchItemCompleteResult.BatchFullyProcessed;
-        }
+        if (!items.IsEmpty)
+            return BatchItemCompleteResult.Completed;
 
-        return BatchItemCompleteResult.Completed;
+        timer?.Change(TimeSpan.Zero, Timeout.InfiniteTimeSpan);
+        return BatchItemCompleteResult.BatchFullyProcessed;
     }
 
     public IEnumerable<BatchItem<T>> Unpack()