Repository: osbuild/osbuild
Pull request #485 (2020-07-24 12:06:20)

use cache in assembler tests
use one osbuild executor for all qemu tests

This way the test can benefit from osbuild's internal cache:

The first subtest builds all the stages and  runs the assembler
The next subtests can reuse the built stages and just run the assembler

Some data from my machine:

Building the manifest takes about 120 seconds
Running just the assembler on the cache's content takes 30 seconds.

Before this change, the whole manifest was built 3 times:
3 * 120 = 360 seconds

After this change, the whole manifest is built once and the cache
is reused 2 times:
1 * 120 + 2 * 30 = 180 seconds

teg: Nice stuff! Though the tar test also needs an `osb` object passed in.
teg: I think this is a clear improvement, but could we take it one step further and share the cache for all the assembler tests? We now checkpoint the same tree for each type of assembler (but no longer for each image type). That could be a follow-up though.

Patch:
diff --git a/osbuild/api.py b/osbuild/api.py
index 4d77d2b2..7b47571b 100644
--- a/osbuild/api.py
+++ b/osbuild/api.py
@@ -1,22 +1,131 @@
+import abc
 import asyncio
 import io
 import json
 import os
+import sys
 import tempfile
 import threading
+from typing import Dict, Optional
+from .util.types import PathLike
 from .util import jsoncomm
 
 
-class API:
-    def __init__(self, socket_address, args, monitor):
+__all__ = [
+    "API"
+]
+
+
+class BaseAPI(abc.ABC):
+    """Base class for all API providers
+
+    This base class provides the basic scaffolding for setting
+    up API endpoints, normally to be used for bi-directional
+    communication from and to the sandbox. It is to be used as
+    a context manger. The communication channel will only be
+    established on entering the context and will be shut down
+    when the context is left.
+
+    New messages are delivered via the `_message` method, that
+    needs to be implemented by deriving classes.
+
+    Optionally, the `_cleanup` method can be implemented, to
+    clean up resources after the context is left and the
+    communication channel shut down.
+
+    On incoming messages, first the `_dispatch` method will be
+    called; the default implementation will receive the message
+    call `_message.`
+    """
+    def __init__(self, socket_address: Optional[PathLike] = None):
         self.socket_address = socket_address
+        self.barrier = threading.Barrier(2)
+        self.event_loop = None
+        self.thread = None
+        self._socketdir = None
+
+    @property
+    @classmethod
+    @abc.abstractmethod
+    def endpoint(cls):
+        """The name of the API endpoint"""
+
+    @abc.abstractmethod
+    def _message(self, msg: Dict, fds: jsoncomm.FdSet, sock: jsoncomm.Socket, addr: str):
+        """Called for a new incoming message
+
+        The file descriptor set `fds` will be closed after the call.
+        Use the `FdSet.steal()` method to extract file descriptors.
+        """
+
+    def _cleanup(self):
+        """Called after the event loop is shut down"""
+
+    @classmethod
+    def _make_socket_dir(cls):
+        """Called to create the temporary socket dir"""
+        return tempfile.TemporaryDirectory(prefix="api-", dir="/run/osbuild")
+
+    def _dispatch(self, sock: jsoncomm.Socket):
+        """Called when data is available on the socket"""
+        msg, fds, addr = sock.recv()
+        self._message(msg, fds, sock, addr)
+        fds.close()
+
+    def _run_event_loop(self):
+        with jsoncomm.Socket.new_server(self.socket_address) as server:
+            self.barrier.wait()
+            self.event_loop.add_reader(server, self._dispatch, server)
+            asyncio.set_event_loop(self.event_loop)
+            self.event_loop.run_forever()
+            self.event_loop.remove_reader(server)
+
+    def __enter__(self):
+        # We are not re-entrant, so complain if re-entered.
+        assert self.event_loop is None
+
+        if not self.socket_address:
+            self._socketdir = self._make_socket_dir()
+            address = os.path.join(self._socketdir.name, self.endpoint)
+            self.socket_address = address
+
+        self.event_loop = asyncio.new_event_loop()
+        self.thread = threading.Thread(target=self._run_event_loop)
+
+        self.barrier.reset()
+        self.thread.start()
+        self.barrier.wait()
+
+        return self
+
+    def __exit__(self, *args):
+        self.event_loop.call_soon_threadsafe(self.event_loop.stop)
+        self.thread.join()
+        self.event_loop.close()
+
+        # Give deriving classes a chance to clean themselves up
+        self._cleanup()
+
+        self.thread = None
+        self.event_loop = None
+
+        if self._socketdir:
+            self._socketdir.cleanup()
+            self._socketdir = None
+            self.socket_address = None
+
+
+class API(BaseAPI):
+    """The main OSBuild API"""
+
+    endpoint = "osbuild"
+
+    def __init__(self, args, monitor, *, socket_address=None):
+        super().__init__(socket_address)
         self.input = args
         self._output_data = io.StringIO()
         self._output_pipe = None
         self.monitor = monitor
-        self.event_loop = asyncio.new_event_loop()
-        self.thread = threading.Thread(target=self._run_event_loop)
-        self.barrier = threading.Barrier(2)
 
     @property
     def output(self):
@@ -56,28 +165,24 @@ class API:
 
             server.send(msg, fds=fds, destination=addr)
 
-    def _dispatch(self, server):
-        msg, _, addr = server.recv()
+    def _message(self, msg, fds, sock, addr):
         if msg["method"] == 'setup-stdio':
