sources/azure: refactor _should_reprovision[_after_nic_attach]() logic

cloudinit/sources/DataSourceAzure.py

@@ -55,7 +55,6 @@
 # DMI chassis-asset-tag is set static for all azure instances
 AZURE_CHASSIS_ASSET_TAG = "7783-7084-3265-9085-8269-3286-77"
 REPROVISION_MARKER_FILE = "/var/lib/cloud/data/poll_imds"
-REPROVISION_NIC_ATTACH_MARKER_FILE = "/var/lib/cloud/data/wait_for_nic_attach"
 REPROVISION_NIC_DETACHED_MARKER_FILE = "/var/lib/cloud/data/nic_detached"
 REPORTED_READY_MARKER_FILE = "/var/lib/cloud/data/reported_ready"
 AGENT_SEED_DIR = "/var/lib/waagent"
@@ -80,6 +79,13 @@ class MetadataType(Enum):
     REPROVISION_DATA = "{}/reprovisiondata".format(IMDS_URL)
 
 
+class PPSType(Enum):
+    NONE = "None"
+    RUNNING = "Running"
+    SAVABLE = "Savable"
+    UNKNOWN = "Unknown"
+
+
 PLATFORM_ENTROPY_SOURCE = "/sys/firmware/acpi/tables/OEM0"
 
 # List of static scripts and network config artifacts created by
@@ -337,31 +343,19 @@ def crawl_metadata(self):
         # it determines the value of ret. More specifically, the first one in
         # the candidate list determines the path to take in order to get the
         # metadata we need.
-        reprovision = False
         ovf_is_accessible = False
-        reprovision_after_nic_attach = False
         metadata_source = None
         md = {}
         userdata_raw = ""
         cfg = {}
         files = {}
         if os.path.isfile(REPROVISION_MARKER_FILE):
-            reprovision = True
             metadata_source = "IMDS"
             report_diagnostic_event(
                 "Reprovision marker file already present "
                 "before crawling Azure metadata: %s" % REPROVISION_MARKER_FILE,
                 logger_func=LOG.debug,
             )
-        elif os.path.isfile(REPROVISION_NIC_ATTACH_MARKER_FILE):
-            reprovision_after_nic_attach = True
-            metadata_source = "NIC_ATTACH_MARKER_PRESENT"
-            report_diagnostic_event(
-                "Reprovision nic attach marker file "
-                "already present before crawling Azure "
-                "metadata: %s" % REPROVISION_NIC_ATTACH_MARKER_FILE,
-                logger_func=LOG.debug,
-            )
         else:
             for src in list_possible_azure_ds(self.seed_dir, ddir):
                 try:
@@ -418,21 +412,18 @@ def crawl_metadata(self):
             report_diagnostic_event(msg)
             raise sources.InvalidMetaDataException(msg)
 
-        perform_reprovision = reprovision or self._should_reprovision(
-            cfg, imds_md
-        )
-        perform_reprovision_after_nic_attach = (
-            reprovision_after_nic_attach
-            or self._should_reprovision_after_nic_attach(cfg, imds_md)
-        )
-
-        if perform_reprovision or perform_reprovision_after_nic_attach:
+        pps_type = self._determine_pps_type(cfg, imds_md)
+        if pps_type != PPSType.NONE:
             if util.is_FreeBSD():
                 msg = "Free BSD is not supported for PPS VMs"
                 report_diagnostic_event(msg, logger_func=LOG.error)
                 raise sources.InvalidMetaDataException(msg)
-            if perform_reprovision_after_nic_attach:
+
+            self._write_reprovision_marker()
+
+            if pps_type == PPSType.SAVABLE:
                 self._wait_for_all_nics_ready()
