Repository: canonical/operator

Issue:
No interface to `juju-reboot` hook tool in ops
Hello there! I am currently working through a predicament in the [SLURM charms](https://charmhub.io/?q=slurm) where the machines need to be rebooted after package installation to apply security updates. Originally, the SLURM charms were deployed manually so the human operator would issue a machine reboot command based on the substrate being used (e.g. `openstack server reboot <machine_id>`). We are now working to make deploying SLURM an automatic process, however, we still have the issue where the machines need to be rebooted to apply updates.

My idea was to use the `juju-reboot` hook tool located in `/var/lib/juju/tools/machine-<#>/`, however, it does not seem that ops has a direct method for invoking `juju-reboot`. It would be nice if ops had a defined method for scheduling machine reboots at the end of hook execution:

```python3
# Assume charm is fully defined

def _on_install(self, _: InstallEvent) -> None:
    # Install some packages, update base image, etc.
    ...

    if pathlib.Path("/var/run/reboot-required").exists():
        logger.info("Scheduling machine for reboot...")
        self.unit.reboot()
```

Currently I can circumvent my reboot issue by directly invoking the `juju-reboot` hook tool, but I thought that this might be a nice edition to ops. Let me know what your thoughts are!

Discussion:
benhoyt: Oh, looks like `tox -e docs` is failing right now.
jameinel: > When it's run for a K8s charm, Juju logs an error message but doesn't return an response with the error (or have a non-zero exit), which means that we have to fail silently.

If we can tell that you're in k8s, can't we fail inside `ops` itself? 
I do think Juju should be failing more appropriately (and I don't think there is a reason why we couldn't force the pod to terminate and restart). I don't think people have been asking for that ability, but it would be more consistent to have it do something.

benhoyt: > If we can tell that you're in k8s, can't we fail inside ops itself?

We don't have an easy way to do that right now. You can look for `containers` in `metadata.yaml`, but charms with no containers can be deployed on K8s too. Harry [mentioned](https://chat.charmhub.io/charmhub/pl/wkxduuqszt8cjdcxq7ihoe8mwr) that `JUJU_CONTAINER_NAMES` will be set on K8s (it may be empty, but it'll be set).

So maybe it's time to add a helper method `is_k8s_sidecar`. We've needed this before too. Or maybe `charm.type` (though something more specific than "type") or something which is an enum, one of `machine`, `k8s-podspec`, `k8s-sidecar`.
tonyandrewmeyer: > I think just having reboot do something on k8s is also fine, and we wouldn't want to work around the workaround in the future.

If Juju adds support for doing something to k8s when `juju-reboot` is run, and we just have the note in the docs, then updating your Juju version but without ops changing `reboot` will just start working.

If we have something like:

```python
if container_type == "machine":
    raise ModelError("reboot() is only available when running on machine controllers")
```

Then you'd need to update ops as well (to the version where we presumably strip that out). That would be less surprising, which seems good.

> So maybe it's time to add a helper method is_k8s_sidecar. We've needed this before too. Or maybe charm.type (though something more specific than "type") or something which is an enum, one of machine, k8s-podspec, k8s-sidecar.

