public async Task <TaskAgentMessage> GetNextMessageAsync(CancellationToken token)
{
Trace.Entering();
ArgUtil.NotNull(Session, nameof(Session));
ArgUtil.NotNull(_settings, nameof(_settings));
var agentServer = HostContext.GetService <IAgentServer>();
while (true)
{
token.ThrowIfCancellationRequested();
TaskAgentMessage message = null;
try
{
message = await agentServer.GetAgentMessageAsync(_settings.PoolId,
Session.SessionId,
_lastMessageId,
token);
if (message != null)
{
_lastMessageId = message.MessageId;
}
}
catch (TimeoutException)
{
Trace.Verbose($"{nameof(TimeoutException)} received.");
}
catch (TaskCanceledException)
{
Trace.Verbose($"{nameof(TaskCanceledException)} received.");
}
catch (TaskAgentSessionExpiredException)
{
Trace.Verbose($"{nameof(TaskAgentSessionExpiredException)} received.");
// TODO: Throw a specific exception so the caller can control the flow appropriately.
throw;
}
catch (OperationCanceledException)
{
throw;
}
catch (Exception ex)
{
Trace.Error(ex);
// TODO: Throw a specific exception so the caller can control the flow appropriately.
throw;
}
if (message == null)
{
Trace.Verbose($"No message retrieved from session '{Session.SessionId}'.");
continue;
}
Trace.Verbose($"Message '{message.MessageId}' received from session '{Session.SessionId}'.");
return(message);
}
}
private async Task DeleteMessageAsync(TaskAgentMessage message)
{
if (message != null && _listener.Session.SessionId != Guid.Empty)
{
var agentServer = HostContext.GetService <IAgentServer>();
using (var cs = new CancellationTokenSource(TimeSpan.FromSeconds(30)))
{
await agentServer.DeleteAgentMessageAsync(_poolId, message.MessageId, _listener.Session.SessionId, cs.Token);
}
}
}
public async Task DeleteMessageAsync(TaskAgentMessage message)
{
Trace.Entering();
ArgUtil.NotNull(_session, nameof(_session));
if (message != null && _session.SessionId != Guid.Empty)
{
using (var cs = new CancellationTokenSource(TimeSpan.FromSeconds(30)))
{
await _runnerServer.DeleteAgentMessageAsync(_settings.PoolId, message.MessageId, _session.SessionId, cs.Token);
}
}
}
private ICryptoTransform GetMessageDecryptor(
Aes aes,
TaskAgentMessage message)
{
if (_session.EncryptionKey.Encrypted)
{
// The agent session encryption key uses the AES symmetric algorithm
var keyManager = HostContext.GetService <IRSAKeyManager>();
using (var rsa = keyManager.GetKey())
{
return(aes.CreateDecryptor(rsa.Decrypt(_session.EncryptionKey.Value, RSAEncryptionPadding.OaepSHA1), message.IV));
}
}
else
{
return(aes.CreateDecryptor(_session.EncryptionKey.Value, message.IV));
}
}
private TaskAgentMessage DecryptMessage(TaskAgentMessage message)
{
if (_session.EncryptionKey == null ||
_session.EncryptionKey.Value.Length == 0 ||
message == null ||
message.IV == null ||
message.IV.Length == 0)
{
return(message);
}
using (var aes = Aes.Create())
using (var decryptor = GetMessageDecryptor(aes, message))
using (var body = new MemoryStream(Convert.FromBase64String(message.Body)))
using (var cryptoStream = new CryptoStream(body, decryptor, CryptoStreamMode.Read))
using (var bodyReader = new StreamReader(cryptoStream, Encoding.UTF8))
{
message.Body = bodyReader.ReadToEnd();
}
return(message);
}
//create worker manager, create message listener and start listening to the queue
private async Task <int> RunAsync(CancellationToken token, AgentSettings settings)
{
Trace.Info(nameof(RunAsync));
// Load the settings.
_poolId = settings.PoolId;
var listener = HostContext.GetService <IMessageListener>();
if (!await listener.CreateSessionAsync(token))
{
return(1);
}
_term.WriteLine(StringUtil.Loc("ListenForJobs"));
_sessionId = listener.Session.SessionId;
TaskAgentMessage message = null;
try
{
using (var workerManager = HostContext.GetService <IWorkerManager>())
{
while (!token.IsCancellationRequested)
{
try
{
message = await listener.GetNextMessageAsync(token); //get next message
if (string.Equals(message.MessageType, AgentRefreshMessage.MessageType, StringComparison.OrdinalIgnoreCase))
{
Trace.Warning("Referesh message received, but not yet handled by agent implementation.");
}
else if (string.Equals(message.MessageType, JobRequestMessage.MessageType, StringComparison.OrdinalIgnoreCase))
{
var newJobMessage = JsonUtility.FromString <JobRequestMessage>(message.Body);
workerManager.Run(newJobMessage);
}
else if (string.Equals(message.MessageType, JobCancelMessage.MessageType, StringComparison.OrdinalIgnoreCase))
{
var cancelJobMessage = JsonUtility.FromString <JobCancelMessage>(message.Body);
workerManager.Cancel(cancelJobMessage);
}
}
finally
{
if (message != null)
{
//TODO: make sure we don't mask more important exception
await DeleteMessageAsync(message);
message = null;
}
}
}
}
}
finally
{
//TODO: make sure we don't mask more important exception
await listener.DeleteSessionAsync();
}
return(0);
}
public async Task <TaskAgentMessage> GetNextMessageAsync(CancellationToken token)
{
Trace.Entering();
ArgUtil.NotNull(_session, nameof(_session));
ArgUtil.NotNull(_settings, nameof(_settings));
bool encounteringError = false;
int continuousError = 0;
string errorMessage = string.Empty;
Stopwatch heartbeat = new Stopwatch();
heartbeat.Restart();
while (true)
{
token.ThrowIfCancellationRequested();
TaskAgentMessage message = null;
try
{
message = await _runnerServer.GetAgentMessageAsync(_settings.PoolId,
_session.SessionId,
_lastMessageId,
token);
// Decrypt the message body if the session is using encryption
message = DecryptMessage(message);
if (message != null)
{
_lastMessageId = message.MessageId;
}
if (encounteringError) //print the message once only if there was an error
{
_term.WriteLine($"{DateTime.UtcNow:u}: Runner reconnected.");
encounteringError = false;
continuousError = 0;
}
}
catch (OperationCanceledException) when(token.IsCancellationRequested)
{
Trace.Info("Get next message has been cancelled.");
throw;
}
catch (TaskAgentAccessTokenExpiredException)
{
Trace.Info("Runner OAuth token has been revoked. Unable to pull message.");
throw;
}
catch (Exception ex)
{
Trace.Error("Catch exception during get next message.");
Trace.Error(ex);
// don't retry if SkipSessionRecover = true, DT service will delete agent session to stop agent from taking more jobs.
if (ex is TaskAgentSessionExpiredException && !_settings.SkipSessionRecover && await CreateSessionAsync(token))
{
Trace.Info($"{nameof(TaskAgentSessionExpiredException)} received, recovered by recreate session.");
}
else if (!IsGetNextMessageExceptionRetriable(ex))
{
throw;
}
else
{
continuousError++;
//retry after a random backoff to avoid service throttling
//in case of there is a service error happened and all agents get kicked off of the long poll and all agent try to reconnect back at the same time.
if (continuousError <= 5)
{
// random backoff [15, 30]
_getNextMessageRetryInterval = BackoffTimerHelper.GetRandomBackoff(TimeSpan.FromSeconds(15), TimeSpan.FromSeconds(30), _getNextMessageRetryInterval);
}
else
{
// more aggressive backoff [30, 60]
_getNextMessageRetryInterval = BackoffTimerHelper.GetRandomBackoff(TimeSpan.FromSeconds(30), TimeSpan.FromSeconds(60), _getNextMessageRetryInterval);
}
if (!encounteringError)
{
//print error only on the first consecutive error
_term.WriteError($"{DateTime.UtcNow:u}: Runner connect error: {ex.Message}. Retrying until reconnected.");
encounteringError = true;
}
// re-create VssConnection before next retry
await _runnerServer.RefreshConnectionAsync(RunnerConnectionType.MessageQueue, TimeSpan.FromSeconds(60));
Trace.Info("Sleeping for {0} seconds before retrying.", _getNextMessageRetryInterval.TotalSeconds);
await HostContext.Delay(_getNextMessageRetryInterval, token);
}
}
if (message == null)
{
if (heartbeat.Elapsed > TimeSpan.FromMinutes(30))
{
Trace.Info($"No message retrieved from session '{_session.SessionId}' within last 30 minutes.");
heartbeat.Restart();
}
else
{
Trace.Verbose($"No message retrieved from session '{_session.SessionId}'.");
}
continue;
}
Trace.Info($"Message '{message.MessageId}' received from session '{_session.SessionId}'.");
return(message);
}
}
private TaskAgentMessage DecryptMessage(TaskAgentMessage message)
{
if (Session.EncryptionKey == null ||
Session.EncryptionKey.Value.Length == 0 ||
message == null ||
message.IV == null ||
message.IV.Length == 0)
{
return message;
}
using (var aes = Aes.Create())
using (var decryptor = GetMessageDecryptor(aes, message))
using (var body = new MemoryStream(Convert.FromBase64String(message.Body)))
using (var cryptoStream = new CryptoStream(body, decryptor, CryptoStreamMode.Read))
using (var bodyReader = new StreamReader(cryptoStream, Encoding.UTF8))
{
message.Body = bodyReader.ReadToEnd();
}
return message;
}
public async Task <TaskAgentMessage> GetNextMessageAsync(CancellationToken token)
{
Trace.Entering();
ArgUtil.NotNull(_session, nameof(_session));
ArgUtil.NotNull(_settings, nameof(_settings));
bool encounteringError = false;
string errorMessage = string.Empty;
Stopwatch heartbeat = new Stopwatch();
heartbeat.Restart();
while (true)
{
token.ThrowIfCancellationRequested();
TaskAgentMessage message = null;
try
{
message = await _agentServer.GetAgentMessageAsync(_settings.PoolId,
_session.SessionId,
_lastMessageId,
token);
// Decrypt the message body if the session is using encryption
message = DecryptMessage(message);
if (message != null)
{
_lastMessageId = message.MessageId;
}
if (encounteringError) //print the message once only if there was an error
{
_term.WriteLine(StringUtil.Loc("QueueConnected", DateTime.UtcNow));
encounteringError = false;
}
}
catch (OperationCanceledException) when(token.IsCancellationRequested)
{
Trace.Info("Get next message has been cancelled.");
throw;
}
catch (Exception ex)
{
Trace.Error("Catch exception during get next message.");
Trace.Error(ex);
if (ex is TaskAgentSessionExpiredException && await CreateSessionAsync(token))
{
Trace.Info($"{nameof(TaskAgentSessionExpiredException)} received, recoverd by recreate session.");
}
else if (!IsGetNextMessageExceptionRetriable(ex))
{
throw;
}
else
{
//retry after a delay
if (!encounteringError)
{
//print error only on the first consecutive error
_term.WriteError(StringUtil.Loc("QueueConError", DateTime.UtcNow, ex.Message, _getNextMessageRetryInterval.TotalSeconds));
encounteringError = true;
}
Trace.Info("Sleeping for {0} seconds before retrying.", _getNextMessageRetryInterval.TotalSeconds);
await HostContext.Delay(_getNextMessageRetryInterval, token);
}
}
if (message == null)
{
if (heartbeat.Elapsed > TimeSpan.FromMinutes(30))
{
Trace.Info($"No message retrieved from session '{_session.SessionId}' within last 30 minutes.");
heartbeat.Restart();
}
else
{
Trace.Verbose($"No message retrieved from session '{_session.SessionId}'.");
}
continue;
}
Trace.Info($"Message '{message.MessageId}' received from session '{_session.SessionId}'.");
return(message);
}
}
//create worker manager, create message listener and start listening to the queue
private async Task <int> RunAsync(CancellationToken token, AgentSettings settings, bool runAsService)
{
Trace.Info(nameof(RunAsync));
// Load the settings.
_poolId = settings.PoolId;
var listener = HostContext.GetService <IMessageListener>();
if (!await listener.CreateSessionAsync(token))
{
return(Constants.Agent.ReturnCode.TerminatedError);
}
_term.WriteLine(StringUtil.Loc("ListenForJobs", DateTime.UtcNow));
_sessionId = listener.Session.SessionId;
IJobDispatcher jobDispatcher = null;
try
{
bool disableAutoUpdate = !string.IsNullOrEmpty(Environment.GetEnvironmentVariable("agent.disableupdate"));
bool autoUpdateInProgress = false;
Task <bool> selfUpdateTask = null;
jobDispatcher = HostContext.CreateService <IJobDispatcher>();
while (!token.IsCancellationRequested)
{
TaskAgentMessage message = null;
bool skipMessageDeletion = false;
try
{
Task <TaskAgentMessage> getNextMessage = listener.GetNextMessageAsync(token);
if (autoUpdateInProgress)
{
Trace.Verbose("Auto update task running at backend, waiting for getNextMessage or selfUpdateTask to finish.");
Task completeTask = await Task.WhenAny(getNextMessage, selfUpdateTask);
if (completeTask == selfUpdateTask)
{
autoUpdateInProgress = false;
if (await selfUpdateTask)
{
Trace.Info("Auto update task finished at backend, an agent update is ready to apply exit the current agent instance.");
return(Constants.Agent.ReturnCode.AgentUpdating);
}
else
{
Trace.Info("Auto update task finished at backend, there is no available agent update needs to apply, continue message queue looping.");
}
}
}
message = await getNextMessage; //get next message
if (string.Equals(message.MessageType, AgentRefreshMessage.MessageType, StringComparison.OrdinalIgnoreCase))
{
if (disableAutoUpdate)
{
Trace.Info("Refresh message received, skip autoupdate since environment variable agent.disableupdate is set.");
}
else
{
if (autoUpdateInProgress == false)
{
autoUpdateInProgress = true;
var selfUpdater = HostContext.GetService <ISelfUpdater>();
selfUpdateTask = selfUpdater.SelfUpdate(jobDispatcher, !runAsService, token);
Trace.Info("Refresh message received, kick-off selfupdate background process.");
}
else
{
Trace.Info("Refresh message received, skip autoupdate since a previous autoupdate is already running.");
}
}
}
else if (string.Equals(message.MessageType, JobRequestMessage.MessageType, StringComparison.OrdinalIgnoreCase))
{
if (autoUpdateInProgress)
{
skipMessageDeletion = true;
}
else
{
var newJobMessage = JsonUtility.FromString <JobRequestMessage>(message.Body);
jobDispatcher.Run(newJobMessage);
}
}
else if (string.Equals(message.MessageType, JobCancelMessage.MessageType, StringComparison.OrdinalIgnoreCase))
{
var cancelJobMessage = JsonUtility.FromString <JobCancelMessage>(message.Body);
bool jobCancelled = jobDispatcher.Cancel(cancelJobMessage);
skipMessageDeletion = autoUpdateInProgress && !jobCancelled;
}
}
finally
{
if (!skipMessageDeletion && message != null)
{
try
{
await DeleteMessageAsync(message);
}
catch (Exception ex)
{
Trace.Error($"Catch exception during delete message from message queue. message id: {message.MessageId}");
Trace.Error(ex);
}
finally
{
message = null;
}
}
}
}
}
finally
{
if (jobDispatcher != null)
{
await jobDispatcher.ShutdownAsync();
}
//TODO: make sure we don't mask more important exception
await listener.DeleteSessionAsync();
}
return(Constants.Agent.ReturnCode.Success);
}
public async Task <TaskAgentMessage> GetNextMessageAsync(CancellationToken token)
{
Trace.Entering();
ArgUtil.NotNull(_session, nameof(_session));
ArgUtil.NotNull(_settings, nameof(_settings));
bool encounteringError = false;
int continuousError = 0;
string errorMessage = string.Empty;
Stopwatch heartbeat = new Stopwatch();
heartbeat.Restart();
while (true)
{
token.ThrowIfCancellationRequested();
TaskAgentMessage message = null;
try
{
message = await _runnerServer.GetAgentMessageAsync(_settings.PoolId,
_session.SessionId,
_lastMessageId,
token);
// Decrypt the message body if the session is using encryption
message = DecryptMessage(message);
if (message != null)
{
_lastMessageId = message.MessageId;
}
if (encounteringError) //print the message once only if there was an error
{
_term.WriteLine($"{DateTime.UtcNow:u}: Runner reconnected.");
encounteringError = false;
continuousError = 0;
}
if (_needToCheckAuthorizationUrlUpdate &&
_authorizationUrlMigrationBackgroundTask?.IsCompleted == true)
{
if (HostContext.GetService <IJobDispatcher>().Busy ||
HostContext.GetService <ISelfUpdater>().Busy)
{
Trace.Info("Job or runner updates in progress, update credentials next time.");
}
else
{
try
{
var newCred = await _authorizationUrlMigrationBackgroundTask;
await _runnerServer.ConnectAsync(new Uri(_settings.ServerUrl), newCred);
Trace.Info("Updated connection to use migrated credential for next GetMessage call.");
_useMigratedCredentials = true;
_authorizationUrlMigrationBackgroundTask = null;
_needToCheckAuthorizationUrlUpdate = false;
}
catch (Exception ex)
{
Trace.Error("Fail to refresh connection with new authorization url.");
Trace.Error(ex);
}
}
}
if (_authorizationUrlRollbackReattemptDelayBackgroundTask?.IsCompleted == true)
{
try
{
// we rolled back to use original creds about 2 days before, now it's a good time to try migrated creds again.
Trace.Info("Re-attempt to use migrated credential");
var migratedCreds = _credMgr.LoadCredentials();
await _runnerServer.ConnectAsync(new Uri(_settings.ServerUrl), migratedCreds);
_useMigratedCredentials = true;
_authorizationUrlRollbackReattemptDelayBackgroundTask = null;
}
catch (Exception ex)
{
Trace.Error("Fail to refresh connection with new authorization url on rollback reattempt.");
Trace.Error(ex);
}
}
}
catch (OperationCanceledException) when(token.IsCancellationRequested)
{
Trace.Info("Get next message has been cancelled.");
throw;
}
catch (TaskAgentAccessTokenExpiredException)
{
Trace.Info("Runner OAuth token has been revoked. Unable to pull message.");
throw;
}
catch (Exception ex)
{
Trace.Error("Catch exception during get next message.");
Trace.Error(ex);
// don't retry if SkipSessionRecover = true, DT service will delete agent session to stop agent from taking more jobs.
if (ex is TaskAgentSessionExpiredException && !_settings.SkipSessionRecover && await CreateSessionAsync(token))
{
Trace.Info($"{nameof(TaskAgentSessionExpiredException)} received, recovered by recreate session.");
}
else if (!IsGetNextMessageExceptionRetriable(ex))
{
if (_useMigratedCredentials)
{
// migrated credentials might cause lose permission during permission check,
// we will force to use original credential and try again
_useMigratedCredentials = false;
var reattemptBackoff = BackoffTimerHelper.GetRandomBackoff(TimeSpan.FromHours(24), TimeSpan.FromHours(36));
_authorizationUrlRollbackReattemptDelayBackgroundTask = HostContext.Delay(reattemptBackoff, token); // retry migrated creds in 24-36 hours.
var originalCreds = _credMgr.LoadCredentials(false);
await _runnerServer.ConnectAsync(new Uri(_settings.ServerUrl), originalCreds);
Trace.Error("Fallback to original credentials and try again.");
}
else
{
throw;
}
}
else
{
continuousError++;
//retry after a random backoff to avoid service throttling
//in case of there is a service error happened and all agents get kicked off of the long poll and all agent try to reconnect back at the same time.
if (continuousError <= 5)
{
// random backoff [15, 30]
_getNextMessageRetryInterval = BackoffTimerHelper.GetRandomBackoff(TimeSpan.FromSeconds(15), TimeSpan.FromSeconds(30), _getNextMessageRetryInterval);
}
else
{
// more aggressive backoff [30, 60]
_getNextMessageRetryInterval = BackoffTimerHelper.GetRandomBackoff(TimeSpan.FromSeconds(30), TimeSpan.FromSeconds(60), _getNextMessageRetryInterval);
}
if (!encounteringError)
{
//print error only on the first consecutive error
_term.WriteError($"{DateTime.UtcNow:u}: Runner connect error: {ex.Message}. Retrying until reconnected.");
encounteringError = true;
}
// re-create VssConnection before next retry
await _runnerServer.RefreshConnectionAsync(RunnerConnectionType.MessageQueue, TimeSpan.FromSeconds(60));
Trace.Info("Sleeping for {0} seconds before retrying.", _getNextMessageRetryInterval.TotalSeconds);
await HostContext.Delay(_getNextMessageRetryInterval, token);
}
}
if (message == null)
{
if (heartbeat.Elapsed > TimeSpan.FromMinutes(30))
{
Trace.Info($"No message retrieved from session '{_session.SessionId}' within last 30 minutes.");
heartbeat.Restart();
}
else
{
Trace.Verbose($"No message retrieved from session '{_session.SessionId}'.");
}
continue;
}
Trace.Info($"Message '{message.MessageId}' received from session '{_session.SessionId}'.");
return(message);
}
}
//create worker manager, create message listener and start listening to the queue
private async Task <int> RunAsync(AgentSettings settings)
{
Trace.Info(nameof(RunAsync));
_listener = HostContext.GetService <IMessageListener>();
if (!await _listener.CreateSessionAsync(HostContext.AgentShutdownToken))
{
return(Constants.Agent.ReturnCode.TerminatedError);
}
_term.WriteLine(StringUtil.Loc("ListenForJobs", DateTime.UtcNow));
IJobDispatcher jobDispatcher = null;
CancellationTokenSource messageQueueLoopTokenSource = CancellationTokenSource.CreateLinkedTokenSource(HostContext.AgentShutdownToken);
try
{
var notification = HostContext.GetService <IJobNotification>();
if (!String.IsNullOrEmpty(settings.NotificationSocketAddress))
{
notification.StartClient(settings.NotificationSocketAddress);
}
else
{
notification.StartClient(settings.NotificationPipeName, HostContext.AgentShutdownToken);
}
// this is not a reliable way to disable auto update.
// we need server side work to really enable the feature
// https://github.com/Microsoft/vsts-agent/issues/446 (Feature: Allow agent / pool to opt out of automatic updates)
bool disableAutoUpdate = !string.IsNullOrEmpty(Environment.GetEnvironmentVariable("agent.disableupdate"));
bool autoUpdateInProgress = false;
Task <bool> selfUpdateTask = null;
jobDispatcher = HostContext.CreateService <IJobDispatcher>();
while (!HostContext.AgentShutdownToken.IsCancellationRequested)
{
TaskAgentMessage message = null;
bool skipMessageDeletion = false;
try
{
Task <TaskAgentMessage> getNextMessage = _listener.GetNextMessageAsync(messageQueueLoopTokenSource.Token);
if (autoUpdateInProgress)
{
Trace.Verbose("Auto update task running at backend, waiting for getNextMessage or selfUpdateTask to finish.");
Task completeTask = await Task.WhenAny(getNextMessage, selfUpdateTask);
if (completeTask == selfUpdateTask)
{
autoUpdateInProgress = false;
if (await selfUpdateTask)
{
Trace.Info("Auto update task finished at backend, an agent update is ready to apply exit the current agent instance.");
Trace.Info("Stop message queue looping.");
messageQueueLoopTokenSource.Cancel();
try
{
await getNextMessage;
}
catch (Exception ex)
{
Trace.Info($"Ignore any exception after cancel message loop. {ex}");
}
return(Constants.Agent.ReturnCode.AgentUpdating);
}
else
{
Trace.Info("Auto update task finished at backend, there is no available agent update needs to apply, continue message queue looping.");
}
}
}
message = await getNextMessage; //get next message
if (string.Equals(message.MessageType, AgentRefreshMessage.MessageType, StringComparison.OrdinalIgnoreCase))
{
if (disableAutoUpdate)
{
Trace.Info("Refresh message received, skip autoupdate since environment variable agent.disableupdate is set.");
}
else
{
if (autoUpdateInProgress == false)
{
autoUpdateInProgress = true;
var agentUpdateMessage = JsonUtility.FromString <AgentRefreshMessage>(message.Body);
var selfUpdater = HostContext.GetService <ISelfUpdater>();
selfUpdateTask = selfUpdater.SelfUpdate(agentUpdateMessage, jobDispatcher, HostContext.StartupType != StartupType.Service, HostContext.AgentShutdownToken);
Trace.Info("Refresh message received, kick-off selfupdate background process.");
}
else
{
Trace.Info("Refresh message received, skip autoupdate since a previous autoupdate is already running.");
}
}
}
else if (string.Equals(message.MessageType, JobRequestMessageTypes.AgentJobRequest, StringComparison.OrdinalIgnoreCase))
{
if (autoUpdateInProgress)
{
skipMessageDeletion = true;
}
else
{
var newJobMessage = JsonUtility.FromString <AgentJobRequestMessage>(message.Body);
jobDispatcher.Run(newJobMessage);
}
}
else if (string.Equals(message.MessageType, JobCancelMessage.MessageType, StringComparison.OrdinalIgnoreCase))
{
var cancelJobMessage = JsonUtility.FromString <JobCancelMessage>(message.Body);
bool jobCancelled = jobDispatcher.Cancel(cancelJobMessage);
skipMessageDeletion = autoUpdateInProgress && !jobCancelled;
}
else
{
Trace.Error($"Received message {message.MessageId} with unsupported message type {message.MessageType}.");
}
}
finally
{
if (!skipMessageDeletion && message != null)
{
try
{
await _listener.DeleteMessageAsync(message);
}
catch (Exception ex)
{
Trace.Error($"Catch exception during delete message from message queue. message id: {message.MessageId}");
Trace.Error(ex);
}
finally
{
message = null;
}
}
}
}
}
finally
{
if (jobDispatcher != null)
{
await jobDispatcher.ShutdownAsync();
}
//TODO: make sure we don't mask more important exception
await _listener.DeleteSessionAsync();
messageQueueLoopTokenSource.Dispose();
}
return(Constants.Agent.ReturnCode.Success);
}
//process 2 new job messages, and one cancel message
public async void TestRunAsync()
{
using (var hc = new TestHostContext(this))
using (var tokenSource = new CancellationTokenSource())
{
//Arrange
var agent = new Agent.Listener.Agent();
agent.TokenSource = tokenSource;
hc.SetSingleton <IConfigurationManager>(_configurationManager.Object);
hc.SetSingleton <IJobNotification>(_jobNotification.Object);
hc.SetSingleton <IMessageListener>(_messageListener.Object);
hc.SetSingleton <IPromptManager>(_promptManager.Object);
hc.SetSingleton <IAgentServer>(_agentServer.Object);
hc.SetSingleton <IVstsAgentWebProxy>(_proxy.Object);
agent.Initialize(hc);
var settings = new AgentSettings
{
PoolId = 43242
};
var message = new TaskAgentMessage()
{
Body = JsonUtility.ToString(CreateJobRequestMessage("job1")),
MessageId = 4234,
MessageType = JobRequestMessageTypes.AgentJobRequest
};
var messages = new Queue <TaskAgentMessage>();
messages.Enqueue(message);
var signalWorkerComplete = new SemaphoreSlim(0, 1);
_configurationManager.Setup(x => x.LoadSettings())
.Returns(settings);
_configurationManager.Setup(x => x.IsConfigured())
.Returns(true);
_messageListener.Setup(x => x.CreateSessionAsync(It.IsAny <CancellationToken>()))
.Returns(Task.FromResult <bool>(true));
_messageListener.Setup(x => x.GetNextMessageAsync(It.IsAny <CancellationToken>()))
.Returns(async() =>
{
if (0 == messages.Count)
{
signalWorkerComplete.Release();
await Task.Delay(2000, tokenSource.Token);
}
return(messages.Dequeue());
});
_messageListener.Setup(x => x.DeleteSessionAsync())
.Returns(Task.CompletedTask);
_messageListener.Setup(x => x.DeleteMessageAsync(It.IsAny <TaskAgentMessage>()))
.Returns(Task.CompletedTask);
_jobDispatcher.Setup(x => x.Run(It.IsAny <AgentJobRequestMessage>()))
.Callback(() =>
{
});
_jobNotification.Setup(x => x.StartClient(It.IsAny <String>(), It.IsAny <CancellationToken>()))
.Callback(() =>
{
});
_jobNotification.Setup(x => x.StartClient(It.IsAny <String>()))
.Callback(() =>
{
});
hc.EnqueueInstance <IJobDispatcher>(_jobDispatcher.Object);
//Act
var command = new CommandSettings(hc, new string[] { "run" });
Task agentTask = agent.ExecuteCommand(command);
//Assert
//wait for the agent to run one job
if (!await signalWorkerComplete.WaitAsync(2000))
{
Assert.True(false, $"{nameof(_messageListener.Object.GetNextMessageAsync)} was not invoked.");
}
else
{
//Act
tokenSource.Cancel(); //stop Agent
//Assert
Task[] taskToWait2 = { agentTask, Task.Delay(2000) };
//wait for the Agent to exit
await Task.WhenAny(taskToWait2);
Assert.True(agentTask.IsCompleted, $"{nameof(agent.ExecuteCommand)} timed out.");
Assert.True(!agentTask.IsFaulted, agentTask.Exception?.ToString());
Assert.True(agentTask.IsCanceled);
_jobDispatcher.Verify(x => x.Run(It.IsAny <AgentJobRequestMessage>()), Times.Once(),
$"{nameof(_jobDispatcher.Object.Run)} was not invoked.");
_messageListener.Verify(x => x.GetNextMessageAsync(It.IsAny <CancellationToken>()), Times.AtLeastOnce());
_messageListener.Verify(x => x.CreateSessionAsync(It.IsAny <CancellationToken>()), Times.Once());
_messageListener.Verify(x => x.DeleteSessionAsync(), Times.Once());
_messageListener.Verify(x => x.DeleteMessageAsync(It.IsAny <TaskAgentMessage>()), Times.AtLeastOnce());
}
}
}
private ICryptoTransform GetMessageDecryptor(
Aes aes,
TaskAgentMessage message)
{
if (Session.EncryptionKey.Encrypted)
{
// The agent session encryption key uses the AES symmetric algorithm
var keyManager = HostContext.GetService<IRSAKeyManager>();
using (var rsa = keyManager.GetKey())
{
return aes.CreateDecryptor(rsa.Decrypt(Session.EncryptionKey.Value, RSAEncryptionPadding.OaepSHA1), message.IV);
}
}
else
{
return aes.CreateDecryptor(Session.EncryptionKey.Value, message.IV);
}
}
//process 2 new job messages, and one cancel message
public async void TestRunAsync()
{
using (var hc = new TestHostContext(this))
{
//Arrange
var runner = new Runner.Listener.Runner();
hc.SetSingleton <IConfigurationManager>(_configurationManager.Object);
hc.SetSingleton <IJobNotification>(_jobNotification.Object);
hc.SetSingleton <IMessageListener>(_messageListener.Object);
hc.SetSingleton <IPromptManager>(_promptManager.Object);
hc.SetSingleton <IRunnerServer>(_runnerServer.Object);
hc.SetSingleton <IConfigurationStore>(_configStore.Object);
runner.Initialize(hc);
var settings = new RunnerSettings
{
PoolId = 43242
};
var message = new TaskAgentMessage()
{
Body = JsonUtility.ToString(CreateJobRequestMessage("job1")),
MessageId = 4234,
MessageType = JobRequestMessageTypes.PipelineAgentJobRequest
};
var messages = new Queue <TaskAgentMessage>();
messages.Enqueue(message);
var signalWorkerComplete = new SemaphoreSlim(0, 1);
_configurationManager.Setup(x => x.LoadSettings())
.Returns(settings);
_configurationManager.Setup(x => x.IsConfigured())
.Returns(true);
_messageListener.Setup(x => x.CreateSessionAsync(It.IsAny <CancellationToken>()))
.Returns(Task.FromResult <bool>(true));
_messageListener.Setup(x => x.GetNextMessageAsync(It.IsAny <CancellationToken>()))
.Returns(async() =>
{
if (0 == messages.Count)
{
signalWorkerComplete.Release();
await Task.Delay(2000, hc.RunnerShutdownToken);
}
return(messages.Dequeue());
});
_messageListener.Setup(x => x.DeleteSessionAsync())
.Returns(Task.CompletedTask);
_messageListener.Setup(x => x.DeleteMessageAsync(It.IsAny <TaskAgentMessage>()))
.Returns(Task.CompletedTask);
_jobDispatcher.Setup(x => x.Run(It.IsAny <Pipelines.AgentJobRequestMessage>(), It.IsAny <bool>()))
.Callback(() =>
{
});
_jobNotification.Setup(x => x.StartClient(It.IsAny <String>()))
.Callback(() =>
{
});
hc.EnqueueInstance <IJobDispatcher>(_jobDispatcher.Object);
_configStore.Setup(x => x.IsServiceConfigured()).Returns(false);
//Act
var command = new CommandSettings(hc, new string[] { "run" });
Task runnerTask = runner.ExecuteCommand(command);
//Assert
//wait for the runner to run one job
if (!await signalWorkerComplete.WaitAsync(2000))
{
Assert.True(false, $"{nameof(_messageListener.Object.GetNextMessageAsync)} was not invoked.");
}
else
{
//Act
hc.ShutdownRunner(ShutdownReason.UserCancelled); //stop Runner
//Assert
Task[] taskToWait2 = { runnerTask, Task.Delay(2000) };
//wait for the runner to exit
await Task.WhenAny(taskToWait2);
Assert.True(runnerTask.IsCompleted, $"{nameof(runner.ExecuteCommand)} timed out.");
Assert.True(!runnerTask.IsFaulted, runnerTask.Exception?.ToString());
Assert.True(runnerTask.IsCanceled);
_jobDispatcher.Verify(x => x.Run(It.IsAny <Pipelines.AgentJobRequestMessage>(), It.IsAny <bool>()), Times.Once(),
$"{nameof(_jobDispatcher.Object.Run)} was not invoked.");
_messageListener.Verify(x => x.GetNextMessageAsync(It.IsAny <CancellationToken>()), Times.AtLeastOnce());
_messageListener.Verify(x => x.CreateSessionAsync(It.IsAny <CancellationToken>()), Times.Once());
_messageListener.Verify(x => x.DeleteSessionAsync(), Times.Once());
_messageListener.Verify(x => x.DeleteMessageAsync(It.IsAny <TaskAgentMessage>()), Times.AtLeastOnce());
// verify that we didn't try to delete local settings file (since we're not ephemeral)
_configurationManager.Verify(x => x.DeleteLocalRunnerConfig(), Times.Never());
}
}
}
public async void TestRunOnceOnlyTakeOneJobMessage()
{
using (var hc = new TestHostContext(this))
{
//Arrange
var runner = new Runner.Listener.Runner();
hc.SetSingleton <IConfigurationManager>(_configurationManager.Object);
hc.SetSingleton <IJobNotification>(_jobNotification.Object);
hc.SetSingleton <IMessageListener>(_messageListener.Object);
hc.SetSingleton <IPromptManager>(_promptManager.Object);
hc.SetSingleton <IRunnerServer>(_runnerServer.Object);
hc.SetSingleton <IConfigurationStore>(_configStore.Object);
runner.Initialize(hc);
var settings = new RunnerSettings
{
PoolId = 43242,
Ephemeral = true
};
var message1 = new TaskAgentMessage()
{
Body = JsonUtility.ToString(CreateJobRequestMessage("job1")),
MessageId = 4234,
MessageType = JobRequestMessageTypes.PipelineAgentJobRequest
};
var message2 = new TaskAgentMessage()
{
Body = JsonUtility.ToString(CreateJobRequestMessage("job1")),
MessageId = 4235,
MessageType = JobRequestMessageTypes.PipelineAgentJobRequest
};
var messages = new Queue <TaskAgentMessage>();
messages.Enqueue(message1);
messages.Enqueue(message2);
_configurationManager.Setup(x => x.LoadSettings())
.Returns(settings);
_configurationManager.Setup(x => x.IsConfigured())
.Returns(true);
_messageListener.Setup(x => x.CreateSessionAsync(It.IsAny <CancellationToken>()))
.Returns(Task.FromResult <bool>(true));
_messageListener.Setup(x => x.GetNextMessageAsync(It.IsAny <CancellationToken>()))
.Returns(async() =>
{
if (0 == messages.Count)
{
await Task.Delay(2000);
}
return(messages.Dequeue());
});
_messageListener.Setup(x => x.DeleteSessionAsync())
.Returns(Task.CompletedTask);
_messageListener.Setup(x => x.DeleteMessageAsync(It.IsAny <TaskAgentMessage>()))
.Returns(Task.CompletedTask);
var runOnceJobCompleted = new TaskCompletionSource <bool>();
_jobDispatcher.Setup(x => x.RunOnceJobCompleted)
.Returns(runOnceJobCompleted);
_jobDispatcher.Setup(x => x.Run(It.IsAny <Pipelines.AgentJobRequestMessage>(), It.IsAny <bool>()))
.Callback(() =>
{
runOnceJobCompleted.TrySetResult(true);
});
_jobNotification.Setup(x => x.StartClient(It.IsAny <String>()))
.Callback(() =>
{
});
hc.EnqueueInstance <IJobDispatcher>(_jobDispatcher.Object);
_configStore.Setup(x => x.IsServiceConfigured()).Returns(false);
//Act
var command = new CommandSettings(hc, new string[] { "run" });
Task <int> runnerTask = runner.ExecuteCommand(command);
//Assert
//wait for the runner to run one job and exit
await Task.WhenAny(runnerTask, Task.Delay(30000));
Assert.True(runnerTask.IsCompleted, $"{nameof(runner.ExecuteCommand)} timed out.");
Assert.True(!runnerTask.IsFaulted, runnerTask.Exception?.ToString());
Assert.True(runnerTask.Result == Constants.Runner.ReturnCode.Success);
_jobDispatcher.Verify(x => x.Run(It.IsAny <Pipelines.AgentJobRequestMessage>(), true), Times.Once(),
$"{nameof(_jobDispatcher.Object.Run)} was not invoked.");
_messageListener.Verify(x => x.GetNextMessageAsync(It.IsAny <CancellationToken>()), Times.AtLeastOnce());
_messageListener.Verify(x => x.CreateSessionAsync(It.IsAny <CancellationToken>()), Times.Once());
_messageListener.Verify(x => x.DeleteSessionAsync(), Times.Once());
_messageListener.Verify(x => x.DeleteMessageAsync(It.IsAny <TaskAgentMessage>()), Times.Once());
}
}
public async Task <TaskAgentMessage> GetNextMessageAsync(CancellationToken token)
{
Trace.Entering();
ArgUtil.NotNull(Session, nameof(Session));
ArgUtil.NotNull(_settings, nameof(_settings));
var agentServer = HostContext.GetService <IAgentServer>();
int consecutiveErrors = 0; //number of consecutive exceptions thrown by GetAgentMessageAsync
string errorMessage = string.Empty;
while (true)
{
token.ThrowIfCancellationRequested();
TaskAgentMessage message = null;
try
{
consecutiveErrors++;
message = await agentServer.GetAgentMessageAsync(_settings.PoolId,
Session.SessionId,
_lastMessageId,
token);
if (message != null)
{
_lastMessageId = message.MessageId;
}
if (consecutiveErrors > 1) //print the message once only if there was an error
{
_term.WriteLine(StringUtil.Loc("QueueConnected", DateTime.UtcNow));
}
consecutiveErrors = 0;
}
catch (TimeoutException ex)
{
Trace.Verbose($"{nameof(TimeoutException)} received.");
//retry after a delay
errorMessage = ex.Message;
}
catch (TaskAgentSessionExpiredException)
{
Trace.Verbose($"{nameof(TaskAgentSessionExpiredException)} received.");
if (!await CreateSessionAsync(token))
{
throw;
}
consecutiveErrors = 0;
}
catch (OperationCanceledException ex)
{
Trace.Verbose($"{nameof(OperationCanceledException)} received.");
//we get here when the agent is stopped with CTRL-C or service is stopped or HttpClient has timed out
if (token.IsCancellationRequested) //Distinguish timeout from user cancellation
{
throw;
}
//retry after a delay
errorMessage = ex.Message;
}
catch (Exception ex)
{
Trace.Error(ex);
if (IsFatalException(ex))
{
throw;
}
//retry after a delay
errorMessage = ex.Message;
}
//print an error and add a delay
if (consecutiveErrors > 0)
{
TimeSpan interval = TimeSpan.FromSeconds(15);
if (consecutiveErrors == 1)
{
//print error only on the first consecutive error
_term.WriteError(StringUtil.Loc("QueueConError", DateTime.UtcNow, errorMessage, interval.TotalSeconds));
}
Trace.Info("Sleeping for {0} seconds before retrying.", interval.TotalSeconds);
await HostContext.Delay(interval, token);
}
if (message == null)
{
Trace.Verbose($"No message retrieved from session '{Session.SessionId}'.");
continue;
}
Trace.Verbose($"Message '{message.MessageId}' received from session '{Session.SessionId}'.");
return(message);
}
}
public async void GetNextMessage()
{
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());
_agentServer
.Setup(x => x.CreateAgentSessionAsync(
_settings.PoolId,
It.Is<TaskAgentSession>(y => y != null),
tokenSource.Token))
.Returns(Task.FromResult(expectedSession));
_capabilitiesManager.Setup(x => x.GetCapabilitiesAsync(_settings, It.IsAny<CancellationToken>())).Returns(Task.FromResult(new Dictionary<string, string>()));
_credMgr.Setup(x => x.LoadCredentials()).Returns(new Common.VssCredentials());
// Act.
MessageListener listener = new MessageListener();
listener.Initialize(tc);
bool result = await listener.CreateSessionAsync(tokenSource.Token);
Assert.True(result);
Assert.Equal(expectedSession, listener.Session);
var arMessages = new TaskAgentMessage[]
{
new TaskAgentMessage
{
Body = "somebody1",
MessageId = 4234,
MessageType = JobRequestMessage.MessageType
},
new TaskAgentMessage
{
Body = "somebody2",
MessageId = 4235,
MessageType = JobCancelMessage.MessageType
},
null, //should be skipped by GetNextMessageAsync implementation
null,
new TaskAgentMessage
{
Body = "somebody3",
MessageId = 4236,
MessageType = JobRequestMessage.MessageType
}
};
var messages = new Queue<TaskAgentMessage>(arMessages);
_agentServer
.Setup(x => x.GetAgentMessageAsync(
_settings.PoolId, expectedSession.SessionId, It.IsAny<long?>(), tokenSource.Token))
.Returns(async (Int32 poolId, Guid sessionId, Int64? lastMessageId, CancellationToken cancellationToken) =>
{
await Task.Yield();
return messages.Dequeue();
});
TaskAgentMessage message1 = await listener.GetNextMessageAsync(tokenSource.Token);
TaskAgentMessage message2 = await listener.GetNextMessageAsync(tokenSource.Token);
TaskAgentMessage message3 = await listener.GetNextMessageAsync(tokenSource.Token);
Assert.Equal(arMessages[0], message1);
Assert.Equal(arMessages[1], message2);
Assert.Equal(arMessages[4], message3);
//Assert
_agentServer
.Verify(x => x.GetAgentMessageAsync(
_settings.PoolId, expectedSession.SessionId, It.IsAny<long?>(), tokenSource.Token), Times.Exactly(arMessages.Length));
}
}
private async Task DeleteMessageAsync(TaskAgentMessage message)
{
if (message != null && _listener.Session.SessionId != Guid.Empty)
{
var agentServer = HostContext.GetService<IAgentServer>();
using (var cs = new CancellationTokenSource(TimeSpan.FromSeconds(30)))
{
await agentServer.DeleteAgentMessageAsync(_poolId, message.MessageId, _listener.Session.SessionId, cs.Token);
}
}
}
//process 2 new job messages, and one cancel message
public async void TestRunAsync()
{
using (var hc = new TestHostContext(this))
using (var tokenSource = new CancellationTokenSource())
{
//Arrange
var agent = new Agent.Listener.Agent();
agent.TokenSource = tokenSource;
hc.SetSingleton <IConfigurationManager>(_configurationManager.Object);
hc.SetSingleton <IMessageListener>(_messageListener.Object);
hc.SetSingleton <IPromptManager>(_promptManager.Object);
hc.SetSingleton <IAgentServer>(_agentServer.Object);
agent.Initialize(hc);
var settings = new AgentSettings
{
PoolId = 43242
};
var taskAgentSession = new TaskAgentSession();
//we use reflection to achieve this, because "set" is internal
PropertyInfo sessionIdProperty = taskAgentSession.GetType().GetProperty("SessionId", BindingFlags.Instance | BindingFlags.NonPublic | BindingFlags.Public);
Assert.NotNull(sessionIdProperty);
sessionIdProperty.SetValue(taskAgentSession, Guid.NewGuid());
var message = new TaskAgentMessage()
{
Body = JsonUtility.ToString(CreateJobRequestMessage("job1")),
MessageId = 4234,
MessageType = JobRequestMessage.MessageType
};
var messages = new Queue <TaskAgentMessage>();
messages.Enqueue(message);
var signalWorkerComplete = new SemaphoreSlim(0, 1);
_configurationManager.Setup(x => x.LoadSettings())
.Returns(settings);
_configurationManager.Setup(x => x.IsConfigured())
.Returns(true);
_configurationManager.Setup(x => x.EnsureConfiguredAsync(It.IsAny <CommandSettings>()))
.Returns(Task.CompletedTask);
_messageListener.Setup(x => x.CreateSessionAsync(It.IsAny <CancellationToken>()))
.Returns(Task.FromResult <bool>(true));
_messageListener.Setup(x => x.Session)
.Returns(taskAgentSession);
_messageListener.Setup(x => x.GetNextMessageAsync(It.IsAny <CancellationToken>()))
.Returns(async() =>
{
if (0 == messages.Count)
{
signalWorkerComplete.Release();
await Task.Delay(2000, tokenSource.Token);
}
return(messages.Dequeue());
});
_messageListener.Setup(x => x.DeleteSessionAsync())
.Returns(Task.CompletedTask);
_jobDispatcher.Setup(x => x.Run(It.IsAny <JobRequestMessage>()))
.Callback(() =>
{
});
_agentServer.Setup(x => x.DeleteAgentMessageAsync(settings.PoolId, message.MessageId, taskAgentSession.SessionId, It.IsAny <CancellationToken>()))
.Returns((Int32 poolId, Int64 messageId, Guid sessionId, CancellationToken cancellationToken) =>
{
return(Task.CompletedTask);
});
hc.EnqueueInstance <IJobDispatcher>(_jobDispatcher.Object);
//Act
var command = new CommandSettings(hc, new string[0]);
Task agentTask = agent.ExecuteCommand(command);
//Assert
//wait for the agent to run one job
if (!await signalWorkerComplete.WaitAsync(2000))
{
Assert.True(false, $"{nameof(_messageListener.Object.GetNextMessageAsync)} was not invoked.");
}
else
{
//Act
tokenSource.Cancel(); //stop Agent
//Assert
Task[] taskToWait2 = { agentTask, Task.Delay(2000) };
//wait for the Agent to exit
await Task.WhenAny(taskToWait2);
Assert.True(agentTask.IsCompleted, $"{nameof(agent.ExecuteCommand)} timed out.");
Assert.True(!agentTask.IsFaulted, agentTask.Exception?.ToString());
Assert.True(agentTask.IsCanceled);
_jobDispatcher.Verify(x => x.Run(It.IsAny <JobRequestMessage>()), Times.Once(),
$"{nameof(_jobDispatcher.Object.Run)} was not invoked.");
_messageListener.Verify(x => x.GetNextMessageAsync(It.IsAny <CancellationToken>()), Times.AtLeastOnce());
_messageListener.Verify(x => x.CreateSessionAsync(It.IsAny <CancellationToken>()), Times.Once());
_messageListener.Verify(x => x.DeleteSessionAsync(), Times.Once());
_agentServer.Verify(x => x.DeleteAgentMessageAsync(settings.PoolId, message.MessageId, taskAgentSession.SessionId, It.IsAny <CancellationToken>()), Times.AtLeastOnce());
}
}
}
public async void GetNextMessage()
{
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());
_agentServer
.Setup(x => x.CreateAgentSessionAsync(
_settings.PoolId,
It.Is <TaskAgentSession>(y => y != null),
tokenSource.Token))
.Returns(Task.FromResult(expectedSession));
_capabilitiesManager.Setup(x => x.GetCapabilitiesAsync(_settings, It.IsAny <CancellationToken>())).Returns(Task.FromResult(new Dictionary <string, string>()));
_credMgr.Setup(x => x.LoadCredentials()).Returns(new Common.VssCredentials());
// Act.
MessageListener listener = new MessageListener();
listener.Initialize(tc);
bool result = await listener.CreateSessionAsync(tokenSource.Token);
Assert.True(result);
Assert.Equal(expectedSession, listener.Session);
var arMessages = new TaskAgentMessage[]
{
new TaskAgentMessage
{
Body = "somebody1",
MessageId = 4234,
MessageType = JobRequestMessageTypes.AgentJobRequest
},
new TaskAgentMessage
{
Body = "somebody2",
MessageId = 4235,
MessageType = JobCancelMessage.MessageType
},
null, //should be skipped by GetNextMessageAsync implementation
null,
new TaskAgentMessage
{
Body = "somebody3",
MessageId = 4236,
MessageType = JobRequestMessageTypes.AgentJobRequest
}
};
var messages = new Queue <TaskAgentMessage>(arMessages);
_agentServer
.Setup(x => x.GetAgentMessageAsync(
_settings.PoolId, expectedSession.SessionId, It.IsAny <long?>(), tokenSource.Token))
.Returns(async(Int32 poolId, Guid sessionId, Int64? lastMessageId, CancellationToken cancellationToken) =>
{
await Task.Yield();
return(messages.Dequeue());
});
TaskAgentMessage message1 = await listener.GetNextMessageAsync(tokenSource.Token);
TaskAgentMessage message2 = await listener.GetNextMessageAsync(tokenSource.Token);
TaskAgentMessage message3 = await listener.GetNextMessageAsync(tokenSource.Token);
Assert.Equal(arMessages[0], message1);
Assert.Equal(arMessages[1], message2);
Assert.Equal(arMessages[4], message3);
//Assert
_agentServer
.Verify(x => x.GetAgentMessageAsync(
_settings.PoolId, expectedSession.SessionId, It.IsAny <long?>(), tokenSource.Token), Times.Exactly(arMessages.Length));
}
}
public async void TestRunOnceHandleUpdateMessage()
{
using (var hc = new TestHostContext(this))
{
//Arrange
var agent = new Agent.Listener.Agent();
hc.SetSingleton <IConfigurationManager>(_configurationManager.Object);
hc.SetSingleton <IJobNotification>(_jobNotification.Object);
hc.SetSingleton <IMessageListener>(_messageListener.Object);
hc.SetSingleton <IPromptManager>(_promptManager.Object);
hc.SetSingleton <IAgentServer>(_agentServer.Object);
hc.SetSingleton <IVstsAgentWebProxy>(_proxy.Object);
hc.SetSingleton <IAgentCertificateManager>(_cert.Object);
hc.SetSingleton <IConfigurationStore>(_configStore.Object);
hc.SetSingleton <ISelfUpdater>(_updater.Object);
agent.Initialize(hc);
var settings = new AgentSettings
{
PoolId = 43242,
AgentId = 5678
};
var message1 = new TaskAgentMessage()
{
Body = JsonUtility.ToString(new AgentRefreshMessage(settings.AgentId, "2.123.0")),
MessageId = 4234,
MessageType = AgentRefreshMessage.MessageType
};
var messages = new Queue <TaskAgentMessage>();
messages.Enqueue(message1);
_updater.Setup(x => x.SelfUpdate(It.IsAny <AgentRefreshMessage>(), It.IsAny <IJobDispatcher>(), It.IsAny <bool>(), It.IsAny <CancellationToken>()))
.Returns(Task.FromResult(true));
_configurationManager.Setup(x => x.LoadSettings())
.Returns(settings);
_configurationManager.Setup(x => x.IsConfigured())
.Returns(true);
_messageListener.Setup(x => x.CreateSessionAsync(It.IsAny <CancellationToken>()))
.Returns(Task.FromResult <bool>(true));
_messageListener.Setup(x => x.GetNextMessageAsync(It.IsAny <CancellationToken>()))
.Returns(async() =>
{
if (0 == messages.Count)
{
await Task.Delay(2000);
}
return(messages.Dequeue());
});
_messageListener.Setup(x => x.DeleteSessionAsync())
.Returns(Task.CompletedTask);
_messageListener.Setup(x => x.DeleteMessageAsync(It.IsAny <TaskAgentMessage>()))
.Returns(Task.CompletedTask);
_jobNotification.Setup(x => x.StartClient(It.IsAny <String>(), It.IsAny <String>(), It.IsAny <CancellationToken>()))
.Callback(() =>
{
});
_jobNotification.Setup(x => x.StartClient(It.IsAny <String>(), It.IsAny <String>()))
.Callback(() =>
{
});
hc.EnqueueInstance <IJobDispatcher>(_jobDispatcher.Object);
_configStore.Setup(x => x.IsServiceConfigured()).Returns(false);
//Act
var command = new CommandSettings(hc, new string[] { "run", "--once" });
Task <int> agentTask = agent.ExecuteCommand(command);
//Assert
//wait for the agent to exit with right return code
await Task.WhenAny(agentTask, Task.Delay(30000));
Assert.True(agentTask.IsCompleted, $"{nameof(agent.ExecuteCommand)} timed out.");
Assert.True(!agentTask.IsFaulted, agentTask.Exception?.ToString());
Assert.True(agentTask.Result == Constants.Agent.ReturnCode.RunOnceAgentUpdating);
_updater.Verify(x => x.SelfUpdate(It.IsAny <AgentRefreshMessage>(), It.IsAny <IJobDispatcher>(), false, It.IsAny <CancellationToken>()), Times.Once);
_jobDispatcher.Verify(x => x.Run(It.IsAny <Pipelines.AgentJobRequestMessage>(), true), Times.Never());
_messageListener.Verify(x => x.GetNextMessageAsync(It.IsAny <CancellationToken>()), Times.AtLeastOnce());
_messageListener.Verify(x => x.CreateSessionAsync(It.IsAny <CancellationToken>()), Times.Once());
_messageListener.Verify(x => x.DeleteSessionAsync(), Times.Once());
_messageListener.Verify(x => x.DeleteMessageAsync(It.IsAny <TaskAgentMessage>()), Times.Once());
}
}
//create worker manager, create message listener and start listening to the queue
private async Task <int> RunAsync(RunnerSettings settings, bool runOnce = false)
{
try
{
Trace.Info(nameof(RunAsync));
_listener = HostContext.GetService <IMessageListener>();
if (!await _listener.CreateSessionAsync(HostContext.RunnerShutdownToken))
{
return(Constants.Runner.ReturnCode.TerminatedError);
}
HostContext.WritePerfCounter("SessionCreated");
_term.WriteLine($"{DateTime.UtcNow:u}: Listening for Jobs");
IJobDispatcher jobDispatcher = null;
CancellationTokenSource messageQueueLoopTokenSource = CancellationTokenSource.CreateLinkedTokenSource(HostContext.RunnerShutdownToken);
try
{
var notification = HostContext.GetService <IJobNotification>();
notification.StartClient(settings.MonitorSocketAddress);
bool autoUpdateInProgress = false;
Task <bool> selfUpdateTask = null;
bool runOnceJobReceived = false;
jobDispatcher = HostContext.CreateService <IJobDispatcher>();
while (!HostContext.RunnerShutdownToken.IsCancellationRequested)
{
TaskAgentMessage message = null;
bool skipMessageDeletion = false;
try
{
Task <TaskAgentMessage> getNextMessage = _listener.GetNextMessageAsync(messageQueueLoopTokenSource.Token);
if (autoUpdateInProgress)
{
Trace.Verbose("Auto update task running at backend, waiting for getNextMessage or selfUpdateTask to finish.");
Task completeTask = await Task.WhenAny(getNextMessage, selfUpdateTask);
if (completeTask == selfUpdateTask)
{
autoUpdateInProgress = false;
if (await selfUpdateTask)
{
Trace.Info("Auto update task finished at backend, an runner update is ready to apply exit the current runner instance.");
Trace.Info("Stop message queue looping.");
messageQueueLoopTokenSource.Cancel();
try
{
await getNextMessage;
}
catch (Exception ex)
{
Trace.Info($"Ignore any exception after cancel message loop. {ex}");
}
if (runOnce)
{
return(Constants.Runner.ReturnCode.RunOnceRunnerUpdating);
}
else
{
return(Constants.Runner.ReturnCode.RunnerUpdating);
}
}
else
{
Trace.Info("Auto update task finished at backend, there is no available runner update needs to apply, continue message queue looping.");
}
}
}
if (runOnceJobReceived)
{
Trace.Verbose("One time used runner has start running its job, waiting for getNextMessage or the job to finish.");
Task completeTask = await Task.WhenAny(getNextMessage, jobDispatcher.RunOnceJobCompleted.Task);
if (completeTask == jobDispatcher.RunOnceJobCompleted.Task)
{
Trace.Info("Job has finished at backend, the runner will exit since it is running under onetime use mode.");
Trace.Info("Stop message queue looping.");
messageQueueLoopTokenSource.Cancel();
try
{
await getNextMessage;
}
catch (Exception ex)
{
Trace.Info($"Ignore any exception after cancel message loop. {ex}");
}
return(Constants.Runner.ReturnCode.Success);
}
}
message = await getNextMessage; //get next message
HostContext.WritePerfCounter($"MessageReceived_{message.MessageType}");
if (string.Equals(message.MessageType, AgentRefreshMessage.MessageType, StringComparison.OrdinalIgnoreCase))
{
if (autoUpdateInProgress == false)
{
autoUpdateInProgress = true;
var runnerUpdateMessage = JsonUtility.FromString <AgentRefreshMessage>(message.Body);
var selfUpdater = HostContext.GetService <ISelfUpdater>();
selfUpdateTask = selfUpdater.SelfUpdate(runnerUpdateMessage, jobDispatcher, !runOnce && HostContext.StartupType != StartupType.Service, HostContext.RunnerShutdownToken);
Trace.Info("Refresh message received, kick-off selfupdate background process.");
}
else
{
Trace.Info("Refresh message received, skip autoupdate since a previous autoupdate is already running.");
}
}
else if (string.Equals(message.MessageType, JobRequestMessageTypes.PipelineAgentJobRequest, StringComparison.OrdinalIgnoreCase))
{
if (autoUpdateInProgress || runOnceJobReceived)
{
skipMessageDeletion = true;
Trace.Info($"Skip message deletion for job request message '{message.MessageId}'.");
}
else
{
var jobMessage = StringUtil.ConvertFromJson <Pipelines.AgentJobRequestMessage>(message.Body);
jobDispatcher.Run(jobMessage, runOnce);
if (runOnce)
{
Trace.Info("One time used runner received job message.");
runOnceJobReceived = true;
}
}
}
else if (string.Equals(message.MessageType, JobCancelMessage.MessageType, StringComparison.OrdinalIgnoreCase))
{
var cancelJobMessage = JsonUtility.FromString <JobCancelMessage>(message.Body);
bool jobCancelled = jobDispatcher.Cancel(cancelJobMessage);
skipMessageDeletion = (autoUpdateInProgress || runOnceJobReceived) && !jobCancelled;
if (skipMessageDeletion)
{
Trace.Info($"Skip message deletion for cancellation message '{message.MessageId}'.");
}
}
else
{
Trace.Error($"Received message {message.MessageId} with unsupported message type {message.MessageType}.");
}
}
finally
{
if (!skipMessageDeletion && message != null)
{
try
{
await _listener.DeleteMessageAsync(message);
}
catch (Exception ex)
{
Trace.Error($"Catch exception during delete message from message queue. message id: {message.MessageId}");
Trace.Error(ex);
}
finally
{
message = null;
}
}
}
}
}
finally
{
if (jobDispatcher != null)
{
await jobDispatcher.ShutdownAsync();
}
//TODO: make sure we don't mask more important exception
await _listener.DeleteSessionAsync();
messageQueueLoopTokenSource.Dispose();
}
}
catch (TaskAgentAccessTokenExpiredException)
{
Trace.Info("Runner OAuth token has been revoked. Shutting down.");
}
return(Constants.Runner.ReturnCode.Success);
}
//process 2 new job messages, and one cancel message
public async void TestRunAsync()
{
using (var hc = new TestHostContext(this))
using (var tokenSource = new CancellationTokenSource())
{
//Arrange
var agent = new Agent.Listener.Agent();
agent.TokenSource = tokenSource;
hc.SetSingleton<IConfigurationManager>(_configurationManager.Object);
hc.SetSingleton<IJobNotification>(_jobNotification.Object);
hc.SetSingleton<IMessageListener>(_messageListener.Object);
hc.SetSingleton<IPromptManager>(_promptManager.Object);
hc.SetSingleton<IAgentServer>(_agentServer.Object);
agent.Initialize(hc);
var settings = new AgentSettings
{
PoolId = 43242
};
var taskAgentSession = new TaskAgentSession();
//we use reflection to achieve this, because "set" is internal
PropertyInfo sessionIdProperty = taskAgentSession.GetType().GetProperty("SessionId", BindingFlags.Instance | BindingFlags.NonPublic | BindingFlags.Public);
Assert.NotNull(sessionIdProperty);
sessionIdProperty.SetValue(taskAgentSession, Guid.NewGuid());
var message = new TaskAgentMessage()
{
Body = JsonUtility.ToString(CreateJobRequestMessage("job1")),
MessageId = 4234,
MessageType = JobRequestMessage.MessageType
};
var messages = new Queue<TaskAgentMessage>();
messages.Enqueue(message);
var signalWorkerComplete = new SemaphoreSlim(0, 1);
_configurationManager.Setup(x => x.LoadSettings())
.Returns(settings);
_configurationManager.Setup(x => x.IsConfigured())
.Returns(true);
_configurationManager.Setup(x => x.EnsureConfiguredAsync(It.IsAny<CommandSettings>()))
.Returns(Task.CompletedTask);
_messageListener.Setup(x => x.CreateSessionAsync(It.IsAny<CancellationToken>()))
.Returns(Task.FromResult<bool>(true));
_messageListener.Setup(x => x.Session)
.Returns(taskAgentSession);
_messageListener.Setup(x => x.GetNextMessageAsync(It.IsAny<CancellationToken>()))
.Returns(async () =>
{
if (0 == messages.Count)
{
signalWorkerComplete.Release();
await Task.Delay(2000, tokenSource.Token);
}
return messages.Dequeue();
});
_messageListener.Setup(x => x.DeleteSessionAsync())
.Returns(Task.CompletedTask);
_jobDispatcher.Setup(x => x.Run(It.IsAny<JobRequestMessage>()))
.Callback(()=>
{
});
_agentServer.Setup(x => x.DeleteAgentMessageAsync(settings.PoolId, message.MessageId, taskAgentSession.SessionId, It.IsAny<CancellationToken>()))
.Returns((Int32 poolId, Int64 messageId, Guid sessionId, CancellationToken cancellationToken) =>
{
return Task.CompletedTask;
});
_jobNotification.Setup(x => x.StartClient(It.IsAny<String>(), It.IsAny<CancellationToken>()))
.Callback(() =>
{
});
hc.EnqueueInstance<IJobDispatcher>(_jobDispatcher.Object);
//Act
var command = new CommandSettings(hc, new string[0]);
Task agentTask = agent.ExecuteCommand(command);
//Assert
//wait for the agent to run one job
if (!await signalWorkerComplete.WaitAsync(2000))
{
Assert.True(false, $"{nameof(_messageListener.Object.GetNextMessageAsync)} was not invoked.");
}
else
{
//Act
tokenSource.Cancel(); //stop Agent
//Assert
Task[] taskToWait2 = { agentTask, Task.Delay(2000) };
//wait for the Agent to exit
await Task.WhenAny(taskToWait2);
Assert.True(agentTask.IsCompleted, $"{nameof(agent.ExecuteCommand)} timed out.");
Assert.True(!agentTask.IsFaulted, agentTask.Exception?.ToString());
Assert.True(agentTask.IsCanceled);
_jobDispatcher.Verify(x => x.Run(It.IsAny<JobRequestMessage>()), Times.Once(),
$"{nameof(_jobDispatcher.Object.Run)} was not invoked.");
_messageListener.Verify(x => x.GetNextMessageAsync(It.IsAny<CancellationToken>()), Times.AtLeastOnce());
_messageListener.Verify(x => x.CreateSessionAsync(It.IsAny<CancellationToken>()), Times.Once());
_messageListener.Verify(x => x.DeleteSessionAsync(), Times.Once());
_agentServer.Verify(x => x.DeleteAgentMessageAsync(settings.PoolId, message.MessageId, taskAgentSession.SessionId, It.IsAny<CancellationToken>()), Times.AtLeastOnce());
}
}
}