-            self._setup_stdio(server, addr)
+            self._setup_stdio(sock, addr)
 
-    def _run_event_loop(self):
-        with jsoncomm.Socket.new_server(self.socket_address) as server:
-            self.barrier.wait()
-            self.event_loop.add_reader(server, self._dispatch, server)
-            asyncio.set_event_loop(self.event_loop)
-            self.event_loop.run_forever()
-            self.event_loop.remove_reader(server)
-
-    def __enter__(self):
-        self.thread.start()
-        self.barrier.wait()
-        return self
-
-    def __exit__(self, *args):
-        self.event_loop.call_soon_threadsafe(self.event_loop.stop)
-        self.thread.join()
-        self.event_loop.close()
+    def _cleanup(self):
         if self._output_pipe:
             os.close(self._output_pipe)
             self._output_pipe = None
+
+
+def setup_stdio(path="/run/osbuild/api/osbuild"):
+    """Replace standard i/o with the ones provided by the API"""
+    with jsoncomm.Socket.new_client(path) as client:
+        req = {"method": "setup-stdio"}
+        client.send(req)
+        msg, fds, _ = client.recv()
+        for sio in ["stdin", "stdout", "stderr"]:
+            target = getattr(sys, sio)
+            source = fds[msg[sio]]
+            os.dup2(source, target.fileno())
+        fds.close()
diff --git a/osbuild/buildroot.py b/osbuild/buildroot.py
index decf9906..486d8e73 100644
--- a/osbuild/buildroot.py
+++ b/osbuild/buildroot.py
@@ -40,7 +40,7 @@ class BuildRoot(contextlib.AbstractContextManager):
         self._vardir = var
         self._libdir = libdir
         self._runner = runner
-        self.api = None
+        self._apis = []
         self.dev = None
         self.var = None
 
@@ -60,10 +60,6 @@ class BuildRoot(contextlib.AbstractContextManager):
             #
             # For now, this includes:
             #
-            #   * We create an API directory where the caller can place sockets
-            #     before we bind-mount it into the container on
-            #     `/run/osbuild/api`.
-            #
             #   * We create a tmpfs instance *without* `nodev` which we then use
             #     as `/dev` in the container. This is required for the container
             #     to create device nodes for loop-devices.
@@ -73,12 +69,11 @@ class BuildRoot(contextlib.AbstractContextManager):
             #     create throw-away data that it does not want to put into a
             #     tmpfs.
 
-            api = tempfile.TemporaryDirectory(prefix="osbuild-api-", dir=self._rundir)
-            self.api = self._exitstack.enter_context(api)
-
+            os.makedirs(self._rundir, exist_ok=True)
             dev = tempfile.TemporaryDirectory(prefix="osbuild-dev-", dir=self._rundir)
             self.dev = self._exitstack.enter_context(dev)
 
+            os.makedirs(self._vardir, exist_ok=True)
             var = tempfile.TemporaryDirectory(prefix="osbuild-var-", dir=self._vardir)
             self.var = self._exitstack.enter_context(var)
 
@@ -92,6 +87,10 @@ class BuildRoot(contextlib.AbstractContextManager):
             self._mknod(self.dev, "tty", 0o666, 5, 0)
             self._mknod(self.dev, "zero", 0o666, 1, 5)
 
+            # Prepare all registered API endpoints
+            for api in self._apis:
+                self._exitstack.enter_context(api)
+
             self._exitstack = self._exitstack.pop_all()
 
         return self
@@ -100,6 +99,17 @@ class BuildRoot(contextlib.AbstractContextManager):
         self._exitstack.close()
         self._exitstack = None
 
