From 812dfb97d3f2dd049f67348320802eade6f4480f Mon Sep 17 00:00:00 2001
From: David Fowler <davidfowl@gmail.com>
Date: Fri, 17 Apr 2020 19:05:44 -0700
Subject: [PATCH] Fix TimerAwaitable rooting (#20937)

* Fix TimerAwaitable rooting
- This fixes an issue where components that use timer awaitable currently need to be explicitly stopped and disposed for it to be unrooted. This makes it possible to write a finalizer.
---
 .../test/UnitTests/TimerAwaitableTests.cs     | 69 +++++++++++++++++++
 src/SignalR/common/Shared/TimerAwaitable.cs   | 17 ++++-
 2 files changed, 84 insertions(+), 2 deletions(-)
 create mode 100644 src/SignalR/clients/csharp/Client/test/UnitTests/TimerAwaitableTests.cs

diff --git a/src/SignalR/clients/csharp/Client/test/UnitTests/TimerAwaitableTests.cs b/src/SignalR/clients/csharp/Client/test/UnitTests/TimerAwaitableTests.cs
new file mode 100644
index 0000000000..2c005ebe0f
--- /dev/null
+++ b/src/SignalR/clients/csharp/Client/test/UnitTests/TimerAwaitableTests.cs
@@ -0,0 +1,69 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Runtime.CompilerServices;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Internal;
+using Xunit;
+
+namespace Microsoft.AspNetCore.SignalR.Client.Tests
+{
+    public class TimerAwaitableTests
+    {
+        [Fact]
+        public void FinalizerRunsIfTimerAwaitableReferencesObject()
+        {
+            var tcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
+            UseTimerAwaitableAndUnref(tcs);
+
+            // Make sure it *really* cleans up
+            for (int i = 0; i < 5 && !tcs.Task.IsCompleted; i++)
+            {
+                GC.Collect();
+                GC.WaitForPendingFinalizers();
+            }
+
+            // Make sure the finalizer runs
+            Assert.True(tcs.Task.IsCompleted);
+        }
+
+        [MethodImpl(MethodImplOptions.NoInlining)]
+        private void UseTimerAwaitableAndUnref(TaskCompletionSource<object> tcs)
+        {
+            _ = new ObjectWithTimerAwaitable(tcs).Start();
+        }
+    }
+
+    // This object holds onto a TimerAwaitable referencing the callback (the async continuation is the callback)
+    // it also has a finalizer that triggers a tcs so callers can be notified when this object is being cleaned up.
+    public class ObjectWithTimerAwaitable
+    {
+        private readonly TimerAwaitable _timer;
+        private readonly TaskCompletionSource<object> _tcs;
+        private int _count;
+
+        public ObjectWithTimerAwaitable(TaskCompletionSource<object> tcs)
+        {
+            _tcs = tcs;
+            _timer = new TimerAwaitable(TimeSpan.Zero, TimeSpan.FromSeconds(1));
+            _timer.Start();
+        }
+
+        public async Task Start()
+        {
+            using (_timer)
+            {
+                while (await _timer)
+                {
+                    _count++;
+                }
+            }
+        }
+
+        ~ObjectWithTimerAwaitable()
+        {
+            _tcs.TrySetResult(null);
+        }
+    }
+}
diff --git a/src/SignalR/common/Shared/TimerAwaitable.cs b/src/SignalR/common/Shared/TimerAwaitable.cs
index ee6cb0fd23..c40d55b42d 100644
--- a/src/SignalR/common/Shared/TimerAwaitable.cs
+++ b/src/SignalR/common/Shared/TimerAwaitable.cs
@@ -50,7 +50,20 @@ namespace Microsoft.AspNetCore.Internal
                                 restoreFlow = true;
                             }
 
-                            _timer = new Timer(state => ((TimerAwaitable)state).Tick(), this, _dueTime, _period);
+                            // This fixes the cycle by using a WeakReference to the state object. The object graph now looks like this:
+                            // Timer -> TimerHolder -> TimerQueueTimer -> WeakReference<TimerAwaitable> -> Timer -> ...
+                            // If TimerAwaitable falls out of scope, the timer should be released.
+                            _timer = new Timer(state =>
+                            {
+                                var weakRef = (WeakReference<TimerAwaitable>)state;
+                                if (weakRef.TryGetTarget(out var thisRef))
+                                {
+                                    thisRef.Tick();
+                                }
+                            },
+                            new WeakReference<TimerAwaitable>(this),
+                            _dueTime,
+                            _period);
                         }
                         finally
                         {
@@ -125,4 +138,4 @@ namespace Microsoft.AspNetCore.Internal
             }
         }
     }
-}
\ No newline at end of file
+}