Persist daemon resource runtime state

daemon/registry.py

@@ -1,9 +1,12 @@
 #!/usr/bin/env python3
 
+import copy
+import hashlib
 import importlib.util
 import inspect
 import json
 import os
+import tempfile
 from pathlib import Path
 
 
@@ -122,10 +125,21 @@ class Plugin:
 
 
 class Resource:
-    def __init__(self, resource_id, plugin, config):
+    def __init__(
+        self,
+        resource_id,
+        plugin,
+        config,
+        state_path=None,
+    ):
         self.resource_id = resource_id
         self.plugin = plugin
         self.config = config
+        self.state_path = (
+            Path(state_path)
+            if state_path is not None
+            else None
+        )
 
         # Generic daemon-lifetime state storage. The daemon does not interpret
         # plugin keys or values; applications may use this for transition
@@ -133,6 +147,101 @@ class Resource:
         self.runtime_state = {}
 
         self.plugin.validate_config(config)
+        self.load_state()
+
+    def load_state(self):
+        if self.state_path is None:
+            return
+
+        if not self.state_path.exists():
+            return
+
+        try:
+            with self.state_path.open(
+                "r",
+                encoding="utf-8",
+            ) as stream:
+                state = json.load(stream)
+
+        except (OSError, json.JSONDecodeError) as exc:
+            raise RegistryError(
+                f"unable to load resource state: {self.resource_id}"
+            ) from exc
+
+        if not isinstance(state, dict):
+            raise RegistryError(
+                f"resource state is not an object: {self.resource_id}"
+            )
+
+        self.runtime_state = state
+
+        # Monotonic timestamps cannot survive a daemon restart.
+        postgresql_state = self.runtime_state.get("postgresql")
+
+        if isinstance(postgresql_state, dict):
+            postgresql_state.pop(
+                "reconnect_started_at",
+                None,
+            )
+            postgresql_state.pop(
+                "state_changed_at",
+                None,
+            )
+
+    def persistent_state(self):
+        state = copy.deepcopy(self.runtime_state)
+        postgresql_state = state.get("postgresql")
+
+        if isinstance(postgresql_state, dict):
+            postgresql_state.pop(
+                "reconnect_started_at",
+                None,
+            )
+            postgresql_state.pop(
+                "state_changed_at",
+                None,
+            )
+
+        return state
+
+    def save_state(self):
+        if self.state_path is None:
+            return
+
+        self.state_path.parent.mkdir(
+            parents=True,
+            exist_ok=True,
+        )
+
+        temporary_path = None
+
+        try:
+            with tempfile.NamedTemporaryFile(
+                mode="w",
+                encoding="utf-8",
+                dir=self.state_path.parent,
+                prefix=f".{self.state_path.name}.",
+                delete=False,
+            ) as stream:
+                temporary_path = Path(stream.name)
+                json.dump(
+                    self.persistent_state(),
+                    stream,
+                    separators=(",", ":"),
+                    sort_keys=True,
+                )
+                stream.write("\n")
+
+            os.chmod(temporary_path, 0o600)
+            os.replace(temporary_path, self.state_path)
+
+        except OSError as exc:
+            if temporary_path is not None:
+                temporary_path.unlink(missing_ok=True)
+
+            raise RegistryError(
+                f"unable to save resource state: {self.resource_id}"
+            ) from exc
 
     def describe(self):
         return {
@@ -154,6 +263,9 @@ class Registry:
         )
 
         self.resources = {}
+        self.state_dir = os.environ.get(
+            "MANAGED_CONTAINERD_STATE_DIR"
+        )
         self._load_environment_resource()
 
     def _load_environment_resource(self):
@@ -180,10 +292,21 @@ class Registry:
             for key in plugin.config_keys
         }
 
+        state_path = None
+
+        if self.state_dir:
+            state_name = hashlib.sha256(
+                resource_id.encode("utf-8")
+            ).hexdigest() + ".json"
+            state_path = (
+                Path(self.state_dir) / state_name
+            )
+
         self.resources[resource_id] = Resource(
             resource_id,
             plugin,
             config,
+            state_path,
         )
 
     def list_resources(self):
@@ -213,6 +336,8 @@ class Registry:
             request_parameters,
         )
 
+        resource.save_state()
+
         return {
             "ok": True,
             "resource": resource_id,

tests/test-state-persistence.py

@@ -0,0 +1,122 @@
+#!/usr/bin/env python3
+
+import json
+import os
+import tempfile
+from pathlib import Path
+from unittest.mock import patch
+
+from registry import Registry
+
+
+PLUGIN_SOURCE = r"""
+def notify(config, runtime_state, request_parameters):
+    runtime_state["postgresql"] = {
+        "notification": dict(request_parameters),
+        "reconnect_started_at": 123.4,
+        "state_changed_at": 567.8,
+    }
+
+    return {
+        "ocf_rc": 0,
+    }
+"""
+
+
+def require(condition, message):
+    if not condition:
+        raise SystemExit("FAIL: " + message)
+
+
+print("===== DAEMON STATE PERSISTENCE =====")
+
+with tempfile.TemporaryDirectory() as temporary_root:
+    root = Path(temporary_root)
+    plugin_root = root / "plugins" / "fake"
+    plugin_root.mkdir(parents=True)
+
+    (plugin_root / "manifest.json").write_text(
+        json.dumps({
+            "name": "fake",
+            "version": 1,
+            "module": "plugin.py",
+            "actions": ["notify"],
+            "required_config": [],
+        })
+    )
+    (plugin_root / "plugin.py").write_text(PLUGIN_SOURCE)
+
+    state_dir = root / "state"
+
+    environment = {
+        "MANAGED_CONTAINERD_PLUGIN_ROOT": str(
+            root / "plugins"
+        ),
+        "MANAGED_CONTAINERD_STATE_DIR": str(state_dir),
+        "RESOURCE_ID": "fake-resource",
+        "PLUGIN": "fake",
+    }
+
+    parameters = {
+        "type": "pre",
+        "operation": "demote",
+        "promote_uname": "david-eve",
+    }
+
+    with patch.dict(os.environ, environment, clear=False):
+        first = Registry()
+        result = first.execute(
+            "fake-resource",
+            "notify",
+            parameters,
+        )
+
+        require(
+            result["result"]["ocf_rc"] == 0,
+            "initial notify failed",
+        )
+
+        files = list(state_dir.glob("*.json"))
+        require(
+            len(files) == 1,
+            "state file was not created",
+        )
+
+        mode = files[0].stat().st_mode & 0o777
+        require(
+            mode == 0o600,
+            "state file is not mode 0600",
+        )
+
+        saved = json.loads(files[0].read_text())
+        require(
+            saved["postgresql"]["notification"]
+            == parameters,
+            "notification was not persisted",
+        )
+        require(
+            "reconnect_started_at"
+            not in saved["postgresql"],
+            "volatile reconnect state was persisted",
+        )
+        require(
+            "state_changed_at"
+            not in saved["postgresql"],
+            "volatile monotonic state was persisted",
+        )
+        print("atomic_state_write      persisted")
+
+        second = Registry()
+        restored = second.resources[
+            "fake-resource"
+        ].runtime_state
+
+        require(
+            restored["postgresql"]["notification"]
+            == parameters,
+            "notification was not restored",
+        )
+        print("restart_state_load       restored")
+
+print()
+print("STATE_PERSISTENCE_TEST_PASS=1")