diff --git a/daemon/registry.py b/daemon/registry.py index 84f4a56..e05840e 100644 --- a/daemon/registry.py +++ b/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, diff --git a/tests/test-state-persistence.py b/tests/test-state-persistence.py new file mode 100755 index 0000000..0dc6de5 --- /dev/null +++ b/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")