Repository: Azure/WALinuxAgent
Pull request #3033 (2024-01-31 01:02:31)

don't allow agent update if attempts reached max limit
<!-- DO NOT DELETE THIS TEMPLATE -->

## Description

Current problem: 
agent upgrade from v2 to v3
we clean up all agents except v2
download v3 and exit v2
deamon unable to start v3 and eventually start v2
v2 repeats v3 and it's in loop

Solution:
Keep track number of updates attempts. If attempts >3 and still agent is blacklisted, we don't proceed with update.

Issue # <!-- if any -->
<!--
Please add an informative description that covers that changes made by the pull request. 
This checklist is used to make sure that common issues in a pull request are addressed.
This will expedite the process of getting your pull request merged and avoid extra work on your part to fix issues discovered during the review process.
-->

---

### PR information
- [ ] The title of the PR is clear and informative.
- [ ] There are a small number of commits, each of which has an informative message. This means that previously merged commits do not appear in the history of the PR. For information on cleaning up the commits in your pull request, [see this page](https://github.com/Azure/azure-powershell/blob/master/documentation/development-docs/cleaning-up-commits.md).
- [ ] If applicable, the PR references the bug/issue that it fixes in the description.
- [ ] New Unit tests were added for the changes made

### Quality of Code and Contribution Guidelines
- [ ] I have read the [contribution guidelines](https://github.com/Azure/WALinuxAgent/blob/master/.github/CONTRIBUTING.md).

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/3033?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) Report
Attention: `8 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`bf3738b`)](https://app.codecov.io/gh/Azure/WALinuxAgent/commit/bf3738b0c18a9c86af1389f222758d78471df92e?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) 71.90% compared to head [(`ca6d050`)](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/3033?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) 71.93%.
> Report is 1 commits behind head on develop.

| [Files](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/3033?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) | Patch % | Lines |
|---|---|---|
| [azurelinuxagent/ga/guestagent.py](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/3033?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2dhL2d1ZXN0YWdlbnQucHk=) | 81.81% | [8 Missing :warning: ](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/3033?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) |

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##           develop    #3033      +/-   ##
===========================================
+ Coverage    71.90%   71.93%   +0.03%     
===========================================
  Files          109      109              
  Lines        16253    16292      +39     
  Branches      2328     2330       +2     
===========================================
+ Hits         11686    11720      +34     
- Misses        4024     4029       +5     
  Partials       543      543              
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/3033?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure).


Patch:
diff --git a/azurelinuxagent/ga/agent_update_handler.py b/azurelinuxagent/ga/agent_update_handler.py
index 02549fa5..8caec108 100644
--- a/azurelinuxagent/ga/agent_update_handler.py
+++ b/azurelinuxagent/ga/agent_update_handler.py
@@ -182,12 +182,23 @@ class AgentUpdateHandler(object):
             if not self._updater.is_retrieved_version_allowed_to_update(agent_family):
                 return
             self._updater.log_new_agent_update_message()
-            self._updater.purge_extra_agents_from_disk()
             agent = self._updater.download_and_get_new_agent(self._protocol, agent_family, goal_state)
-            if agent.is_blacklisted or not agent.is_downloaded:
-                msg = "Downloaded agent version is in bad state : {0} , skipping agent update".format(
-                    str(agent.version))
+
+            # Below condition is to break the update loop if new agent is in bad state in previous attempts
+            # If the bad agent update already attempted 3 times, we don't want to continue with update anymore.
+            # Otherewise we allow the update by increment the update attempt count and clear the bad state to make good agent
+            # [Note: As a result, it is breaking contract between RSM and agent, we may NOT honor the RSM retries for that version]
+            if agent.get_update_attempt_count() >= 3:
+                msg = "Attempted enough update retries for version: {0} but still agent not recovered from bad state. So, we stop updating to this version".format(str(agent.version))
                 raise AgentUpdateError(msg)
+            else:
+                agent.clear_error()
+                agent.inc_update_attempt_count()
+                msg = "Agent update attempt count: {0} for version: {1}".format(agent.get_update_attempt_count(), str(agent.version))
+                logger.info(msg)
+                add_event(op=WALAEventOperation.AgentUpgrade, message=msg, log_event=False)
+
+            self._updater.purge_extra_agents_from_disk()
             self._updater.proceed_with_update()
 
         except Exception as err:
diff --git a/azurelinuxagent/ga/ga_version_updater.py b/azurelinuxagent/ga/ga_version_updater.py
index 46ae1f31..82a621ea 100644
--- a/azurelinuxagent/ga/ga_version_updater.py
+++ b/azurelinuxagent/ga/ga_version_updater.py
@@ -26,8 +26,8 @@ from azurelinuxagent.common.future import ustr
 from azurelinuxagent.common.protocol.extensions_goal_state import GoalStateSource
 from azurelinuxagent.common.utils import fileutil
 from azurelinuxagent.common.utils.flexible_version import FlexibleVersion
-from azurelinuxagent.common.version import AGENT_NAME, AGENT_DIR_PATTERN
-from azurelinuxagent.ga.guestagent import GuestAgent
+from azurelinuxagent.common.version import AGENT_NAME, AGENT_DIR_PATTERN, CURRENT_VERSION
+from azurelinuxagent.ga.guestagent import GuestAgent, AGENT_MANIFEST_FILE
 
 
 class GAVersionUpdater(object):
@@ -77,12 +77,6 @@ class GAVersionUpdater(object):
         """
         raise NotImplementedError
 
-    def purge_extra_agents_from_disk(self):
-        """
-        Method remove the extra agents from disk.
-        """
-        raise NotImplementedError
-
     def proceed_with_update(self):
         """
         performs upgrade/downgrade
@@ -104,6 +98,33 @@ class GAVersionUpdater(object):
         """
         self._gs_id = gs_id
 
+    @staticmethod
+    def download_new_agent_pkg(package_to_download, protocol, is_fast_track_goal_state):
+        """
+        Function downloads the new agent.
+        @param package_to_download: package to download
+        @param protocol: protocol object
+        @param is_fast_track_goal_state: True if goal state is fast track else False
+        """
+        agent_name = "{0}-{1}".format(AGENT_NAME, package_to_download.version)
+        agent_dir = os.path.join(conf.get_lib_dir(), agent_name)
+        agent_pkg_path = ".".join((os.path.join(conf.get_lib_dir(), agent_name), "zip"))
+        agent_handler_manifest_file = os.path.join(agent_dir, AGENT_MANIFEST_FILE)
+        if not os.path.exists(agent_dir) or not os.path.isfile(agent_handler_manifest_file):
+            protocol.client.download_zip_package("agent package", package_to_download.uris, agent_pkg_path, agent_dir, use_verify_header=is_fast_track_goal_state)
+        else:
+            logger.info("Agent {0} was previously downloaded - skipping download", agent_name)
+
+        if not os.path.isfile(agent_handler_manifest_file):
+            try:
+                # Clean up the agent directory if the manifest file is missing
+                logger.info("Agent handler manifest file is missing, cleaning up the agent directory: {0}".format(agent_dir))
+                if os.path.isdir(agent_dir):
+                    shutil.rmtree(agent_dir, ignore_errors=True)
+            except Exception as err:
+                logger.warn("Unable to delete Agent directory: {0}".format(err))
+            raise AgentUpdateError("Downloaded agent package: {0} is missing agent handler manifest file: {1}".format(agent_name, agent_handler_manifest_file))
+
     def download_and_get_new_agent(self, protocol, agent_family, goal_state):
         """
         Function downloads the new agent and returns the downloaded version.
@@ -116,9 +137,17 @@ class GAVersionUpdater(object):
             self._agent_manifest = goal_state.fetch_agent_manifest(agent_family.name, agent_family.uris)
         package_to_download = self._get_agent_package_to_download(self._agent_manifest, self._version)
         is_fast_track_goal_state = goal_state.extensions_goal_state.source == GoalStateSource.FastTrack
-        agent = GuestAgent.from_agent_package(package_to_download, protocol, is_fast_track_goal_state)
+        self.download_new_agent_pkg(package_to_download, protocol, is_fast_track_goal_state)
+        agent = GuestAgent.from_agent_package(package_to_download)
         return agent
 
+    def purge_extra_agents_from_disk(self):
+        """
+        Remove the agents from disk except current version and new agent version
+        """
+        known_agents = [CURRENT_VERSION, self._version]
+        self._purge_unknown_agents_from_disk(known_agents)
+
     def _get_agent_package_to_download(self, agent_manifest, version):
         """
         Returns the package of the given Version found in the manifest. If not found, returns exception
diff --git a/azurelinuxagent/ga/guestagent.py b/azurelinuxagent/ga/guestagent.py
index 35472c74..b4b2d05b 100644
--- a/azurelinuxagent/ga/guestagent.py
+++ b/azurelinuxagent/ga/guestagent.py
@@ -10,29 +10,26 @@ from azurelinuxagent.common.utils import textutil
 from azurelinuxagent.common import logger, conf
 from azurelinuxagent.common.exception import UpdateError
 from azurelinuxagent.common.utils.flexible_version import FlexibleVersion
-from azurelinuxagent.common.version import AGENT_DIR_PATTERN, AGENT_NAME, CURRENT_VERSION
+from azurelinuxagent.common.version import AGENT_DIR_PATTERN, AGENT_NAME
 from azurelinuxagent.ga.exthandlers import HandlerManifest
 
 AGENT_ERROR_FILE = "error.json"  # File name for agent error record
 AGENT_MANIFEST_FILE = "HandlerManifest.json"
 MAX_FAILURE = 3  # Max failure allowed for agent before declare bad agent
+AGENT_UPDATE_COUNT_FILE = "update_attempt.json"  # File for tracking agent update attempt count
 
 
 class GuestAgent(object):
-    def __init__(self, path, pkg, protocol, is_fast_track_goal_state):
+    def __init__(self, path, pkg):
         """
         If 'path' is given, the object is initialized to the version installed under that path.
 
         If 'pkg' is given, the version specified in the package information is downloaded and the object is
         initialized to that version.
 
-        'is_fast_track_goal_state' and 'protocol' are used only when a package is downloaded.
-
         NOTE: Prefer using the from_installed_agent and from_agent_package methods instead of calling __init__ directly
         """
-        self._is_fast_track_goal_state = is_fast_track_goal_state
         self.pkg = pkg
-        self._protocol = protocol
         version = None
         if path is not None:
             m = AGENT_DIR_PATTERN.match(path)
@@ -52,11 +49,13 @@ class GuestAgent(object):
         self.error = GuestAgentError(self.get_agent_error_file())
         self.error.load()
 
+        self.update_attempt_data = GuestAgentUpdateAttempt(self.get_agent_update_count_file())
+        self.update_attempt_data.load()
+
         try:
-            self._ensure_downloaded()
             self._ensure_loaded()
         except Exception as e:
-            # If we're unable to download/unpack the agent, delete the Agent directory
+            # If we're unable to unpack the agent, delete the Agent directory
             try:
                 if os.path.isdir(self.get_agent_dir()):
                     shutil.rmtree(self.get_agent_dir(), ignore_errors=True)
@@ -77,14 +76,14 @@ class GuestAgent(object):
         """
         Creates an instance of GuestAgent using the agent installed in the given 'path'.
         """
-        return GuestAgent(path, None, None, False)
+        return GuestAgent(path, None)
 
     @staticmethod
-    def from_agent_package(package, protocol, is_fast_track_goal_state):
+    def from_agent_package(package):
         """
         Creates an instance of GuestAgent using the information provided in the 'package'; if that version of the agent is not installed it, it installs it.
         """
-        return GuestAgent(None, package, protocol, is_fast_track_goal_state)
+        return GuestAgent(None, package)
 
     @property
     def name(self):
@@ -99,6 +98,9 @@ class GuestAgent(object):
     def get_agent_error_file(self):
         return os.path.join(conf.get_lib_dir(), self.name, AGENT_ERROR_FILE)
 
+    def get_agent_update_count_file(self):
+        return os.path.join(conf.get_lib_dir(), self.name, AGENT_UPDATE_COUNT_FILE)
+
     def get_agent_manifest_path(self):
         return os.path.join(self.get_agent_dir(), AGENT_MANIFEST_FILE)
 
@@ -136,45 +138,20 @@ class GuestAgent(object):
         except Exception as e:
             logger.warn(u"Agent {0} failed recording error state: {1}", self.name, ustr(e))
 
-    def _ensure_downloaded(self):
-        logger.verbose(u"Ensuring Agent {0} is downloaded", self.name)
-
-        if self.is_downloaded:
-            logger.verbose(u"Agent {0} was previously downloaded - skipping download", self.name)
-            return
-
-        if self.pkg is None:
-            raise UpdateError(u"Agent {0} is missing package and download URIs".format(
-                self.name))
-
-        self._download()
+    def inc_update_attempt_count(self):
+        try:
+            self.update_attempt_data.inc_count()
+            self.update_attempt_data.save()
+        except Exception as e:
+            logger.warn(u"Agent {0} failed recording update attempt: {1}", self.name, ustr(e))
 
-        msg = u"Agent {0} downloaded successfully".format(self.name)
-        logger.verbose(msg)
-        add_event(
-            AGENT_NAME,
-            version=self.version,
-            op=WALAEventOperation.Install,
-            is_success=True,
-            message=msg)
+    def get_update_attempt_count(self):
+        return self.update_attempt_data.count
 
     def _ensure_loaded(self):
         self._load_manifest()
         self._load_error()
 
-    def _download(self):
-        try:
-            self._protocol.client.download_zip_package("agent package", self.pkg.uris, self.get_agent_pkg_path(), self.get_agent_dir(), use_verify_header=self._is_fast_track_goal_state)
-        except Exception as exception:
-            msg = "Unable to download Agent {0}: {1}".format(self.name, ustr(exception))
-            add_event(
-                AGENT_NAME,
-                op=WALAEventOperation.Download,
-                version=CURRENT_VERSION,
-                is_success=False,
-                message=msg)
-            raise UpdateError(msg)
-
     def _load_error(self):
         try:
             self.error = GuestAgentError(self.get_agent_error_file())
@@ -303,3 +280,52 @@ class GuestAgentError(object):
             self.failure_count,
             self.was_fatal,
             self.reason)
+
+
+class GuestAgentUpdateAttempt(object):
+    def __init__(self, path):
+        self.count = 0
+        if path is None:
+            raise UpdateError(u"GuestAgentUpdateAttempt requires a path")
+        self.path = path
+
+        self.clear()
+
+    def inc_count(self):
+        self.count += 1
+
+    def clear(self):
+        self.count = 0
+
+    def load(self):
+        if self.path is not None and os.path.isfile(self.path):
+            try:
+                with open(self.path, 'r') as f:
+                    self.from_json(json.load(f))
+            except Exception as error:
+                # The update_attempt.json file is only supposed to be written only by the agent.
+                # If for whatever reason the file is malformed, just delete it to reset state of the errors.
+                logger.warn(
+                    "Ran into error when trying to load error file {0}, deleting it to clean state. Error: {1}".format(
+                        self.path, textutil.format_exception(error)))
+                try:
+                    os.remove(self.path)
+                except Exception:
+                    # We try best case efforts to delete the file, ignore error if we're unable to do so
+                    pass
+
+    def save(self):
+        if os.path.isdir(os.path.dirname(self.path)):
+            with open(self.path, 'w') as f:
+                json.dump(self.to_json(), f)
+
+    def from_json(self, data):
+        self.count = data.get(u"count", 0)
+
+    def to_json(self):
+        data = {
+            u"count": self.count
+        }
+        return data
+
+
diff --git a/azurelinuxagent/ga/rsm_version_updater.py b/azurelinuxagent/ga/rsm_version_updater.py
index a7a8bd97..366f1d70 100644
--- a/azurelinuxagent/ga/rsm_version_updater.py
+++ b/azurelinuxagent/ga/rsm_version_updater.py
@@ -104,14 +104,6 @@ class RSMVersionUpdater(GAVersionUpdater):
         logger.info(msg)
         add_event(op=WALAEventOperation.AgentUpgrade, message=msg, log_event=False)
 
-    def purge_extra_agents_from_disk(self):
-        """
-        Remove the agents( including rsm version if exists) from disk except current version. There is a chance that rsm version could exist and/or blacklisted
-        on previous update attempts. So we should remove it from disk in order to honor current rsm version update.
-        """
-        known_agents = [CURRENT_VERSION]
-        self._purge_unknown_agents_from_disk(known_agents)
-
     def proceed_with_update(self):
         """
         upgrade/downgrade to the new version.
diff --git a/azurelinuxagent/ga/self_update_version_updater.py b/azurelinuxagent/ga/self_update_version_updater.py
index 6605a28e..5a839851 100644
--- a/azurelinuxagent/ga/self_update_version_updater.py
+++ b/azurelinuxagent/ga/self_update_version_updater.py
@@ -171,13 +171,6 @@ class SelfUpdateVersionUpdater(GAVersionUpdater):
         logger.info(msg)
         add_event(op=WALAEventOperation.AgentUpgrade, message=msg, log_event=False)
 
-    def purge_extra_agents_from_disk(self):
-        """
-        Remove the agents from disk except current version and new agent version if exists
-        """
-        known_agents = [CURRENT_VERSION, self._version]
-        self._purge_unknown_agents_from_disk(known_agents)
-
     def proceed_with_update(self):
         """
         upgrade to largest version. Downgrade is not supported.
diff --git a/azurelinuxagent/ga/update.py b/azurelinuxagent/ga/update.py
index f2c20b0f..a8d34f7c 100644
--- a/azurelinuxagent/ga/update.py
+++ b/azurelinuxagent/ga/update.py
@@ -372,7 +372,6 @@ class UpdateHandler(object):
             self._ensure_extension_telemetry_state_configured_properly(protocol)
             self._ensure_firewall_rules_persisted(dst_ip=protocol.get_endpoint())
             self._add_accept_tcp_firewall_rule_if_not_enabled(dst_ip=protocol.get_endpoint())
-            self._reset_legacy_blacklisted_agents()
             self._cleanup_legacy_goal_state_history()
 
             # Get all thread handlers
@@ -1199,16 +1198,3 @@ class UpdateHandler(object):
         except Exception as e:
             msg = "Error while checking ip table rules:{0}".format(ustr(e))
             logger.error(msg)
-
-    def _reset_legacy_blacklisted_agents(self):
-        # Reset the state of all blacklisted agents that were blacklisted by legacy agents (i.e. not during auto-update)
-
-        # Filter legacy agents which are blacklisted but do not contain a `reason` in their error.json files
-        # (this flag signifies that this agent was blacklisted by the newer agents).
-        try:
-            legacy_blacklisted_agents = [agent for agent in self._load_agents() if
-                                         agent.is_blacklisted and agent.error.reason == '']
-            for agent in legacy_blacklisted_agents:
-                agent.clear_error()
-        except Exception as err:
-            logger.warn("Unable to reset legacy blacklisted agents due to: {0}".format(err))


Test patch:
diff --git a/tests/data/ga/WALinuxAgent-9.9.9.9-no_manifest.zip b/tests/data/ga/WALinuxAgent-9.9.9.10-no_manifest.zip
similarity index 100%
rename from tests/data/ga/WALinuxAgent-9.9.9.9-no_manifest.zip
rename to tests/data/ga/WALinuxAgent-9.9.9.10-no_manifest.zip
diff --git a/tests/data/wire/ga_manifest.xml b/tests/data/wire/ga_manifest.xml
index 799e1f11..c51bdbbc 100644
--- a/tests/data/wire/ga_manifest.xml
+++ b/tests/data/wire/ga_manifest.xml
@@ -31,7 +31,7 @@
         <Plugin>
             <Version>9.9.9.10</Version>
             <Uris>
-                <Uri>http://mock-goal-state/ga-manifests/OSTCExtensions.WALinuxAgent__99999.0.0.0</Uri>
+                <Uri>http://mock-goal-state/ga-manifests/OSTCExtensions.WALinuxAgent__9.9.9.10</Uri>
             </Uris>
         </Plugin>
         <Plugin>
diff --git a/tests/ga/test_agent_update_handler.py b/tests/ga/test_agent_update_handler.py
index fc58a6ee..82fcd362 100644
--- a/tests/ga/test_agent_update_handler.py
+++ b/tests/ga/test_agent_update_handler.py
@@ -9,8 +9,9 @@ from azurelinuxagent.common.future import ustr, httpclient
 from azurelinuxagent.common.protocol.restapi import VMAgentUpdateStatuses
 
 from azurelinuxagent.common.protocol.util import ProtocolUtil
-from azurelinuxagent.common.version import CURRENT_VERSION
+from azurelinuxagent.common.version import CURRENT_VERSION, AGENT_NAME
 from azurelinuxagent.ga.agent_update_handler import get_agent_update_handler
+from azurelinuxagent.ga.guestagent import GuestAgent
 from tests.ga.test_update import UpdateTestCase
 from tests.lib.http_request_predicates import HttpRequestPredicates
 from tests.lib.mock_wire_protocol import mock_wire_protocol, MockHttpResponse
@@ -27,7 +28,7 @@ class TestAgentUpdate(UpdateTestCase):
         clear_singleton_instances(ProtocolUtil)
 
     @contextlib.contextmanager
-    def _get_agent_update_handler(self, test_data=None, autoupdate_frequency=0.001, autoupdate_enabled=True, protocol_get_error=False):
+    def _get_agent_update_handler(self, test_data=None, autoupdate_frequency=0.001, autoupdate_enabled=True, protocol_get_error=False, mock_get_header=None, mock_put_header=None):
         # Default to DATA_FILE of test_data parameter raises the pylint warning
         # W0102: Dangerous default value DATA_FILE (builtins.dict) as argument (dangerous-default-value)
         test_data = DATA_FILE if test_data is None else test_data
@@ -51,7 +52,10 @@ class TestAgentUpdate(UpdateTestCase):
                 protocol.aggregate_status = json.loads(args[0])
                 return MockHttpResponse(status=201)
 
-            protocol.set_http_handlers(http_get_handler=get_handler, http_put_handler=put_handler)
+            http_get_handler = mock_get_header if mock_get_header else get_handler
+            http_put_handler = mock_put_header if mock_put_header else put_handler
+
+            protocol.set_http_handlers(http_get_handler=http_get_handler, http_put_handler=http_put_handler)
 
             with patch("azurelinuxagent.common.conf.get_autoupdate_enabled", return_value=autoupdate_enabled):
                 with patch("azurelinuxagent.common.conf.get_autoupdate_frequency", return_value=autoupdate_frequency):
@@ -401,7 +405,7 @@ class TestAgentUpdate(UpdateTestCase):
             self.assertEqual(VMAgentUpdateStatuses.Error, vm_agent_update_status.status)
             self.assertEqual(1, vm_agent_update_status.code)
             self.assertEqual("9.9.9.10", vm_agent_update_status.expected_version)
-            self.assertIn("Downloaded agent version is in bad state", vm_agent_update_status.message)
+            self.assertIn("Failed to download agent package from all URIs", vm_agent_update_status.message)
 
     def test_it_should_report_update_status_with_missing_rsm_version_error(self):
         data_file = DATA_FILE.copy()
@@ -470,3 +474,63 @@ class TestAgentUpdate(UpdateTestCase):
                 agent_update_handler.run(agent_update_handler._protocol.get_goal_state(), True)
 
             self._assert_agent_directories_exist_and_others_dont_exist(versions=[str(CURRENT_VERSION)])
+
+    def test_it_should_continue_with_update_if_number_of_update_attempts_less_than_3(self):
+        data_file = DATA_FILE.copy()
+        data_file["ext_conf"] = "wire/ext_conf_rsm_version.xml"
+
+        latest_version = self.prepare_agents(count=2)
+        self.expand_agents()
+        latest_path = os.path.join(self.tmp_dir, "{0}-{1}".format(AGENT_NAME, latest_version))
+        agent = GuestAgent.from_installed_agent(latest_path)
+        # marking agent as bad agent on first attempt
+        agent.mark_failure(is_fatal=True)
+        agent.inc_update_attempt_count()
+        self.assertTrue(agent.is_blacklisted, "Agent should be blacklisted")
+        self.assertEqual(1, agent.get_update_attempt_count(), "Agent update attempts should be 1")
+        with self._get_agent_update_handler(test_data=data_file) as (agent_update_handler, mock_telemetry):
+            # Rest 2 attempts it should continue with update even agent is marked as bad agent in first attempt
+            for i in range(2):
+                with self.assertRaises(AgentUpgradeExitException):
+                    agent_update_handler._protocol.mock_wire_data.set_version_in_agent_family(
+                        str(latest_version))
+                    agent_update_handler._protocol.mock_wire_data.set_incarnation(i+2)
+                    agent_update_handler._protocol.client.update_goal_state()
+                    agent_update_handler.run(agent_update_handler._protocol.get_goal_state(), True)
+                self._assert_agent_directories_exist_and_others_dont_exist(versions=[str(CURRENT_VERSION), str(latest_version)])
+                agent = GuestAgent.from_installed_agent(latest_path)
+                self.assertFalse(agent.is_blacklisted, "Agent should not be blacklisted")
+                self.assertEqual(i+2, agent.get_update_attempt_count(), "Agent update attempts should be {0}".format(i+2))
+
+            # check if next update is not attempted
+            agent.mark_failure(is_fatal=True)
+            agent_update_handler.run(agent_update_handler._protocol.get_goal_state(), True)
+            agent = GuestAgent.from_installed_agent(latest_path)
+            self.assertTrue(agent.is_blacklisted, "Agent should be blacklisted")
+            self.assertEqual(3, agent.get_update_attempt_count(), "Agent update attempts should be 3")
+            self.assertEqual(1, len([kwarg['message'] for _, kwarg in mock_telemetry.call_args_list if
+                                     "Attempted enough update retries for version: {0} but still agent not recovered from bad state".format(latest_version) in kwarg[
+                                         'message'] and kwarg[
+                                         'op'] == WALAEventOperation.AgentUpgrade]),
+                             "Update is not allowed after 3 attempts")
+
+    def test_it_should_fail_the_update_if_agent_pkg_is_invalid(self):
+        agent_uri = 'https://foo.blob.core.windows.net/bar/OSTCExtensions.WALinuxAgent__9.9.9.10'
+
+        def http_get_handler(uri, *_, **__):
+            if uri in (agent_uri, 'http://168.63.129.16:32526/extensionArtifact'):
+                response = load_bin_data("ga/WALinuxAgent-9.9.9.10-no_manifest.zip")
+                return MockHttpResponse(status=httpclient.OK, body=response)
+            return None
+        self.prepare_agents(count=1)
+        data_file = DATA_FILE.copy()
+        data_file["ext_conf"] = "wire/ext_conf_rsm_version.xml"
+        with self._get_agent_update_handler(test_data=data_file, mock_get_header=http_get_handler) as (agent_update_handler, mock_telemetry):
+            agent_update_handler._protocol.mock_wire_data.set_version_in_agent_family("9.9.9.10")
+            agent_update_handler._protocol.mock_wire_data.set_incarnation(2)
+            agent_update_handler._protocol.client.update_goal_state()
+            agent_update_handler.run(agent_update_handler._protocol.get_goal_state(), True)
+            self._assert_agent_directories_exist_and_others_dont_exist(versions=[str(CURRENT_VERSION)])
+            self.assertEqual(1, len([kwarg['message'] for _, kwarg in mock_telemetry.call_args_list if
+                                     "Downloaded agent package: WALinuxAgent-9.9.9.10 is missing agent handler manifest file" in kwarg['message'] and kwarg[
+                                         'op'] == WALAEventOperation.AgentUpgrade]), "Agent update should fail")
diff --git a/tests/ga/test_guestagent.py b/tests/ga/test_guestagent.py
index 3d25ff79..972e603c 100644
--- a/tests/ga/test_guestagent.py
+++ b/tests/ga/test_guestagent.py
@@ -1,17 +1,14 @@
+import contextlib
 import json
 import os
+import tempfile
 
 from azurelinuxagent.common import conf
 from azurelinuxagent.common.exception import UpdateError
 from azurelinuxagent.ga.guestagent import GuestAgent, AGENT_MANIFEST_FILE, AGENT_ERROR_FILE, GuestAgentError, \
-    MAX_FAILURE
-from azurelinuxagent.common.future import httpclient
-from azurelinuxagent.common.protocol.restapi import ExtHandlerPackage
+    MAX_FAILURE, GuestAgentUpdateAttempt
 from azurelinuxagent.common.version import AGENT_NAME
 from tests.ga.test_update import UpdateTestCase, EMPTY_MANIFEST, WITH_ERROR, NO_ERROR
-from tests.lib import wire_protocol_data
-from tests.lib.mock_wire_protocol import MockHttpResponse, mock_wire_protocol
-from tests.lib.tools import load_bin_data, patch
 
 
 class TestGuestAgent(UpdateTestCase):
@@ -102,6 +99,22 @@ class TestGuestAgent(UpdateTestCase):
         self.assertEqual(2, agent.error.failure_count)
         self.assertTrue(agent.is_blacklisted)
 
+    def test_inc_update_attempt_count(self):
+        agent = GuestAgent.from_installed_agent(self.agent_path)
+        agent.inc_update_attempt_count()
+        self.assertEqual(1, agent.update_attempt_data.count)
+
+        agent.inc_update_attempt_count()
+        self.assertEqual(2, agent.update_attempt_data.count)
+
+    def test_get_update_count(self):
+        agent = GuestAgent.from_installed_agent(self.agent_path)
+        agent.inc_update_attempt_count()
+        self.assertEqual(1, agent.get_update_attempt_count())
+
+        agent.inc_update_attempt_count()
+        self.assertEqual(2, agent.get_update_attempt_count())
+
     def test_load_manifest(self):
         self.expand_agents()
         agent = GuestAgent.from_installed_agent(self.agent_path)
@@ -140,94 +153,6 @@ class TestGuestAgent(UpdateTestCase):
         agent._load_error()
         self.assertTrue(agent.error is not None)
 
-    def test_download(self):
-        self.remove_agents()
-        self.assertFalse(os.path.isdir(self.agent_path))
-
-        agent_uri = 'https://foo.blob.core.windows.net/bar/OSTCExtensions.WALinuxAgent__1.0.0'
-
-        def http_get_handler(uri, *_, **__):
-            if uri == agent_uri:
-                response = load_bin_data(self._get_agent_file_name(), self._agent_zip_dir)
-                return MockHttpResponse(status=httpclient.OK, body=response)
-            return None
-
-        pkg = ExtHandlerPackage(version=str(self._get_agent_version()))
-        pkg.uris.append(agent_uri)
-
-        with mock_wire_protocol(wire_protocol_data.DATA_FILE) as protocol:
-            protocol.set_http_handlers(http_get_handler=http_get_handler)
-            agent = GuestAgent.from_agent_package(pkg, protocol, False)
-
-        self.assertTrue(os.path.isdir(agent.get_agent_dir()))
-        self.assertTrue(agent.is_downloaded)
-
-    def test_download_fail(self):
-        self.remove_agents()
-        self.assertFalse(os.path.isdir(self.agent_path))
-
-        agent_uri = 'https://foo.blob.core.windows.net/bar/OSTCExtensions.WALinuxAgent__1.0.0'
-
-        def http_get_handler(uri, *_, **__):
-            if uri in (agent_uri, 'http://168.63.129.16:32526/extensionArtifact'):
-                return MockHttpResponse(status=httpclient.SERVICE_UNAVAILABLE)
-            return None
-
-        agent_version = self._get_agent_version()
-        pkg = ExtHandlerPackage(version=str(self._get_agent_version()))
-        pkg.uris.append(agent_uri)
-
-        with mock_wire_protocol(wire_protocol_data.DATA_FILE) as protocol:
-            protocol.set_http_handlers(http_get_handler=http_get_handler)
-            with patch("azurelinuxagent.ga.guestagent.add_event") as add_event:
-                agent = GuestAgent.from_agent_package(pkg, protocol, False)
-
-        self.assertFalse(os.path.isfile(self.agent_path))
-
-        messages = [kwargs['message'] for _, kwargs in add_event.call_args_list if kwargs['op'] == 'Install' and kwargs['is_success'] == False]
-        self.assertEqual(1, len(messages), "Expected exactly 1 install error/ Got: {0}".format(add_event.call_args_list))
-        self.assertIn(str.format('[UpdateError] Unable to download Agent WALinuxAgent-{0}', agent_version), messages[0], "The install error does not include the expected message")
-
-        self.assertFalse(agent.is_blacklisted, "Download failures should not blacklist the Agent")
-
-    def test_invalid_agent_package_does_not_blacklist_the_agent(self):
-        agent_uri = 'https://foo.blob.core.windows.net/bar/OSTCExtensions.WALinuxAgent__9.9.9.9'
-
-        def http_get_handler(uri, *_, **__):
-            if uri in (agent_uri, 'http://168.63.129.16:32526/extensionArtifact'):
-                response = load_bin_data("ga/WALinuxAgent-9.9.9.9-no_manifest.zip")
-                return MockHttpResponse(status=httpclient.OK, body=response)
-            return None
-
-        pkg = ExtHandlerPackage(version="9.9.9.9")
-        pkg.uris.append(agent_uri)
-
-        with mock_wire_protocol(wire_protocol_data.DATA_FILE) as protocol:
-            protocol.set_http_handlers(http_get_handler=http_get_handler)
-            agent = GuestAgent.from_agent_package(pkg, protocol, False)
-
-        self.assertFalse(agent.is_blacklisted, "The agent should not be blacklisted if unable to unpack/download")
-        self.assertFalse(os.path.exists(agent.get_agent_dir()), "Agent directory should be cleaned up")
-
-    @patch("azurelinuxagent.ga.update.GuestAgent._download")
-    def test_ensure_download_skips_blacklisted(self, mock_download):
-        agent = GuestAgent.from_installed_agent(self.agent_path)
-        self.assertEqual(0, mock_download.call_count)
-
-        agent.clear_error()
-        agent.mark_failure(is_fatal=True)
-        self.assertTrue(agent.is_blacklisted)
-
-        pkg = ExtHandlerPackage(version=str(self._get_agent_version()))
-        pkg.uris.append(None)
-        # _download is mocked so there will be no http request; passing a None protocol
-        agent = GuestAgent.from_agent_package(pkg, None, False)
-
-        self.assertEqual(1, agent.error.failure_count)
-        self.assertTrue(agent.error.was_fatal)
-        self.assertTrue(agent.is_blacklisted)
-        self.assertEqual(0, mock_download.call_count)
-
 
 class TestGuestAgentError(UpdateTestCase):
     def test_creation(self):
@@ -308,3 +233,69 @@ class TestGuestAgentError(UpdateTestCase):
             WITH_ERROR["reason"])
         self.assertEqual(s, str(err))
         return
+
+
+UPDATE_ATTEMPT = {
+    "count": 2
+}
+
+NO_ATTEMPT = {
+    "count": 0
+}
+
+
+class TestGuestAgentUpdateAttempt(UpdateTestCase):
+    @contextlib.contextmanager
+    def get_attempt_count_file(self, attempt_count=None):
+        if attempt_count is None:
+            attempt_count = NO_ATTEMPT
+        with tempfile.NamedTemporaryFile(mode="w") as fp:
+            json.dump(attempt_count, fp)
+            fp.seek(0)
+            yield fp
+
+    def test_creation(self):
+        self.assertRaises(TypeError, GuestAgentUpdateAttempt)
+        self.assertRaises(UpdateError, GuestAgentUpdateAttempt, None)
+
+        with self.get_attempt_count_file(UPDATE_ATTEMPT) as path:
+            update_data = GuestAgentUpdateAttempt(path.name)
+            update_data.load()
+            self.assertEqual(path.name, update_data.path)
+        self.assertNotEqual(None, update_data)
+
+        self.assertEqual(UPDATE_ATTEMPT["count"], update_data.count)
+
+    def test_clear(self):
+        with self.get_attempt_count_file(UPDATE_ATTEMPT) as path:
+            update_data = GuestAgentUpdateAttempt(path.name)
+            update_data.load()
+            self.assertEqual(path.name, update_data.path)
+        self.assertNotEqual(None, update_data)
+
+        update_data.clear()
+        self.assertEqual(NO_ATTEMPT["count"], update_data.count)
+
+    def test_save(self):
+        with self.get_attempt_count_file(UPDATE_ATTEMPT) as path:
+            update_data = GuestAgentUpdateAttempt(path.name)
+            update_data.load()
+        update_data.inc_count()
+        update_data.save()
+
+        with self.get_attempt_count_file(update_data.to_json()) as path:
+            new_data = GuestAgentUpdateAttempt(path.name)
+            new_data.load()
+
+        self.assertEqual(update_data.count, new_data.count)
+
+    def test_inc_count(self):
+        with self.get_attempt_count_file() as path:
+            update_data = GuestAgentUpdateAttempt(path.name)
+            update_data.load()
+
+        self.assertEqual(0, update_data.count)
+        update_data.inc_count()
+        self.assertEqual(1, update_data.count)
+        update_data.inc_count()
+        self.assertEqual(2, update_data.count)
diff --git a/tests/ga/test_update.py b/tests/ga/test_update.py
index 079f896b..58732de5 100644
--- a/tests/ga/test_update.py
+++ b/tests/ga/test_update.py
@@ -1340,27 +1340,6 @@ class TestUpdate(UpdateTestCase):
                      "Fetching the goal state recovered from previous errors." in args[0]]
         self.assertTrue(len(info_msgs) > 0, "Agent should've logged a message when recovered from GS errors")
 
-    def test_it_should_reset_legacy_blacklisted_agents_on_process_start(self):
-        # Add some good agents
-        self.prepare_agents(count=10)
-        good_agents = [agent.name for agent in self.agents()]
-
-        # Add a set of blacklisted agents
-        self.prepare_agents(count=20, is_available=False)
-        for agent in self.agents():
-            # Assert the test environment is correctly set
-            if agent.name not in good_agents:
-                self.assertTrue(agent.is_blacklisted, "Agent {0} should be blacklisted".format(agent.name))
-            else:
-                self.assertFalse(agent.is_blacklisted, "Agent {0} should not be blacklisted".format(agent.name))
-
-        with _get_update_handler(autoupdate_enabled=False) as (update_handler, _):
-            update_handler.run(debug=True)
-            self.assertEqual(20, self.agent_count(), "All agents should be available on disk")
-            # Ensure none of the agents are blacklisted
-            for agent in self.agents():
-                self.assertFalse(agent.is_blacklisted, "Legacy Agent should not be blacklisted")
-
 
 class TestUpdateWaitForCloudInit(AgentTestCase):
     @staticmethod