From 631f54fbe414ccb5965475f5f15afc694e2ded63 Mon Sep 17 00:00:00 2001
From: Sebastian Stehle <sebastian@squidex.io>
Date: Sun, 23 Jul 2017 17:51:23 +0200
Subject: [PATCH] Timer implementation improved.

---
 .../Timers/CompletionTimer.cs                 | 69 +++++++++++--------
 .../Timers/CompletionTimerTests.cs            | 34 +++++++++
 2 files changed, 76 insertions(+), 27 deletions(-)
 create mode 100644 tests/Squidex.Infrastructure.Tests/Timers/CompletionTimerTests.cs

diff --git a/src/Squidex.Infrastructure/Timers/CompletionTimer.cs b/src/Squidex.Infrastructure/Timers/CompletionTimer.cs
index 16563d938..eec58576c 100644
--- a/src/Squidex.Infrastructure/Timers/CompletionTimer.cs
+++ b/src/Squidex.Infrastructure/Timers/CompletionTimer.cs
@@ -7,7 +7,6 @@
 // ==========================================================================
 
 using System;
-using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -19,38 +18,17 @@ namespace Squidex.Infrastructure.Timers
     {
         private readonly CancellationTokenSource disposeToken = new CancellationTokenSource();
         private readonly Task runTask;
-        private CancellationTokenSource delayToken;
+        private int requiresAtLeastOne;
+        private CancellationTokenSource wakeupToken;
         
-        public CompletionTimer(int delayInMs, Func<CancellationToken, Task> callback)
+        public CompletionTimer(int delayInMs, Func<CancellationToken, Task> callback, int initialDelay = 0)
         {
             Guard.NotNull(callback, nameof(callback));
             Guard.GreaterThan(delayInMs, 0, nameof(delayInMs));
 
-            runTask = RunInternal(delayInMs, callback);
+            runTask = RunInternal(delayInMs, initialDelay, callback);
         }
 
-        private async Task RunInternal(int delay, Func<CancellationToken, Task> callback)
-        {
-            while (!disposeToken.IsCancellationRequested)
-            {
-                try
-                {
-                    await callback(disposeToken.Token).ConfigureAwait(false);
-
-                    delayToken = new CancellationTokenSource();
-
-                    using (var cts = CancellationTokenSource.CreateLinkedTokenSource(disposeToken.Token, delayToken.Token))
-                    {
-                        await Task.Delay(delay, cts.Token).ConfigureAwait(false);
-                    }
-                }
-                catch (TaskCanceledException)
-                {
-                    Debug.WriteLine("Task in TriggerTimer has been cancelled.");
-                }
-            }
-        }
-        
         protected override void DisposeObject(bool disposing)
         {
             if (disposing)
@@ -65,7 +43,44 @@ namespace Squidex.Infrastructure.Timers
         {
             ThrowIfDisposed();
 
-            delayToken?.Cancel();
+            Interlocked.CompareExchange(ref requiresAtLeastOne, 2, 0);
+
+            wakeupToken?.Cancel();
+        }
+
+        private async Task RunInternal(int delay, int initialDelay, Func<CancellationToken, Task> callback)
+        {
+            if (initialDelay > 0)
+            {
+                await WaitAsync(initialDelay);
+            }
+
+            while (requiresAtLeastOne == 2 || !disposeToken.IsCancellationRequested)
+            {
+                try
+                {
+                    await callback(disposeToken.Token).ConfigureAwait(false);
+                }
+                catch (OperationCanceledException) { }
+
+                requiresAtLeastOne = 1;
+
+                await WaitAsync(delay);
+            }
+        }
+
+        private async Task WaitAsync(int intervall)
+        {
+            try
+            {
+                wakeupToken = new CancellationTokenSource();
+
+                using (var cts = CancellationTokenSource.CreateLinkedTokenSource(disposeToken.Token, wakeupToken.Token))
+                {
+                    await Task.Delay(intervall, cts.Token).ConfigureAwait(false);
+                }
+            }
+            catch (OperationCanceledException) { }
         }
     }
 }
diff --git a/tests/Squidex.Infrastructure.Tests/Timers/CompletionTimerTests.cs b/tests/Squidex.Infrastructure.Tests/Timers/CompletionTimerTests.cs
new file mode 100644
index 000000000..aa7f228c0
--- /dev/null
+++ b/tests/Squidex.Infrastructure.Tests/Timers/CompletionTimerTests.cs
@@ -0,0 +1,34 @@
+// ==========================================================================
+//  CompletionTimerTests.cs
+//  Squidex Headless CMS
+// ==========================================================================
+//  Copyright (c) Squidex Group
+//  All rights reserved.
+// ==========================================================================
+
+using Squidex.Infrastructure.Tasks;
+using Xunit;
+
+namespace Squidex.Infrastructure.Timers
+{
+    public class CompletionTimerTests
+    {
+        [Fact]
+        public void Should_invoke_once_even_with_delay()
+        {
+            var called = false;
+
+            var timer = new CompletionTimer(20000, ct =>
+            {
+                called = true;
+
+                return TaskHelper.Done;
+            }, 20000);
+
+            timer.Wakeup();
+            timer.Dispose();
+
+            Assert.True(called);
+        }
+    }
+}