Add health-check timeout to gRPC workitem streaming (#390)

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,18 @@ 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.
+            // The timeout token is used to detect when we are no longer getting any messages, including health checks.
+            // If this is the case, it signifies the connection has been dropped silently and we need to reconnect.
+            using var timeoutSource = new CancellationTokenSource();
+            timeoutSource.CancelAfter(TimeSpan.FromSeconds(60));
+            using 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 +243,20 @@ 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.
+                    if (!cancellation.IsCancellationRequested)
+                    {
+                        // Since the cancellation came from the timeout, log a warning.
+                        this.Logger.ConnectionTimeout();
+                    }
+
+                    return;
+                }
             }
         }
 

src/Worker/Grpc/Logs.cs

@@ -48,5 +48,8 @@ static partial class Logs
 
         [LoggerMessage(EventId = 55, Level = LogLevel.Information, Message = "{instanceId}: Evaluating custom retry handler for failed '{name}' task. Attempt = {attempt}.")]
         public static partial void RetryingTask(this ILogger logger, string instanceId, string name, int attempt);
+
+        [LoggerMessage(EventId = 56, Level = LogLevel.Warning, Message = "Channel to backend has stopped receiving traffic, will attempt to reconnect.")]
+        public static partial void ConnectionTimeout(this ILogger logger);
     }
 }