[PATCH net-next 2/4] selftests: drv-net: add a way to wait for a local process

[Jakub Kicinski <kuba@kernel.org>]

We use wait_port_listen() extensively to wait for a process
we spawned to be ready. Not all processes will open listening
sockets. Add a method of explicitly waiting for a child to
be ready. Pass a FD to the spawned process and wait for it
to write a message to us. FD number is passed via KSFT_READY_FD
env variable.

Make use of this method in the queues test to make it less flaky.

Signed-off-by: Jakub Kicinski <kuba@kernel.org>
---
 .../selftests/drivers/net/xdp_helper.c        | 22 ++++++-
 tools/testing/selftests/drivers/net/queues.py | 46 ++++++---------
 tools/testing/selftests/net/lib/py/utils.py   | 58 +++++++++++++++++--
 3 files changed, 93 insertions(+), 33 deletions(-)

diff --git a/tools/testing/selftests/drivers/net/xdp_helper.c b/tools/testing/selftests/drivers/net/xdp_helper.c
index cf06a88b830b..8f77da4f798f 100644
--- a/tools/testing/selftests/drivers/net/xdp_helper.c
+++ b/tools/testing/selftests/drivers/net/xdp_helper.c
@@ -14,6 +14,25 @@
 #define UMEM_SZ (1U << 16)
 #define NUM_DESC (UMEM_SZ / 2048)
 
