From 0b7379426745d6bc297879d78087801af27c63a8 Mon Sep 17 00:00:00 2001
From: Tingluo Huang <tingluohuang@github.com>
Date: Fri, 20 May 2022 21:31:21 -0400
Subject: [PATCH] Set timeout on sending live console log. (#1903)

---
 src/Runner.Common/JobServer.cs      | 24 ++++++++++++------------
 src/Runner.Common/JobServerQueue.cs |  6 +++++-
 src/dir.proj                        |  2 +-
 3 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/src/Runner.Common/JobServer.cs b/src/Runner.Common/JobServer.cs
index ffa70a48e93..1cac64f65b6 100644
--- a/src/Runner.Common/JobServer.cs
+++ b/src/Runner.Common/JobServer.cs
@@ -1,5 +1,4 @@
-using GitHub.DistributedTask.WebApi;
-using System;
+using System;
 using System.Collections.Generic;
 using System.IO;
 using System.Linq;
@@ -9,6 +8,7 @@
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
+using GitHub.DistributedTask.WebApi;
 using GitHub.Runner.Sdk;
 using GitHub.Services.Common;
 using GitHub.Services.WebApi;
@@ -140,8 +140,8 @@ private async Task CheckNetworkEndpointsAsync(int attemptsLeft)
 
         public void InitializeWebsocketClient(ServiceEndpoint serviceEndpoint)
         {
-             this._serviceEndpoint = serviceEndpoint;
-             InitializeWebsocketClient(TimeSpan.Zero);
+            this._serviceEndpoint = serviceEndpoint;
+            InitializeWebsocketClient(TimeSpan.Zero);
         }
 
         public ValueTask DisposeAsync()
@@ -163,9 +163,9 @@ private void CheckConnection()
 
         private void InitializeWebsocketClient(TimeSpan delay)
         {
-             if (_serviceEndpoint.Authorization != null &&
-                 _serviceEndpoint.Authorization.Parameters.TryGetValue(EndpointAuthorizationParameters.AccessToken, out var accessToken) &&
-                 !string.IsNullOrEmpty(accessToken))
+            if (_serviceEndpoint.Authorization != null &&
+                _serviceEndpoint.Authorization.Parameters.TryGetValue(EndpointAuthorizationParameters.AccessToken, out var accessToken) &&
+                !string.IsNullOrEmpty(accessToken))
             {
                 if (_serviceEndpoint.Data.TryGetValue("FeedStreamUrl", out var feedStreamUrl) && !string.IsNullOrEmpty(feedStreamUrl))
                 {
@@ -177,7 +177,7 @@ private void InitializeWebsocketClient(TimeSpan delay)
                     var userAgentValues = new List<ProductInfoHeaderValue>();
                     userAgentValues.AddRange(UserAgentUtility.GetDefaultRestUserAgent());
                     userAgentValues.AddRange(HostContext.UserAgents);
-                    this._websocketClient.Options.SetRequestHeader("User-Agent", string.Join(" ", userAgentValues.Select(x=>x.ToString())));
+                    this._websocketClient.Options.SetRequestHeader("User-Agent", string.Join(" ", userAgentValues.Select(x => x.ToString())));
 
                     this._websocketConnectTask = ConnectWebSocketClient(feedStreamUrl, delay);
                 }
@@ -201,7 +201,7 @@ private async Task ConnectWebSocketClient(string feedStreamUrl, TimeSpan delay)
                 await this._websocketClient.ConnectAsync(new Uri(feedStreamUrl), default(CancellationToken));
                 Trace.Info($"Successfully started websocket client.");
             }
-            catch(Exception ex)
+            catch (Exception ex)
             {
                 Trace.Info("Exception caught during websocket client connect, fallback of HTTP would be used now instead of websocket.");
                 Trace.Error(ex);
@@ -231,7 +231,7 @@ public async Task AppendTimelineRecordFeedAsync(Guid scopeIdentifier, string hub
             // ...in other words, if websocket client is null, we will skip sending to websocket and just use rest api calls to send data
             if (_websocketClient != null)
             {
-                var linesWrapper =  startLine.HasValue? new TimelineRecordFeedLinesWrapper(stepId, lines, startLine.Value):  new TimelineRecordFeedLinesWrapper(stepId, lines);
+                var linesWrapper = startLine.HasValue ? new TimelineRecordFeedLinesWrapper(stepId, lines, startLine.Value) : new TimelineRecordFeedLinesWrapper(stepId, lines);
                 var jsonData = StringUtil.ConvertToJson(linesWrapper);
                 try
                 {
@@ -242,7 +242,7 @@ public async Task AppendTimelineRecordFeedAsync(Guid scopeIdentifier, string hub
                     {
                         var lastChunk = i + (1 * 1024) >= jsonDataBytes.Length;
                         var chunk = new ArraySegment<byte>(jsonDataBytes, i, Math.Min(1 * 1024, jsonDataBytes.Length - i));
-                        await _websocketClient.SendAsync(chunk, WebSocketMessageType.Text, endOfMessage:lastChunk, cancellationToken);
+                        await _websocketClient.SendAsync(chunk, WebSocketMessageType.Text, endOfMessage: lastChunk, cancellationToken);
                     }
 
                     pushedLinesViaWebsocket = true;
@@ -274,7 +274,7 @@ public async Task AppendTimelineRecordFeedAsync(Guid scopeIdentifier, string hub
                 }
             }
 
-            if (!pushedLinesViaWebsocket)
+            if (!pushedLinesViaWebsocket && !cancellationToken.IsCancellationRequested)
             {
                 if (startLine.HasValue)
                 {
diff --git a/src/Runner.Common/JobServerQueue.cs b/src/Runner.Common/JobServerQueue.cs
index 822f8c2c098..3da4400a9a6 100644
--- a/src/Runner.Common/JobServerQueue.cs
+++ b/src/Runner.Common/JobServerQueue.cs
@@ -299,7 +299,11 @@ private async Task ProcessWebConsoleLinesQueueAsync(bool runOnce = false)
                         {
                             try
                             {
-                                await _jobServer.AppendTimelineRecordFeedAsync(_scopeIdentifier, _hubName, _planId, _jobTimelineId, _jobTimelineRecordId, stepRecordId, batch.Select(logLine => logLine.Line).ToList(), batch[0].LineNumber, default(CancellationToken));
+                                // Give at most 60s for each request. 
+                                using (var timeoutTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(60)))
+                                {
+                                    await _jobServer.AppendTimelineRecordFeedAsync(_scopeIdentifier, _hubName, _planId, _jobTimelineId, _jobTimelineRecordId, stepRecordId, batch.Select(logLine => logLine.Line).ToList(), batch[0].LineNumber, timeoutTokenSource.Token);
+                                }
 
                                 if (_firstConsoleOutputs)
                                 {
diff --git a/src/dir.proj b/src/dir.proj
index cc0b796ebc1..fdd1a5cee76 100644
--- a/src/dir.proj
+++ b/src/dir.proj
@@ -48,7 +48,7 @@
 
     <Target Name="Test" DependsOnTargets="GenerateConstant">
         <Exec Command="dotnet build Test/Test.csproj -c $(BUILDCONFIG) /p:PackageRuntime=$(PackageRuntime)" ConsoleToMSBuild="true" />
-        <Exec Command="dotnet test Test/Test.csproj --no-build --logger:trx" ConsoleToMSBuild="true" />
+        <Exec Command="dotnet test Test/Test.csproj -c $(BUILDCONFIG) --no-build --logger:trx" ConsoleToMSBuild="true" />
     </Target>
 
     <Target Name="Layout" DependsOnTargets="Clean;Build">