Add timeout to gRPC workitem streaming

This commit adds a timeout to the gRPC stream used to communicate
with the backend. This was done because the backend could restart
and drop the connection and the worker would not know. This causes
the worker to hang and not receive any new work items. The fix is
to reset the connection if a long enough period of time has passed
between receiving anything on the stream.

Signed-off-by: halspang <[email protected]>

Author: halspang

src/Worker/Grpc/GrpcDurableTaskWorker.Processor.cs

@@ -1,13 +1,11 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-using System.Diagnostics.CodeAnalysis;
 using System.Text;
 using DurableTask.Core;
 using DurableTask.Core.Entities;
 using DurableTask.Core.Entities.OperationFormat;
 using DurableTask.Core.History;
-using Grpc.Core;
 using Microsoft.DurableTask.Entities;
 using Microsoft.DurableTask.Worker.Shims;
 using Microsoft.Extensions.DependencyInjection;
@@ -186,10 +184,16 @@ async ValueTask<OrchestrationRuntimeState> BuildRuntimeStateAsync(
 
         async Task ProcessWorkItemsAsync(AsyncServerStreamingCall<P.WorkItem> stream, CancellationToken cancellation)
         {
+            // Create a new token source for timing out and a final token source that keys off of them both.
+            var timeoutSource = new CancellationTokenSource();
+            timeoutSource.CancelAfter(TimeSpan.FromSeconds(60));
+            var tokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellation, timeoutSource.Token);
+
             while (!cancellation.IsCancellationRequested)
             {
-                await foreach (P.WorkItem workItem in stream.ResponseStream.ReadAllAsync(cancellation))
+                await foreach (P.WorkItem workItem in stream.ResponseStream.ReadAllAsync(tokenSource.Token))
                 {
+                    timeoutSource.CancelAfter(TimeSpan.FromSeconds(60));
                     if (workItem.RequestCase == P.WorkItem.RequestOneofCase.OrchestratorRequest)
                     {
                         this.RunBackgroundTask(
@@ -237,6 +241,14 @@ async Task ProcessWorkItemsAsync(AsyncServerStreamingCall<P.WorkItem> stream, Ca
                         this.Logger.UnexpectedWorkItemType(workItem.RequestCase.ToString());
                     }
                 }
+
+                if (tokenSource.IsCancellationRequested || tokenSource.Token.IsCancellationRequested)
+                {
+                    // The token has cancelled, this means either:
+                    // 1. The broader 'cancellation' was triggered, return here to start a graceful shutdown.
+                    // 2. The timeoutSource was triggered, return here to trigger a reconnect to the backend.
+                    return;
+                }
             }
         }