Check connectivity for endpoints requested by service. (#2691)

src/Runner.Worker/JobExtension.cs

@@ -4,6 +4,7 @@ using System.Diagnostics;
 using System.Globalization;
 using System.IO;
 using System.Linq;
+using System.Net.Http;
 using System.Runtime.Serialization;
 using System.Threading;
 using System.Threading.Tasks;
@@ -15,6 +16,7 @@ using GitHub.DistributedTask.WebApi;
 using GitHub.Runner.Common;
 using GitHub.Runner.Common.Util;
 using GitHub.Runner.Sdk;
+using GitHub.Services.Common;
 using Pipelines = GitHub.DistributedTask.Pipelines;
 
 namespace GitHub.Runner.Worker
@@ -34,12 +36,13 @@ namespace GitHub.Runner.Worker
     public interface IJobExtension : IRunnerService
     {
         Task<List<IStep>> InitializeJob(IExecutionContext jobContext, Pipelines.AgentJobRequestMessage message);
-        void FinalizeJob(IExecutionContext jobContext, Pipelines.AgentJobRequestMessage message, DateTime jobStartTimeUtc);
+        Task FinalizeJob(IExecutionContext jobContext, Pipelines.AgentJobRequestMessage message, DateTime jobStartTimeUtc);
     }
 
     public sealed class JobExtension : RunnerService, IJobExtension
     {
         private readonly HashSet<string> _existingProcesses = new(StringComparer.OrdinalIgnoreCase);
+        private readonly List<Task<string>> _connectivityCheckTasks = new();
         private bool _processCleanup;
         private string _processLookupId = $"github_{Guid.NewGuid()}";
         private CancellationTokenSource _diskSpaceCheckToken = new();
@@ -428,6 +431,22 @@ namespace GitHub.Runner.Worker
                         _diskSpaceCheckTask = CheckDiskSpaceAsync(context, _diskSpaceCheckToken.Token);
                     }
 
+                    // Check server connectivity in background
+                    ServiceEndpoint systemConnection = message.Resources.Endpoints.Single(x => string.Equals(x.Name, WellKnownServiceEndpointNames.SystemVssConnection, StringComparison.OrdinalIgnoreCase));
+                    if (systemConnection.Data.TryGetValue("ConnectivityChecks", out var connectivityChecksPayload) &&
+                        !string.IsNullOrEmpty(connectivityChecksPayload))
+                    {
+                        Trace.Info($"Start checking server connectivity.");
+                        var checkUrls = StringUtil.ConvertFromJson<List<string>>(connectivityChecksPayload);
+                        if (checkUrls?.Count > 0)
+                        {
+                            foreach (var checkUrl in checkUrls)
+                            {
+                                _connectivityCheckTasks.Add(CheckConnectivity(checkUrl));
+                            }
+                        }
+                    }
+
                     return steps;
                 }
                 catch (OperationCanceledException ex) when (jobContext.CancellationToken.IsCancellationRequested)
@@ -472,7 +491,7 @@ namespace GitHub.Runner.Worker
             return reference;
         }
 
-        public void FinalizeJob(IExecutionContext jobContext, Pipelines.AgentJobRequestMessage message, DateTime jobStartTimeUtc)
+        public async Task FinalizeJob(IExecutionContext jobContext, Pipelines.AgentJobRequestMessage message, DateTime jobStartTimeUtc)
         {
             Trace.Entering();
             ArgUtil.NotNull(jobContext, nameof(jobContext));
@@ -649,6 +668,28 @@ namespace GitHub.Runner.Worker
                     {
                         _diskSpaceCheckToken.Cancel();
                     }
+
+                    // Collect server connectivity check result
+                    if (_connectivityCheckTasks.Count > 0)
+                    {
+                        try
+                        {
+                            Trace.Info($"Wait for all connectivity checks to finish.");
+                            await Task.WhenAll(_connectivityCheckTasks);
+                            foreach (var check in _connectivityCheckTasks)
+                            {
+                                var result = await check;
+                                Trace.Info($"Connectivity check result: {result}");
+                                context.Global.JobTelemetry.Add(new JobTelemetry() { Type = JobTelemetryType.ConnectivityCheck, Message = result });
+                            }
+                        }
+                        catch (Exception ex)
+                        {
+                            Trace.Error($"Fail to check server connectivity.");
+                            Trace.Error(ex);
+                            context.Global.JobTelemetry.Add(new JobTelemetry() { Type = JobTelemetryType.ConnectivityCheck, Message = $"Fail to check server connectivity. {ex.Message}" });
+                        }
+                    }
                 }
                 catch (Exception ex)
                 {
@@ -664,6 +705,37 @@ namespace GitHub.Runner.Worker
             }
         }
 