+
             md, userdata_raw, cfg, files = self._reprovision()
             # fetch metadata again as it has changed after reprovisioning
             imds_md = self.get_imds_data_with_api_fallback(
@@ -515,7 +506,7 @@ def crawl_metadata(self):
             crawled_data["metadata"]["random_seed"] = seed
         crawled_data["metadata"]["instance-id"] = self._iid()
 
-        if perform_reprovision or perform_reprovision_after_nic_attach:
+        if pps_type != PPSType.NONE:
             LOG.info("Reporting ready to Azure after getting ReprovisionData")
             use_cached_ephemeral = (
                 self.distro.networking.is_up(self.fallback_interface)
@@ -1398,67 +1389,39 @@ def _ppstype_from_imds(self, imds_md: dict = None) -> str:
             )
             return None
 
-    def _should_reprovision_after_nic_attach(
-        self, cfg: dict, imds_md=None
-    ) -> bool:
-        """Whether or not we should wait for nic attach and then poll
-        IMDS for reprovisioning data. Also sets a marker file to poll IMDS.
-
-        The marker file is used for the following scenario: the VM boots into
-        wait for nic attach, which we expect to be proceeding infinitely until
-        the nic is attached. If for whatever reason the platform moves us to a
-        new host (for instance a hardware issue), we need to keep waiting.
-        However, since the VM reports ready to the Fabric, we will not attach
-        the ISO, thus cloud-init needs to have a way of knowing that it should
-        jump back into the waiting mode in order to retrieve the ovf_env.
-
-        @param cfg: OVF cfg.
-        @param imds_md: Metadata obtained from IMDS
-        @return: Whether to reprovision after waiting for nics to be attached.
-        """
-        path = REPROVISION_NIC_ATTACH_MARKER_FILE
-        if (
-            cfg.get("PreprovisionedVMType", None) == "Savable"
-            or self._ppstype_from_imds(imds_md) == "Savable"
-            or os.path.isfile(path)
+    def _determine_pps_type(self, ovf_cfg: dict, imds_md: dict) -> PPSType:
+        """Determine PPS type using OVF, IMDS data, and reprovision marker."""
+        if os.path.isfile(REPROVISION_MARKER_FILE):
+            pps_type = PPSType.UNKNOWN
+        elif (
+            ovf_cfg.get("PreprovisionedVMType", None) == PPSType.SAVABLE.value
+            or self._ppstype_from_imds(imds_md) == PPSType.SAVABLE.value
         ):
-            if not os.path.isfile(path):
-                LOG.info(
-                    "Creating a marker file to wait for nic attach: %s", path
-                )
-                util.write_file(
-                    path,
-                    "{pid}: {time}\n".format(pid=os.getpid(), time=time()),
-                )
-            return True
-        return False
-
-    def _should_reprovision(self, cfg: dict, imds_md=None):
-        """Whether or not we should poll IMDS for reprovisioning data.
-        Also sets a marker file to poll IMDS.
-
-        The marker file is used for the following scenario: the VM boots into
-        this polling loop, which we expect to be proceeding infinitely until
-        the VM is picked. If for whatever reason the platform moves us to a
-        new host (for instance a hardware issue), we need to keep polling.
-        However, since the VM reports ready to the Fabric, we will not attach
-        the ISO, thus cloud-init needs to have a way of knowing that it should
-        jump back into the polling loop in order to retrieve the ovf_env."""
-        path = REPROVISION_MARKER_FILE
-        if (
-            cfg.get("PreprovisionedVm") is True
-            or cfg.get("PreprovisionedVMType", None) == "Running"
-            or self._ppstype_from_imds(imds_md) == "Running"
-            or os.path.isfile(path)
+            pps_type = PPSType.SAVABLE
+        elif (
+            ovf_cfg.get("PreprovisionedVm") is True
+            or ovf_cfg.get("PreprovisionedVMType", None)
+            == PPSType.RUNNING.value
+            or self._ppstype_from_imds(imds_md) == PPSType.RUNNING.value
         ):
-            if not os.path.isfile(path):
-                LOG.info("Creating a marker file to poll imds: %s", path)
-                util.write_file(
-                    path,
-                    "{pid}: {time}\n".format(pid=os.getpid(), time=time()),
-                )
-            return True
-        return False
+            pps_type = PPSType.RUNNING
+        else:
+            pps_type = PPSType.NONE
+
+        report_diagnostic_event(
+            "PPS type: %s" % pps_type.value, logger_func=LOG.info
+        )
+        return pps_type
+
+    def _write_reprovision_marker(self):
+        """Write reprovision marker file in case system is rebooted."""
+        LOG.info(
+            "Creating a marker file to poll imds: %s", REPROVISION_MARKER_FILE
+        )
+        util.write_file(
+            REPROVISION_MARKER_FILE,
+            "{pid}: {time}\n".format(pid=os.getpid(), time=time()),
+        )
 
     def _reprovision(self):
         """Initiate the reprovisioning workflow."""
@@ -1507,7 +1470,6 @@ def _negotiate(self):
 
         util.del_file(REPORTED_READY_MARKER_FILE)
         util.del_file(REPROVISION_MARKER_FILE)
-        util.del_file(REPROVISION_NIC_ATTACH_MARKER_FILE)
         util.del_file(REPROVISION_NIC_DETACHED_MARKER_FILE)
         return fabric_data