Handle non success raw http responses (#2541)

* handle non success responses

* fix format

* nits

Co-authored-by: Tingluo Huang <tingluohuang@github.com>

---------

Co-authored-by: Tingluo Huang <tingluohuang@github.com>

src/Runner.Common/RunServer.cs

@@ -51,14 +51,8 @@ namespace GitHub.Runner.Common
         public Task<AgentJobRequestMessage> GetJobMessageAsync(string id, CancellationToken cancellationToken)
         {
             CheckConnection();
-            var jobMessage = RetryRequest<AgentJobRequestMessage>(
+            return RetryRequest<AgentJobRequestMessage>(
                 async () => await _runServiceHttpClient.GetJobMessageAsync(requestUri, id, cancellationToken), cancellationToken);
-            if (jobMessage == null)
-            {
-                throw new TaskOrchestrationJobNotFoundException(id);
-            }
-
-            return jobMessage;
         }
 
         public Task CompleteJobAsync(Guid planId, Guid jobId, TaskResult result, Dictionary<String, VariableValue> outputs, IList<StepResult> stepResults, CancellationToken cancellationToken)
@@ -71,14 +65,8 @@ namespace GitHub.Runner.Common
         public Task<RenewJobResponse> RenewJobAsync(Guid planId, Guid jobId, CancellationToken cancellationToken)
         {
             CheckConnection();
-            var renewJobResponse = RetryRequest<RenewJobResponse>(
+            return RetryRequest<RenewJobResponse>(
                 async () => await _runServiceHttpClient.RenewJobAsync(requestUri, planId, jobId, cancellationToken), cancellationToken);
-            if (renewJobResponse == null)
-            {
-                throw new TaskOrchestrationJobNotFoundException(jobId.ToString());
-            }
-
-            return renewJobResponse;
         }
     }
 }

src/Sdk/RSWebApi/RunServiceHttpClient.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Net;
 using System.Net.Http;
 using System.Threading;
 using System.Threading.Tasks;
@@ -55,7 +56,7 @@ namespace GitHub.Actions.RunService.WebApi
         {
         }
 
-        public Task<AgentJobRequestMessage> GetJobMessageAsync(
+        public async Task<AgentJobRequestMessage> GetJobMessageAsync(
             Uri requestUri,
             string messageId,
             CancellationToken cancellationToken = default)
@@ -69,14 +70,27 @@ namespace GitHub.Actions.RunService.WebApi
             requestUri = new Uri(requestUri, "acquirejob");
 
             var requestContent = new ObjectContent<AcquireJobRequest>(payload, new VssJsonMediaTypeFormatter(true));
-            return SendAsync<AgentJobRequestMessage>(
+            var result = await SendAsync<AgentJobRequestMessage>(
                 httpMethod,
                 requestUri: requestUri,
                 content: requestContent,
                 cancellationToken: cancellationToken);
+
+            if (result.IsSuccess)
+            {
+                return result.Value;
+            }
+
+            switch (result.StatusCode)
+            {
+                case HttpStatusCode.NotFound:
+                    throw new TaskOrchestrationJobNotFoundException($"Job message not found: {messageId}");
+                default:
+                    throw new Exception($"Failed to get job message: {result.Error}");
+            }
         }
 
-        public Task CompleteJobAsync(
+        public async Task CompleteJobAsync(
             Uri requestUri,
             Guid planId,
             Guid jobId,
@@ -98,14 +112,26 @@ namespace GitHub.Actions.RunService.WebApi
             requestUri = new Uri(requestUri, "completejob");
 
             var requestContent = new ObjectContent<CompleteJobRequest>(payload, new VssJsonMediaTypeFormatter(true));
-            return SendAsync(
+            var response = await SendAsync(
                     httpMethod,
                     requestUri,
                     content: requestContent,
                     cancellationToken: cancellationToken);
+            if (response.IsSuccessStatusCode)
+            {
+                return;
+            }
+
+            switch (response.StatusCode)
+            {
+                case HttpStatusCode.NotFound:
+                    throw new TaskOrchestrationJobNotFoundException($"Job not found: {jobId}");
+                default:
+                    throw new Exception($"Failed to complete job: {response.ReasonPhrase}");
+            }
         }
 
-        public Task<RenewJobResponse> RenewJobAsync(
+        public async Task<RenewJobResponse> RenewJobAsync(
             Uri requestUri,
             Guid planId,
             Guid jobId,
@@ -121,11 +147,24 @@ namespace GitHub.Actions.RunService.WebApi
             requestUri = new Uri(requestUri, "renewjob");
 
             var requestContent = new ObjectContent<RenewJobRequest>(payload, new VssJsonMediaTypeFormatter(true));
-            return SendAsync<RenewJobResponse>(
+            var result = await SendAsync<RenewJobResponse>(
                 httpMethod,
                 requestUri,
                 content: requestContent,
                 cancellationToken: cancellationToken);
+
+            if (result.IsSuccess)
+            {
+                return result.Value;
+            }
+
+            switch (result.StatusCode)
+            {
+                case HttpStatusCode.NotFound:
+                    throw new TaskOrchestrationJobNotFoundException($"Job not found: {jobId}");
+                default:
+                    throw new Exception($"Failed to renew job: {result.Error}");
+            }
         }
     }
 }

src/Sdk/WebApi/WebApi/BrokerHttpClient.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.Net;
 using System.Net.Http;
 using System.Threading;
 using System.Threading.Tasks;
@@ -55,7 +56,7 @@ namespace GitHub.Actions.RunService.WebApi
         {
         }
 
-        public Task<TaskAgentMessage> GetRunnerMessageAsync(
+        public async Task<TaskAgentMessage> GetRunnerMessageAsync(
             string runnerVersion,
             TaskAgentStatus? status,
             CancellationToken cancellationToken = default
@@ -74,11 +75,18 @@ namespace GitHub.Actions.RunService.WebApi
                 queryParams.Add("runnerVersion", runnerVersion);
             }
 
-            return SendAsync<TaskAgentMessage>(
+            var result = await SendAsync<TaskAgentMessage>(
                 new HttpMethod("GET"),
                 requestUri: requestUri,
                 queryParameters: queryParams,
                 cancellationToken: cancellationToken);
+
+            if (result.IsSuccess)
+            {
+                return result.Value;
+            }
+
+            throw new Exception($"Failed to get job message: {result.Error}");
         }
     }
 }

src/Sdk/WebApi/WebApi/RawHttpClientBase.cs

@@ -21,7 +21,7 @@ using Newtonsoft.Json.Linq;
 
 namespace Sdk.WebApi.WebApi
 {
-    public class RawHttpClientBase: IDisposable
+    public class RawHttpClientBase : IDisposable
     {
         protected RawHttpClientBase(
             Uri baseUrl,
@@ -101,7 +101,7 @@ namespace Sdk.WebApi.WebApi
             }
         }
 
-        protected Task<T> SendAsync<T>(
+        protected Task<RawHttpClientResult<T>> SendAsync<T>(
             HttpMethod method,
             Uri requestUri,
             HttpContent content = null,
@@ -112,7 +112,7 @@ namespace Sdk.WebApi.WebApi
             return SendAsync<T>(method, null, requestUri, content, queryParameters, userState, cancellationToken);
         }
 
-        protected async Task<T> SendAsync<T>(
+        protected async Task<RawHttpClientResult<T>> SendAsync<T>(
             HttpMethod method,
             IEnumerable<KeyValuePair<String, String>> additionalHeaders,
             Uri requestUri,
@@ -128,7 +128,7 @@ namespace Sdk.WebApi.WebApi
             }
         }
 
-        protected async Task<T> SendAsync<T>(
+        protected async Task<RawHttpClientResult<T>> SendAsync<T>(
             HttpRequestMessage message,
             Object userState = null,
             CancellationToken cancellationToken = default(CancellationToken))
@@ -138,7 +138,16 @@ namespace Sdk.WebApi.WebApi
             //from deadlocking...
             using (HttpResponseMessage response = await this.SendAsync(message, userState, cancellationToken).ConfigureAwait(false))
             {
-                return await ReadContentAsAsync<T>(response, cancellationToken).ConfigureAwait(false);
+                if (response.IsSuccessStatusCode)
+                {
+                    T data = await ReadContentAsAsync<T>(response, cancellationToken).ConfigureAwait(false);
+                    return RawHttpClientResult<T>.Ok(data);
+                }
+                else
+                {
+                    string errorMessage = $"Error: {response.ReasonPhrase}";
+                    return RawHttpClientResult<T>.Fail(errorMessage, response.StatusCode);
+                }
             }
         }
 

src/Sdk/WebApi/WebApi/RawHttpClientResult.cs

@@ -0,0 +1,33 @@
+using System.Net;
+
+namespace Sdk.WebApi.WebApi
+{
+    public class RawHttpClientResult
+    {
+        public bool IsSuccess { get; protected set; }
+        public string Error { get; protected set; }
+        public HttpStatusCode StatusCode { get; protected set; }
+        public bool IsFailure => !IsSuccess;
+
+        protected RawHttpClientResult(bool isSuccess, string error, HttpStatusCode statusCode)
+        {
+            IsSuccess = isSuccess;
+            Error = error;
+            StatusCode = statusCode;
+        }
+    }
+
+    public class RawHttpClientResult<T> : RawHttpClientResult
+    {
+        public T Value { get; private set; }
+
+        protected internal RawHttpClientResult(T value, bool isSuccess, string error, HttpStatusCode statusCode)
+            : base(isSuccess, error, statusCode)
+        {
+            Value = value;
+        }
+
+        public static RawHttpClientResult<T> Fail(string message, HttpStatusCode statusCode) => new RawHttpClientResult<T>(default(T), false, message, statusCode);
+        public static RawHttpClientResult<T> Ok(T value) => new RawHttpClientResult<T>(value, true, string.Empty, HttpStatusCode.OK);
+    }
+}