Merge pull request #9676 from abpframework/liangshiwei/cancellation-token

Allowed to override CancellationToken for a specific call

framework/Volo.Abp.sln

@@ -383,6 +383,8 @@ Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Volo.Abp.TextTemplating.Scr
 EndProject
 Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Volo.Abp.MongoDB.Tests.SecondContext", "test\Volo.Abp.MongoDB.Tests.SecondContext\Volo.Abp.MongoDB.Tests.SecondContext.csproj", "{90B1866A-EF99-40B9-970E-B898E5AA523F}"
 EndProject
+Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Volo.Abp.Threading.Tests", "test\Volo.Abp.Threading.Tests\Volo.Abp.Threading.Tests.csproj", "{7B2FCAD6-86E6-49C8-ADBE-A61B4F4B101B}"
+EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution
 		Debug|Any CPU = Debug|Any CPU
@@ -1141,6 +1143,10 @@ Global
 		{90B1866A-EF99-40B9-970E-B898E5AA523F}.Debug|Any CPU.Build.0 = Debug|Any CPU
 		{90B1866A-EF99-40B9-970E-B898E5AA523F}.Release|Any CPU.ActiveCfg = Release|Any CPU
 		{90B1866A-EF99-40B9-970E-B898E5AA523F}.Release|Any CPU.Build.0 = Release|Any CPU
+		{7B2FCAD6-86E6-49C8-ADBE-A61B4F4B101B}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{7B2FCAD6-86E6-49C8-ADBE-A61B4F4B101B}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{7B2FCAD6-86E6-49C8-ADBE-A61B4F4B101B}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{7B2FCAD6-86E6-49C8-ADBE-A61B4F4B101B}.Release|Any CPU.Build.0 = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(SolutionProperties) = preSolution
 		HideSolutionNode = FALSE
@@ -1334,6 +1340,7 @@ Global
 		{C996F458-98FB-483D-9306-4701290E2FC1} = {447C8A77-E5F0-4538-8687-7383196D04EA}
 		{75D8DADB-3FA9-4C1D-B23A-DBFD08133B7C} = {447C8A77-E5F0-4538-8687-7383196D04EA}
 		{90B1866A-EF99-40B9-970E-B898E5AA523F} = {447C8A77-E5F0-4538-8687-7383196D04EA}
+		{7B2FCAD6-86E6-49C8-ADBE-A61B4F4B101B} = {447C8A77-E5F0-4538-8687-7383196D04EA}
 	EndGlobalSection
 	GlobalSection(ExtensibilityGlobals) = postSolution
 		SolutionGuid = {BB97ECF4-9A84-433F-A80B-2A3285BDD1D5}

framework/src/Volo.Abp.AspNetCore/Volo/Abp/AspNetCore/Threading/HttpContextCancellationTokenProvider.cs