+    def register_api(self, api: "BaseAPI"):
+        """Register an API endpoint.
+
+        The context of the API endpoint will be bound to the context of
+        this `BuildRoot`.
+        """
+        self._apis.append(api)
+
+        if self._exitstack:
+            self._exitstack.enter_context(api)
+
     def run(self, argv, binds=None, readonly_binds=None):
         """Runs a command in the buildroot.
 
@@ -141,9 +151,6 @@ class BuildRoot(contextlib.AbstractContextManager):
         mounts += ["--proc", "/proc"]
         mounts += ["--bind", "/sys", "/sys"]
 
-        # Make osbuild API-calls accessible to the container.
-        mounts += ["--ro-bind", f"{self.api}", "/run/osbuild/api"]
-
         # We execute our own modules by bind-mounting them from the host into
         # the build-root. We have minimal requirements on the build-root, so
         # these modules can be executed. Everything else we provide ourselves.
@@ -164,6 +171,13 @@ class BuildRoot(contextlib.AbstractContextManager):
         for b in readonly_binds or []:
             mounts += ["--ro-bind"] + b.split(":")
 
+        # Prepare all registered API endpoints: bind mount the address with
+        # the `endpoint` name, provided by the API, into the well known path
+        mounts += ["--dir", "/run/osbuild/api"]
+        for api in self._apis:
+            api_path = "/run/osbuild/api/" + api.endpoint
+            mounts += ["--bind", api.socket_address, api_path]
+
         cmd = [
             "bwrap",
             "--cap-add", "CAP_MAC_ADMIN",
diff --git a/osbuild/objectstore.py b/osbuild/objectstore.py
index 4c055cef..03ff7d91 100644
--- a/osbuild/objectstore.py
+++ b/osbuild/objectstore.py
@@ -5,8 +5,9 @@ import os
 import subprocess
 import tempfile
 import weakref
-from typing import Optional, Union
+from typing import Optional
 
+from osbuild.util.types import PathLike
 from osbuild.util import ctx, rmrf
 from . import treesum
 
@@ -16,10 +17,6 @@ __all__ = [
 ]
 
 
-# Type for valid inputs to `os.fspath`
-PathLike = Union[str, bytes, os.PathLike]
-
-
 def mount(source, target, bind=True, ro=True, private=True, mode="0755"):
     options = []
     if bind:
diff --git a/osbuild/pipeline.py b/osbuild/pipeline.py
index ad834be0..45889c45 100644
--- a/osbuild/pipeline.py
+++ b/osbuild/pipeline.py
@@ -78,18 +78,20 @@ class Stage:
 
             ro_binds = [f"{sources_output}:/run/osbuild/sources"]
 
-            with API(f"{build_root.api}/osbuild", args, monitor) as api, \
-                sources.SourcesServer(f"{build_root.api}/sources",
-                                      libdir or "/usr/lib/osbuild",
-                                      self.sources,
-                                      os.path.join(cache, "sources"),
-                                      sources_output):
-                r = build_root.run(
-                    [f"/run/osbuild/lib/stages/{self.name}"],
-                    binds=[os.fspath(tree) + ":/run/osbuild/tree"],
-                    readonly_binds=ro_binds,
-                )
-                return BuildResult(self, r.returncode, api.output)
+            api = API(args, monitor)
+            build_root.register_api(api)
+
+            src = sources.SourcesServer(libdir or "/usr/lib/osbuild",
+                                        self.sources,
+                                        os.path.join(cache, "sources"),
+                                        sources_output)
+            build_root.register_api(src)
+
+            r = build_root.run([f"/run/osbuild/lib/stages/{self.name}"],
+                               binds=[os.fspath(tree) + ":/run/osbuild/tree"],
+                               readonly_binds=ro_binds)
+
+            return BuildResult(self, r.returncode, api.output)
 
 
 class Assembler:
@@ -137,14 +139,17 @@ class Assembler:
 
             ro_binds = [os.fspath(tree) + ":/run/osbuild/tree"]
 
-            with remoteloop.LoopServer(f"{build_root.api}/remoteloop"), \
-                API(f"{build_root.api}/osbuild", args, monitor) as api:
-                r = build_root.run(
-                    [f"/run/osbuild/lib/assemblers/{self.name}"],
-                    binds=binds,
-                    readonly_binds=ro_binds,
-                )
-                return BuildResult(self, r.returncode, api.output)
+            api = API(args, monitor)
+            build_root.register_api(api)
+
+            rls = remoteloop.LoopServer()
+            build_root.register_api(rls)
+
+            r = build_root.run([f"/run/osbuild/lib/assemblers/{self.name}"],
+                               binds=binds,
+                               readonly_binds=ro_binds)
+
+            return BuildResult(self, r.returncode, api.output)
 
 
 class Pipeline:
diff --git a/osbuild/remoteloop.py b/osbuild/remoteloop.py
index 75b0eea9..0306d891 100644
--- a/osbuild/remoteloop.py
+++ b/osbuild/remoteloop.py
@@ -1,19 +1,17 @@
-import asyncio
 import contextlib
 import errno
 import os
-import threading
+from . import api
 from . import loop
 from .util import jsoncomm
 
-
 __all__ = [
     "LoopClient",
     "LoopServer"
 ]
 
 
-class LoopServer:
+class LoopServer(api.BaseAPI):
     """Server for creating loopback devices
 
     The server listens for requests on a AF_UNIX/SOCK_DRGAM sockets.
