Fix null scope for task activity

jviau · Apr 29, 2020 · e07fc62 · e07fc62
1 parent af82cf6
commit e07fc62
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 3 deletions.
diff --git a/eng/Common.props b/eng/Common.props
@@ -25,8 +25,7 @@
     <SymbolPackageFormat>snupkg</SymbolPackageFormat>
     <EmbedUntrackedSources>true</EmbedUntrackedSources>
     <PackageReleaseNotes>
-      - Fix orchestration and activity incorrect version registration.
-      - Fix orchestrations never completing due to extra async work.
+      - Fix task activity null scope.
     </PackageReleaseNotes>
   </PropertyGroup>
 

diff --git a/src/DurableTask.DependencyInjection/src/DefaultTaskHubWorkerBuilder.cs b/src/DurableTask.DependencyInjection/src/DefaultTaskHubWorkerBuilder.cs
@@ -108,6 +108,7 @@ public TaskHubWorker Build(IServiceProvider serviceProvider)
                 worker.AddOrchestrationDispatcherMiddleware(WrapMiddleware(middlewareType));
             }
 
+            worker.AddActivityDispatcherMiddleware(BeginMiddlewareScope(serviceProvider));
             worker.AddActivityDispatcherMiddleware(WrapMiddleware(typeof(ServiceProviderActivityMiddleware)));
             foreach (Type middlewareType in _activitiesMiddleware)
             {
@@ -134,7 +135,7 @@ private static Func<DispatchMiddlewareContext, Func<Task>, Task> BeginMiddleware
                 IOrchestrationScope scope = null;
                 try
                 {
-                    scope = OrchestrationScope.CreateScope(context.GetProperty<OrchestrationInstance>(), serviceProvider);
+                    scope = OrchestrationScope.GetOrCreateScope(context.GetProperty<OrchestrationInstance>(), serviceProvider);
                     await next().ConfigureAwait(false);
                 }
                 finally

diff --git a/src/DurableTask.DependencyInjection/src/OrchestrationScope.cs b/src/DurableTask.DependencyInjection/src/OrchestrationScope.cs
@@ -44,6 +44,26 @@ public static IOrchestrationScope GetScope(OrchestrationInstance orchestrationIn
             }
         }
 
+        /// <summary>
+        /// Gets or creates a new <see cref="IOrchestrationScope"/> for the orchestration instance.
+        /// </summary>
+        /// <param name="orchestrationInstance">The orchestration instance. Not null.</param>
+        /// <param name="serviceProvider">The service provider. Not null.</param>
+        /// <returns>The newly created scope.</returns>
+        public static IOrchestrationScope GetOrCreateScope(
+            OrchestrationInstance orchestrationInstance, IServiceProvider serviceProvider)
+        {
+            Check.NotNull(orchestrationInstance, nameof(orchestrationInstance));
+            Check.NotNull(serviceProvider, nameof(serviceProvider));
+
+            lock (s_scopes)
+            {
+                return s_scopes.ContainsKey(orchestrationInstance)
+                    ? GetScope(orchestrationInstance)
+                    : CreateScope(orchestrationInstance, serviceProvider);
+            }
+        }
+
         /// <summary>
         /// Creates a new <see cref="IOrchestrationScope"/> for the orchestration instance.
         /// </summary>

diff --git a/src/DurableTask.DependencyInjection/test/OrchestrationScopeTests.cs b/src/DurableTask.DependencyInjection/test/OrchestrationScopeTests.cs
@@ -100,6 +100,36 @@ public void CreateScope_Created()
             scope.ServiceProvider.Should().NotBeNull();
         }
 
+        [Fact]
+        public void GetOrCreateScope_Created()
+        {
+            // arrange
+            var instance = new OrchestrationInstance();
+
+            // act
+            IOrchestrationScope scope = OrchestrationScope.GetOrCreateScope(instance, GetServiceProvider());
+
+            // assert
+            scope.Should().NotBeNull();
+            scope.ServiceProvider.Should().NotBeNull();
+        }
+
+        [Fact]
+        public void GetOrCreateScope_Found()
+        {
+            // arrange
+            var instance = new OrchestrationInstance();
+
+            // act
+            IOrchestrationScope first = OrchestrationScope.CreateScope(instance, GetServiceProvider());
+            IOrchestrationScope second = OrchestrationScope.GetOrCreateScope(instance, GetServiceProvider());
+
+            // assert
+            second.Should().NotBeNull();
+            second.ServiceProvider.Should().NotBeNull();
+            second.Should().BeSameAs(first);
+        }
+
         [Fact]
         public async Task SafeDisposeScopeAsync_ArgumentNull()
         {