+/* Move this to a common header when reused! */
+static void ksft_ready(void)
+{
+	const char msg[7] = "ready\n";
+	char *env_str;
+	int fd;
+
+	env_str = getenv("KSFT_READY_FD");
+	if (!env_str)
+		return;
+
+	fd = atoi(env_str);
+	if (!fd)
+		return;
+
+	write(fd, msg, sizeof(msg));
+	close(fd);
+}
+
 /* this is a simple helper program that creates an XDP socket and does the
  * minimum necessary to get bind() to succeed.
  *
@@ -85,8 +104,7 @@ int main(int argc, char **argv)
 		return 1;
 	}
 
-	/* give the parent program some data when the socket is ready*/
-	fprintf(stdout, "%d\n", sock_fd);
+	ksft_ready();
 
 	/* parent program will write a byte to stdin when its ready for this
 	 * helper to exit
diff --git a/tools/testing/selftests/drivers/net/queues.py b/tools/testing/selftests/drivers/net/queues.py
index b6896a57a5fd..91e344d108ee 100755
--- a/tools/testing/selftests/drivers/net/queues.py
+++ b/tools/testing/selftests/drivers/net/queues.py
@@ -5,13 +5,12 @@ from lib.py import ksft_disruptive, ksft_exit, ksft_run
 from lib.py import ksft_eq, ksft_raises, KsftSkipEx, KsftFailEx
 from lib.py import EthtoolFamily, NetdevFamily, NlError
 from lib.py import NetDrvEnv
-from lib.py import cmd, defer, ip
+from lib.py import bkg, cmd, defer, ip
 import errno
 import glob
 import os
 import socket
 import struct
-import subprocess
 
 def sys_get_queues(ifname, qtype='rx') -> int:
     folders = glob.glob(f'/sys/class/net/{ifname}/queues/{qtype}-*')
@@ -25,37 +24,30 @@ import subprocess
     return None
 
 def check_xdp(cfg, nl, xdp_queue_id=0) -> None:
-    xdp = subprocess.Popen([cfg.rpath("xdp_helper"), f"{cfg.ifindex}", f"{xdp_queue_id}"],
-                           stdin=subprocess.PIPE, stdout=subprocess.PIPE, bufsize=1,
-                           text=True)
-    defer(xdp.kill)
+    with bkg(f'{cfg.rpath("xdp_helper")} {cfg.ifindex} {xdp_queue_id}',
+             wait_init=3):
 
-    stdout, stderr = xdp.communicate(timeout=10)
-    rx = tx = False
+        rx = tx = False
 
-    if xdp.returncode == 255:
-        raise KsftSkipEx('AF_XDP unsupported')
-    elif xdp.returncode > 0:
-        raise KsftFailEx('unable to create AF_XDP socket')
+        queues = nl.queue_get({'ifindex': cfg.ifindex}, dump=True)
+        if not queues:
+            raise KsftSkipEx("Netlink reports no queues")
 
-    queues = nl.queue_get({'ifindex': cfg.ifindex}, dump=True)
-    if not queues:
-        raise KsftSkipEx("Netlink reports no queues")
+        for q in queues:
+            if q['id'] == 0:
+                if q['type'] == 'rx':
+                    rx = True
+                if q['type'] == 'tx':
+                    tx = True
 
-    for q in queues:
-        if q['id'] == 0:
-            if q['type'] == 'rx':
-                rx = True
-            if q['type'] == 'tx':
-                tx = True
+                ksft_eq(q['xsk'], {})
+            else:
+                if 'xsk' in q:
+                    _fail("Check failed: xsk attribute set.")
 
-            ksft_eq(q['xsk'], {})
-        else:
-            if 'xsk' in q:
-                _fail("Check failed: xsk attribute set.")
+        ksft_eq(rx, True)
+        ksft_eq(tx, True)
 
-    ksft_eq(rx, True)
-    ksft_eq(tx, True)
 
 def get_queues(cfg, nl) -> None:
     snl = NetdevFamily(recv_size=4096)
diff --git a/tools/testing/selftests/net/lib/py/utils.py b/tools/testing/selftests/net/lib/py/utils.py
index 9e3bcddcf3e8..efb9b8fc1447 100644
--- a/tools/testing/selftests/net/lib/py/utils.py
+++ b/tools/testing/selftests/net/lib/py/utils.py
@@ -2,8 +2,10 @@
 
 import errno
 import json as _json
+import os
 import random
 import re
+import select
 import socket
 import subprocess
 import time
@@ -15,8 +17,22 @@ import time
         self.cmd = cmd_obj
 
 
+def fd_read_timeout(fd, timeout):
+    rlist, _, _ = select.select([fd], [], [], timeout)
+    if rlist:
+        return os.read(fd, 1024)
+    else:
+        raise TimeoutError("Timeout waiting for fd read")
+
+
 class cmd:
-    def __init__(self, comm, shell=True, fail=True, ns=None, background=False, host=None, timeout=5):
+    """
+    Execute a command on local or remote host.
+
+    Use bkg() instead to run a command in the background.
+    """
+    def __init__(self, comm, shell=True, fail=True, ns=None, background=False,
+                 host=None, timeout=5, wait_init=None):
         if ns:
             comm = f'ip netns exec {ns} ' + comm
 
@@ -28,8 +44,23 @@ import time
         if host:
             self.proc = host.cmd(comm)
         else:
+            # wait_init lets us wait for the background process to fully start,
+            # we pass an FD to the child process, and wait for it to write back
+            pass_fds = ()
+            env = os.environ.copy()
+            if wait_init is not None:
+                rfd, wfd = os.pipe()
+                pass_fds = (wfd, )
+                env["KSFT_READY_FD"] = str(wfd)
             self.proc = subprocess.Popen(comm, shell=shell, stdout=subprocess.PIPE,
-                                         stderr=subprocess.PIPE)
+                                         stderr=subprocess.PIPE, pass_fds=pass_fds,
+                                         env=env)
+            if wait_init is not None:
+                os.close(wfd)
+                msg = fd_read_timeout(rfd, wait_init)
+                os.close(rfd)
+                if not msg:
+                    raise Exception("Did not receive ready message")
         if not background:
             self.process(terminate=False, fail=fail, timeout=timeout)
 
@@ -54,10 +85,29 @@ import time
 
 
 class bkg(cmd):
+    """
+    Run a command in the background.
+
+    Examples usage:
+
+    Run a command on remote host, and wait for it to finish.
+    This is usually paired with wait_port_listen() to make sure
+    the command has initialized:
+
+        with bkg("socat ...", exit_wait=True, host=cfg.remote) as nc:
+            ...
+
+    Run a command and expect it to let us know that it's ready
+    by writing to a special file decriptor passed via KSFT_READY_FD.
+    Command will be terminated when we exit the context manager:
+
+        with bkg("my_binary", wait_init=5):
+    """
     def __init__(self, comm, shell=True, fail=None, ns=None, host=None,
-                 exit_wait=False):
+                 exit_wait=False, wait_init=None):
         super().__init__(comm, background=True,
-                         shell=shell, fail=fail, ns=ns, host=host)
+                         shell=shell, fail=fail, ns=ns, host=host,
+                         wait_init=wait_init)
         self.terminate = not exit_wait
         self.check_fail = fail
 
-- 
2.48.1