@@ -37,13 +35,12 @@ class LoopServer:
     object.
     """
 
-    def __init__(self, socket_address):
-        self.socket_address = socket_address
+    endpoint = "remoteloop"
+
+    def __init__(self, *, socket_address=None):
+        super().__init__(socket_address)
         self.devs = []
         self.ctl = loop.LoopControl()
-        self.event_loop = asyncio.new_event_loop()
-        self.thread = threading.Thread(target=self._run_event_loop)
-        self.barrier = threading.Barrier(2)
 
     def _create_device(self, fd, dir_fd, offset=None, sizelimit=None):
         while True:
@@ -73,35 +70,16 @@ class LoopServer:
         self.devs.append(lo)
         return lo.devname
 
-    def _dispatch(self, server):
-        args, fds, addr = server.recv()
-
-        fd = fds[args["fd"]]
-        dir_fd = fds[args["dir_fd"]]
-        offset = args.get("offset")
-        sizelimit = args.get("sizelimit")
+    def _message(self, msg, fds, sock, addr):
+        fd = fds[msg["fd"]]
+        dir_fd = fds[msg["dir_fd"]]
+        offset = msg.get("offset")
+        sizelimit = msg.get("sizelimit")
 
         devname = self._create_device(fd, dir_fd, offset, sizelimit)
-        server.send({"devname": devname}, destination=addr)
-        fds.close()
-
-    def _run_event_loop(self):
-        with jsoncomm.Socket.new_server(self.socket_address) as server:
-            self.barrier.wait()
-            self.event_loop.add_reader(server, self._dispatch, server)
-            asyncio.set_event_loop(self.event_loop)
-            self.event_loop.run_forever()
-            self.event_loop.remove_reader(server)
-
-    def __enter__(self):
-        self.thread.start()
-        self.barrier.wait()
-        return self
-
-    def __exit__(self, *args):
-        self.event_loop.call_soon_threadsafe(self.event_loop.stop)
-        self.thread.join()
-        self.event_loop.close()
+        sock.send({"devname": devname}, destination=addr)
+
+    def _cleanup(self):
         for lo in self.devs:
             lo.close()
 
diff --git a/osbuild/sources.py b/osbuild/sources.py
index 459188fb..2adf9e5b 100644
--- a/osbuild/sources.py
+++ b/osbuild/sources.py
@@ -1,20 +1,19 @@
-import asyncio
 import json
 import subprocess
-import threading
+from . import api
 from .util import jsoncomm
 
 
-class SourcesServer:
-    def __init__(self, socket_address, libdir, options, cache, output):
-        self.socket_address = socket_address
+class SourcesServer(api.BaseAPI):
+
+    endpoint = "sources"
+
+    def __init__(self, libdir, options, cache, output, *, socket_address=None):
+        super().__init__(socket_address)
         self.libdir = libdir
         self.cache = cache
         self.output = output
         self.options = options or {}
-        self.barrier = threading.Barrier(2)
-        self.event_loop = None
-        self.thread = None
 
     def _run_source(self, source, checksums):
         msg = {
@@ -37,39 +36,9 @@ class SourcesServer:
         except ValueError:
             return {"error": f"source returned malformed json: {r.stdout}"}
 
-    def _dispatch(self, server):
-        request, _, addr = server.recv()
-        reply = self._run_source(request["source"], request["checksums"])
-        server.send(reply, destination=addr)
-
-    def _run_event_loop(self):
-        with jsoncomm.Socket.new_server(self.socket_address) as server:
-            self.barrier.wait()
-            self.event_loop.add_reader(server, self._dispatch, server)
-            asyncio.set_event_loop(self.event_loop)
-            self.event_loop.run_forever()
-            self.event_loop.remove_reader(server)
-
-    def __enter__(self):
-        # We are not re-entrant, so complain if re-entered.
-        assert self.event_loop is None
-
-        self.event_loop = asyncio.new_event_loop()
-        self.thread = threading.Thread(target=self._run_event_loop)
-
-        self.barrier.reset()
-        self.thread.start()
-        self.barrier.wait()
-
-        return self
-
-    def __exit__(self, *args):
-        self.event_loop.call_soon_threadsafe(self.event_loop.stop)
-        self.thread.join()
-        self.event_loop.close()
-
-        self.thread = None
-        self.event_loop = None
+    def _message(self, msg, fds, sock, addr):
+        reply = self._run_source(msg["source"], msg["checksums"])
+        sock.send(reply, destination=addr)
 
 
 def get(source, checksums, api_path="/run/osbuild/api/sources"):
diff --git a/osbuild/util/jsoncomm.py b/osbuild/util/jsoncomm.py
index 4042c726..f2c77ff5 100644
--- a/osbuild/util/jsoncomm.py
+++ b/osbuild/util/jsoncomm.py
@@ -14,9 +14,10 @@ import os
 import socket
 from typing import Any
 from typing import Optional
+from .types import PathLike
 
 
-class FdSet():
+class FdSet:
     """File-Descriptor Set
 
     This object wraps an array of file-descriptors. Unlike a normal integer
@@ -137,7 +138,7 @@ class Socket(contextlib.AbstractContextManager):
             self._unlink = None
 
     @classmethod
-    def new_client(cls, connect_to: Optional[str] = None):
+    def new_client(cls, connect_to: Optional[PathLike] = None):
         """Create Client
 
         Create a new client socket.
@@ -163,7 +164,7 @@ class Socket(contextlib.AbstractContextManager):
             # Connect the socket. This has no effect other than specifying the
             # default destination for send operations.
             if connect_to is not None:
-                sock.connect(connect_to)
+                sock.connect(os.fspath(connect_to))
         except:
             if sock is not None:
                 sock.close()
@@ -172,7 +173,7 @@ class Socket(contextlib.AbstractContextManager):
         return cls(sock, None)
 
     @classmethod
-    def new_server(cls, bind_to: str):
+    def new_server(cls, bind_to: PathLike):
         """Create Server
 
         Create a new listener socket.
@@ -199,7 +200,7 @@ class Socket(contextlib.AbstractContextManager):
             # as well. We do not guarantee atomicity, so you better make sure
             # you do not rely on it.
             sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
-            sock.bind(bind_to)
+            sock.bind(os.fspath(bind_to))
             unlink = os.open(os.path.join(".", path[0]), os.O_CLOEXEC | os.O_PATH)
         except:
             if unlink is not None:
diff --git a/osbuild/util/types.py b/osbuild/util/types.py
new file mode 100644
index 00000000..74629012
--- /dev/null
+++ b/osbuild/util/types.py
@@ -0,0 +1,11 @@
+#
+# Define some useful typing abbreviations
+#
+
+import os
+
+from typing import Union
+
+
+#: Represents a file system path. See also `os.fspath`.
+PathLike = Union[str, bytes, os.PathLike]
diff --git a/runners/org.osbuild.fedora30 b/runners/org.osbuild.fedora30
index 3768dbe9..6c234d60 100755
--- a/runners/org.osbuild.fedora30
+++ b/runners/org.osbuild.fedora30
@@ -3,7 +3,7 @@
 import os
 import subprocess
 import sys
-from osbuild.util import jsoncomm
+import osbuild.api
 
 
 def ldconfig():
@@ -36,20 +36,8 @@ def nsswitch():
         pass
 
 
