Harden ExecutionContext SuppressFlow

dotnet · Dec 27, 2017 · afba626 · afba626
1 parent e20cb38
commit afba626
Show file tree

Hide file tree

Showing 5 changed files with 100 additions and 15 deletions.
diff --git a/src/System.Data.SqlClient/src/System/Data/ProviderBase/DbConnectionFactory.cs b/src/System.Data.SqlClient/src/System/Data/ProviderBase/DbConnectionFactory.cs
@@ -125,12 +125,27 @@ virtual internal DbConnectionPoolGroupProviderInfo CreateConnectionPoolGroupProv
 
         private Timer CreatePruningTimer()
         {
+            Timer timer;
             // Don't capture the current ExecutionContext and its AsyncLocals onto the timer causing them to live forever
-            ExecutionContext.SuppressFlow();
-            TimerCallback callback = new TimerCallback(PruneConnectionPoolGroups);
-            Timer timer = new Timer(callback, null, PruningDueTime, PruningPeriod);
-            // Restore the current ExecutionContext
-            ExecutionContext.RestoreFlow();
+            bool restoreFlow = false;
+            try
+            {
+                if (!ExecutionContext.IsFlowSuppressed())
+                {
+                    ExecutionContext.SuppressFlow();
+                    restoreFlow = true;
+                }
+
+                TimerCallback callback = new TimerCallback(PruneConnectionPoolGroups);
+                timer = new Timer(callback, null, PruningDueTime, PruningPeriod);
+            }
+            finally
+            {
+                // Restore the current ExecutionContext
+                if (restoreFlow)
+                    ExecutionContext.RestoreFlow();
+            }
+
             return timer;
         }
 

diff --git a/src/System.Data.SqlClient/src/System/Data/ProviderBase/DbConnectionPool.cs b/src/System.Data.SqlClient/src/System/Data/ProviderBase/DbConnectionPool.cs
@@ -666,11 +666,26 @@ internal void Clear()
 
         private Timer CreateCleanupTimer()
         {
+            Timer timer;
             // Don't capture the current ExecutionContext and its AsyncLocals onto the timer causing them to live forever
-            ExecutionContext.SuppressFlow();
-            Timer timer = (new Timer(new TimerCallback(this.CleanupCallback), null, _cleanupWait, _cleanupWait));
-            // Restore the current ExecutionContext
-            ExecutionContext.RestoreFlow();
+            bool restoreFlow = false;
+            try
+            {
+                if (!ExecutionContext.IsFlowSuppressed())
+                {
+                    ExecutionContext.SuppressFlow();
+                    restoreFlow = true;
+                }
+
+                timer = (new Timer(new TimerCallback(this.CleanupCallback), null, _cleanupWait, _cleanupWait));
+            }
+            finally
+            {
+                // Restore the current ExecutionContext
+                if (restoreFlow)
+                    ExecutionContext.RestoreFlow();
+            }
+
             return timer;
         }
 