+        private async Task<string> CheckConnectivity(string endpointUrl)
+        {
+            Trace.Info($"Check server connectivity for {endpointUrl}.");
+            string result = string.Empty;
+            using (var timeoutTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(5)))
+            {
+                try
+                {
+                    using (var httpClientHandler = HostContext.CreateHttpClientHandler())
+                    using (var httpClient = new HttpClient(httpClientHandler))
+                    {
+                        httpClient.DefaultRequestHeaders.UserAgent.AddRange(HostContext.UserAgents);
+                        var response = await httpClient.GetAsync(endpointUrl, timeoutTokenSource.Token);
+                        result = $"{endpointUrl}: {response.StatusCode}";
+                    }
+                }
+                catch (Exception ex) when (ex is OperationCanceledException && timeoutTokenSource.IsCancellationRequested)
+                {
+                    Trace.Error($"Request timeout during connectivity check: {ex}");
+                    result = $"{endpointUrl}: timeout";
+                }
+                catch (Exception ex)
+                {
+                    Trace.Error($"Catch exception during connectivity check: {ex}");
+                    result = $"{endpointUrl}: {ex.Message}";
+                }
+            }
+
+            return result;
+        }
+
         private async Task CheckDiskSpaceAsync(IExecutionContext context, CancellationToken token)
         {
             while (!token.IsCancellationRequested)

src/Runner.Worker/JobRunner.cs

@@ -229,7 +229,7 @@ namespace GitHub.Runner.Worker
                 finally
                 {
                     Trace.Info("Finalize job.");
-                    jobExtension.FinalizeJob(jobContext, message, jobStartTimeUtc);
+                    await jobExtension.FinalizeJob(jobContext, message, jobStartTimeUtc);
                 }
 
                 Trace.Info($"Job result after all job steps finish: {jobContext.Result ?? TaskResult.Succeeded}");

src/Sdk/DTWebApi/WebApi/JobTelemetryType.cs

@@ -9,5 +9,8 @@ namespace GitHub.DistributedTask.WebApi
 
         [EnumMember]
         ActionCommand = 1,
+
+        [EnumMember]
+        ConnectivityCheck = 2,
     }
 }

src/Test/L0/Worker/JobExtensionL0.cs

@@ -1,13 +1,13 @@
-using GitHub.DistributedTask.WebApi;
+using System;
-using GitHub.Runner.Worker;
-using Moq;
-using System;
 using System.Collections.Generic;
 using System.Linq;
 using System.Runtime.CompilerServices;
-using System.Threading.Tasks;
-using Xunit;
 using System.Threading;
+using System.Threading.Tasks;
+using GitHub.DistributedTask.WebApi;
+using GitHub.Runner.Worker;
+using Moq;
+using Xunit;
 using Pipelines = GitHub.DistributedTask.Pipelines;
 
 namespace GitHub.Runner.Common.Tests.Worker
@@ -105,6 +105,18 @@ namespace GitHub.Runner.Common.Tests.Worker
             github["repository"] = new Pipelines.ContextData.StringContextData("actions/runner");
             github["secret_source"] = new Pipelines.ContextData.StringContextData("Actions");
             _message.ContextData.Add("github", github);
+            _message.Resources.Endpoints.Add(new ServiceEndpoint()
+            {
+                Name = WellKnownServiceEndpointNames.SystemVssConnection,
+                Url = new Uri("https://pipelines.actions.githubusercontent.com"),
+                Authorization = new EndpointAuthorization()
+                {
+                    Scheme = "Test",
+                    Parameters = {
+                        {"AccessToken", "token"}
+                    }
+                },
+            });
 
             hc.SetSingleton(_actionManager.Object);
             hc.SetSingleton(_config.Object);
@@ -231,7 +243,7 @@ namespace GitHub.Runner.Common.Tests.Worker
         [Fact]
         [Trait("Level", "L0")]
         [Trait("Category", "Worker")]