-def setup_stdio():
-    with jsoncomm.Socket.new_client("/run/osbuild/api/osbuild") as client:
-        req = {'method': 'setup-stdio'}
-        client.send(req)
-        msg, fds, _ = client.recv()
-        for io in ['stdin', 'stdout', 'stderr']:
-            target = getattr(sys, io)
-            source = fds[msg[io]]
-            os.dup2(source, target.fileno())
-        fds.close()
-
-
 if __name__ == "__main__":
-    setup_stdio()
+    osbuild.api.setup_stdio()
     ldconfig()
     sysusers()
     tmpfiles()
diff --git a/runners/org.osbuild.linux b/runners/org.osbuild.linux
index ffa1c5b1..7bc3706b 100755
--- a/runners/org.osbuild.linux
+++ b/runners/org.osbuild.linux
@@ -1,25 +1,12 @@
 #!/usr/bin/python3
 
-import os
 import subprocess
 import sys
-from osbuild.util import jsoncomm
-
-
-def setup_stdio():
-    with jsoncomm.Socket.new_client("/run/osbuild/api/osbuild") as client:
-        req = {'method': 'setup-stdio'}
-        client.send(req)
-        msg, fds, _ = client.recv()
-        for io in ['stdin', 'stdout', 'stderr']:
-            target = getattr(sys, io)
-            source = fds[msg[io]]
-            os.dup2(source, target.fileno())
-        fds.close()
+import osbuild.api
 
 
 if __name__ == "__main__":
-    setup_stdio()
+    osbuild.api.setup_stdio()
 
     r = subprocess.run(sys.argv[1:], check=False)
     sys.exit(r.returncode)
diff --git a/runners/org.osbuild.rhel81 b/runners/org.osbuild.rhel81
index 66e63561..933268a3 100755
--- a/runners/org.osbuild.rhel81
+++ b/runners/org.osbuild.rhel81
@@ -3,7 +3,7 @@
 import os
 import subprocess
 import sys
-from osbuild.util import jsoncomm
+import osbuild.api
 
 
 def ldconfig():
@@ -36,17 +36,6 @@ def nsswitch():
         pass
 
 
-def setup_stdio():
-    with jsoncomm.Socket.new_client("/run/osbuild/api/osbuild") as client:
-        req = {'method': 'setup-stdio'}
-        client.send(req)
-        msg, fds, _ = client.recv()
-        for io in ['stdin', 'stdout', 'stderr']:
-            target = getattr(sys, io)
-            source = fds[msg[io]]
-            os.dup2(source, target.fileno())
-        fds.close()
-
 def os_release():
     """/usr/lib/os-release doesn't exist. The `redhat-release` package
     generates `/etc/os-release directly. To work around this, do the same here.
@@ -82,7 +71,7 @@ def python_alternatives():
         pass
 
 if __name__ == "__main__":
-    setup_stdio()
+    osbuild.api.setup_stdio()
     ldconfig()
     sysusers()
     tmpfiles()
diff --git a/runners/org.osbuild.rhel82 b/runners/org.osbuild.rhel82
index 1ab092bc..94fe5d0b 100755
--- a/runners/org.osbuild.rhel82
+++ b/runners/org.osbuild.rhel82
@@ -3,7 +3,7 @@
 import os
 import subprocess
 import sys
-from osbuild.util import jsoncomm
+import osbuild.api
 
 
 def ldconfig():
@@ -36,17 +36,6 @@ def nsswitch():
         pass
 
 
-def setup_stdio():
-    with jsoncomm.Socket.new_client("/run/osbuild/api/osbuild") as client:
-        req = {'method': 'setup-stdio'}
-        client.send(req)
-        msg, fds, _ = client.recv()
-        for io in ['stdin', 'stdout', 'stderr']:
-            target = getattr(sys, io)
-            source = fds[msg[io]]
-            os.dup2(source, target.fileno())
-        fds.close()
-
 
 def python_alternatives():
     """/usr/bin/python3 is a symlink to /etc/alternatives/python3, which points
@@ -60,7 +49,7 @@ def python_alternatives():
         pass
 
 if __name__ == "__main__":
-    setup_stdio()
+    osbuild.api.setup_stdio()
     ldconfig()
     sysusers()
     tmpfiles()
diff --git a/runners/org.osbuild.ubuntu1804 b/runners/org.osbuild.ubuntu1804
index 3768dbe9..aed4be83 100755
--- a/runners/org.osbuild.ubuntu1804
+++ b/runners/org.osbuild.ubuntu1804
@@ -3,7 +3,8 @@
 import os
 import subprocess
 import sys
-from osbuild.util import jsoncomm
+
+import osbuild.api
 
 
 def ldconfig():
@@ -36,20 +37,8 @@ def nsswitch():
         pass
 
 
-def setup_stdio():
-    with jsoncomm.Socket.new_client("/run/osbuild/api/osbuild") as client:
-        req = {'method': 'setup-stdio'}
-        client.send(req)
-        msg, fds, _ = client.recv()
-        for io in ['stdin', 'stdout', 'stderr']:
-            target = getattr(sys, io)
-            source = fds[msg[io]]
-            os.dup2(source, target.fileno())
-        fds.close()
-
-
 if __name__ == "__main__":
-    setup_stdio()
+    osbuild.api.setup_stdio()
     ldconfig()
     sysusers()
     tmpfiles()
