sources: obj.pkl cache should be written anyime get_data is run (#1669)

When metadata update events trigger a new datasource.get_data run
ensure we are syncing the cached obj.pkl to disk so subsequent
boot stages can leverage the updated metadata.

Add write_cache param to persist_instance_data to avoid
persisting instance data when init.ds_restored from cache.

This avoids a race on clouds where network config is updated
per boot in init-local timeframe but init-network uses stale network
metadata from cache because updated metadata was not
persisted.

Migate _pkl_load and _pkl_store out of stages module and into
sources as it really is only applicable to datasource serialization.

5 files changed, 98 insertions, 45 deletions

diff --git a/cloudinit/cmd/main.py b/cloudinit/cmd/main.py
index 2860126a..6134d7c4 100755
--- a/cloudinit/cmd/main.py
+++ b/cloudinit/cmd/main.py
@@ -809,7 +809,7 @@ def _maybe_persist_instance_data(init):
             init.paths.run_dir, sources.INSTANCE_JSON_FILE
         )
         if not os.path.exists(instance_data_file):
-            init.datasource.persist_instance_data()
+            init.datasource.persist_instance_data(write_cache=False)
 
 
 def _maybe_set_hostname(init, stage, retry_stage):
diff --git a/cloudinit/sources/__init__.py b/cloudinit/sources/__init__.py
index b621fb6e..c399beb6 100644
--- a/cloudinit/sources/__init__.py
+++ b/cloudinit/sources/__init__.py
@@ -12,9 +12,10 @@ import abc
 import copy
 import json
 import os
+import pickle
 from collections import namedtuple
 from enum import Enum, unique
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 from cloudinit import dmi, importer
 from cloudinit import log as logging
@@ -373,14 +374,19 @@ class DataSource(CloudInitPickleMixin, metaclass=abc.ABCMeta):
         self.persist_instance_data()
         return return_value
 
-    def persist_instance_data(self):
+    def persist_instance_data(self, write_cache=True):
         """Process and write INSTANCE_JSON_FILE with all instance metadata.
 
         Replace any hyphens with underscores in key names for use in template
         processing.
 
+        :param write_cache: boolean set True to persist obj.pkl when
+            instance_link exists.
+
         @return True on successful write, False otherwise.
         """
+        if write_cache and os.path.lexists(self.paths.instance_link):
+            pkl_store(self, self.paths.get_ipath_cur("obj_pkl"))
         if hasattr(self, "_crawled_metadata"):
             # Any datasource with _crawled_metadata will best represent
             # most recent, 'raw' metadata
@@ -1063,4 +1069,43 @@ def list_from_depends(depends, ds_list):
     return ret_list
 
 
+def pkl_store(obj: DataSource, fname: str) -> bool:
+    """Use pickle to serialize Datasource to a file as a cache.
+
+    :return: True on success
+    """
+    try:
+        pk_contents = pickle.dumps(obj)
+    except Exception:
+        util.logexc(LOG, "Failed pickling datasource %s", obj)
+        return False
+    try:
+        util.write_file(fname, pk_contents, omode="wb", mode=0o400)
+    except Exception:
+        util.logexc(LOG, "Failed pickling datasource to %s", fname)
+        return False
+    return True
+
+
+def pkl_load(fname: str) -> Optional[DataSource]:
+    """Use pickle to deserialize a instance Datasource from a cache file."""
+    pickle_contents = None
+    try:
+        pickle_contents = util.load_file(fname, decode=False)
+    except Exception as e:
+        if os.path.isfile(fname):
+            LOG.warning("failed loading pickle in %s: %s", fname, e)
+
+    # This is allowed so just return nothing successfully loaded...
+    if not pickle_contents:
+        return None
+    try:
+        return pickle.loads(pickle_contents)
+    except DatasourceUnpickleUserDataError:
+        return None
+    except Exception:
+        util.logexc(LOG, "Failed loading pickled blob from %s", fname)
+        return None
+
+
 # vi: ts=4 expandtab
