Add ability to filter orchestrations at worker (#443)

Add ability to filter orchestrations at worker

This change allows a user to implement an interface that allows
for programmatic orchestration filtering. As it is an interface,
it can be plugged into DI. Returning true allows the
orchestration to proceed, false will block the orchestration.

In this initial commit, the filter can work off of the orchestration
name and the orchestration tags.

Signed-off-by: Hal Spang <[email protected]>

Author: halspang

src/Worker/Core/DependencyInjection/DurableTaskWorkerBuilderExtensions.cs

@@ -109,4 +109,32 @@ public static IDurableTaskWorkerBuilder UseVersioning(this IDurableTaskWorkerBui
         });
         return builder;
     }
+
+    /// <summary>
+    /// Adds an orchestration filter to the specified <see cref="IDurableTaskWorkerBuilder"/>.
+    /// </summary>
+    /// <param name="builder">The builder to set the builder target for.</param>
+    /// <typeparam name="TOrchestrationFilter">The implementation of a <see cref="IOrchestrationFilter"/> that will be bound.</typeparam>
+    /// <returns>The same <see cref="IDurableTaskWorkerBuilder"/> instance, allowing for method chaining.</returns>
+    [Obsolete("Experimental")]
+    public static IDurableTaskWorkerBuilder UseOrchestrationFilter<TOrchestrationFilter>(this IDurableTaskWorkerBuilder builder) where TOrchestrationFilter : class, IOrchestrationFilter
+    {
+        Check.NotNull(builder);
+        builder.Services.AddSingleton<IOrchestrationFilter, TOrchestrationFilter>();
+        return builder;
+    }
+
+    /// <summary>
+    /// Adds an orchestration filter to the specified <see cref="IDurableTaskWorkerBuilder"/>.
+    /// </summary>
+    /// <param name="builder">The builder to set the builder target for.</param>
+    /// <param name="filter">The instance of an <see cref="IOrchestrationFilter"/> to use.</param>
+    /// <returns>The same <see cref="IDurableTaskWorkerBuilder"/> instance, allowing for method chaining.</returns>
+    [Obsolete("Experimental")]
+    public static IDurableTaskWorkerBuilder UseOrchestrationFilter(this IDurableTaskWorkerBuilder builder, IOrchestrationFilter filter)
+    {
+        Check.NotNull(builder);
+        builder.Services.AddSingleton(filter);
+        return builder;
+    }
 }

src/Worker/Core/DurableTaskWorkerOptions.cs

@@ -145,6 +145,12 @@ public DataConverter DataConverter
     /// </summary>
     public bool IsVersioningSet { get; internal set; }
 
+    /// <summary>
+    /// Gets or sets a callback function that determines whether an orchestration should be accepted for work.
+    /// </summary>
+    [Obsolete("Experimental")]
+    public IOrchestrationFilter? OrchestrationFilter { get; set; }
+
     /// <summary>
     /// Gets a value indicating whether <see cref="DataConverter" /> was explicitly set or not.
     /// </summary>
@@ -156,6 +162,7 @@ public DataConverter DataConverter
     /// </remarks>
     internal bool DataConverterExplicitlySet { get; private set; }
 
+
     /// <summary>
     /// Applies these option values to another.
     /// </summary>
@@ -169,6 +176,7 @@ internal void ApplyTo(DurableTaskWorkerOptions other)
             other.MaximumTimerInterval = this.MaximumTimerInterval;
             other.EnableEntitySupport = this.EnableEntitySupport;
             other.Versioning = this.Versioning;
+            other.OrchestrationFilter = this.OrchestrationFilter;
         }
     }
 

src/Worker/Core/IOrchestrationFilter.cs

@@ -0,0 +1,35 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+namespace Microsoft.DurableTask.Worker;
+
+/// <summary>
+/// Defines a filter for validating orchestrations.
+/// </summary>
+[Obsolete("Experimental")]
+public interface IOrchestrationFilter
+{
+    /// <summary>
+    /// Validate the orchestration against the filter represented by this interface.
+    /// </summary>
+    /// <param name="info">The information on the orchestration to validate.</param>
+    /// <param name="cancellationToken">The cancellation token for the request to timeout.</param>
+    /// <returns><code>true</code> if the orchestration is valid <code>false</code> otherwise.</returns>
+    ValueTask<bool> IsOrchestrationValidAsync(OrchestrationFilterParameters info, CancellationToken cancellationToken = default);
+}
+
+/// <summary>
+/// Struct representation of orchestration information.
+/// </summary>
+public struct OrchestrationFilterParameters
+{
+    /// <summary>
+    /// Gets the name of the orchestration.
+    /// </summary>
+    public string? Name { get; init; }
+
+    /// <summary>
+    /// Gets the tags associated with the orchestration.
+    /// </summary>
+    public IReadOnlyDictionary<string, string>? Tags { get; init; }
+}

src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs

@@ -30,13 +30,16 @@ class Processor
         readonly TaskHubSidecarServiceClient client;
         readonly DurableTaskShimFactory shimFactory;
         readonly GrpcDurableTaskWorkerOptions.InternalOptions internalOptions;
+        [Obsolete("Experimental")]
+        readonly IOrchestrationFilter? orchestrationFilter;
 
-        public Processor(GrpcDurableTaskWorker worker, TaskHubSidecarServiceClient client)
+        public Processor(GrpcDurableTaskWorker worker, TaskHubSidecarServiceClient client, IOrchestrationFilter? orchestrationFilter = null)
         {
             this.worker = worker;
             this.client = client;
             this.shimFactory = new DurableTaskShimFactory(this.worker.grpcOptions, this.worker.loggerFactory);
             this.internalOptions = this.worker.grpcOptions.Internal;
+            this.orchestrationFilter = orchestrationFilter;
         }
 
         ILogger Logger => this.worker.logger;
