Use ICancellationTokenProvider

framework/src/Volo.Abp.EntityFrameworkCore/Volo/Abp/Uow/EntityFrameworkCore/UnitOfWorkDbContextProvider.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.EntityFrameworkCore;
 using Microsoft.EntityFrameworkCore.Storage;
@@ -8,6 +9,7 @@ using Microsoft.Extensions.Logging.Abstractions;
 using Volo.Abp.Data;
 using Volo.Abp.EntityFrameworkCore;
 using Volo.Abp.EntityFrameworkCore.DependencyInjection;
+using Volo.Abp.Threading;
 
 namespace Volo.Abp.Uow.EntityFrameworkCore
 {
@@ -20,13 +22,16 @@ namespace Volo.Abp.Uow.EntityFrameworkCore
 
         private readonly IUnitOfWorkManager _unitOfWorkManager;
         private readonly IConnectionStringResolver _connectionStringResolver;
+        private readonly ICancellationTokenProvider _cancellationTokenProvider;
 
         public UnitOfWorkDbContextProvider(
             IUnitOfWorkManager unitOfWorkManager,
-            IConnectionStringResolver connectionStringResolver)
+            IConnectionStringResolver connectionStringResolver,
+            ICancellationTokenProvider cancellationTokenProvider)
         {
             _unitOfWorkManager = unitOfWorkManager;
             _connectionStringResolver = connectionStringResolver;
+            _cancellationTokenProvider = cancellationTokenProvider;
 
             Logger = NullLogger<UnitOfWorkDbContextProvider<TDbContext>>.Instance;
         }
@@ -194,8 +199,8 @@ namespace Volo.Abp.Uow.EntityFrameworkCore
                 var dbContext = unitOfWork.ServiceProvider.GetRequiredService<TDbContext>();
 
                 var dbTransaction = unitOfWork.Options.IsolationLevel.HasValue
-                    ? await dbContext.Database.BeginTransactionAsync(unitOfWork.Options.IsolationLevel.Value)
-                    : await dbContext.Database.BeginTransactionAsync();
+                    ? await dbContext.Database.BeginTransactionAsync(unitOfWork.Options.IsolationLevel.Value, GetCancellationToken())
+                    : await dbContext.Database.BeginTransactionAsync(GetCancellationToken());
 
                 unitOfWork.AddTransactionApi(
                     transactionApiKey,
@@ -215,11 +220,11 @@ namespace Volo.Abp.Uow.EntityFrameworkCore
 
                 if (dbContext.As<DbContext>().HasRelationalTransactionManager())
                 {
-                    await dbContext.Database.UseTransactionAsync(activeTransaction.DbContextTransaction.GetDbTransaction());
+                    await dbContext.Database.UseTransactionAsync(activeTransaction.DbContextTransaction.GetDbTransaction(), GetCancellationToken());
                 }
                 else
                 {
-                    await dbContext.Database.BeginTransactionAsync(); //TODO: Why not using the new created transaction?
+                    await dbContext.Database.BeginTransactionAsync(GetCancellationToken()); //TODO: Why not using the new created transaction?
                 }
 
                 activeTransaction.AttendedDbContexts.Add(dbContext);
@@ -227,5 +232,10 @@ namespace Volo.Abp.Uow.EntityFrameworkCore
                 return dbContext;
             }
         }
+
+        protected virtual CancellationToken GetCancellationToken(CancellationToken preferredValue = default)
+        {
+            return _cancellationTokenProvider.FallbackToProvider(preferredValue);
+        }
     }
 }

framework/src/Volo.Abp.MongoDB/Volo/Abp/Uow/MongoDB/UnitOfWorkMongoDbContextProvider.cs

@@ -6,6 +6,7 @@ using MongoDB.Bson;
 using MongoDB.Driver;
 using Volo.Abp.Data;
 using Volo.Abp.MongoDB;
+using Volo.Abp.Threading;
 
 namespace Volo.Abp.Uow.MongoDB
 {
@@ -14,13 +15,16 @@ namespace Volo.Abp.Uow.MongoDB
     {
         private readonly IUnitOfWorkManager _unitOfWorkManager;
         private readonly IConnectionStringResolver _connectionStringResolver;
+        private readonly ICancellationTokenProvider _cancellationTokenProvider;
 
         public UnitOfWorkMongoDbContextProvider(
             IUnitOfWorkManager unitOfWorkManager,
-            IConnectionStringResolver connectionStringResolver)
+            IConnectionStringResolver connectionStringResolver,
+            ICancellationTokenProvider cancellationTokenProvider)
         {
             _unitOfWorkManager = unitOfWorkManager;
             _connectionStringResolver = connectionStringResolver;
+            _cancellationTokenProvider = cancellationTokenProvider;
         }
 
         [Obsolete("Use CreateDbContextAsync")]
@@ -183,7 +187,7 @@ namespace Volo.Abp.Uow.MongoDB
 
             if (activeTransaction?.SessionHandle == null)
             {
-                var session = await client.StartSessionAsync(cancellationToken: cancellationToken);
+                var session = await client.StartSessionAsync(cancellationToken: GetCancellationToken(cancellationToken));
 
                 if (unitOfWork.Options.Timeout.HasValue)
                 {
@@ -206,5 +210,10 @@ namespace Volo.Abp.Uow.MongoDB
 
             return dbContext;
         }
+
+        protected virtual CancellationToken GetCancellationToken(CancellationToken preferredValue = default)
+        {
+            return _cancellationTokenProvider.FallbackToProvider(preferredValue);
+        }
     }
 }

framework/src/Volo.Abp.Uow/Volo/Abp/Uow/UnitOfWorkInterceptor.cs

@@ -23,7 +23,7 @@ namespace Volo.Abp.Uow
             _defaultOptions = options.Value;
         }
 
-        public async override Task InterceptAsync(IAbpMethodInvocation invocation)
+        public override async Task InterceptAsync(IAbpMethodInvocation invocation)
         {
             if (!UnitOfWorkHelper.IsUnitOfWorkMethod(invocation.Method, out var unitOfWorkAttribute))
             {