fix: Add CancellationToken to job worker and handler

Allows passing a CancellationToken to JobWorker Open() method to inject a host cancellationToken and expose token to AsyncJobHandler delegate.

Refs: #519

Author: ofreundAdvz

Client/Api/Worker/IJobWorkerBuilderStep1.cs

@@ -15,6 +15,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Threading;
 using System.Threading.Tasks;
 using Zeebe.Client.Api.Commands;
 using Zeebe.Client.Api.Responses;
@@ -43,8 +44,9 @@ public interface IJobWorkerBuilderStep1
 /// </summary>
 /// <param name="client">the job client to complete or fail the job.</param>
 /// <param name="activatedJob">the job, which was activated by the worker.</param>
+/// <param name="cancellationToken">the token to cancel the job handler.</param>
 /// <returns>A <see cref="Task" /> representing the asynchronous operation.</returns>
-public delegate Task AsyncJobHandler(IJobClient client, IJob activatedJob);
+public delegate Task AsyncJobHandler(IJobClient client, IJob activatedJob, CancellationToken cancellationToken = default);
 
 public interface IJobWorkerBuilderStep2
 {
@@ -247,6 +249,7 @@ public interface IJobWorkerBuilderStep3 : ITenantIdsCommandStep<IJobWorkerBuilde
     /// <summary>
     ///     Open the worker and start to work on available tasks.
     /// </summary>
+    /// <param name="cancellationToken">the token to cancel the job worker.</param>
     /// <returns>the worker.</returns>
-    IJobWorker Open();
+    IJobWorker Open(CancellationToken cancellationToken = default);
 }

Client/Impl/Worker/JobWorker.cs

@@ -45,7 +45,8 @@ public sealed class JobWorker : IJobWorker
     private readonly int maxJobsActive;
     private readonly TimeSpan pollInterval;
 
-    private readonly CancellationTokenSource source;
+    private readonly CancellationTokenSource localCts;
+    private CancellationTokenSource? linkedCts;
     private readonly double thresholdJobsActivation;
 
     private int currentJobsActive;