diff --git a/stages/org.osbuild.script b/stages/org.osbuild.script
deleted file mode 100755
index 9eeab56b..00000000
--- a/stages/org.osbuild.script
+++ /dev/null
@@ -1,58 +0,0 @@
-#!/usr/bin/python3
-"""
-Run an arbitrary script inside the target tree
-
-Runs an arbitrary script inside the target tree.
-
-Writes the contents of the `script` item to `/osbuild-script`, sets the
-permissions of the script to 0550 (-r-xr-x---), then uses the host's `chroot`
-binary to chroot into the tree and execute the script. The script is removed
-after it completes.
-
-WARNING: running code inside the tree is unsafe, unreliable, and generally
-discouraged. Using this stage may result in unexplained failures or other
-undefined behavior, and should only be done as a last resort.
-
-NOTE: if `script` does not start with a line like '#!/bin/bash
-', executing
-it will fail with ENOEXEC. Some `chroot` binaries will try to run the script
-through `/bin/sh` in that case, so it might still work, but that behavior is
-not guaranteed.
-"""
-
-
-import atexit
-import json
-import os
-import subprocess
-import sys
-
-SCHEMA = """
-"additionalProperties": false,
-"required": ["script"],
-"properties": {
-  "script": {
-    "type": "string",
-    "description": "contents of the script file to be executed."
-  }
-}
-"""
-
-def main(tree, options):
-    script = options["script"]
-
-    scriptfile = f"{tree}/osbuild-script"
-
-    with open(scriptfile, "w") as f:
-        f.write(script)
-
-    os.chmod(scriptfile, 0o550)
-    atexit.register(lambda: os.unlink(scriptfile))
-
-    return subprocess.run(["chroot", tree, "/osbuild-script"], check=False).returncode
-
-
-if __name__ == '__main__':
-    args = json.load(sys.stdin)
-    r = main(args["tree"], args["options"])
-    sys.exit(r)


Test patch:
diff --git a/test/mod/test_api.py b/test/mod/test_api.py
new file mode 100644
index 00000000..014c5c59
--- /dev/null
+++ b/test/mod/test_api.py
@@ -0,0 +1,95 @@
+#
+# Test for API infrastructure
+#
+
+import pathlib
+import os
+import sys
+import tempfile
+import unittest
+
+import osbuild
+from osbuild.buildroot import BuildRoot
+from osbuild.monitor import NullMonitor
+from osbuild.util import jsoncomm
+
+
+class APITester(osbuild.api.BaseAPI):
+    """Records the number of messages and if it got cleaned up"""
+    def __init__(self, sockaddr):
+        super().__init__(sockaddr)
+        self.clean = False
+        self.messages = 0
+
+    endpoint = "test-api"
+
+    def _message(self, msg, _fds, sock, addr):
+        self.messages += 1
+
+        if msg["method"] == "echo":
+            msg["method"] = "reply"
+            sock.send(msg, destination=addr)
+
+    def _cleanup(self):
+        self.clean = True
+
+
+class TestAPI(unittest.TestCase):
+    """Check API infrastructure"""
+    def setUp(self):
+        self.tmp = tempfile.TemporaryDirectory()
+
+    def tearDown(self):
+        self.tmp.cleanup()
+
+    def test_basic(self):
+        # Basic API communication and cleanup checks
+
+        socket = os.path.join(self.tmp.name, "socket")
+        api = APITester(socket)
+        with api:
+            with jsoncomm.Socket.new_client(socket) as client:
+                req = {'method': 'echo', 'data': 'Hello'}
+                client.send(req)
+                msg, _, _ = client.recv()
+                self.assertEqual(msg["method"], "reply")
+                self.assertEqual(req["data"], msg["data"])
+
+        self.assertEqual(api.clean, True)
+        self.assertEqual(api.messages, 1)
+
+        # Assert proper cleanup
+        self.assertIsNone(api.thread)
+        self.assertIsNone(api.event_loop)
+
+    def test_reentrancy_guard(self):
+        socket = os.path.join(self.tmp.name, "socket")
+        api = APITester(socket)
+        with api:
+            with self.assertRaises(AssertionError):
+                with api:
+                    pass
+
+    def test_buildroot(self):
+        # Check API and BuildRoot integration: the runner will call
+        # api.setup_stdio and thus check that connecting to the api
+        # works correctly
+        runner = "org.osbuild.linux"
+        libdir = os.path.abspath(os.curdir)
+        var = pathlib.Path(self.tmp.name, "var")
+        var.mkdir()
+
+        monitor = NullMonitor(sys.stderr.fileno())
+        with BuildRoot("/", runner, libdir=libdir, var=var) as root:
+            api = osbuild.api.API({}, monitor)
+            root.register_api(api)
+
+            r = root.run(["/usr/bin/true"])
+            self.assertEqual(r.returncode, 0)
+
+            # Test we can use `.run` multiple times
+            r = root.run(["/usr/bin/true"])
+            self.assertEqual(r.returncode, 0)
+
+            r = root.run(["/usr/bin/false"])
+            self.assertNotEqual(r.returncode, 0)
diff --git a/test/mod/test_monitor.py b/test/mod/test_monitor.py
index 36a2ddd3..8004a76b 100644
--- a/test/mod/test_monitor.py
+++ b/test/mod/test_monitor.py
@@ -15,23 +15,9 @@ import osbuild
 import osbuild.meta
 from osbuild.api import API
 from osbuild.monitor import LogMonitor
-from osbuild.util import jsoncomm
 from .. import test
 
 
