Fixed CancellationTokenSource Issue

agoda-com · Feb 27, 2024 · 18f1b8d · 18f1b8d
1 parent 6f1c823
commit 18f1b8d
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 25 deletions.
diff --git a/Agoda.Frameworks.DB.Tests/DbRepositoryTest.cs b/Agoda.Frameworks.DB.Tests/DbRepositoryTest.cs
@@ -666,28 +666,6 @@ public void QueryMultipleAsync_Retry_Failure()
             Assert.IsNotNull(_onQueryCompleteEvents[0].Error);
             Assert.IsNotNull(_onQueryCompleteEvents[1].Error);
         }
-
-        [Test]
-        public async Task ExecuteReaderAsync_CancellationToken_Cancelled()
-        {
-            var cancellationToken = new CancellationTokenSource(TimeSpan.FromDays(1));
-            cancellationToken.Cancel();
-            var maxAttemptCount = 2;
-            Assert.ThrowsAsync<TaskCanceledException>(async () => await _db.ExecuteReaderAsync<string>("mobile_ro", "db.v1.sp_foo", 1,
-                maxAttemptCount,
-                cancellationToken.Token,
-                new IDbDataParameter[]
-                {
-                    new SqlParameter("@param1", "value1"),
-                    new SqlParameter("@param2", "value2")
-
-                }, reader => { return Task.FromResult("");})
-            );
-
-            _dbResources.Verify(x => x.ChooseDb("mobile_ro").UpdateWeight(It.IsAny<string>(), false), Times.Exactly(maxAttemptCount));
-        }
-
-
         protected class FakeStoredProc : IStoredProc<string, string>
         {
             public string DbName => "mobile_ro";

diff --git a/Agoda.Frameworks.DB/DbRepository.cs b/Agoda.Frameworks.DB/DbRepository.cs
@@ -303,14 +303,16 @@ public Task<T> ExecuteReaderAsync<T>(
             string storedProc,
             int timeoutSecs,
             int maxAttemptCount,
-            CancellationToken token,
+            int taskCancellationTimeOutInMilliSecs,
             IDbDataParameter[] parameters,
             Func<SqlDataReader, Task<T>> callback)
         {
             return _dbResources.ChooseDb(database).ExecuteAsync(async (connectionStr, _) =>
             {
                 var stopwatch = Stopwatch.StartNew();
                 Exception error = null;
+                var cancellationTokenSource = new CancellationTokenSource();
+                cancellationTokenSource.CancelAfter(taskCancellationTimeOutInMilliSecs);
                 try
                 {
                     using (var connection = _generateConnection(connectionStr))
@@ -332,7 +334,7 @@ public Task<T> ExecuteReaderAsync<T>(
                                 CommandTimeout = timeoutSecs
                             };
                             sqlCommand.Parameters.AddRange(parameters);
-                            using (var reader = await sqlCommand.ExecuteReaderAsync(token))
+                            using (var reader = await sqlCommand.ExecuteReaderAsync(cancellationTokenSource.Token))
                             {
                                 return await callback(reader);
                             }

diff --git a/Agoda.Frameworks.DB/IDbRepository.cs b/Agoda.Frameworks.DB/IDbRepository.cs
@@ -141,7 +141,7 @@ Task<T> ExecuteReaderAsync<T>(
             string storedProc,
             int timeoutSecs,
             int maxAttemptCount,
-            CancellationToken token,
+            int taskCancellationTimeOutInMilliSecs,
             IDbDataParameter[] parameters,
             Func<SqlDataReader, Task<T>> callback);