Stop leaking JoinableTask via ExecutionContext

Fixes microsoft#274 and gets the recently added test to pass
AArnott · Aug 12, 2018 · ff42b9d · ff42b9d
1 parent aa254f1
commit ff42b9d
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 3 deletions.
diff --git 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,24 @@ internal SynchronizationContext ApplicableJobSyncContext
             }
         }
 
+        /// <summary>
+        /// Gets a weak reference to this object.
+        /// </summary>
+        internal WeakReference<JoinableTask> WeakSelf
+        {
+            get
+            {
+                if (this.weakSelf == null)
+                {
+                    // We take the trouble to be thread-safe about it so that we might optimize by comparing
+                    // the WeakReference objects themselves rather than have to access their values.
+                    Interlocked.CompareExchange(ref this.weakSelf, new WeakReference<JoinableTask>(this), null);
+                }
+
+                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
@@ -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>