-def setup_stdio(path):
-    """Copy of what the osbuild runners do to setup stdio"""
-    with jsoncomm.Socket.new_client(path) as client:
-        req = {'method': 'setup-stdio'}
-        client.send(req)
-        msg, fds, _ = client.recv()
-        for sio in ['stdin', 'stdout', 'stderr']:
-            target = getattr(sys, sio)
-            source = fds[msg[sio]]
-            os.dup2(source, target.fileno())
-        fds.close()
-
-
 def echo(path):
     """echo stdin to stdout after setting stdio up via API
 
@@ -39,7 +25,7 @@ def echo(path):
     simulating an osbuild runner and a stage run which does
     nothing but returns the supplied options to stdout again.
     """
-    setup_stdio(path)
+    osbuild.api.setup_stdio(path)
     data = json.load(sys.stdin)
     json.dump(data, sys.stdout)
     sys.exit(0)
@@ -88,12 +74,12 @@ class TestMonitor(unittest.TestCase):
             logfile = os.path.join(tmpdir, "log.txt")
 
             with open(logfile, "w") as log, \
-                 API(path, args, LogMonitor(log.fileno())) as api:
+                 API(args, LogMonitor(log.fileno()), socket_address=path) as api:
                 p = mp.Process(target=echo, args=(path, ))
                 p.start()
                 p.join()
                 self.assertEqual(p.exitcode, 0)
-                output = api.output
+                output = api.output  # pylint: disable=no-member
                 assert output
 
             self.assertEqual(json.dumps(args), output)
diff --git a/test/mod/test_util_jsoncomm.py b/test/mod/test_util_jsoncomm.py
index 30bcf0e7..96e58184 100644
--- a/test/mod/test_util_jsoncomm.py
+++ b/test/mod/test_util_jsoncomm.py
@@ -4,6 +4,7 @@
 
 import asyncio
 import os
+import pathlib
 import tempfile
 import unittest
 
@@ -13,7 +14,7 @@ from osbuild.util import jsoncomm
 class TestUtilJsonComm(unittest.TestCase):
     def setUp(self):
         self.dir = tempfile.TemporaryDirectory()
-        self.address = os.path.join(self.dir.name, "listener")
+        self.address = pathlib.Path(self.dir.name, "listener")
         self.server = jsoncomm.Socket.new_server(self.address)
         self.client = jsoncomm.Socket.new_client(self.address)
 
diff --git a/test/run/test_assemblers.py b/test/run/test_assemblers.py
index 3d470b4b..988c3614 100644
--- a/test/run/test_assemblers.py
+++ b/test/run/test_assemblers.py
@@ -25,24 +25,23 @@ class TestAssemblers(test.TestBase):
         self.osbuild = test.OSBuild(self)
 
     @contextlib.contextmanager
-    def run_assembler(self, name, options, output_path):
-        with self.osbuild as osb:
-            with open(os.path.join(self.locate_test_data(),
-                                   "manifests/filesystem.json")) as f:
-                manifest = json.load(f)
-            manifest["pipeline"] = dict(
-                manifest["pipeline"],
-                assembler={"name": name, "options": options}
-            )
-            data = json.dumps(manifest)
-
-            treeid = osb.treeid_from_manifest(data)
-            assert treeid
-
-            osb.compile(data, checkpoints=[treeid])
-            with osb.map_object(treeid) as tree, \
-                 osb.map_output(output_path) as output:
-                yield tree, output
+    def run_assembler(self, osb, name, options, output_path):
+        with open(os.path.join(self.locate_test_data(),
+                               "manifests/filesystem.json")) as f:
+            manifest = json.load(f)
+        manifest["pipeline"] = dict(
+            manifest["pipeline"],
+            assembler={"name": name, "options": options}
+        )
+        data = json.dumps(manifest)
+
+        treeid = osb.treeid_from_manifest(data)
+        assert treeid
+
+        osb.compile(data, checkpoints=[treeid])
+        with osb.map_object(treeid) as tree, \
+             osb.map_output(output_path) as output:
+            yield tree, output
 
     def assertImageFile(self, filename, fmt, expected_size=None):
         info = json.loads(subprocess.check_output(["qemu-img", "info", "--output", "json", filename]))
@@ -95,9 +94,10 @@ class TestAssemblers(test.TestBase):
             "root_fs_uuid": "016a1cda-5182-4ab3-bf97-426b00b74eb0",
             "size": 2 * 1024 * 1024 * 1024
         }
-        with self.run_assembler("org.osbuild.rawfs", options, "image.raw") as (tree, image):
-            self.assertImageFile(image, "raw", options["size"])
-            self.assertFilesystem(image, options["root_fs_uuid"], "ext4", tree)
+        with self.osbuild as osb:
+            with self.run_assembler(osb, "org.osbuild.rawfs", options, "image.raw") as (tree, image):
+                self.assertImageFile(image, "raw", options["size"])
+                self.assertFilesystem(image, options["root_fs_uuid"], "ext4", tree)
 
     @unittest.skipUnless(test.TestBase.have_tree_diff(), "tree-diff missing")
     def test_ostree(self):
@@ -132,57 +132,61 @@ class TestAssemblers(test.TestBase):
     @unittest.skipUnless(test.TestBase.have_tree_diff(), "tree-diff missing")
     def test_qemu(self):
         loctl = loop.LoopControl()
