Wire up deadline on Nexus context (#636)

Author: Quinn-With-Two-Ns

src/Temporalio/Worker/NexusWorker.cs

@@ -99,9 +99,10 @@ public async Task ExecuteAsync()
                             // we're late-binding it here.
                             var running = new RunningTask();
                             runningTasks[task.Task.TaskToken] = running;
+                            var requestDeadline = task.RequestDeadline?.ToDateTime();
 #pragma warning disable CA2008 // We don't have to pass a scheduler, factory already implies one
                             running.Task = worker.Options.NexusTaskFactory.StartNew(
-                                () => HandlePollTaskAsync(running, task.Task)).Unwrap();
+                                () => HandlePollTaskAsync(running, task.Task, requestDeadline)).Unwrap();
 #pragma warning restore CA2008
                             break;
                         case NexusTask.VariantOneofCase.CancelTask:
@@ -141,12 +142,12 @@ private static void RemoveInvalidHeaders(MapField<string, string> headers)
             headers.Remove("Request-Timeout");
         }
 
-        private async Task HandlePollTaskAsync(RunningTask running, PollNexusTaskQueueResponse task)
+        private async Task HandlePollTaskAsync(RunningTask running, PollNexusTaskQueueResponse task, DateTime? requestDeadline)
         {
             try
             {
                 // Handle poll and post back to Core
-                var completion = await HandlePollTaskInternalAsync(running, task).ConfigureAwait(false);
+                var completion = await HandlePollTaskInternalAsync(running, task, requestDeadline).ConfigureAwait(false);
                 logger.LogTrace("Sending Nexus completion: {Completion}", completion);
                 await worker.BridgeWorker.CompleteNexusTaskAsync(completion).ConfigureAwait(false);
             }
@@ -165,7 +166,7 @@ private async Task HandlePollTaskAsync(RunningTask running, PollNexusTaskQueueRe
         }
 
         private async Task<StartOperationResponse> HandleStartOperationAsync(
-            RunningTask running, PollNexusTaskQueueResponse task)
+            RunningTask running, PollNexusTaskQueueResponse task, DateTime? requestDeadline)
         {
             // Create context
             RemoveInvalidHeaders(task.Request.Header);
@@ -195,6 +196,7 @@ private async Task<StartOperationResponse> HandleStartOperationAsync(
                             e);
                     }
                 }).ToList(),
+                RequestDeadline = requestDeadline,
             };
             running.OnCancelReason = reason => context.CancellationReason = reason;
 
@@ -253,7 +255,7 @@ private async Task<StartOperationResponse> HandleStartOperationAsync(
         }
 
         private async Task<CancelOperationResponse> HandleCancelOperationAsync(
-            RunningTask running, PollNexusTaskQueueResponse task)
+            RunningTask running, PollNexusTaskQueueResponse task, DateTime? requestDeadline)
         {
             // Create context
             RemoveInvalidHeaders(task.Request.Header);
@@ -266,6 +268,7 @@ private async Task<CancelOperationResponse> HandleCancelOperationAsync(
             {
                 Headers = task.Request.Header.Count == 0 ? null :
                     new Dictionary<string, string>(task.Request.Header, StringComparer.OrdinalIgnoreCase),
+                RequestDeadline = requestDeadline,
             };
             running.OnCancelReason = reason => context.CancellationReason = reason;
 
@@ -287,22 +290,22 @@ private async Task<CancelOperationResponse> HandleCancelOperationAsync(
         }
 
         private async Task<NexusTaskCompletion> HandlePollTaskInternalAsync(
-            RunningTask running, PollNexusTaskQueueResponse task)
+            RunningTask running, PollNexusTaskQueueResponse task, DateTime? requestDeadline)
         {
             try
             {
                 // Handle each case
                 switch (task.Request.VariantCase)
                 {
                     case Request.VariantOneofCase.StartOperation:
-                        var startResp = await HandleStartOperationAsync(running, task).ConfigureAwait(false);
+                        var startResp = await HandleStartOperationAsync(running, task, requestDeadline).ConfigureAwait(false);
                         return new()
                         {
                             TaskToken = task.TaskToken,
                             Completed = new() { StartOperation = startResp },
                         };
                     case Request.VariantOneofCase.CancelOperation:
-                        var cancelResp = await HandleCancelOperationAsync(running, task).ConfigureAwait(false);
+                        var cancelResp = await HandleCancelOperationAsync(running, task, requestDeadline).ConfigureAwait(false);
                         return new()
                         {
                             TaskToken = task.TaskToken,

tests/Temporalio.Tests/Worker/NexusWorkerTests.cs

@@ -301,6 +301,29 @@ await Workflow.CreateNexusWorkflowClient<IStringService>(endpoint).
         Assert.Equal("timed out", ctx.CancellationReason);
     }
 
+    [Fact]
+    public async Task ExecuteNexusOperationAsync_RequestDeadline_SetOnContext()
+    {
+        var contextSource = new TaskCompletionSource<OperationStartContext>();
+        var workerOptions = new TemporalWorkerOptions($"tq-{Guid.NewGuid()}").
+            AddNexusService(new AsyncFuncStringService(
+                start: (ctx, input) =>
+                {
+                    contextSource.SetResult(ctx);
+                    return Task.FromResult(OperationStartResult.SyncResult($"Hello, {input}"));
+                }));
+        var endpoint = await CreateNexusEndpointAsync(workerOptions.TaskQueue!);
+        await RunInWorkflowAsync(workerOptions, async () =>
+        {
+            var result = await Workflow.CreateNexusWorkflowClient<IStringService>(endpoint).
+                ExecuteNexusOperationAsync(svc => svc.DoSomething("some-name"));
+            Assert.Equal("Hello, some-name", result);
+        });
+        var ctx = await contextSource.Task;
+        // Server always sends a request timeout header, so the deadline should be set
+        Assert.NotNull(ctx.RequestDeadline);
+    }
+
     [Fact]
     public async Task ExecuteNexusOperationAsync_OperationSummary_FoundInHistory()
     {