diff --git a/cloudinit/stages.py b/cloudinit/stages.py
index 66e12eed..132dd83b 100644
--- a/cloudinit/stages.py
+++ b/cloudinit/stages.py
@@ -6,7 +6,6 @@
 
 import copy
 import os
-import pickle
 import sys
 from collections import namedtuple
 from typing import Dict, Iterable, List, Optional, Set
@@ -247,7 +246,7 @@ class Init(object):
         # We try to restore from a current link and static path
         # by using the instance link, if purge_cache was called
         # the file wont exist.
-        return _pkl_load(self.paths.get_ipath_cur("obj_pkl"))
+        return sources.pkl_load(self.paths.get_ipath_cur("obj_pkl"))
 
     def _write_to_cache(self):
         if self.datasource is None:
@@ -260,7 +259,9 @@ class Init(object):
                 omode="w",
                 content="",
             )
-        return _pkl_store(self.datasource, self.paths.get_ipath_cur("obj_pkl"))
+        return sources.pkl_store(
+            self.datasource, self.paths.get_ipath_cur("obj_pkl")
+        )
 
     def _get_datasources(self):
         # Any config provided???
@@ -973,38 +974,4 @@ def fetch_base_config():
     )
 
 
-def _pkl_store(obj, fname):
-    try:
-        pk_contents = pickle.dumps(obj)
-    except Exception:
-        util.logexc(LOG, "Failed pickling datasource %s", obj)
-        return False
-    try:
-        util.write_file(fname, pk_contents, omode="wb", mode=0o400)
-    except Exception:
-        util.logexc(LOG, "Failed pickling datasource to %s", fname)
-        return False
-    return True
-
-
-def _pkl_load(fname):
-    pickle_contents = None
-    try:
-        pickle_contents = util.load_file(fname, decode=False)
-    except Exception as e:
-        if os.path.isfile(fname):
-            LOG.warning("failed loading pickle in %s: %s", fname, e)
-
-    # This is allowed so just return nothing successfully loaded...
-    if not pickle_contents:
-        return None
-    try:
-        return pickle.loads(pickle_contents)
-    except sources.DatasourceUnpickleUserDataError:
-        return None
-    except Exception:
-        util.logexc(LOG, "Failed loading pickled blob from %s", fname)
-        return None
-
-
 # vi: ts=4 expandtab
diff --git a/tests/unittests/sources/test_init.py b/tests/unittests/sources/test_init.py
index a42c6a72..52f6cbfc 100644
--- a/tests/unittests/sources/test_init.py
+++ b/tests/unittests/sources/test_init.py
@@ -17,6 +17,7 @@ from cloudinit.sources import (
     UNSET,
     DataSource,
     canonical_cloud_id,
+    pkl_load,
     redact_sensitive_keys,
 )
 from cloudinit.user_data import UserDataProcessor
@@ -672,8 +673,12 @@ class TestDataSource(CiTestCase):
     def test_persist_instance_data_writes_ec2_metadata_when_set(self):
         """When ec2_metadata class attribute is set, persist to json."""
         tmp = self.tmp_dir()
+        cloud_dir = os.path.join(tmp, "cloud")
+        util.ensure_dir(cloud_dir)
         datasource = DataSourceTestSubclassNet(
-            self.sys_cfg, self.distro, Paths({"run_dir": tmp})
+            self.sys_cfg,
+            self.distro,
+            Paths({"run_dir": tmp, "cloud_dir": cloud_dir}),
         )
         datasource.ec2_metadata = UNSET
         datasource.get_data()
@@ -690,8 +695,12 @@ class TestDataSource(CiTestCase):
     def test_persist_instance_data_writes_canonical_cloud_id_and_symlink(self):
         """canonical-cloud-id class attribute is set, persist to json."""
         tmp = self.tmp_dir()
