From 7f94248314408f5cfd914452b5599e3f8a439724 Mon Sep 17 00:00:00 2001
From: Wraith2 <wraith2@gmail.com>
Date: Thu, 6 Mar 2025 01:31:03 +0000
Subject: [PATCH] add try-catch around getting open connection rework retry to
 avoid lambda capture

---
 .../Microsoft/Data/SqlClient/SqlCommand.cs    | 44 ++++++++++++-------
 .../Microsoft/Data/SqlClient/SqlCommand.cs    | 44 ++++++++++++-------
 2 files changed, 56 insertions(+), 32 deletions(-)

diff --git a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlCommand.cs b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlCommand.cs
index e64d81da40..13317a80e4 100644
--- a/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlCommand.cs
+++ b/src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlCommand.cs
@@ -2445,9 +2445,16 @@ long firstAttemptStart
                                 _cachedAsyncState.ResetAsyncState();
                             }
 
-                            _activeConnection.GetOpenTdsConnection().DecrementAsyncCount();
+                            try
+                            {
+                                _activeConnection.GetOpenTdsConnection().DecrementAsyncCount();
 
-                            globalCompletion.TrySetException(e);
+                                globalCompletion.TrySetException(e);
+                            }
+                            catch (Exception e2)
+                            {
+                                globalCompletion.TrySetException(e2);
+                            }
                         }
                         else
                         {
@@ -2464,21 +2471,26 @@ long firstAttemptStart
                                     TdsParserStaticMethods.GetRemainingTimeout(timeout, firstAttemptStart), true /*inRetry*/,
                                     asyncWrite);
 
-                                retryTask.ContinueWith(retryTsk =>
-                                {
-                                    if (retryTsk.IsFaulted)
-                                    {
-                                        globalCompletion.TrySetException(retryTsk.Exception.InnerException);
-                                    }
-                                    else if (retryTsk.IsCanceled)
+                                retryTask.ContinueWith(
+                                    static (Task<object> retryTask, object state) =>
                                     {
-                                        globalCompletion.TrySetCanceled();
-                                    }
-                                    else
-                                    {
-                                        globalCompletion.TrySetResult(retryTsk.Result);
-                                    }
-                                }, TaskScheduler.Default);
+                                        TaskCompletionSource<object> completion = (TaskCompletionSource<object>)state;
+                                        if (retryTask.IsFaulted)
+                                        {
+                                            completion.TrySetException(retryTask.Exception.InnerException);
+                                        }
+                                        else if (retryTask.IsCanceled)
+                                        {
+                                            completion.TrySetCanceled();
+                                        }
+                                        else
+                                        {
+                                            completion.TrySetResult(retryTask.Result);
+                                        }
+                                    }, 
+                                    state: globalCompletion,
+                                    TaskScheduler.Default
+                                );
                             }
                             catch (Exception e2)
                             {
diff --git a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlCommand.cs b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlCommand.cs
index 579c64181a..4e6d7d0099 100644
--- a/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlCommand.cs
+++ b/src/Microsoft.Data.SqlClient/netfx/src/Microsoft/Data/SqlClient/SqlCommand.cs
@@ -2718,9 +2718,16 @@ private bool TriggerInternalEndAndRetryIfNecessary(CommandBehavior behavior, obj
                                 {
                                     _cachedAsyncState.ResetAsyncState();
                                 }
-                                _activeConnection.GetOpenTdsConnection().DecrementAsyncCount();
+                                try
+                                {
+                                    _activeConnection.GetOpenTdsConnection().DecrementAsyncCount();
 
-                                globalCompletion.TrySetException(e);
+                                    globalCompletion.TrySetException(e);
+                                }
+                                catch (Exception e2)
+                                {
+                                    globalCompletion.TrySetException(e2);
+                                }
                             }
                             else
                             {
@@ -2735,21 +2742,26 @@ private bool TriggerInternalEndAndRetryIfNecessary(CommandBehavior behavior, obj
                                     _internalEndExecuteInitiated = false;
                                     Task<object> retryTask = (Task<object>)retryFunc(behavior, null, stateObject, TdsParserStaticMethods.GetRemainingTimeout(timeout, firstAttemptStart), true/*inRetry*/, asyncWrite);
 
-                                    retryTask.ContinueWith(retryTsk =>
-                                    {
-                                        if (retryTsk.IsFaulted)
-                                        {
-                                            globalCompletion.TrySetException(retryTsk.Exception.InnerException);
-                                        }
-                                        else if (retryTsk.IsCanceled)
+                                    retryTask.ContinueWith(
+                                        static (Task<object> retryTask, object state) =>
                                         {
-                                            globalCompletion.TrySetCanceled();
-                                        }
-                                        else
-                                        {
-                                            globalCompletion.TrySetResult(retryTsk.Result);
-                                        }
-                                    }, TaskScheduler.Default);
+                                            TaskCompletionSource<object> completion = (TaskCompletionSource<object>)state;
+                                            if (retryTask.IsFaulted)
+                                            {
+                                                completion.TrySetException(retryTask.Exception.InnerException);
+                                            }
+                                            else if (retryTask.IsCanceled)
+                                            {
+                                                completion.TrySetCanceled();
+                                            }
+                                            else
+                                            {
+                                                completion.TrySetResult(retryTask.Result);
+                                            }
+                                        },
+                                        state: globalCompletion,
+                                        TaskScheduler.Default
+                                    );
                                 }
                                 catch (Exception e2)
                                 {