Fix leak in DbCommand's use of CancellationToken

We're only disposing of the registration on error; we should also dispose of it on success. While doing this, I've also avoided a delegate allocation per operation.
dotnet · Feb 24, 2017 · 297fcc3 · 297fcc3
1 parent 391964b
commit 297fcc3
Showing 1 changed file with 15 additions and 6 deletions.
diff --git a/src/System.Data.Common/src/System/Data/Common/DbCommand.cs b/src/System.Data.Common/src/System/Data/Common/DbCommand.cs
@@ -128,7 +128,7 @@ public virtual Task<int> ExecuteNonQueryAsync(CancellationToken cancellationToke
                 CancellationTokenRegistration registration = new CancellationTokenRegistration();
                 if (cancellationToken.CanBeCanceled)
                 {
-                    registration = cancellationToken.Register(CancelIgnoreFailure);
+                    registration = cancellationToken.Register(s => ((DbCommand)s).CancelIgnoreFailure(), this);
                 }
 
                 try
@@ -137,9 +137,12 @@ public virtual Task<int> ExecuteNonQueryAsync(CancellationToken cancellationToke
                 }
                 catch (Exception e)
                 {
-                    registration.Dispose();
                     return Task.FromException<int>(e);
                 }
+                finally
+                {
+                    registration.Dispose();
+                }
             }
         }
 
@@ -166,7 +169,7 @@ protected virtual Task<DbDataReader> ExecuteDbDataReaderAsync(CommandBehavior be
                 CancellationTokenRegistration registration = new CancellationTokenRegistration();
                 if (cancellationToken.CanBeCanceled)
                 {
-                    registration = cancellationToken.Register(CancelIgnoreFailure);
+                    registration = cancellationToken.Register(s => ((DbCommand)s).CancelIgnoreFailure(), this);
                 }
 
                 try
@@ -175,9 +178,12 @@ protected virtual Task<DbDataReader> ExecuteDbDataReaderAsync(CommandBehavior be
                 }
                 catch (Exception e)
                 {
-                    registration.Dispose();
                     return Task.FromException<DbDataReader>(e);
                 }
+                finally
+                {
+                    registration.Dispose();
+                }
             }
         }
 
@@ -195,7 +201,7 @@ public virtual Task<object> ExecuteScalarAsync(CancellationToken cancellationTok
                 CancellationTokenRegistration registration = new CancellationTokenRegistration();
                 if (cancellationToken.CanBeCanceled)
                 {
-                    registration = cancellationToken.Register(CancelIgnoreFailure);
+                    registration = cancellationToken.Register(s => ((DbCommand)s).CancelIgnoreFailure(), this);
                 }
 
                 try
@@ -204,9 +210,12 @@ public virtual Task<object> ExecuteScalarAsync(CancellationToken cancellationTok
                 }
                 catch (Exception e)
                 {
-                    registration.Dispose();
                     return Task.FromException<object>(e);
                 }
+                finally
+                {
+                    registration.Dispose();
+                }
             }
         }