@@ -54,7 +55,7 @@ public sealed class JobWorker : IJobWorker
     internal JobWorker(JobWorkerBuilder builder)
     {
         jobWorkerBuilder = builder;
-        source = new CancellationTokenSource();
+        localCts = new CancellationTokenSource();
         logger = builder.LoggerFactory?.CreateLogger<JobWorker>();
         jobHandler = jobWorkerBuilder.Handler();
         autoCompletion = builder.AutoCompletionEnabled();
@@ -69,13 +70,15 @@ internal JobWorker(JobWorkerBuilder builder)
     /// <inheritdoc />
     public void Dispose()
     {
-        source.Cancel();
+        localCts.Cancel();
         // delay disposing, since poll and handler take some time to close
         _ = Task.Delay(TimeSpan.FromMilliseconds(pollInterval.TotalMilliseconds * 2))
         .ContinueWith(t =>
         {
             logger?.LogError("Dispose source");
-            source.Dispose();
+            localCts.Dispose();
+            linkedCts?.Dispose();
+            linkedCts = null;
         });
         isRunning = false;
     }
@@ -96,10 +99,12 @@ public bool IsClosed()
     ///     Opens the configured JobWorker to activate jobs in the given poll interval
     ///     and handle with the given handler.
     /// </summary>
-    internal void Open()
+    /// <param name="stoppingToken">The host cancellation token.</param>
+    internal void Open(CancellationToken stoppingToken)
     {
         isRunning = true;
-        var cancellationToken = source.Token;
+        this.linkedCts = CancellationTokenSource.CreateLinkedTokenSource(stoppingToken, localCts.Token);
+        var cancellationToken = linkedCts.Token;
         var bufferOptions = CreateBufferOptions(cancellationToken);
         var executionOptions = CreateExecutionOptions(cancellationToken);
 
@@ -135,7 +140,7 @@ internal void Open()
 
     private async Task StreamJobs(ITargetBlock<IJob> input, CancellationToken cancellationToken)
     {
-        while (!source.IsCancellationRequested)
+        while (!cancellationToken.IsCancellationRequested)
         {
             try
             {
@@ -151,7 +156,7 @@ private async Task StreamJobs(ITargetBlock<IJob> input, CancellationToken cancel
                     activateJobsRequest.Worker,
                     activateJobsRequest.Type);
 
-                while (!source.IsCancellationRequested)
+                while (!cancellationToken.IsCancellationRequested)
                 {
                     var currentJobs = Thread.VolatileRead(ref currentJobsActive);
 
@@ -178,7 +183,7 @@ private async Task StreamJobs(ITargetBlock<IJob> input, CancellationToken cancel
                         grpcActivatedJob.Worker,
                         grpcActivatedJob.Key);
 
-                    await HandleActivationResponse(input, response, null);
+                    await HandleActivationResponse(input, response, null, cancellationToken);
                 }
             }
             catch (RpcException rpcException)
@@ -210,7 +215,7 @@ private static DataflowBlockOptions CreateBufferOptions(CancellationToken cancel
 
     private async Task PollJobs(ITargetBlock<IJob> input, CancellationToken cancellationToken)
     {
-        while (!source.IsCancellationRequested)
+        while (!cancellationToken.IsCancellationRequested)
         {
             var currentJobs = Thread.VolatileRead(ref currentJobsActive);
             if (currentJobs < thresholdJobsActivation)
@@ -221,7 +226,7 @@ private async Task PollJobs(ITargetBlock<IJob> input, CancellationToken cancella
                 try
                 {
                     await jobActivator.SendActivateRequest(activateJobsRequest,
-                        async jobsResponse => await HandleActivationResponse(input, jobsResponse, jobCount),
+                        async jobsResponse => await HandleActivationResponse(input, jobsResponse, jobCount, cancellationToken),
                         null,
                         cancellationToken);
                 }
@@ -238,7 +243,11 @@ await jobActivator.SendActivateRequest(activateJobsRequest,
         }
     }
 
-    private async Task HandleActivationResponse(ITargetBlock<IJob> input, IActivateJobsResponse response, int? jobCount)
+    private async Task HandleActivationResponse(
+        ITargetBlock<IJob> input,
+        IActivateJobsResponse response,
+        int? jobCount,
+        CancellationToken cancellationToken)
     {
         if (jobCount.HasValue)
         {
@@ -258,7 +267,7 @@ private async Task HandleActivationResponse(ITargetBlock<IJob> input, IActivateJ
 
         foreach (var job in response.Jobs)
         {
-            _ = await input.SendAsync(job);
+            _ = await input.SendAsync(job, cancellationToken);
             _ = Interlocked.Increment(ref currentJobsActive);
         }
     }
@@ -269,7 +278,7 @@ private async Task<IJob> HandleActivatedJob(IJob activatedJob, CancellationToken
 
         try
         {
-            await jobHandler(jobClient, activatedJob);
+            await jobHandler(jobClient, activatedJob, cancellationToken);
             await TryToAutoCompleteJob(jobClient, activatedJob, cancellationToken);
         }
         catch (Exception exception)

Client/Impl/Worker/JobWorkerBuilder.cs

@@ -16,6 +16,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using GatewayProtocol;
 using Microsoft.Extensions.Logging;
@@ -50,7 +51,7 @@ public IJobWorkerBuilderStep2 JobType(string type)
 
     public IJobWorkerBuilderStep3 Handler(JobHandler handler)
     {
-        asyncJobHandler = (c, j) => Task.Run(() => handler.Invoke(c, j));
+        asyncJobHandler = (c, j, cts) => Task.Run(() => handler.Invoke(c, j), cts);
         return this;
     }
 
@@ -147,11 +148,11 @@ public IJobWorkerBuilderStep3 HandlerThreads(byte threadCount)
         return this;
     }
 
-    public IJobWorker Open()
+    public IJobWorker Open(CancellationToken cancellationToken = default)
     {
         var worker = new JobWorker(this);
 
-        worker.Open();
+        worker.Open(cancellationToken);
 
         return worker;
     }