[PATCH 23/27] tasst: Helpers for running daemons with a pidfile

[David Gibson <david@gibson.dropbear.id.au>]

When running commands through nstool, stopping them directly with the
stop() method isn't reliable, because that just kills the nstool exec
process not the actual process.  At least for programs that create their
own pidfile we can do better by using that to kill off the process when
we're done, including during cleanup on failure.  This extends the
Site.bg() method with a pidfile parameter which will allow cleanup by this
means.

There's also the case of processes that start in the foreground then
daemonize themselves.  We'd like to clean those up via pidfile as well, so
add a Site.daemon() method to do exactly that.

Signed-off-by: David Gibson <david@gibson.dropbear.id.au>
---
 test/tasst/exesite.py | 31 +++++++++++++++++++++++++++----
 1 file changed, 27 insertions(+), 4 deletions(-)

diff --git a/test/tasst/exesite.py b/test/tasst/exesite.py
index 1f34eee9..6414c73f 100644
--- a/test/tasst/exesite.py
+++ b/test/tasst/exesite.py
@@ -22,7 +22,7 @@ from avocado_classless.test import (
     assert_eq, assert_eq_unordered, assert_in, assert_raises, test_output
 )
 
-from tasst.typecheck import typecheck
+from tasst.typecheck import typecheck, typecheck_default
 
 
 class SiteProcess(contextlib.AbstractContextManager):
@@ -31,18 +31,32 @@ class SiteProcess(contextlib.AbstractContextManager):
     """
 
     def __init__(self, site, cmd, subp, *,
-                 ignore_status, context_timeout):
+                 ignore_status, context_timeout, pidfile):
         self.site = typecheck(site, Site)
         self.cmd = typecheck(cmd, str)
         self.subproc = typecheck(subp, avocado.utils.process.SubProcess)
         self.ignore_status = typecheck(ignore_status, bool)
         self.context_timeout = float(context_timeout)
+        self.pidfile = typecheck_default(pidfile, str, None)
+        self.pid = None
+
+        if pidfile is not None:
+            site.require_cmds('cat', 'kill')
 
     def __enter__(self):
         self.subproc.start()
+        if self.pidfile is not None:
+            # Wait for the PID file to be written
+            pidstr = None
+            while not pidstr:
+                pidstr = self.site.output(f'cat {self.pidfile}',
+                                          ignore_status=True)
+            self.pid = int(pidstr)
         return self
 
     def __exit__(self, *exc_details):
+        if self.pid is not None and self.subproc.poll() is None:
+            self.site.fg(f'kill -TERM {self.pid}')
         result = self.subproc.run(timeout=self.context_timeout)
         if not self.ignore_status and result.exit_status != 0:
             siteinfo = f'[{self.site.name} site]'
@@ -84,11 +98,20 @@ class Site(contextlib.AbstractContextManager):
         cmd, kwargs = self.hostify(cmd, **kwargs)
         return avocado.utils.process.SubProcess(cmd, **kwargs)
 
-    def bg(self, cmd, context_timeout=1.0, ignore_status=False, **kwargs):
+    def bg(self, cmd, *,
+           context_timeout=1.0, ignore_status=False, pidfile=None, **kwargs):
         subproc = self.subprocess(cmd, **kwargs)
         return SiteProcess(self, cmd, subproc,
                            context_timeout=context_timeout,
-                           ignore_status=ignore_status)
+                           ignore_status=ignore_status, pidfile=pidfile)
+
+    @contextlib.contextmanager
+    def daemon(self, cmd, *, pidfile, **kwargs):
+        self.require_cmds('cat', 'kill')
+        self.fg(cmd, **kwargs)
+        yield
+        pid = int(self.output(f'cat {pidfile}'))
+        self.fg(f'kill -TERM {pid}')
 
     def require_cmds(self, *cmds):
         missing = [c for c in cmds
-- 
@@ -22,7 +22,7 @@ from avocado_classless.test import (
     assert_eq, assert_eq_unordered, assert_in, assert_raises, test_output
 )
 
-from tasst.typecheck import typecheck
+from tasst.typecheck import typecheck, typecheck_default
 
 
 class SiteProcess(contextlib.AbstractContextManager):
@@ -31,18 +31,32 @@ class SiteProcess(contextlib.AbstractContextManager):
     """
 
     def __init__(self, site, cmd, subp, *,
-                 ignore_status, context_timeout):
+                 ignore_status, context_timeout, pidfile):
         self.site = typecheck(site, Site)
         self.cmd = typecheck(cmd, str)
         self.subproc = typecheck(subp, avocado.utils.process.SubProcess)
         self.ignore_status = typecheck(ignore_status, bool)
         self.context_timeout = float(context_timeout)
+        self.pidfile = typecheck_default(pidfile, str, None)
+        self.pid = None
+
+        if pidfile is not None:
+            site.require_cmds('cat', 'kill')
 
     def __enter__(self):
         self.subproc.start()
+        if self.pidfile is not None:
+            # Wait for the PID file to be written
+            pidstr = None
+            while not pidstr:
+                pidstr = self.site.output(f'cat {self.pidfile}',
+                                          ignore_status=True)
+            self.pid = int(pidstr)
         return self
 
     def __exit__(self, *exc_details):
+        if self.pid is not None and self.subproc.poll() is None:
+            self.site.fg(f'kill -TERM {self.pid}')
         result = self.subproc.run(timeout=self.context_timeout)
         if not self.ignore_status and result.exit_status != 0:
             siteinfo = f'[{self.site.name} site]'
@@ -84,11 +98,20 @@ class Site(contextlib.AbstractContextManager):
         cmd, kwargs = self.hostify(cmd, **kwargs)
         return avocado.utils.process.SubProcess(cmd, **kwargs)
 
-    def bg(self, cmd, context_timeout=1.0, ignore_status=False, **kwargs):
+    def bg(self, cmd, *,
+           context_timeout=1.0, ignore_status=False, pidfile=None, **kwargs):
         subproc = self.subprocess(cmd, **kwargs)
         return SiteProcess(self, cmd, subproc,
                            context_timeout=context_timeout,
-                           ignore_status=ignore_status)
+                           ignore_status=ignore_status, pidfile=pidfile)
+
+    @contextlib.contextmanager
+    def daemon(self, cmd, *, pidfile, **kwargs):
+        self.require_cmds('cat', 'kill')
+        self.fg(cmd, **kwargs)
+        yield
+        pid = int(self.output(f'cat {pidfile}'))
+        self.fg(f'kill -TERM {pid}')
 
     def require_cmds(self, *cmds):
         missing = [c for c in cmds
-- 
2.41.0