From 79668d45288e7a17bc627fdc05a609423161a44c Mon Sep 17 00:00:00 2001
From: PengpengSun <40026211+PengpengSun@users.noreply.github.com>
Date: Fri, 29 Mar 2024 22:39:13 +0800
Subject: [PATCH 2/2] fix: Fall back to cached local ds if no valid ds found
 (#4997)

Rebooting an instance which has finished VMware guest
customization with DataSourceVMware will load
DataSourceNone due to metadata is NOT available.

This is mostly a re-post of PR#229, few differences are:
1. Let ds decide if fallback is allowed, not always fall back
   to previous cached LOCAL ds.
2. No comparing instance-id of cached ds with previous instance-id
   due to I think they are always identical.

Fixes GH-3402

(cherry picked from commit 9929a00580d50afc60bf4e0fb9f2f39d4f797b4b)
Signed-off-by: Ani Sinha <anisinha@redhat.com>

Conflicts:
	cloudinit/sources/__init__.py
	Conflicts because of changes in upstream source coming from
	30d5e9a3358f4cbaced ("refactor: Use _unpickle rather than hasattr() in sources")
---
 cloudinit/sources/DataSourceVMware.py | 14 +++++++++-
 cloudinit/sources/__init__.py         | 14 ++++++++++
 cloudinit/stages.py                   | 40 +++++++++++++++++----------
 3 files changed, 53 insertions(+), 15 deletions(-)

diff --git a/cloudinit/sources/DataSourceVMware.py b/cloudinit/sources/DataSourceVMware.py
index 77a2de6c..2a91a307 100644
--- a/cloudinit/sources/DataSourceVMware.py
+++ b/cloudinit/sources/DataSourceVMware.py
@@ -223,7 +223,7 @@ class DataSourceVMware(sources.DataSource):
                 break
 
         if not self.data_access_method:
-            LOG.error("failed to find a valid data access method")
+            LOG.debug("failed to find a valid data access method")
             return False
 
         LOG.info("using data access method %s", self._get_subplatform())
@@ -317,6 +317,18 @@ class DataSourceVMware(sources.DataSource):
             self.metadata["instance-id"] = str(id_file.read()).rstrip().lower()
             return self.metadata["instance-id"]
 
+    def check_if_fallback_is_allowed(self):
+        if (
+            self.data_access_method
+            and self.data_access_method == DATA_ACCESS_METHOD_IMC
+            and is_vmware_platform()
+        ):
+            LOG.debug(
+                "Cache fallback is allowed for : %s", self._get_subplatform()
+            )
+            return True
+        return False
+
     def get_public_ssh_keys(self):
         for key_name in (
             "public-keys-data",
diff --git a/cloudinit/sources/__init__.py b/cloudinit/sources/__init__.py
index 4ea1fc56..e2b3029e 100644
--- a/cloudinit/sources/__init__.py
+++ b/cloudinit/sources/__init__.py
@@ -332,6 +332,10 @@ class DataSource(CloudInitPickleMixin, metaclass=abc.ABCMeta):
             self.vendordata2_raw = None
         if not hasattr(self, "skip_hotplug_detect"):
             self.skip_hotplug_detect = False
+
+        if not hasattr(self, "check_if_fallback_is_allowed"):
+            setattr(self, "check_if_fallback_is_allowed", lambda: False)
+
         if hasattr(self, "userdata") and self.userdata is not None:
             # If userdata stores MIME data, on < python3.6 it will be
             # missing the 'policy' attribute that exists on >=python3.6.
@@ -925,6 +929,16 @@ class DataSource(CloudInitPickleMixin, metaclass=abc.ABCMeta):
         # quickly (local check only) if self.instance_id is still
         return False
 
+    def check_if_fallback_is_allowed(self):
+        """check_if_fallback_is_allowed()
+        Checks if a cached ds is allowed to be restored when no valid ds is
+        found in local mode by checking instance-id and searching valid data
+        through ds list.
+
+        @return True if a ds allows fallback, False otherwise.
+        """
+        return False
+
     @staticmethod
     def _determine_dsmode(candidates, default=None, valid=None):
         # return the first candidate that is non None, warn if not valid
diff --git a/cloudinit/stages.py b/cloudinit/stages.py
index e95bb76f..894eeac5 100644
--- a/cloudinit/stages.py
+++ b/cloudinit/stages.py
@@ -380,20 +380,32 @@ class Init:
             LOG.debug(myrep.description)
 
         if not ds:
-            util.del_file(self.paths.instance_link)
-            (cfg_list, pkg_list) = self._get_datasources()
-            # Deep copy so that user-data handlers can not modify
-            # (which will affect user-data handlers down the line...)
-            (ds, dsname) = sources.find_source(
-                self.cfg,
-                self.distro,
-                self.paths,
-                copy.deepcopy(self.ds_deps),
-                cfg_list,
-                pkg_list,
-                self.reporter,
-            )
-            LOG.info("Loaded datasource %s - %s", dsname, ds)
+            try:
+                cfg_list, pkg_list = self._get_datasources()
+                # Deep copy so that user-data handlers can not modify
+                # (which will affect user-data handlers down the line...)
+                ds, dsname = sources.find_source(
+                    self.cfg,
+                    self.distro,
+                    self.paths,
+                    copy.deepcopy(self.ds_deps),
+                    cfg_list,
+                    pkg_list,
+                    self.reporter,
+                )
+                util.del_file(self.paths.instance_link)
+                LOG.info("Loaded datasource %s - %s", dsname, ds)
+            except sources.DataSourceNotFoundException as e:
+                if existing != "check":
+                    raise e
+                ds = self._restore_from_cache()
+                if ds and ds.check_if_fallback_is_allowed():
+                    LOG.info(
+                        "Restored fallback datasource from checked cache: %s",
+                        ds,
+                    )
+                else:
+                    raise e
         self.datasource = ds
         # Ensure we adjust our path members datasource
         # now that we have one (thus allowing ipath to be used)
-- 
2.39.3