@@ -374,6 +377,31 @@ async Task OnRunOrchestratorAsync(
                     entityConversionState,
                     cancellationToken);
 
+                bool filterPassed = true;
+                if (this.orchestrationFilter != null)
+                {
+                    filterPassed = await this.orchestrationFilter.IsOrchestrationValidAsync(
+                        new OrchestrationFilterParameters
+                        {
+                            Name = runtimeState.Name,
+                            Tags = runtimeState.Tags != null ? new Dictionary<string, string>(runtimeState.Tags) : null,
+                        },
+                        cancellationToken);
+                }
+
+                if (!filterPassed)
+                {
+                    this.Logger.AbandoningOrchestrationDueToOrchestrationFilter(request.InstanceId, completionToken);
+                    await this.client.AbandonTaskOrchestratorWorkItemAsync(
+                        new P.AbandonOrchestrationTaskRequest
+                        {
+                            CompletionToken = completionToken,
+                        },
+                        cancellationToken: cancellationToken);
+
+                    return;
+                }
+
                 // If versioning has been explicitly set, we attempt to follow that pattern. If it is not set, we don't compare versions here.
                 failureDetails = EvaluateOrchestrationVersioning(versioning, runtimeState.Version, out versionFailure);
 

src/Worker/Grpc/GrpcDurableTaskWorker.cs

@@ -17,6 +17,7 @@ sealed partial class GrpcDurableTaskWorker : DurableTaskWorker
     readonly IServiceProvider services;
     readonly ILoggerFactory loggerFactory;
     readonly ILogger logger;
+    readonly IOrchestrationFilter? orchestrationFilter;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="GrpcDurableTaskWorker" /> class.
@@ -27,28 +28,31 @@ sealed partial class GrpcDurableTaskWorker : DurableTaskWorker
     /// <param name="workerOptions">The generic worker options.</param>
     /// <param name="services">The service provider.</param>
     /// <param name="loggerFactory">The logger.</param>
+    /// <param name="orchestrationFilter">The optional <see cref="IOrchestrationFilter"/> used to filter orchestration execution.</param>
     public GrpcDurableTaskWorker(
         string name,
         IDurableTaskFactory factory,
         IOptionsMonitor<GrpcDurableTaskWorkerOptions> grpcOptions,
         IOptionsMonitor<DurableTaskWorkerOptions> workerOptions,
         IServiceProvider services,
-        ILoggerFactory loggerFactory)
+        ILoggerFactory loggerFactory,
+        IOrchestrationFilter? orchestrationFilter = null)
         : base(name, factory)
     {
         this.grpcOptions = Check.NotNull(grpcOptions).Get(name);
         this.workerOptions = Check.NotNull(workerOptions).Get(name);
         this.services = Check.NotNull(services);
         this.loggerFactory = Check.NotNull(loggerFactory);
         this.logger = loggerFactory.CreateLogger("Microsoft.DurableTask"); // TODO: use better category name.
+        this.orchestrationFilter = orchestrationFilter;
     }
 
     /// <inheritdoc />
     protected override async Task ExecuteAsync(CancellationToken stoppingToken)
     {
         await using AsyncDisposable disposable = this.GetCallInvoker(out CallInvoker callInvoker, out string address);
         this.logger.StartingTaskHubWorker(address);
-        await new Processor(this, new(callInvoker)).ExecuteAsync(stoppingToken);
+        await new Processor(this, new(callInvoker), this.orchestrationFilter).ExecuteAsync(stoppingToken);
     }
 
 #if NET6_0_OR_GREATER

src/Worker/Grpc/Logs.cs

@@ -57,5 +57,8 @@ static partial class Logs
 
         [LoggerMessage(EventId = 58, Level = LogLevel.Information, Message = "Abandoning orchestration. InstanceId = '{instanceId}'. Completion token = '{completionToken}'")]
         public static partial void AbandoningOrchestrationDueToVersioning(this ILogger logger, string instanceId, string completionToken);
+
+        [LoggerMessage(EventId = 59, Level = LogLevel.Information, Message = "Abandoning orchestration due to filtering. InstanceId = '{instanceId}'. Completion token = '{completionToken}'")]
+        public static partial void AbandoningOrchestrationDueToOrchestrationFilter(this ILogger logger, string instanceId, string completionToken);
     }
 }

test/Grpc.IntegrationTests/GrpcSidecar/Grpc/TaskHubGrpcServer.cs

@@ -6,9 +6,9 @@
 using DurableTask.Core;
 using DurableTask.Core.History;
 using DurableTask.Core.Query;
-using Microsoft.DurableTask.Sidecar.Dispatcher;
 using Google.Protobuf.WellKnownTypes;
 using Grpc.Core;
+using Microsoft.DurableTask.Sidecar.Dispatcher;
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
@@ -589,6 +589,16 @@ static string GetTaskIdKey(string instanceId, int taskId)
         return string.Concat(instanceId, "__", taskId.ToString());
     }
 
+    public override Task<P.AbandonActivityTaskResponse> AbandonTaskActivityWorkItem(P.AbandonActivityTaskRequest request, ServerCallContext context)
+    {
+        return Task.FromResult<P.AbandonActivityTaskResponse>(new());
+    }
+
+    public override Task<P.AbandonOrchestrationTaskResponse> AbandonTaskOrchestratorWorkItem(P.AbandonOrchestrationTaskRequest request, ServerCallContext context)
+    {
+        return Task.FromResult<P.AbandonOrchestrationTaskResponse>(new());
+    }
+
     /// <summary>
     /// A <see cref="ITrafficSignal"/> implementation that is used to control whether the task hub
     /// dispatcher can fetch new work-items, based on whether a client is currently connected.