Skip DeleteAgentSession when the acess token has been revoked. (#1692)

src/Runner.Listener/MessageListener.cs

@@ -1,18 +1,18 @@
-using GitHub.DistributedTask.WebApi;
-using GitHub.Runner.Listener.Configuration;
-using GitHub.Services.Common;
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
+using System.IO;
+using System.Runtime.InteropServices;
+using System.Security.Cryptography;
+using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
-using System.Security.Cryptography;
+using GitHub.DistributedTask.WebApi;
-using System.IO;
-using System.Text;
-using GitHub.Services.OAuth;
-using System.Diagnostics;
-using System.Runtime.InteropServices;
 using GitHub.Runner.Common;
+using GitHub.Runner.Listener.Configuration;
 using GitHub.Runner.Sdk;
+using GitHub.Services.Common;
+using GitHub.Services.OAuth;
 
 namespace GitHub.Runner.Listener
 {
@@ -33,6 +33,7 @@ namespace GitHub.Runner.Listener
         private IRunnerServer _runnerServer;
         private TaskAgentSession _session;
         private TimeSpan _getNextMessageRetryInterval;
+        private bool _accessTokenRevoked = false;
         private readonly TimeSpan _sessionCreationRetryInterval = TimeSpan.FromSeconds(30);
         private readonly TimeSpan _sessionConflictRetryLimit = TimeSpan.FromMinutes(4);
         private readonly TimeSpan _clockSkewRetryLimit = TimeSpan.FromMinutes(30);
@@ -111,6 +112,7 @@ namespace GitHub.Runner.Listener
                 catch (TaskAgentAccessTokenExpiredException)
                 {
                     Trace.Info("Runner OAuth token has been revoked. Session creation failed.");
+                    _accessTokenRevoked = true;
                     throw;
                 }
                 catch (Exception ex)
@@ -153,12 +155,19 @@ namespace GitHub.Runner.Listener
         public async Task DeleteSessionAsync()
         {
             if (_session != null && _session.SessionId != Guid.Empty)
+            {
+                if (!_accessTokenRevoked)
                 {
                     using (var ts = new CancellationTokenSource(TimeSpan.FromSeconds(30)))
                     {
                         await _runnerServer.DeleteAgentSessionAsync(_settings.PoolId, _session.SessionId, ts.Token);
                     }
                 }
+                else
+                {
+                    Trace.Warning("Runner OAuth token has been revoked. Skip deleting session.");
+                }
+            }
         }
 
         public async Task<TaskAgentMessage> GetNextMessageAsync(CancellationToken token)
@@ -205,6 +214,7 @@ namespace GitHub.Runner.Listener
                 catch (TaskAgentAccessTokenExpiredException)
                 {
                     Trace.Info("Runner OAuth token has been revoked. Unable to pull message.");
+                    _accessTokenRevoked = true;
                     throw;
                 }
                 catch (Exception ex)

src/Test/L0/Listener/MessageListenerL0.cs

@@ -1,18 +1,18 @@
-using GitHub.DistributedTask.WebApi;
+using System;
-using GitHub.Services.Common;
-using GitHub.Services.WebApi;
-using GitHub.Runner.Listener;
-using GitHub.Runner.Listener.Configuration;
-using Moq;
-using System;
-using System.Runtime.CompilerServices;
-using System.Threading.Tasks;
-using Xunit;
-using System.Threading;
-using System.Reflection;
 using System.Collections.Generic;
 using System.IO;
+using System.Reflection;
+using System.Runtime.CompilerServices;
 using System.Security.Cryptography;
+using System.Threading;
+using System.Threading.Tasks;
+using GitHub.DistributedTask.WebApi;
+using GitHub.Runner.Listener;
+using GitHub.Runner.Listener.Configuration;
+using GitHub.Services.Common;
+using GitHub.Services.WebApi;
+using Moq;
+using Xunit;
 
 namespace GitHub.Runner.Common.Tests.Listener
 {
@@ -256,5 +256,67 @@ namespace GitHub.Runner.Common.Tests.Listener
                         tokenSource.Token), Times.Once());
             }
         }
+
+        [Fact]
+        [Trait("Level", "L0")]
+        [Trait("Category", "Runner")]
+        public async void SkipDeleteSession_WhenGetNextMessageGetTaskAgentAccessTokenExpiredException()
+        {
+            using (TestHostContext tc = CreateTestContext())
+            using (var tokenSource = new CancellationTokenSource())
+            {
+                Tracing trace = tc.GetTrace();
+
+                // Arrange.
+                var expectedSession = new TaskAgentSession();
+                PropertyInfo sessionIdProperty = expectedSession.GetType().GetProperty("SessionId", BindingFlags.Instance | BindingFlags.NonPublic | BindingFlags.Public);
+                Assert.NotNull(sessionIdProperty);
+                sessionIdProperty.SetValue(expectedSession, Guid.NewGuid());
+
+                _runnerServer
+                    .Setup(x => x.CreateAgentSessionAsync(
+                        _settings.PoolId,
+                        It.Is<TaskAgentSession>(y => y != null),
+                        tokenSource.Token))
+                    .Returns(Task.FromResult(expectedSession));
+
+                _credMgr.Setup(x => x.LoadCredentials()).Returns(new VssCredentials());
+                _store.Setup(x => x.GetCredentials()).Returns(new CredentialData() { Scheme = Constants.Configuration.OAuthAccessToken });
+                _store.Setup(x => x.GetMigratedCredentials()).Returns(default(CredentialData));
+
+                // Act.
+                MessageListener listener = new MessageListener();
+                listener.Initialize(tc);
+
+                bool result = await listener.CreateSessionAsync(tokenSource.Token);
+                Assert.True(result);
+
+                _runnerServer
+                    .Setup(x => x.GetAgentMessageAsync(
+                        _settings.PoolId, expectedSession.SessionId, It.IsAny<long?>(), tokenSource.Token))
+                    .Throws(new TaskAgentAccessTokenExpiredException("test"));
+                try
+                {
+                    await listener.GetNextMessageAsync(tokenSource.Token);
+                }
+                catch (TaskAgentAccessTokenExpiredException)
+                {
+                    //expected
+                }
+                finally
+                {
+                    await listener.DeleteSessionAsync();
+                }
+
+                //Assert
+                _runnerServer
+                    .Verify(x => x.GetAgentMessageAsync(
+                        _settings.PoolId, expectedSession.SessionId, It.IsAny<long?>(), tokenSource.Token), Times.Once);
+
+                _runnerServer
+                    .Verify(x => x.DeleteAgentSessionAsync(
+                        _settings.PoolId, expectedSession.SessionId, It.IsAny<CancellationToken>()), Times.Never);
+            }
+        }
     }
 }