+        cloud_dir = os.path.join(tmp, "cloud")
+        util.ensure_dir(cloud_dir)
         datasource = DataSourceTestSubclassNet(
-            self.sys_cfg, self.distro, Paths({"run_dir": tmp})
+            self.sys_cfg,
+            self.distro,
+            Paths({"run_dir": tmp, "cloud_dir": cloud_dir}),
         )
         cloud_id_link = os.path.join(tmp, "cloud-id")
         cloud_id_file = os.path.join(tmp, "cloud-id-my-cloud")
@@ -722,8 +731,12 @@ class TestDataSource(CiTestCase):
     def test_persist_instance_data_writes_network_json_when_set(self):
         """When network_data.json class attribute is set, persist to json."""
         tmp = self.tmp_dir()
+        cloud_dir = os.path.join(tmp, "cloud")
+        util.ensure_dir(cloud_dir)
         datasource = DataSourceTestSubclassNet(
-            self.sys_cfg, self.distro, Paths({"run_dir": tmp})
+            self.sys_cfg,
+            self.distro,
+            Paths({"run_dir": tmp, "cloud_dir": cloud_dir}),
         )
         datasource.get_data()
         json_file = self.tmp_path(INSTANCE_JSON_FILE, tmp)
@@ -736,6 +749,34 @@ class TestDataSource(CiTestCase):
             {"network_json": "is good"}, instance_data["ds"]["network_json"]
         )
 
+    def test_persist_instance_serializes_datasource_pickle(self):
+        """obj.pkl is written when instance link present and write_cache."""
+        tmp = self.tmp_dir()
+        cloud_dir = os.path.join(tmp, "cloud")
+        util.ensure_dir(cloud_dir)
+        datasource = DataSourceTestSubclassNet(
+            self.sys_cfg,
+            self.distro,
+            Paths({"run_dir": tmp, "cloud_dir": cloud_dir}),
+        )
+        pkl_cache_file = os.path.join(cloud_dir, "instance/obj.pkl")
+        self.assertFalse(os.path.exists(pkl_cache_file))
+        datasource.network_json = {"network_json": "is good"}
+        # No /var/lib/cloud/instance symlink
+        datasource.persist_instance_data(write_cache=True)
+        self.assertFalse(os.path.exists(pkl_cache_file))
+
+        # Symlink /var/lib/cloud/instance but write_cache=False
+        util.sym_link(cloud_dir, os.path.join(cloud_dir, "instance"))
+        datasource.persist_instance_data(write_cache=False)
+        self.assertFalse(os.path.exists(pkl_cache_file))
+
+        # Symlink /var/lib/cloud/instance and write_cache=True
+        datasource.persist_instance_data(write_cache=True)
+        self.assertTrue(os.path.exists(pkl_cache_file))
+        ds = pkl_load(pkl_cache_file)
+        self.assertEqual(datasource.network_json, ds.network_json)
+
     def test_get_data_base64encodes_unserializable_bytes(self):
         """On py3, get_data base64encodes any unserializable content."""
         tmp = self.tmp_dir()
diff --git a/tests/unittests/test_upgrade.py b/tests/unittests/test_upgrade.py
index d7a721a2..ed3c7efb 100644
--- a/tests/unittests/test_upgrade.py
+++ b/tests/unittests/test_upgrade.py
@@ -18,7 +18,7 @@ import pathlib
 
 import pytest
 
-from cloudinit.stages import _pkl_load
+from cloudinit.sources import pkl_load
 from tests.unittests.helpers import resourceLocation
 
 
@@ -34,7 +34,7 @@ class TestUpgrade:
         Test implementations _must not_ modify the ``previous_obj_pkl`` which
         they are passed, as that will affect tests that run after them.
         """
-        return _pkl_load(str(request.param))
+        return pkl_load(str(request.param))
 
     def test_networking_set_on_distro(self, previous_obj_pkl):
         """We always expect to have ``.networking`` on ``Distro`` objects."""