xnbd: fix killing of old xnbd-server process

Our LAB run lots of XNBD jobs, and lots of xnbd-server process were
never killed.

The port allocated is not used for killing, eg
self.parameters['protocols']['lava-xnbd']['port'] is still set to auto.

As suggested by Neil Williams, it's better to "create a dedicated
function which is called from the test job"

So this patch adds a "set_port" function which get the right port
number.

Existing jobs should now add this:
   protocols:
      lava-xnbd:
      - action: nbd-deploy
        request: set_port
after
actions:
- deploy:

Change-Id: I6b464314ac2c1a137162b0a7f955c4501f8a7424
Suggested-by: Jan-Simon Möller <jsmoeller@linuxfoundation.org>

3 files changed, 60 insertions, 16 deletions

diff --git a/lava_dispatcher/actions/deploy/nbd.py b/lava_dispatcher/actions/deploy/nbd.py
index 75bd22fb7..3a066f9a4 100644
--- a/lava_dispatcher/actions/deploy/nbd.py
+++ b/lava_dispatcher/actions/deploy/nbd.py
@@ -30,8 +30,6 @@ from lava_dispatcher.actions.deploy import DeployAction
 from lava_dispatcher.actions.deploy.download import DownloaderAction
 from lava_dispatcher.utils.shell import which
 from lava_dispatcher.utils.filesystem import tftpd_dir
-from lava_dispatcher.utils.network import get_free_port
-from lava_dispatcher.utils.network import dispatcher_ip
 from lava_dispatcher.protocols.xnbd import XnbdProtocol
 from lava_dispatcher.actions.deploy.overlay import OverlayAction
 
@@ -132,13 +130,6 @@ class NbdAction(DeployAction):  # pylint:disable=too-many-instance-attributes
         # and store in namespace for boot action
         # ip
         parameters['lava-xnbd'] = {}
-        self.nbd_ip = dispatcher_ip(self.job.parameters['dispatcher'])
-        parameters['lava-xnbd']['ip'] = self.nbd_ip
-        self.set_namespace_data(action=self.name, label='nbd', key='nbd_server_ip', value=self.nbd_ip, parameters=parameters)
-        # port
-        self.nbd_port = get_free_port(self.job.parameters['dispatcher'])
-        parameters['lava-xnbd']['port'] = self.nbd_port
-        self.set_namespace_data(action=self.name, label='nbd', key='nbd_server_port', value=self.nbd_port, parameters=parameters)
         # handle XnbdAction next - bring-up xnbd-server
         self.internal_pipeline.add_action(XnbdAction())
 
@@ -159,9 +150,12 @@ class XnbdAction(DeployAction):
         connection = super(XnbdAction, self).run(connection, max_end_time, args)
         self.logger.debug("%s: starting xnbd-server", self.name)
         # pull from parameters - as previously set
-        self.nbd_server_port = self.parameters['lava-xnbd']['port']
-        self.nbd_server_ip = self.parameters['lava-xnbd']['ip']
         self.nbd_root = self.parameters['lava-xnbd']['nbdroot']
