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()); } } }