`charm.platform`? (To match charmhub.io's Platform filter).
tonyandrewmeyer: > I don't know whether `juju-reboot --now` guarantees you will/won't get your process back. But likely charm authors will assume that if they pass now=True they won't execute any more code.

In my tests, I couldn't get any (Python) code to run after the call to reboot (which was done with `subprocess.run` as in the implementation, so Python would be waiting for the subprocess to exit). I guess we could `sys.exit(0)` in the `now=True` case - the behaviour would essentially be the same as now, but provide more of a guarantee that you didn't get stray lines of code executed.
jameinel: FWIW, I raised the question more so that we wouldn't have a flaky test suite, than because I was worried about the reboot time not being precise enough.
But I'm also happy to go with reboot count instead, as the actual datetime isn't something you could really test.


Patch:
diff --git a/CHANGES.md b/CHANGES.md
index cc12ba4..6b42007 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -1,6 +1,9 @@
 # 2.8.0
 
+* Added `Unit.reboot()` and `Harness.reboot_count``
+* Added `RelationMeta.optional`
 * The type of a `Handle`'s `key` was expanded from `str` to `str|None`
+* Narrowed types of `app` and `unit` in relation events to exclude `None` where applicable
 
 # 2.7.0
 
diff --git a/ops/model.py b/ops/model.py
index 29ec6f5..33f7f12 100644
--- a/ops/model.py
+++ b/ops/model.py
@@ -26,6 +26,7 @@ import re
 import shutil
 import stat
 import subprocess
+import sys
 import tempfile
 import time
 import typing
@@ -686,6 +687,8 @@ class Unit:
         Use :meth:`open_port` and :meth:`close_port` to manage ports
         individually.
 
+        *New in version 2.7*
+
         Args:
             ports: The ports to open. Provide an int to open a TCP port, or
                 a :class:`Port` to open a port for another protocol.
@@ -709,6 +712,32 @@ class Unit:
         for protocol, port in desired - existing:
             self._backend.open_port(protocol, port)
 
+    def reboot(self, now: bool = False) -> None:
+        """Reboot the host machine.
+
+        Normally, the reboot will only take place after the current hook successfully
+        completes. Use ``now=True`` to reboot immediately without waiting for the
+        hook to complete; this is useful when multiple restarts are required (Juju
+        will re-run the hook after rebooting).
+
+        This is not supported on Kubernetes charms, can only be called for the current unit,
+        and cannot be used in an action hook.
+
+        *New in version 2.8*
+
+        Args:
+            now: terminate immediately without waiting for the current hook to complete,
+                restarting the hook after reboot.
+
+        Raises:
+            RuntimeError: if called on a remote unit.
+            :class:`ModelError`: if used in an action hook.
+
+        """
+        if not self._is_our_unit:
+            raise RuntimeError(f'cannot reboot a remote unit {self}')
+        self._backend.reboot(now)
+
 
 @dataclasses.dataclass(frozen=True)
 class Port:
@@ -3400,6 +3429,16 @@ class _ModelBackend:
         protocol_lit = typing.cast(typing.Literal['tcp', 'udp'], protocol)
         return Port(protocol_lit, int(port))
 
+    def reboot(self, now: bool = False):
+        if now:
+            self._run("juju-reboot", "--now")
+            # Juju will kill the Charm process, and in testing no code after
+            # this point would execute. However, we want to guarantee that for
+            # Charmers, so we force that to be the case.
+            sys.exit()
+        else:
+            self._run("juju-reboot")
+
 
 class _ModelBackendValidator:
     """Provides facilities for validating inputs and formatting them for model backends."""


Test patch:
diff --git a/ops/testing.py b/ops/testing.py
index c9af783..8d638b9 100755
--- a/ops/testing.py
+++ b/ops/testing.py
@@ -1741,6 +1741,11 @@ class Harness(Generic[CharmType]):
             handler=(lambda _: result) if handler is None else handler  # type: ignore
         )
 
+    @property
+    def reboot_count(self) -> int:
+        """Number of times the charm has called :meth:`ops.Unit.reboot`."""
+        return self._backend._reboot_count
+
 
 def _get_app_or_unit_name(app_or_unit: AppUnitOrName) -> str:
     """Return name of given application or unit (return strings directly)."""
@@ -1957,6 +1962,7 @@ class _TestingModelBackend:
         self._secrets: List[_Secret] = []
         self._opened_ports: Set[model.Port] = set()
         self._networks: Dict[Tuple[Optional[str], Optional[int]], _NetworkDict] = {}
+        self._reboot_count = 0
 
     def _validate_relation_access(self, relation_name: str, relations: List[model.Relation]):
         """Ensures that the named relation exists/has been added.
@@ -2517,6 +2523,14 @@ class _TestingModelBackend:
         else:
             raise model.ModelError(f'ERROR invalid protocol "{protocol}", expected "tcp", "udp", or "icmp"\n')  # NOQA: test_quote_backslashes
 
+    def reboot(self, now: bool = False):
+        self._reboot_count += 1
+        if not now:
+            return
+        # This should exit, reboot, and re-emit the event, but we'll need the caller
+        # to handle everything after the exit.
+        raise SystemExit()
+
 
 @_copy_docstrings(pebble.ExecProcess)
 class _TestingExecProcess:
diff --git a/test/test_model.py b/test/test_model.py
index 365ddff..51b6d20 100755
--- a/test/test_model.py
+++ b/test/test_model.py
@@ -3596,5 +3596,28 @@ class TestPorts(unittest.TestCase):
         ])
 
 
+class TestUnit(unittest.TestCase):
+    def setUp(self):
+        self.model = ops.model.Model(ops.charm.CharmMeta(), ops.model._ModelBackend('myapp/0'))
+        self.unit = self.model.unit
+
+    def test_reboot(self):
+        fake_script(self, 'juju-reboot', 'exit 0')
+        self.unit.reboot()
+        self.assertEqual(fake_script_calls(self, clear=True), [
+            ['juju-reboot', ''],
+        ])
+        with self.assertRaises(SystemExit):
+            self.unit.reboot(now=True)
+        self.assertEqual(fake_script_calls(self, clear=True), [
+            ['juju-reboot', '--now'],
+        ])
+
+        with self.assertRaises(RuntimeError):
+            self.model.get_unit('other').reboot()
+        with self.assertRaises(RuntimeError):
+            self.model.get_unit('other').reboot(now=True)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/test_testing.py b/test/test_testing.py
index 5a0a06c..cf68413 100644
--- a/test/test_testing.py
+++ b/test/test_testing.py
@@ -3345,6 +3345,37 @@ class TestTestingModelBackend(unittest.TestCase):
         client = backend.get_pebble('/custom/socket/path')
         self.assertIsInstance(client, _TestingPebbleClient)
 
+    def test_reboot(self):
+        class RebootingCharm(ops.CharmBase):
+            def __init__(self, framework: ops.Framework):
+                super().__init__(framework)
+                self.framework.observe(self.on.install, self._reboot_now)
+                self.framework.observe(self.on.remove, self._reboot)
+
+            def _reboot_now(self, event: ops.InstallEvent):
+                self.unit.reboot(now=True)
+
+            def _reboot(self, event: ops.RemoveEvent):
+                self.unit.reboot()
+
+        harness = ops.testing.Harness(RebootingCharm, meta='''
+            name: test-app
+            ''')
+        self.addCleanup(harness.cleanup)
+        self.assertEqual(harness.reboot_count, 0)
+        backend = harness._backend
+        backend.reboot()
+        self.assertEqual(harness.reboot_count, 1)
+        with self.assertRaises(SystemExit):
+            backend.reboot(now=True)
+        self.assertEqual(harness.reboot_count, 2)
+        harness.begin()
+        with self.assertRaises(SystemExit):
+            harness.charm.on.install.emit()
+        self.assertEqual(harness.reboot_count, 3)
+        harness.charm.on.remove.emit()
+        self.assertEqual(harness.reboot_count, 4)
+
 
 class _TestingPebbleClientMixin:
     def get_testing_client(self):