Save goal state history explicitly

azurelinuxagent/common/protocol/goal_state.py

@@ -70,7 +70,7 @@ def __init__(self, msg, inner=None):
 
 
 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 @@ def __init__(self, wire_client, goal_state_properties=GoalStateProperties.All, s
         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 @@ def _fetch_manifest(self, manifest_type, name, uris):
         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 @@ def _update(self, force_update):
 
         # 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 @@ def _download_certificates(self, certs_uri):
         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 @@ def _restore_wire_server_goal_state(self, incarnation, xml_text, xml_doc, vm_set
         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 @@ def _restore_wire_server_goal_state(self, incarnation, xml_text, xml_doc, vm_set
             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 @@ def _fetch_full_wire_server_goal_state(self, incarnation, xml_doc):
             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 @@ def _fetch_full_wire_server_goal_state(self, incarnation, xml_doc):
                 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

azurelinuxagent/common/protocol/util.py

@@ -188,7 +188,7 @@ def _clear_wireserver_endpoint(self):
                 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 @@ def _detect_protocol(self, init_goal_state=True):
 
                 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 @@ def clear_protocol(self):
         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 @@ def get_protocol(self, init_goal_state=True):
 
             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)

azurelinuxagent/common/protocol/wire.py

@@ -73,7 +73,7 @@ def __init__(self, endpoint):
             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 @@ def detect(self, init_goal_state=True):
         # 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 @@ def update_host_plugin(self, container_id, role_config_name):
             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,15 +792,15 @@ def update_goal_state(self, silent=False):
         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
         """
         try:
             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 @@ def upload_status_blob(self):
 
         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:

azurelinuxagent/ga/update.py

@@ -339,7 +339,7 @@ def run(self, debug=False):
             # 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 @@ def _try_update_goal_state(self, protocol):
         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()
 

tests/common/protocol/test_goal_state.py

@@ -112,14 +112,18 @@ def test_fetching_the_goal_state_should_save_the_goal_state_to_the_history_direc
             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 @@ def test_update_should_create_new_history_subdirectories(self):
             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 @@ def test_it_should_redact_the_protected_settings_when_saving_to_the_history_dire
             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 @@ def test_it_should_save_vm_settings_on_parse_errors(self):
 
             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():

tests/common/protocol/test_protocol_util.py

@@ -127,14 +127,14 @@ def test_detect_protocol_no_dhcp(self, WireProtocol, mock_get_lib_dir, _):
         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, _):

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 create_vm_status(extension_status):
         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 @@ def test_it_should_process_goal_state_only_on_new_goal_state(self):
             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()

tests/lib/mock_wire_protocol.py

@@ -22,7 +22,7 @@
 
 
 @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 stop():
     # go do it
     try:
         protocol.start()
-        protocol.detect()
+        protocol.detect(save_to_history=save_to_history)
         yield protocol
     finally:
         protocol.stop()