From d2796c6a8751b32bd5952976c1663f969720a482 Mon Sep 17 00:00:00 2001
From: Andrew Arnott <andarno@microsoft.com>
Date: Sun, 12 Aug 2018 11:21:11 -0700
Subject: [PATCH] Stop leaking JoinableTask via ExecutionContext

Fixes #274 and gets the recently added test to pass
---
 .../JoinableTask.cs                           | 21 +++++++++++++++++++
 .../JoinableTaskContext.cs                    | 12 ++++++++---
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/src/Microsoft.VisualStudio.Threading/JoinableTask.cs b/src/Microsoft.VisualStudio.Threading/JoinableTask.cs
index b5116e336..a58849e3c 100644
--- a/src/Microsoft.VisualStudio.Threading/JoinableTask.cs
+++ b/src/Microsoft.VisualStudio.Threading/JoinableTask.cs
@@ -120,6 +120,11 @@ public partial class JoinableTask
         [DebuggerBrowsable(DebuggerBrowsableState.Never)]
         private Delegate initialDelegate;
 
+        /// <summary>
+        /// Backing field for the <see cref="WeakSelf"/> property.
+        /// </summary>
+        private WeakReference<JoinableTask> weakSelf;
+
         /// <summary>
         /// Initializes a new instance of the <see cref="JoinableTask"/> class.
         /// </summary>
@@ -332,6 +337,22 @@ internal SynchronizationContext ApplicableJobSyncContext
             }
         }
 
+        /// <summary>
+        /// Gets a weak reference to this object.
+        /// </summary>
+        internal WeakReference<JoinableTask> WeakSelf
+        {
+            get
+            {
+                if (this.weakSelf == null)
+                {
+                    this.weakSelf = new WeakReference<JoinableTask>(this);
+                }
+
+                return this.weakSelf;
+            }
+        }
+
         /// <summary>
         /// Gets the flags set on this task.
         /// </summary>
diff --git a/src/Microsoft.VisualStudio.Threading/JoinableTaskContext.cs b/src/Microsoft.VisualStudio.Threading/JoinableTaskContext.cs
index c62fe9ca4..c83d0b3e1 100644
--- a/src/Microsoft.VisualStudio.Threading/JoinableTaskContext.cs
+++ b/src/Microsoft.VisualStudio.Threading/JoinableTaskContext.cs
@@ -86,7 +86,7 @@ public partial class JoinableTaskContext : IDisposable
         /// <summary>
         /// An AsyncLocal value that carries the joinable instance associated with an async operation.
         /// </summary>
-        private readonly AsyncLocal<JoinableTask> joinableOperation = new AsyncLocal<JoinableTask>();
+        private readonly AsyncLocal<WeakReference<JoinableTask>> joinableOperation = new AsyncLocal<WeakReference<JoinableTask>>();
 
         /// <summary>
         /// The set of tasks that have started but have not yet completed.
@@ -251,8 +251,14 @@ internal object SyncContextLock
         /// </summary>
         internal JoinableTask AmbientTask
         {
-            get { return this.joinableOperation.Value; }
-            set { this.joinableOperation.Value = value; }
+            get
+            {
+                JoinableTask result = null;
+                this.joinableOperation.Value?.TryGetTarget(out result);
+                return result;
+            }
+
+            set => this.joinableOperation.Value = value?.WeakSelf;
         }
 
         /// <summary>