@@ -732,7 +747,26 @@ private DbConnectionInternal CreateObject(DbConnection owningObject, DbConnectio
                 // Make sure the timer starts even if ThreadAbort occurs after setting the ErrorEvent.
 
                 // timer allocation has to be done out of CER block
-                Timer t = new Timer(new TimerCallback(this.ErrorCallback), null, Timeout.Infinite, Timeout.Infinite);
+                Timer timer;
+                // Don't capture the current ExecutionContext and its AsyncLocals onto the timer causing them to live forever
+                bool restoreFlow = false;
+                try
+                {
+                    if (!ExecutionContext.IsFlowSuppressed())
+                    {
+                        ExecutionContext.SuppressFlow();
+                        restoreFlow = true;
+                    }
+
+                    timer = new Timer(new TimerCallback(this.ErrorCallback), null, Timeout.Infinite, Timeout.Infinite);
+                }
+                finally
+                {
+                    // Restore the current ExecutionContext
+                    if (restoreFlow)
+                        ExecutionContext.RestoreFlow();
+                }
+
                 bool timerIsNotDisposed;
                 try { }
                 finally
@@ -743,8 +777,8 @@ private DbConnectionInternal CreateObject(DbConnection owningObject, DbConnectio
                     // Enable the timer.
                     // Note that the timer is created to allow periodic invocation. If ThreadAbort occurs in the middle of ErrorCallback,
                     // the timer will restart. Otherwise, the timer callback (ErrorCallback) destroys the timer after resetting the error to avoid second callback.
-                    _errorTimer = t;
-                    timerIsNotDisposed = t.Change(_errorWait, _errorWait);
+                    _errorTimer = timer;
+                    timerIsNotDisposed = timer.Change(_errorWait, _errorWait);
                 }
 
                 Debug.Assert(timerIsNotDisposed, "ErrorCallback timer has been disposed");

diff --git a/src/System.Data.SqlClient/src/System/Data/SqlClient/SqlDependencyListener.cs b/src/System.Data.SqlClient/src/System/Data/SqlClient/SqlDependencyListener.cs
@@ -612,7 +612,9 @@ private void Restart(object unused)
                         _timeoutParam.Value = _defaultWaitforTimeout; // If success, reset to default for re-queue.
                         AsynchronouslyQueryServiceBrokerQueue();
                         _errorState = false;
+                        Timer retryTimer = _retryTimer;
                         _retryTimer = null;
+                        retryTimer?.Dispose();
                     }
                 }
 
@@ -1470,4 +1472,4 @@ internal bool Stop(
 
          return stopped;
     }
-}
+}
diff --git a/src/System.Data.SqlClient/src/System/Data/SqlClient/SqlDependencyUtils.cs b/src/System.Data.SqlClient/src/System/Data/SqlClient/SqlDependencyUtils.cs
@@ -80,7 +80,24 @@ private SqlDependencyPerAppDomainDispatcher()
             _notificationIdToDependenciesHash = new Dictionary<string, DependencyList>();
             _commandHashToNotificationId = new Dictionary<string, string>();
 
-            _timeoutTimer = new Timer(new TimerCallback(TimeoutTimerCallback), null, Timeout.Infinite, Timeout.Infinite);
+            // Don't capture the current ExecutionContext and its AsyncLocals onto the timer causing them to live forever
+            bool restoreFlow = false;
+            try
+            {
+                if (!ExecutionContext.IsFlowSuppressed())
+                {
+                    ExecutionContext.SuppressFlow();
+                    restoreFlow = true;
+                }
+
+                _timeoutTimer = new Timer(new TimerCallback(TimeoutTimerCallback), null, Timeout.Infinite, Timeout.Infinite);
+            }
+            finally
+            {
+                // Restore the current ExecutionContext
+                if (restoreFlow)
+                    ExecutionContext.RestoreFlow();
+            }
 
             SubscribeToAppDomainUnload();
         }

diff --git a/src/System.Data.SqlClient/src/System/Data/SqlClient/TdsParserStateObject.cs b/src/System.Data.SqlClient/src/System/Data/SqlClient/TdsParserStateObject.cs
@@ -2295,7 +2295,24 @@ internal void ReadSni(TaskCompletionSource<object> completion)
 
                 if (_networkPacketTimeout == null)
                 {
-                    _networkPacketTimeout = new Timer(OnTimeout, null, Timeout.Infinite, Timeout.Infinite);
+                    // Don't capture the current ExecutionContext and its AsyncLocals onto the timer causing them to live forever
+                    bool restoreFlow = false;
+                    try
+                    {
+                        if (!ExecutionContext.IsFlowSuppressed())
+                        {
+                            ExecutionContext.SuppressFlow();
+                            restoreFlow = true;
+                        }
+
+                        _networkPacketTimeout = new Timer(OnTimeout, null, Timeout.Infinite, Timeout.Infinite);
+                    }
+                    finally
+                    {
+                        // Restore the current ExecutionContext
+                        if (restoreFlow)
+                            ExecutionContext.RestoreFlow();
+                    }
                 }
 
                 // -1 == Infinite