Repository: Azure/WALinuxAgent
Pull request #2977 (2023-11-08 00:45:28)

Save goal state history explicitly
We have multiple threads/processes fetching the goal state and all of those operations save the goal state to the history folder, which can end up messing up the history folder due to concurrency issues. 

We use the history folder to debug issues in extension handling, so we need only the main thread to save the goal state. I added a new parameter, 'save_to_history', which is False by default and only the extension handler thread sets to True. The parameter had to be propagated thru several layers in the code, as well as in the test mocks.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/2977?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) Report
> Merging [#2977](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/2977?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) (7ef29fe) into [develop](https://app.codecov.io/gh/Azure/WALinuxAgent/commit/8f3ee961c2524b2c4e1a54f5823b0891399277f9?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) (8f3ee96) will **decrease** coverage by `0.07%`.
> The diff coverage is `78.94%`.

> :exclamation: Current head 7ef29fe differs from pull request most recent head c2ab99b. Consider uploading reports for the commit c2ab99b to get more accurate results

```diff
@@             Coverage Diff             @@
##           develop    #2977      +/-   ##
===========================================
- Coverage    71.99%   71.93%   -0.07%     
===========================================
  Files          106      106              
  Lines        16086    16096      +10     
  Branches      2300     2309       +9     
===========================================
- Hits         11581    11578       -3     
- Misses        3975     3983       +8     
- Partials       530      535       +5     
```


| [Files](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/2977?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure) | Coverage Δ | |
|---|---|---|
| [azurelinuxagent/common/protocol/util.py](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/2977?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9wcm90b2NvbC91dGlsLnB5) | `80.32% <100.00%> (ø)` | |
| [azurelinuxagent/ga/update.py](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/2977?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2dhL3VwZGF0ZS5weQ==) | `87.44% <100.00%> (ø)` | |
| [azurelinuxagent/common/protocol/wire.py](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/2977?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9wcm90b2NvbC93aXJlLnB5) | `77.37% <85.71%> (ø)` | |
| [azurelinuxagent/common/protocol/goal\_state.py](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/2977?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure#diff-YXp1cmVsaW51eGFnZW50L2NvbW1vbi9wcm90b2NvbC9nb2FsX3N0YXRlLnB5) | `93.42% <72.00%> (-1.33%)` | :arrow_down: |

... and [2 files with indirect coverage changes](https://app.codecov.io/gh/Azure/WALinuxAgent/pull/2977/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Azure)



Patch:
diff --git a/azurelinuxagent/common/protocol/goal_state.py b/azurelinuxagent/common/protocol/goal_state.py
index 1b4bcea8..2eb89c1e 100644
--- a/azurelinuxagent/common/protocol/goal_state.py
+++ b/azurelinuxagent/common/protocol/goal_state.py
@@ -70,7 +70,7 @@ class GoalStateInconsistentError(ProtocolError):
 
 
 class GoalState(object):
-    def __init__(self, wire_client, goal_state_properties=GoalStateProperties.All, silent=False):
+    def __init__(self, wire_client, goal_state_properties=GoalStateProperties.All, silent=False, save_to_history=False):
         """
         Fetches the goal state using the given wire client.
 
@@ -84,6 +84,7 @@ class GoalState(object):
         try:
             self._wire_client = wire_client
             self._history = None
+            self._save_to_history = save_to_history
             self._extensions_goal_state = None  # populated from vmSettings or extensionsConfig
             self._goal_state_properties = goal_state_properties
             self.logger = logger.Logger(logger.DEFAULT_LOGGER)
@@ -186,7 +187,8 @@ class GoalState(object):
         try:
             is_fast_track = self.extensions_goal_state.source == GoalStateSource.FastTrack
             xml_text = self._wire_client.fetch_manifest(manifest_type, uris, use_verify_header=is_fast_track)
-            self._history.save_manifest(name, xml_text)
+            if self._save_to_history:
+                self._history.save_manifest(name, xml_text)
             return ExtensionManifest(xml_text)
         except Exception as e:
             raise ProtocolError("Failed to retrieve {0} manifest. Error: {1}".format(manifest_type, ustr(e)))
@@ -263,11 +265,12 @@ class GoalState(object):
 
         # Start a new history subdirectory and capture the updated goal state
         tag = "{0}".format(incarnation) if vm_settings is None else "{0}-{1}".format(incarnation, vm_settings.etag)
-        self._history = GoalStateHistory(timestamp, tag)
-        if goal_state_updated:
-            self._history.save_goal_state(xml_text)
-        if vm_settings_updated:
-            self._history.save_vm_settings(vm_settings.get_redacted_text())
+        if self._save_to_history:
+            self._history = GoalStateHistory(timestamp, tag)
+            if goal_state_updated:
+                self._history.save_goal_state(xml_text)
+            if vm_settings_updated:
+                self._history.save_vm_settings(vm_settings.get_redacted_text())
 
         #
         # Continue fetching the rest of the goal state
@@ -324,7 +327,8 @@ class GoalState(object):
         if len(certs.warnings) > 0:
             self.logger.warn(certs.warnings)
             add_event(op=WALAEventOperation.GoalState, message=certs.warnings)
-        self._history.save_certificates(json.dumps(certs.summary))
+        if self._save_to_history:
+            self._history.save_certificates(json.dumps(certs.summary))
         return certs
 
     def _check_and_download_missing_certs_on_disk(self):
@@ -357,8 +361,9 @@ class GoalState(object):
         msg = 'The HGAP stopped supporting vmSettings; will fetched the goal state from the WireServer.'
         self.logger.info(msg)
         add_event(op=WALAEventOperation.VmSettings, message=msg)
-        self._history = GoalStateHistory(datetime.datetime.utcnow(), incarnation)
-        self._history.save_goal_state(xml_text)
+        if self._save_to_history:
+            self._history = GoalStateHistory(datetime.datetime.utcnow(), incarnation)
+            self._history.save_goal_state(xml_text)
         self._extensions_goal_state = self._fetch_full_wire_server_goal_state(incarnation, xml_doc)
         if self._extensions_goal_state.created_on_timestamp < vm_settings_support_stopped_error.timestamp:
             self._extensions_goal_state.is_outdated = True
@@ -368,7 +373,8 @@ class GoalState(object):
             add_event(op=WALAEventOperation.VmSettings, message=msg)
 
     def save_to_history(self, data, file_name):
-        self._history.save(data, file_name)
+        if self._save_to_history:
+            self._history.save(data, file_name)
 
     @staticmethod
     def _fetch_goal_state(wire_client):
@@ -463,21 +469,24 @@ class GoalState(object):
             else:
                 xml_text = self._wire_client.fetch_config(extensions_config_uri, self._wire_client.get_header())
                 extensions_config = ExtensionsGoalStateFactory.create_from_extensions_config(incarnation, xml_text, self._wire_client)
-                self._history.save_extensions_config(extensions_config.get_redacted_text())
+                if self._save_to_history:
+                    self._history.save_extensions_config(extensions_config.get_redacted_text())
 
             hosting_env = None
             if GoalStateProperties.HostingEnv & self._goal_state_properties:
                 hosting_env_uri = findtext(xml_doc, "HostingEnvironmentConfig")
                 xml_text = self._wire_client.fetch_config(hosting_env_uri, self._wire_client.get_header())
                 hosting_env = HostingEnv(xml_text)
-                self._history.save_hosting_env(xml_text)
+                if self._save_to_history:
+                    self._history.save_hosting_env(xml_text)
 
             shared_config = None
             if GoalStateProperties.SharedConfig & self._goal_state_properties:
                 shared_conf_uri = findtext(xml_doc, "SharedConfig")
                 xml_text = self._wire_client.fetch_config(shared_conf_uri, self._wire_client.get_header())
                 shared_config = SharedConfig(xml_text)
-                self._history.save_shared_conf(xml_text)
+                if self._save_to_history:
+                    self._history.save_shared_conf(xml_text)
                 # SharedConfig.xml is used by other components (Azsec and Singularity/HPC Infiniband), so save it to the agent's root directory as well
                 shared_config_file = os.path.join(conf.get_lib_dir(), SHARED_CONF_FILE_NAME)
                 try:
@@ -496,7 +505,8 @@ class GoalState(object):
                 if remote_access_uri is not None:
                     xml_text = self._wire_client.fetch_config(remote_access_uri, self._wire_client.get_header_for_cert())
                     remote_access = RemoteAccess(xml_text)
-                    self._history.save_remote_access(xml_text)
+                    if self._save_to_history:
+                        self._history.save_remote_access(xml_text)
 
             self._incarnation = incarnation
             self._role_instance_id = role_instance_id
diff --git a/azurelinuxagent/common/protocol/util.py b/azurelinuxagent/common/protocol/util.py
index 7d7f9016..b90e9542 100644
--- a/azurelinuxagent/common/protocol/util.py
+++ b/azurelinuxagent/common/protocol/util.py
@@ -188,7 +188,7 @@ class ProtocolUtil(SingletonPerThread):
                 return
             logger.error("Failed to clear wiresever endpoint: {0}", e)
 
-    def _detect_protocol(self, init_goal_state=True):
+    def _detect_protocol(self, save_to_history, init_goal_state=True):
         """
         Probe protocol endpoints in turn.
         """
@@ -217,7 +217,7 @@ class ProtocolUtil(SingletonPerThread):
 
                 try:
                     protocol = WireProtocol(endpoint)
-                    protocol.detect(init_goal_state=init_goal_state)
+                    protocol.detect(init_goal_state=init_goal_state, save_to_history=save_to_history)
                     self._set_wireserver_endpoint(endpoint)
                     return protocol
 
@@ -268,7 +268,7 @@ class ProtocolUtil(SingletonPerThread):
         finally:
             self._lock.release()
 
-    def get_protocol(self, init_goal_state=True):
+    def get_protocol(self, init_goal_state=True, save_to_history=False):
         """
         Detect protocol by endpoint.
         :returns: protocol instance
@@ -296,7 +296,7 @@ class ProtocolUtil(SingletonPerThread):
 
             logger.info("Detect protocol endpoint")
 
-            protocol = self._detect_protocol(init_goal_state=init_goal_state)
+            protocol = self._detect_protocol(save_to_history=save_to_history, init_goal_state=init_goal_state)
 
             IOErrorCounter.set_protocol_endpoint(endpoint=protocol.get_endpoint())
             self._save_protocol(WIRE_PROTOCOL_NAME)
diff --git a/azurelinuxagent/common/protocol/wire.py b/azurelinuxagent/common/protocol/wire.py
index ea397f04..c93624cb 100644
--- a/azurelinuxagent/common/protocol/wire.py
+++ b/azurelinuxagent/common/protocol/wire.py
@@ -73,7 +73,7 @@ class WireProtocol(DataContract):
             raise ProtocolError("WireProtocol endpoint is None")
         self.client = WireClient(endpoint)
 
-    def detect(self, init_goal_state=True):
+    def detect(self, init_goal_state=True, save_to_history=False):
         self.client.check_wire_protocol_version()
 
         trans_prv_file = os.path.join(conf.get_lib_dir(),
@@ -86,7 +86,7 @@ class WireProtocol(DataContract):
         # Initialize the goal state, including all the inner properties
         if init_goal_state:
             logger.info('Initializing goal state during protocol detection')
-            self.client.reset_goal_state()
+            self.client.reset_goal_state(save_to_history=save_to_history)
 
     def update_host_plugin_from_goal_state(self):
         self.client.update_host_plugin_from_goal_state()
@@ -777,13 +777,13 @@ class WireClient(object):
             self._host_plugin.update_container_id(container_id)
             self._host_plugin.update_role_config_name(role_config_name)
 
-    def update_goal_state(self, silent=False):
+    def update_goal_state(self, silent=False, save_to_history=False):
         """
         Updates the goal state if the incarnation or etag changed
         """
         try:
             if self._goal_state is None:
-                self._goal_state = GoalState(self, silent=silent)
+                self._goal_state = GoalState(self, silent=silent, save_to_history=save_to_history)
             else:
                 self._goal_state.update(silent=silent)
 
@@ -792,7 +792,7 @@ class WireClient(object):
         except Exception as exception:
             raise ProtocolError("Error fetching goal state: {0}".format(ustr(exception)))
 
-    def reset_goal_state(self, goal_state_properties=GoalStateProperties.All, silent=False):
+    def reset_goal_state(self, goal_state_properties=GoalStateProperties.All, silent=False, save_to_history=False):
         """
         Resets the goal state
         """
@@ -800,7 +800,7 @@ class WireClient(object):
             if not silent:
                 logger.info("Forcing an update of the goal state.")
 
-            self._goal_state = GoalState(self, goal_state_properties=goal_state_properties, silent=silent)
+            self._goal_state = GoalState(self, goal_state_properties=goal_state_properties, silent=silent, save_to_history=save_to_history)
 
         except ProtocolError:
             raise
@@ -936,7 +936,7 @@ class WireClient(object):
 
         if extensions_goal_state.status_upload_blob is None:
             # the status upload blob is in ExtensionsConfig so force a full goal state refresh
-            self.reset_goal_state(silent=True)
+            self.reset_goal_state(silent=True, save_to_history=True)
             extensions_goal_state = self.get_goal_state().extensions_goal_state
 
             if extensions_goal_state.status_upload_blob is None:
diff --git a/azurelinuxagent/ga/update.py b/azurelinuxagent/ga/update.py
index 2862f8e9..6e7b5b91 100644
--- a/azurelinuxagent/ga/update.py
+++ b/azurelinuxagent/ga/update.py
@@ -339,7 +339,7 @@ class UpdateHandler(object):
             # Initialize the goal state; some components depend on information provided by the goal state and this
             # call ensures the required info is initialized (e.g. telemetry depends on the container ID.)
             #
-            protocol = self.protocol_util.get_protocol()
+            protocol = self.protocol_util.get_protocol(save_to_history=True)
 
             self._initialize_goal_state(protocol)
 
@@ -503,7 +503,7 @@ class UpdateHandler(object):
         try:
             max_errors_to_log = 3
 
-            protocol.client.update_goal_state(silent=self._update_goal_state_error_count >= max_errors_to_log)
+            protocol.client.update_goal_state(silent=self._update_goal_state_error_count >= max_errors_to_log, save_to_history=True)
 
             self._goal_state = protocol.get_goal_state()
 


Test patch:
diff --git a/tests/common/protocol/test_goal_state.py b/tests/common/protocol/test_goal_state.py
index c33851e9..5b4a2948 100644
--- a/tests/common/protocol/test_goal_state.py
+++ b/tests/common/protocol/test_goal_state.py
@@ -112,14 +112,18 @@ class GoalStateTestCase(AgentTestCase, HttpRequestPredicates):
             protocol.mock_wire_data.set_incarnation(999)
             protocol.mock_wire_data.set_etag(888)
 
-            _ = GoalState(protocol.client)
+            _ = GoalState(protocol.client, save_to_history=True)
 
             self._assert_directory_contents(
                 self._find_history_subdirectory("999-888"),
                 ["GoalState.xml", "ExtensionsConfig.xml", "VmSettings.json", "Certificates.json", "SharedConfig.xml", "HostingEnvironmentConfig.xml"])
 
+    @staticmethod
+    def _get_history_directory():
+        return os.path.join(conf.get_lib_dir(), ARCHIVE_DIRECTORY_NAME)
+
     def _find_history_subdirectory(self, tag):
-        matches = glob.glob(os.path.join(self.tmp_dir, ARCHIVE_DIRECTORY_NAME, "*_{0}".format(tag)))
+        matches = glob.glob(os.path.join(self._get_history_directory(), "*_{0}".format(tag)))
         self.assertTrue(len(matches) == 1, "Expected one history directory for tag {0}. Got: {1}".format(tag, matches))
         return matches[0]
 
@@ -136,7 +140,7 @@ class GoalStateTestCase(AgentTestCase, HttpRequestPredicates):
             protocol.mock_wire_data.set_incarnation(123)
             protocol.mock_wire_data.set_etag(654)
 
-            goal_state = GoalState(protocol.client)
+            goal_state = GoalState(protocol.client, save_to_history=True)
             self._assert_directory_contents(
                 self._find_history_subdirectory("123-654"),
                 ["GoalState.xml", "ExtensionsConfig.xml", "VmSettings.json",  "Certificates.json", "SharedConfig.xml", "HostingEnvironmentConfig.xml"])
@@ -164,7 +168,7 @@ class GoalStateTestCase(AgentTestCase, HttpRequestPredicates):
             protocol.mock_wire_data.set_incarnation(888)
             protocol.mock_wire_data.set_etag(888)
 
-            goal_state = GoalState(protocol.client)
+            goal_state = GoalState(protocol.client, save_to_history=True)
 
             extensions_goal_state = goal_state.extensions_goal_state
             protected_settings = []
@@ -221,6 +225,12 @@ class GoalStateTestCase(AgentTestCase, HttpRequestPredicates):
 
             self.assertEqual(expected, actual, "The vmSettings were not saved correctly")
 
+    def test_should_not_save_to_the_history_by_default(self):
+        with mock_wire_protocol(wire_protocol_data.DATA_FILE_VM_SETTINGS) as protocol:
+            _ = GoalState(protocol.client)  # omit the save_to_history parameter
+            history = self._get_history_directory()
+            self.assertFalse(os.path.exists(history), "The history directory not should have been created")
+
     @staticmethod
     @contextlib.contextmanager
     def _create_protocol_ws_and_hgap_in_sync():
diff --git a/tests/common/protocol/test_protocol_util.py b/tests/common/protocol/test_protocol_util.py
index 06683cef..b60ca9af 100644
--- a/tests/common/protocol/test_protocol_util.py
+++ b/tests/common/protocol/test_protocol_util.py
@@ -127,14 +127,14 @@ class TestProtocolUtil(AgentTestCase):
         endpoint_file = protocol_util._get_wireserver_endpoint_file_path()  # pylint: disable=unused-variable
 
         # Test wire protocol when no endpoint file has been written
-        protocol_util._detect_protocol()
+        protocol_util._detect_protocol(save_to_history=False)
         self.assertEqual(KNOWN_WIRESERVER_IP, protocol_util.get_wireserver_endpoint())
 
         # Test wire protocol on dhcp failure
         protocol_util.osutil.is_dhcp_available.return_value = True
         protocol_util.dhcp_handler.run.side_effect = DhcpError()
 
-        self.assertRaises(ProtocolError, protocol_util._detect_protocol)
+        self.assertRaises(ProtocolError, lambda: protocol_util._detect_protocol(save_to_history=False))
 
     @patch("azurelinuxagent.common.protocol.util.WireProtocol")
     def test_get_protocol(self, WireProtocol, _):
diff --git a/tests/ga/test_update.py b/tests/ga/test_update.py
index 49a4e0be..8bac6774 100644
--- a/tests/ga/test_update.py
+++ b/tests/ga/test_update.py
@@ -2176,7 +2176,7 @@ def _create_update_handler():
 
 
 @contextlib.contextmanager
-def _mock_exthandlers_handler(extension_statuses=None):
+def _mock_exthandlers_handler(extension_statuses=None, save_to_history=False):
     """
     Creates an ExtHandlersHandler that doesn't actually handle any extensions, but that returns status for 1 extension.
     The returned ExtHandlersHandler uses a mock WireProtocol, and both the run() and report_ext_handlers_status() are
@@ -2191,7 +2191,7 @@ def _mock_exthandlers_handler(extension_statuses=None):
         vm_status.vmAgent.extensionHandlers[0].extension_status.status = extension_status
         return vm_status
 
-    with mock_wire_protocol(DATA_FILE) as protocol:
+    with mock_wire_protocol(DATA_FILE, save_to_history=save_to_history) as protocol:
         exthandlers_handler = ExtHandlersHandler(protocol)
         exthandlers_handler.run = Mock()
         if extension_statuses is None:
@@ -2237,7 +2237,7 @@ class ProcessGoalStateTestCase(AgentTestCase):
             self.assertEqual(3, agent_update_handler.run.call_count, "agent_update_handler.run() should have been called on the new goal state")
 
     def test_it_should_write_the_agent_status_to_the_history_folder(self):
-        with _mock_exthandlers_handler() as exthandlers_handler:
+        with _mock_exthandlers_handler(save_to_history=True) as exthandlers_handler:
             update_handler = _create_update_handler()
             remote_access_handler = Mock()
             remote_access_handler.run = Mock()
diff --git a/tests/lib/mock_wire_protocol.py b/tests/lib/mock_wire_protocol.py
index 4e3521fe..78cbc59e 100644
--- a/tests/lib/mock_wire_protocol.py
+++ b/tests/lib/mock_wire_protocol.py
@@ -22,7 +22,7 @@ from tests.lib import wire_protocol_data
 
 
 @contextlib.contextmanager
-def mock_wire_protocol(mock_wire_data_file, http_get_handler=None, http_post_handler=None, http_put_handler=None, do_not_mock=lambda method, url: False, fail_on_unknown_request=True):
+def mock_wire_protocol(mock_wire_data_file, http_get_handler=None, http_post_handler=None, http_put_handler=None, do_not_mock=lambda method, url: False, fail_on_unknown_request=True, save_to_history=False):
     """
     Creates a WireProtocol object that handles requests to the WireServer, the Host GA Plugin, and some requests to storage (requests that provide mock data
     in wire_protocol_data.py).
@@ -38,6 +38,8 @@ def mock_wire_protocol(mock_wire_data_file, http_get_handler=None, http_post_han
     The 'do_not_mock' lambda can be used to skip the mocks for specific requests; if the lambda returns True, the mocks won't be applied and the
     original common.utils.restutil.http_request will be invoked instead.
 
+    The 'save_to_history' parameter is passed thru in the call to WireProtocol.detect().
+
     The returned protocol object maintains a list of "tracked" urls. When a handler function returns a value than is not None the url for the
     request is automatically added to the tracked list. The handler function can add other items to this list using the track_url() method on
     the mock.
@@ -147,7 +149,7 @@ def mock_wire_protocol(mock_wire_data_file, http_get_handler=None, http_post_han
     # go do it
     try:
         protocol.start()
-        protocol.detect()
+        protocol.detect(save_to_history=save_to_history)
         yield protocol
     finally:
         protocol.stop()