-        for fmt in ["raw", "raw.xz", "qcow2", "vmdk", "vdi"]:
-            with self.subTest(fmt=fmt):
-                print(f"  {fmt}", flush=True)
-                options = {
-                    "format": fmt,
-                    "filename": f"image.{fmt}",
-                    "ptuuid": "b2c09a39-db93-44c5-846a-81e06b1dc162",
-                    "root_fs_uuid": "aff010e9-df95-4f81-be6b-e22317251033",
-                    "size": 2 * 1024 * 1024 * 1024
-                }
-                with self.run_assembler("org.osbuild.qemu",
-                                        options,
-                                        f"image.{fmt}") as (tree, image):
-                    if fmt == "raw.xz":
-                        subprocess.run(["unxz", "--keep", "--force", image], check=True)
-                        image = image[:-3]
-                        fmt = "raw"
-                    self.assertImageFile(image, fmt, options["size"])
-                    with open_image(loctl, image, fmt) as (target, device):
-                        ptable = self.read_partition_table(device)
-                        self.assertPartitionTable(ptable,
-                                                  "dos",
-                                                  options["ptuuid"],
-                                                  1,
-                                                  boot_partition=1)
-                        self.assertGRUB2(device,
-                                         "26e3327c6b5ac9b5e21d8b86f19ff7cb4d12fb2d0406713f936997d9d89de3ee",
-                                         "9b31c8fbc59602a38582988bf91c3948ae9c6f2a231ab505ea63a7005e302147",
-                                         1024 * 1024)
-
-                        p1 = ptable["partitions"][0]
-                        ssize = ptable.get("sectorsize", 512)
-                        start, size = p1["start"] * ssize, p1["size"] * ssize
-                        with loop_open(loctl, target, offset=start, size=size) as dev:
-                            self.assertFilesystem(dev, options["root_fs_uuid"], "ext4", tree)
+        with self.osbuild as osb:
+            for fmt in ["raw", "raw.xz", "qcow2", "vmdk", "vdi"]:
+                with self.subTest(fmt=fmt):
+                    print(f"  {fmt}", flush=True)
+                    options = {
+                        "format": fmt,
+                        "filename": f"image.{fmt}",
+                        "ptuuid": "b2c09a39-db93-44c5-846a-81e06b1dc162",
+                        "root_fs_uuid": "aff010e9-df95-4f81-be6b-e22317251033",
+                        "size": 2 * 1024 * 1024 * 1024
+                    }
+                    with self.run_assembler(osb,
+                                            "org.osbuild.qemu",
+                                            options,
+                                            f"image.{fmt}") as (tree, image):
+                        if fmt == "raw.xz":
+                            subprocess.run(["unxz", "--keep", "--force", image], check=True)
+                            image = image[:-3]
+                            fmt = "raw"
+                        self.assertImageFile(image, fmt, options["size"])
+                        with open_image(loctl, image, fmt) as (target, device):
+                            ptable = self.read_partition_table(device)
+                            self.assertPartitionTable(ptable,
+                                                      "dos",
+                                                      options["ptuuid"],
+                                                      1,
+                                                      boot_partition=1)
+                            self.assertGRUB2(device,
+                                             "26e3327c6b5ac9b5e21d8b86f19ff7cb4d12fb2d0406713f936997d9d89de3ee",
+                                             "9b31c8fbc59602a38582988bf91c3948ae9c6f2a231ab505ea63a7005e302147",
+                                             1024 * 1024)
+
+                            p1 = ptable["partitions"][0]
+                            ssize = ptable.get("sectorsize", 512)
+                            start, size = p1["start"] * ssize, p1["size"] * ssize
+                            with loop_open(loctl, target, offset=start, size=size) as dev:
+                                self.assertFilesystem(dev, options["root_fs_uuid"], "ext4", tree)
 
     def test_tar(self):
         cases = [
             ("tree.tar.gz", None, ["application/x-tar"]),
             ("tree.tar.gz", "gzip", ["application/x-gzip", "application/gzip"])
         ]
-        for filename, compression, expected_mimetypes in cases:
-            options = {"filename": filename}
-            if compression:
-                options["compression"] = compression
-            with self.run_assembler("org.osbuild.tar",
-                                    options,
-                                    filename) as (_, image):
-                output = subprocess.check_output(["file", "--mime-type", image], encoding="utf-8")
-                _, mimetype = output.strip().split(": ") # "filename: mimetype"
-                self.assertIn(mimetype, expected_mimetypes)
+        with self.osbuild as osb:
+            for filename, compression, expected_mimetypes in cases:
+                options = {"filename": filename}
+                if compression:
+                    options["compression"] = compression
+                with self.run_assembler(osb,
+                                        "org.osbuild.tar",
+                                        options,
+                                        filename) as (_, image):
+                    output = subprocess.check_output(["file", "--mime-type", image], encoding="utf-8")
+                    _, mimetype = output.strip().split(": ") # "filename: mimetype"
+                    self.assertIn(mimetype, expected_mimetypes)
 
 
 @contextlib.contextmanager
diff --git a/test/run/test_sources.py b/test/run/test_sources.py
index 2e8830eb..e1f44fc0 100644
--- a/test/run/test_sources.py
+++ b/test/run/test_sources.py
@@ -97,9 +97,9 @@ class TestSources(test.TestBase):
                 with tempfile.TemporaryDirectory() as tmpdir, \
                     fileServer(self.locate_test_data()), \
                     osbuild.sources.SourcesServer(
-                            f"{tmpdir}/sources-api",
                             "./", source_options,
-                            f"{tmpdir}/cache", f"{tmpdir}/dst"):
+                            f"{tmpdir}/cache", f"{tmpdir}/dst",
+                            socket_address=f"{tmpdir}/sources-api"):
                     self.check_case(source, case_options, f"{tmpdir}/sources-api")
                     self.check_case(source, case_options, f"{tmpdir}/sources-api")