From aeb9e7d73de74a4abb68fab6f249614060cc1130 Mon Sep 17 00:00:00 2001
From: Dominic <dominic.couture@gmail.com>
Date: Wed, 24 Oct 2018 01:05:18 -0400
Subject: [PATCH] Fix #779: NRE on Dispose in ExecutionTimer (#782)

* Fix ThreadStatic issue with an object pool
---
 .../Utility/ExecutionTimer.cs                 | 18 +++++++-----
 .../Utility/ObjectPool.cs                     | 29 +++++++++++++++++++
 .../Utility/ExecutionTimerTests.cs            | 25 ++++++++++++++++
 .../Utility/ObjectPoolTests.cs                | 23 +++++++++++++++
 4 files changed, 87 insertions(+), 8 deletions(-)
 create mode 100644 src/PowerShellEditorServices/Utility/ObjectPool.cs
 create mode 100644 test/PowerShellEditorServices.Test/Utility/ExecutionTimerTests.cs
 create mode 100644 test/PowerShellEditorServices.Test/Utility/ObjectPoolTests.cs

diff --git a/src/PowerShellEditorServices/Utility/ExecutionTimer.cs b/src/PowerShellEditorServices/Utility/ExecutionTimer.cs
index e82d5b55c..994f83d27 100644
--- a/src/PowerShellEditorServices/Utility/ExecutionTimer.cs
+++ b/src/PowerShellEditorServices/Utility/ExecutionTimer.cs
@@ -20,8 +20,9 @@ namespace Microsoft.PowerShell.EditorServices.Utility
     /// </example>
     public struct ExecutionTimer : IDisposable
     {
-        [ThreadStatic]
-        private static Stopwatch t_stopwatch;
+        private static readonly ObjectPool<Stopwatch> s_stopwatchPool = new ObjectPool<Stopwatch>();
+
+        private Stopwatch _stopwatch;
 
         private readonly ILogger _logger;
 
@@ -50,7 +51,7 @@ public static ExecutionTimer Start(
             [CallerLineNumber] int callerLineNumber = -1)
         {
             var timer = new ExecutionTimer(logger, message, callerMemberName, callerFilePath, callerLineNumber);
-            Stopwatch.Start();
+            timer._stopwatch.Start();
             return timer;
         }
 
@@ -66,6 +67,7 @@ internal ExecutionTimer(
             _callerMemberName = callerMemberName;
             _callerFilePath = callerFilePath;
             _callerLineNumber = callerLineNumber;
+            _stopwatch = s_stopwatchPool.Rent();
         }
 
         /// <summary>
@@ -74,16 +76,18 @@ internal ExecutionTimer(
         /// </summary>
         public void Dispose()
         {
-            t_stopwatch.Stop();
+            _stopwatch.Stop();
 
             string logMessage = new StringBuilder()
                 .Append(_message)
                 .Append(" [")
-                .Append(t_stopwatch.ElapsedMilliseconds)
+                .Append(_stopwatch.ElapsedMilliseconds)
                 .Append("ms]")
                 .ToString();
 
-            t_stopwatch.Reset();
+            _stopwatch.Reset();
+
+            s_stopwatchPool.Return(_stopwatch);
 
             _logger.Write(
                 LogLevel.Verbose,
@@ -92,7 +96,5 @@ public void Dispose()
                 callerSourceFile: _callerFilePath,
                 callerLineNumber: _callerLineNumber);
         }
-
-        private static Stopwatch Stopwatch => t_stopwatch ?? (t_stopwatch = new Stopwatch());
     }
 }
diff --git a/src/PowerShellEditorServices/Utility/ObjectPool.cs b/src/PowerShellEditorServices/Utility/ObjectPool.cs
new file mode 100644
index 000000000..1a372f02e
--- /dev/null
+++ b/src/PowerShellEditorServices/Utility/ObjectPool.cs
@@ -0,0 +1,29 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using System.Collections.Concurrent;
+
+namespace Microsoft.PowerShell.EditorServices.Utility
+{
+    /// <summary>
+    /// A basic implementation of the object pool pattern.
+    /// </summary>
+    internal class ObjectPool<T>
+        where T : new()
+    {
+        private ConcurrentBag<T> _pool = new ConcurrentBag<T>();
+
+        /// <summary>
+        /// Get an instance of an object, either new or from the pool depending on availability.
+        /// </summary>
+        public T Rent() => _pool.TryTake(out var obj) ? obj : new T();
+
+        /// <summary>
+        /// Return an object to the pool.
+        /// </summary>
+        /// <param name="obj">The object to return to the pool.</param>
+        public void Return(T obj) => _pool.Add(obj);
+    }
+}
diff --git a/test/PowerShellEditorServices.Test/Utility/ExecutionTimerTests.cs b/test/PowerShellEditorServices.Test/Utility/ExecutionTimerTests.cs
new file mode 100644
index 000000000..098f70830
--- /dev/null
+++ b/test/PowerShellEditorServices.Test/Utility/ExecutionTimerTests.cs
@@ -0,0 +1,25 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using Microsoft.PowerShell.EditorServices.Utility;
+using System;
+using System.IO;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+using Xunit;
+
+namespace Microsoft.PowerShell.EditorServices.Test.Utility
+{
+    public class ExecutionTimerTests
+    {
+        [Fact]
+        public async void DoesNotThrowExceptionWhenDisposedOnAnotherThread()
+        {
+            var timer = ExecutionTimer.Start(Logging.CreateLogger().Build(), "Message");
+            await Task.Run(() => timer.Dispose());
+        }
+    }
+}
diff --git a/test/PowerShellEditorServices.Test/Utility/ObjectPoolTests.cs b/test/PowerShellEditorServices.Test/Utility/ObjectPoolTests.cs
new file mode 100644
index 000000000..8b37bce86
--- /dev/null
+++ b/test/PowerShellEditorServices.Test/Utility/ObjectPoolTests.cs
@@ -0,0 +1,23 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using Microsoft.PowerShell.EditorServices.Utility;
+using Xunit;
+
+namespace Microsoft.PowerShell.EditorServices.Test.Utility
+{
+    public class ObjectPoolTests
+    {
+        [Fact]
+        public void DoesNotCreateNewObjects()
+        {
+            var pool = new ObjectPool<object>();
+            var obj = pool.Rent();
+            pool.Return(obj);
+
+            Assert.Same(obj, pool.Rent());
+        }
+    }
+}