-        public void UploadDiganosticLogIfEnvironmentVariableSet()
+        public async Task UploadDiganosticLogIfEnvironmentVariableSet()
         {
             using (TestHostContext hc = CreateTestContext())
             {
@@ -244,7 +256,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 _jobEc.Initialize(hc);
                 _jobEc.InitializeJob(_message, _tokenSource.Token);
 
-                jobExtension.FinalizeJob(_jobEc, _message, DateTime.UtcNow);
+                await jobExtension.FinalizeJob(_jobEc, _message, DateTime.UtcNow);
 
                 _diagnosticLogManager.Verify(x =>
                     x.UploadDiagnosticLogs(
@@ -259,7 +271,7 @@ namespace GitHub.Runner.Common.Tests.Worker
         [Fact]
         [Trait("Level", "L0")]
         [Trait("Category", "Worker")]
-        public void DontUploadDiagnosticLogIfEnvironmentVariableFalse()
+        public async Task DontUploadDiagnosticLogIfEnvironmentVariableFalse()
         {
             using (TestHostContext hc = CreateTestContext())
             {
@@ -272,7 +284,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 _jobEc.Initialize(hc);
                 _jobEc.InitializeJob(_message, _tokenSource.Token);
 
-                jobExtension.FinalizeJob(_jobEc, _message, DateTime.UtcNow);
+                await jobExtension.FinalizeJob(_jobEc, _message, DateTime.UtcNow);
 
                 _diagnosticLogManager.Verify(x =>
                     x.UploadDiagnosticLogs(
@@ -287,14 +299,14 @@ namespace GitHub.Runner.Common.Tests.Worker
         [Fact]
         [Trait("Level", "L0")]
         [Trait("Category", "Worker")]
-        public void DontUploadDiagnosticLogIfEnvironmentVariableMissing()
+        public async Task DontUploadDiagnosticLogIfEnvironmentVariableMissing()
         {
             using (TestHostContext hc = CreateTestContext())
             {
                 var jobExtension = new JobExtension();
                 jobExtension.Initialize(hc);
 
-                jobExtension.FinalizeJob(_jobEc, _message, DateTime.UtcNow);
+                await jobExtension.FinalizeJob(_jobEc, _message, DateTime.UtcNow);
 
                 _diagnosticLogManager.Verify(x =>
                     x.UploadDiagnosticLogs(
@@ -309,7 +321,7 @@ namespace GitHub.Runner.Common.Tests.Worker
         [Fact]
         [Trait("Level", "L0")]
         [Trait("Category", "Worker")]
-        public void EnsureFinalizeJobRunsIfMessageHasNoEnvironmentUrl()
+        public async Task EnsureFinalizeJobRunsIfMessageHasNoEnvironmentUrl()
         {
             using (TestHostContext hc = CreateTestContext())
             {
@@ -322,7 +334,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 _jobEc.Initialize(hc);
                 _jobEc.InitializeJob(_message, _tokenSource.Token);
 
-                jobExtension.FinalizeJob(_jobEc, _message, DateTime.UtcNow);
+                await jobExtension.FinalizeJob(_jobEc, _message, DateTime.UtcNow);
 
                 Assert.Equal(TaskResult.Succeeded, _jobEc.Result);
             }
@@ -331,7 +343,7 @@ namespace GitHub.Runner.Common.Tests.Worker
         [Fact]
         [Trait("Level", "L0")]
         [Trait("Category", "Worker")]
-        public void EnsureFinalizeJobHandlesNullEnvironmentUrl()
+        public async Task EnsureFinalizeJobHandlesNullEnvironmentUrl()
         {
             using (TestHostContext hc = CreateTestContext())
             {
@@ -347,7 +359,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 _jobEc.Initialize(hc);
                 _jobEc.InitializeJob(_message, _tokenSource.Token);
 
-                jobExtension.FinalizeJob(_jobEc, _message, DateTime.UtcNow);
+                await jobExtension.FinalizeJob(_jobEc, _message, DateTime.UtcNow);
 
                 Assert.Equal(TaskResult.Succeeded, _jobEc.Result);
             }
@@ -356,7 +368,7 @@ namespace GitHub.Runner.Common.Tests.Worker
         [Fact]
         [Trait("Level", "L0")]
         [Trait("Category", "Worker")]
-        public void EnsureFinalizeJobHandlesNullEnvironment()
+        public async Task EnsureFinalizeJobHandlesNullEnvironment()
         {
             using (TestHostContext hc = CreateTestContext())
             {
@@ -369,7 +381,7 @@ namespace GitHub.Runner.Common.Tests.Worker
                 _jobEc.Initialize(hc);
                 _jobEc.InitializeJob(_message, _tokenSource.Token);
 
-                jobExtension.FinalizeJob(_jobEc, _message, DateTime.UtcNow);
+                await jobExtension.FinalizeJob(_jobEc, _message, DateTime.UtcNow);
 
                 Assert.Equal(TaskResult.Succeeded, _jobEc.Result);
             }
@@ -397,7 +409,7 @@ namespace GitHub.Runner.Common.Tests.Worker
 
                 var hookStart = result.First() as JobExtensionRunner;
 
-                jobExtension.FinalizeJob(_jobEc, _message, DateTime.UtcNow);
+                await jobExtension.FinalizeJob(_jobEc, _message, DateTime.UtcNow);
 
                 Assert.Equal(Constants.Hooks.JobStartedStepName, hookStart.DisplayName);
                 Assert.Equal(Constants.Hooks.JobCompletedStepName, (_jobEc.PostJobSteps.Last() as JobExtensionRunner).DisplayName);
@@ -410,7 +422,7 @@ namespace GitHub.Runner.Common.Tests.Worker
         [Fact]
         [Trait("Level", "L0")]
         [Trait("Category", "Worker")]
-        public void EnsureNoPreAndPostHookSteps()
+        public async Task EnsureNoPreAndPostHookSteps()
         {
             using (TestHostContext hc = CreateTestContext())
             {
@@ -425,7 +437,7 @@ namespace GitHub.Runner.Common.Tests.Worker
 
                 var x = _jobEc.JobSteps;
 
-                jobExtension.FinalizeJob(_jobEc, _message, DateTime.UtcNow);
+                await jobExtension.FinalizeJob(_jobEc, _message, DateTime.UtcNow);
 
                 Assert.Equal(TaskResult.Succeeded, _jobEc.Result);
                 Assert.Equal(0, _jobEc.PostJobSteps.Count);