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,