i3-benchmark-get-tree.py

#!/usr/bin/env python3

import enum
import fcntl
import json
import contextlib
import signal
import struct
import os
import subprocess
import tempfile
import socket
import time
import typing

DISPLAY = ":3"


class IpcType(enum.Enum):
    RUN_COMMAND = 0
    GET_TREE = 4


def i3_msg(client: socket.socket, message_type: IpcType, payload: str = "") -> bytes:
    payload = payload.encode()
    client.sendall(
        b"i3-ipc" + struct.pack("II", len(payload), message_type.value) + payload
    )

    response = b""
    while len(response) < 14:
        response += client.recv(1024 * 1024)
    assert response[:6] == b"i3-ipc"
    length, message_type = struct.unpack("II", response[6:14])
    response = response[14:]
    while len(response) < length:
        response += client.recv(1024 * 1024)
    assert len(response) == length
    # Do not call `json.loads` here to not interfere with the benchmark.
    return response


@contextlib.contextmanager
def run_i3(
    socket_path: str, args: list[str]
) -> typing.Generator[socket.socket, None, None]:
    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
        sock.bind(socket_path)
        sock.listen(1)
        fcntl.fcntl(sock.fileno(), fcntl.F_SETFD, 0)
        assert sock.fileno() == 3

        i3_pid = os.fork()
        if i3_pid == 0:
            os.execvpe(
                file="i3",
                args=["i3", *args],
                env={k: v for k, v in os.environ.items() if k != "I3SOCK"}
                | {
                    "DISPLAY": DISPLAY,
                    "LISTEN_PID": str(os.getpid()),
                    "LISTEN_FDS": "1",
                },
            )
        try:
            with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client:
                client.connect(socket_path)
                yield client
        finally:
            os.kill(i3_pid, signal.SIGTERM)
            os.waitpid(i3_pid, 0)
            os.remove(socket_path)


def count_nodes(node: dict) -> int:
    return 1 + sum(count_nodes(n) for n in node.get("nodes", []))


def main() -> None:
    subprocess.run(["i3", "-v"], check=True)

    with tempfile.TemporaryDirectory() as tmpdir:
        with subprocess.Popen(["Xephyr", "-screen", "640x480", DISPLAY]) as xephyr:
            with open(f"{tmpdir}/config", "w") as config_file:
                config_file.write("font pango:Sans 10\n")

            layout = {"swallows": [{"class": "^XTerm$"}], "type": "con"}
            for i in range(8):
                with open(f"{tmpdir}/layout.json", "w") as layout_file:
                    json.dump(layout, layout_file)
                with run_i3(f"{tmpdir}/i3-ipc.sock", [f"-c{tmpdir}/config"]) as client:
                    i3_msg(
                        client,
                        IpcType.RUN_COMMAND,
                        f"append_layout {tmpdir}/layout.json",
                    )

                    # Besides counting the number of nodes, the following line is used
                    # as a barrier to ensure that i3 is done running `append_layout`.
                    tree = json.loads(i3_msg(client, IpcType.GET_TREE))
                    nodes = count_nodes(tree)

                    times_ms = []
                    for _ in range(500):
                        start = time.monotonic()
                        i3_msg(client, IpcType.GET_TREE)
                        end = time.monotonic()
                        times_ms.append((end - start) * 1000.0)
                    times_ms.sort()

                    print(
                        f"nodes={nodes:3}",
                        f"min={times_ms[0]:5.2f}",
                        f"median={times_ms[len(times_ms) // 2]:5.2f}",
                        f"max={times_ms[-1]:5.2f}",
                        f"avg={sum(times_ms) / len(times_ms):5.2f}",
                    )

                layout = {
                    "layout": ["splith", "splitv"][i % 2],
                    "nodes": [layout, layout],
                }

            xephyr.kill()


if __name__ == "__main__":
    main()