+        self.nbd_server_port = self.get_namespace_data(action='nbd-deploy', label='nbd', key='nbd_server_port')
+        self.nbd_server_ip = self.get_namespace_data(action='nbd-deploy', label='nbd', key='nbd_server_ip')
+        if self.nbd_server_port is None:
+            self.errors = "NBD server port is unset"
+            return connection
         self.logger.debug("NBD-IP: %s, NBD-PORT: %s, NBD-ROOT: %s",
                           self.nbd_server_ip, self.nbd_server_port, self.nbd_root)
         nbd_cmd = ['xnbd-server', '--logpath', '/tmp/xnbd.log.%s' % self.nbd_server_port,
diff --git a/lava_dispatcher/protocols/xnbd.py b/lava_dispatcher/protocols/xnbd.py
index 28ac1dd1c..b4266687e 100644
--- a/lava_dispatcher/protocols/xnbd.py
+++ b/lava_dispatcher/protocols/xnbd.py
@@ -24,9 +24,12 @@ import logging
 from lava_dispatcher.connection import Protocol
 from lava_dispatcher.action import (
     Timeout,
+    JobError,
 )
 from lava_dispatcher.shell import ShellCommand
 from lava_dispatcher.utils.constants import XNBD_SYSTEM_TIMEOUT
+from lava_dispatcher.utils.network import dispatcher_ip
+from lava_dispatcher.utils.network import get_free_port
 
 
 class XnbdProtocol(Protocol):
@@ -42,6 +45,7 @@ class XnbdProtocol(Protocol):
         self.logger = logging.getLogger('dispatcher')
         self.parameters = parameters
         self.port = None
+        self.ports = []
 
     @classmethod
     def accepts(cls, parameters):
@@ -51,6 +55,9 @@ class XnbdProtocol(Protocol):
             return False
         return True
 
+    def collate(self, reply, params):
+        params.update(reply)
+
     def set_up(self):
         """
         Called from the job at the start of the run step.
@@ -59,6 +66,44 @@ class XnbdProtocol(Protocol):
         if 'port' not in self.parameters['protocols']['lava-xnbd']:
             self.errors = ('No port set in parameters for lava-xnbd protocol!\nE.g.:\n protocols:\n  lava-xnbd:\n    port: auto \n')
 
+    def __call__(self, *args, **kwargs):
+        action = kwargs.get('action', None)
+        self.logger.debug("[%s] Checking protocol data for %s", action.name, self.name)
+        try:
+            return self._api_select(args, action=action)
+        except (ValueError, TypeError) as exc:
+            raise JobError("Invalid call to %s %s" % (self.name, exc))
+
+    def _api_select(self, data, action=None):
+        if not data:
+            raise TestError("Protocol called without any data")
+        for item in data:
+            if 'request' not in item:
+                raise JobError("Bad API call over protocol - missing request")
+            if 'set_port' in item['request']:
+                return self.set_port(action=action)
+            else:
+                raise JobError("Unrecognised API call in request.")
+        return None
+
+    def set_port(self, action):
+        msg = {
+            'data': {
+                'nbd_server_port': 10809
+            }
+        }
+        nbd_port = self.parameters['protocols']['lava-xnbd']['port']
+        if nbd_port == 'auto':
+            self.logger.debug("Get a port from pool")
+            nbd_port = get_free_port(self.parameters['dispatcher'])
+        self.ports.append(nbd_port)
+        msg['data']['nbd_server_port'] = nbd_port
+        action.set_namespace_data('nbd-deploy', label='nbd', key='nbd_server_port', value=nbd_port, parameters=action.parameters)
+        nbd_ip = dispatcher_ip(self.parameters['dispatcher'])
+        action.set_namespace_data('nbd-deploy', label='nbd', key='nbd_server_ip', value=nbd_ip, parameters=action.parameters)
+        self.logger.debug("Set_port %d" % nbd_port)
+        return msg['data']
+
     def finalise_protocol(self, device=None):
         """Called by Finalize action to power down and clean up the assigned
         device.
@@ -66,11 +111,12 @@ class XnbdProtocol(Protocol):
         # shutdown xnbd for the given device/job based in the port-number used
         try:
             self.logger.debug("%s cleanup", self.name)
-            self.port = self.parameters['protocols']['lava-xnbd']['port']
-            nbd_cmd = "pkill -f xnbd-server.*%s" % (self.port)
-            shell = ShellCommand("%s\n" % nbd_cmd, self.system_timeout,
-                                 logger=self.logger)
-            shell.expect(pexpect.EOF)
+            for port in self.ports:
+                self.logger.debug("clean NBD port %s" % port)
+                nbd_cmd = "pkill -f xnbd-server.*%s" % (port)
+                shell = ShellCommand("%s\n" % nbd_cmd, self.system_timeout,
+                                     logger=self.logger)
+                shell.expect(pexpect.EOF)
         except Exception as e:
             self.logger.debug(str(e))
             self.logger.debug("xnbd-finalize-protocol failed, but continuing anyway.")
diff --git a/lava_dispatcher/test/sample_jobs/bbb-initrd-nbd.yaml b/lava_dispatcher/test/sample_jobs/bbb-initrd-nbd.yaml
index 276861f7a..67356cb5b 100644
--- a/lava_dispatcher/test/sample_jobs/bbb-initrd-nbd.yaml
+++ b/lava_dispatcher/test/sample_jobs/bbb-initrd-nbd.yaml
@@ -43,3 +43,7 @@ actions:
     - 'root@beaglebone:'
     timeout:
       minutes: 4
+    protocols:
+      lava-xnbd:
+      - action: nbd-deploy
+        request: set_port