@@ -6,13 +6,26 @@ using Volo.Abp.Threading;
 namespace Volo.Abp.AspNetCore.Threading
 {
     [Dependency(ReplaceServices = true)]
-    public class HttpContextCancellationTokenProvider : ICancellationTokenProvider, ITransientDependency
+    public class HttpContextCancellationTokenProvider : CancellationTokenProviderBase, ITransientDependency
     {
-        public CancellationToken Token => _httpContextAccessor.HttpContext?.RequestAborted ?? CancellationToken.None;
-
         private readonly IHttpContextAccessor _httpContextAccessor;
 
-        public HttpContextCancellationTokenProvider(IHttpContextAccessor httpContextAccessor)
+        public override CancellationToken Token
+        {
+            get
+            {
+                if (OverrideValue != null)
+                {
+                    return OverrideValue.CancellationToken;
+                }
+                return _httpContextAccessor.HttpContext?.RequestAborted ?? CancellationToken.None;
+            }
+        }
+
+        public HttpContextCancellationTokenProvider(
+            IAmbientScopeProvider<CancellationTokenOverride> cancellationTokenOverrideScopeProvider,
+            IHttpContextAccessor httpContextAccessor)
+            : base(cancellationTokenOverrideScopeProvider)
         {
             _httpContextAccessor = httpContextAccessor;
         }

framework/src/Volo.Abp.Threading/Volo/Abp/Threading/CancellationTokenOverride.cs

@@ -0,0 +1,14 @@
+using System.Threading;
+
+namespace Volo.Abp.Threading
+{
+    public class CancellationTokenOverride
+    {
+        public CancellationToken CancellationToken { get; }
+
+        public CancellationTokenOverride(CancellationToken cancellationToken)
+        {
+            CancellationToken = cancellationToken;
+        }
+    }
+}

framework/src/Volo.Abp.Threading/Volo/Abp/Threading/CancellationTokenProviderBase.cs

@@ -0,0 +1,26 @@
+using System;
+using System.Threading;
+
+namespace Volo.Abp.Threading
+{
+    public abstract class CancellationTokenProviderBase : ICancellationTokenProvider
+    {
+        public const string CancellationTokenOverrideContextKey = "Volo.Abp.Threading.CancellationToken.Override";
+
+        public abstract CancellationToken Token { get; }
+
+        protected IAmbientScopeProvider<CancellationTokenOverride> CancellationTokenOverrideScopeProvider { get; }
+
+        protected CancellationTokenOverride OverrideValue => CancellationTokenOverrideScopeProvider.GetValue(CancellationTokenOverrideContextKey);
+
+        protected CancellationTokenProviderBase(IAmbientScopeProvider<CancellationTokenOverride> cancellationTokenOverrideScopeProvider)
+        {
+            CancellationTokenOverrideScopeProvider = cancellationTokenOverrideScopeProvider;
+        }
+
+        public IDisposable Use(CancellationToken cancellationToken)
+        {
+            return CancellationTokenOverrideScopeProvider.BeginScope(CancellationTokenOverrideContextKey, new CancellationTokenOverride(cancellationToken));
+        }
+    }
+}

framework/src/Volo.Abp.Threading/Volo/Abp/Threading/ICancellationTokenProvider.cs

@@ -1,9 +1,12 @@
-using System.Threading;
+using System;
+using System.Threading;
 
 namespace Volo.Abp.Threading
 {
     public interface ICancellationTokenProvider
     {
         CancellationToken Token { get; }
+
+        IDisposable Use(CancellationToken cancellationToken);
     }
 }

framework/src/Volo.Abp.Threading/Volo/Abp/Threading/NullCancellationTokenProvider.cs

@@ -2,15 +2,15 @@
 
 namespace Volo.Abp.Threading
 {
-    public class NullCancellationTokenProvider : ICancellationTokenProvider
+    public class NullCancellationTokenProvider : CancellationTokenProviderBase
     {
         public static NullCancellationTokenProvider Instance { get; } = new NullCancellationTokenProvider();
 
-        public CancellationToken Token { get; } = CancellationToken.None;
+        public override CancellationToken Token => OverrideValue?.CancellationToken ?? CancellationToken.None;
 
         private NullCancellationTokenProvider()
+            : base(new AmbientDataContextAmbientScopeProvider<CancellationTokenOverride>(new AsyncLocalAmbientDataContext()))
         {
-            
         }
     }
-}
+}

framework/test/Volo.Abp.Threading.Tests/Volo.Abp.Threading.Tests.csproj

@@ -0,0 +1,16 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+    <Import Project="..\..\..\common.test.props" />
+
+    <PropertyGroup>
+        <TargetFramework>net5.0</TargetFramework>
+        <RootNamespace />
+    </PropertyGroup>
+
+    <ItemGroup>
+        <ProjectReference Include="..\AbpTestBase\AbpTestBase.csproj" />
+        <ProjectReference Include="..\..\src\Volo.Abp.Threading\Volo.Abp.Threading.csproj" />
+        <PackageReference Include="Microsoft.NET.Test.Sdk" Version="$(MicrosoftNETTestSdkPackageVersion)" />
+    </ItemGroup>
+
+</Project>

framework/test/Volo.Abp.Threading.Tests/Volo/Abp/Threading/AbpThreadingTestModule.cs

@@ -0,0 +1,15 @@
+using Volo.Abp.Modularity;
+
+namespace Volo.Abp.Threading
+{
+    [DependsOn(
+        typeof(AbpThreadingModule),
+        typeof(AbpTestBaseModule)
+    )]
+    public class AbpThreadingTestModule : AbpModule
+    {
+        public override void ConfigureServices(ServiceConfigurationContext context)
+        {
+        }
+    }
+}

framework/test/Volo.Abp.Threading.Tests/Volo/Abp/Threading/NullCancellationTokenProvider_Tests.cs

@@ -0,0 +1,63 @@
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Shouldly;
+using Volo.Abp.Testing;
+using Xunit;
+
+namespace Volo.Abp.Threading
+{
+    public class NullCancellationTokenProvider_Tests : AbpIntegratedTest<AbpThreadingTestModule>
+    {
+        private readonly ICancellationTokenProvider _cancellationTokenProvider;
+
+        public NullCancellationTokenProvider_Tests()
+        {
+            _cancellationTokenProvider = NullCancellationTokenProvider.Instance;
+        }
+
+        [Fact]
+        public void Should_Return_None_Token()
+        {
+            _cancellationTokenProvider.Token.ShouldBe(CancellationToken.None);
+        }
+
+        [Fact]
+        public void Should_Return_Specific_Token()
+        {
+            var cts = new CancellationTokenSource();
+
+            using (_cancellationTokenProvider.Use(cts.Token))
+            {
+                var newCancellationTokenProvider = NullCancellationTokenProvider.Instance;
+
+                newCancellationTokenProvider.Token.ShouldBe(cts.Token);
+            }
+
+            _cancellationTokenProvider.Token.ShouldBe(CancellationToken.None);
+        }
+
+        [Fact]
+        public void Should_Cancel_After_100_Milliseconds()
+        {
+            var cts = new CancellationTokenSource();
+            cts.CancelAfter(TimeSpan.FromMilliseconds(100));
+
+            using (_cancellationTokenProvider.Use(cts.Token))
+            {
+                var newCancellationTokenProvider = NullCancellationTokenProvider.Instance;
+                Should.Throw<OperationCanceledException>(() => LongTask(1000, newCancellationTokenProvider.Token));
+            }
+        }
+
+        private void LongTask(int loopCounter, CancellationToken cancellationToken = default)
+        {
+            for (var i = 0; i < loopCounter; i++)
+            {
+                cancellationToken.ThrowIfCancellationRequested();
+
+                Thread.Sleep(10);
+            }
+        }
+    }
+}