mirror of
https://github.com/daijro/camoufox.git
synced 2026-08-25 00:00:25 +00:00
fix(virtdisplay): atomically claim X11 display via Xvfb -displayfd (#597)
Replace the userspace lock-file scan + random-jitter retry loop with Xvfb's own -displayfd mechanism. Xvfb scans up from :0 and atomically binds the first free X11 socket (kernel-mediated, no userspace race), then writes the chosen display number back through an inherited pipe. This eliminates the duplicate-display race that occurred when many camoufox processes started concurrently and all observed the same set of free display numbers before any of them bound. Adds a 10s read timeout on the displayfd pipe so a hung Xvfb fails fast instead of blocking forever, and adds tests covering single launch, idempotent get(), 50 concurrent reservations with uniqueness, and post-kill display reuse.
This commit is contained in:
Jake Writer
GitHub
@@ -1,10 +1,9 @@
|
||||
import os
|
||||
import select
|
||||
import subprocess # nosec
|
||||
from glob import glob
|
||||
from multiprocessing import Lock
|
||||
from random import randrange
|
||||
import time
|
||||
from shutil import which
|
||||
from typing import List, Optional
|
||||
from typing import Optional
|
||||
|
||||
from camoufox.exceptions import (
|
||||
CannotExecuteXvfb,
|
||||
@@ -13,20 +12,12 @@ from camoufox.exceptions import (
|
||||
)
|
||||
from camoufox.pkgman import OS_NAME
|
||||
|
||||
# Safe timeout for Xvfb writing display num, prevents infinite hang.
|
||||
DISPLAYFD_READ_TIMEOUT_S = 10.0
|
||||
|
||||
|
||||
class VirtualDisplay:
|
||||
"""
|
||||
A minimal virtual display implementation for Linux.
|
||||
"""
|
||||
|
||||
def __init__(self, debug: Optional[bool] = False) -> None:
|
||||
"""
|
||||
Constructor for the VirtualDisplay class (singleton object).
|
||||
"""
|
||||
self.debug = debug
|
||||
self.proc: Optional[subprocess.Popen] = None
|
||||
self._display: Optional[int] = None
|
||||
self._lock = Lock()
|
||||
"""A minimal virtual display implementation for Linux."""
|
||||
|
||||
xvfb_args = (
|
||||
# fmt: off
|
||||
@@ -39,18 +30,19 @@ class VirtualDisplay:
|
||||
"-extension", "XVideo",
|
||||
"-extension", "XVideo-MotionCompensation",
|
||||
"-extension", "XINERAMA",
|
||||
"-shmem",
|
||||
"-fp", "built-ins",
|
||||
"-nocursor",
|
||||
"-br",
|
||||
# fmt: on
|
||||
)
|
||||
|
||||
def __init__(self, debug: bool = False) -> None:
|
||||
self.debug = debug
|
||||
self.proc: Optional[subprocess.Popen] = None
|
||||
self._display: Optional[int] = None
|
||||
|
||||
@property
|
||||
def xvfb_path(self) -> str:
|
||||
"""
|
||||
Get the path to the xvfb executable
|
||||
"""
|
||||
path = which("Xvfb")
|
||||
if not path:
|
||||
raise CannotFindXvfb("Please install Xvfb to use headless mode.")
|
||||
@@ -58,89 +50,74 @@ class VirtualDisplay:
|
||||
raise CannotExecuteXvfb(f"I do not have permission to execute Xvfb: {path}")
|
||||
return path
|
||||
|
||||
@property
|
||||
def xvfb_cmd(self) -> List[str]:
|
||||
"""
|
||||
Get the xvfb command
|
||||
"""
|
||||
return [self.xvfb_path, f':{self.display}', *self.xvfb_args]
|
||||
|
||||
def execute_xvfb(self):
|
||||
"""
|
||||
Spawn a detatched process
|
||||
"""
|
||||
if self.debug:
|
||||
print('Starting virtual display:', ' '.join(self.xvfb_cmd))
|
||||
self.proc = subprocess.Popen( # nosec
|
||||
self.xvfb_cmd,
|
||||
stdout=None if self.debug else subprocess.DEVNULL,
|
||||
stderr=None if self.debug else subprocess.DEVNULL,
|
||||
)
|
||||
|
||||
def get(self) -> str:
|
||||
"""
|
||||
Get the display number
|
||||
"""
|
||||
self.assert_linux()
|
||||
self._assert_linux()
|
||||
|
||||
with self._lock:
|
||||
if self.proc is None:
|
||||
self.execute_xvfb()
|
||||
elif self.debug:
|
||||
print(f'Using virtual display: {self.display}')
|
||||
return f':{self.display}'
|
||||
if self.proc is None:
|
||||
# Launch Xvfb with -displayfd so Xvfb itself picks a free display
|
||||
# number atomically and reports it back. Avoids userspace races.
|
||||
# subprocess.Popen's pass_fds keeps an fd at its parent number in
|
||||
# the child (unlike Node's `stdio: [..., 'pipe']` which renumbers
|
||||
# to 3), so we tell Xvfb that exact number.
|
||||
read_fd, write_fd = os.pipe()
|
||||
cmd = [self.xvfb_path, "-displayfd", str(write_fd), *self.xvfb_args]
|
||||
if self.debug:
|
||||
print("Starting virtual display:", " ".join(cmd))
|
||||
self.proc = subprocess.Popen( # nosec
|
||||
cmd,
|
||||
stdin=subprocess.DEVNULL,
|
||||
stdout=None if self.debug else subprocess.DEVNULL,
|
||||
stderr=None if self.debug else subprocess.DEVNULL,
|
||||
start_new_session=True,
|
||||
pass_fds=(write_fd,),
|
||||
env={
|
||||
**os.environ,
|
||||
# Force Mesa software GLX; we don't use the GPU anyway.
|
||||
"__GLX_VENDOR_LIBRARY_NAME": "mesa",
|
||||
"LIBGL_ALWAYS_SOFTWARE": "1",
|
||||
},
|
||||
)
|
||||
os.close(write_fd) # so the read end EOFs when Xvfb closes its end
|
||||
|
||||
def kill(self):
|
||||
"""
|
||||
Terminate the xvfb process
|
||||
"""
|
||||
with self._lock:
|
||||
if self.proc and self.proc.poll() is None:
|
||||
if self.debug:
|
||||
print('Terminating virtual display:', self.display)
|
||||
self.proc.terminate()
|
||||
buf = b""
|
||||
deadline = time.monotonic() + DISPLAYFD_READ_TIMEOUT_S
|
||||
try:
|
||||
while b"\n" not in buf:
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0 or not select.select([read_fd], [], [], remaining)[0]:
|
||||
self.kill()
|
||||
raise CannotExecuteXvfb(
|
||||
f"Xvfb did not report a display within "
|
||||
f"{int(DISPLAYFD_READ_TIMEOUT_S * 1000)}ms"
|
||||
)
|
||||
chunk = os.read(read_fd, 64)
|
||||
if not chunk:
|
||||
self.kill()
|
||||
raise CannotExecuteXvfb(
|
||||
f"Xvfb did not report a display "
|
||||
f"(got {buf!r}, exit={self.proc.poll()})"
|
||||
)
|
||||
buf += chunk
|
||||
finally:
|
||||
os.close(read_fd)
|
||||
|
||||
def __del__(self):
|
||||
"""
|
||||
Kill and delete the VirtualDisplay object
|
||||
"""
|
||||
self.kill()
|
||||
try:
|
||||
self._display = int(buf.strip())
|
||||
except ValueError:
|
||||
self.kill()
|
||||
raise CannotExecuteXvfb(f"Xvfb wrote non-integer display: {buf!r}")
|
||||
elif self.debug:
|
||||
print(f"Using virtual display: {self._display}")
|
||||
|
||||
return f":{self._display}"
|
||||
|
||||
def kill(self) -> None:
|
||||
if self.proc and self.proc.poll() is None:
|
||||
if self.debug:
|
||||
print("Terminating virtual display:", self._display)
|
||||
self.proc.terminate()
|
||||
|
||||
@staticmethod
|
||||
def _get_lock_files() -> List[str]:
|
||||
"""
|
||||
Get list of lock files in /tmp
|
||||
"""
|
||||
tmpd = os.environ.get('TMPDIR', '/tmp') # nosec
|
||||
try:
|
||||
lock_files = glob(os.path.join(tmpd, ".X*-lock"))
|
||||
except FileNotFoundError:
|
||||
return []
|
||||
return [p for p in lock_files if os.path.isfile(p)]
|
||||
|
||||
@staticmethod
|
||||
def _free_display() -> int:
|
||||
"""
|
||||
Search for free display
|
||||
"""
|
||||
ls = list(
|
||||
map(lambda x: int(x.split("X")[1].split("-")[0]), VirtualDisplay._get_lock_files())
|
||||
)
|
||||
return max(99, max(ls) + randrange(3, 20)) if ls else 99 # nosec
|
||||
|
||||
@property
|
||||
def display(self) -> int:
|
||||
"""
|
||||
Get the display number
|
||||
"""
|
||||
if self._display is None:
|
||||
self._display = self._free_display()
|
||||
return self._display
|
||||
|
||||
@staticmethod
|
||||
def assert_linux():
|
||||
"""
|
||||
Assert that the current OS is Linux
|
||||
"""
|
||||
if OS_NAME != 'lin':
|
||||
def _assert_linux() -> None:
|
||||
if OS_NAME != "lin":
|
||||
raise VirtualDisplayNotSupported("Virtual display is only supported on Linux.")
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
"""
|
||||
Tests for camoufox.virtdisplay.
|
||||
|
||||
Mirrors camoufox-js/test/virtdisplay.test.ts.
|
||||
|
||||
Run with:
|
||||
cd pythonlib && python -m pytest tests/test_virtdisplay.py -v
|
||||
|
||||
VIRTDISPLAY_TEST_N controls the concurrent-launch count. Default is kept
|
||||
low so the test passes on any developer box; set to e.g. 1000 to exercise
|
||||
real scaling. At high N you will need `ulimit -n` headroom -- each Xvfb
|
||||
takes one X11 socket plus our -displayfd pipe.
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from typing import List, Set
|
||||
|
||||
import pytest
|
||||
|
||||
# Make `import camoufox` resolve to the in-tree pythonlib without an install.
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
||||
|
||||
from camoufox.virtdisplay import VirtualDisplay # noqa: E402
|
||||
|
||||
DISPLAY_RE = re.compile(r"^:\d+$")
|
||||
N = int(os.environ.get("VIRTDISPLAY_TEST_N", "50"))
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
sys.platform != "linux", reason="VirtualDisplay is Linux-only"
|
||||
)
|
||||
|
||||
|
||||
# Track every VirtualDisplay we spawn so the fixture can guarantee
|
||||
# cleanup even if an assertion fails mid-test.
|
||||
@pytest.fixture
|
||||
def tracked() -> List[VirtualDisplay]:
|
||||
items: List[VirtualDisplay] = []
|
||||
yield items
|
||||
for vd in items:
|
||||
try:
|
||||
vd.kill()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _track(items: List[VirtualDisplay], vd: VirtualDisplay) -> VirtualDisplay:
|
||||
items.append(vd)
|
||||
return vd
|
||||
|
||||
|
||||
def _wait_for_exit(vd: VirtualDisplay, timeout_s: float = 5.0) -> None:
|
||||
deadline = time.monotonic() + timeout_s
|
||||
while time.monotonic() < deadline:
|
||||
if vd.proc is None or vd.proc.poll() is not None:
|
||||
return
|
||||
time.sleep(0.025)
|
||||
|
||||
|
||||
def test_single_launch_returns_valid_display_and_kill_terminates_xvfb(tracked):
|
||||
vd = _track(tracked, VirtualDisplay())
|
||||
display = vd.get()
|
||||
assert DISPLAY_RE.match(display), display
|
||||
assert vd.proc is not None
|
||||
assert vd.proc.poll() is None # alive
|
||||
|
||||
vd.kill()
|
||||
_wait_for_exit(vd)
|
||||
assert vd.proc.poll() is not None # exited
|
||||
|
||||
|
||||
def test_get_is_idempotent_within_one_virtual_display(tracked):
|
||||
vd = _track(tracked, VirtualDisplay())
|
||||
a = vd.get()
|
||||
b = vd.get()
|
||||
assert a == b
|
||||
|
||||
|
||||
def test_concurrent_reservations_all_get_unique_displays(tracked):
|
||||
# Every VirtualDisplay spawns its own Xvfb. Each Xvfb scans up from :0
|
||||
# and atomically claims the first free X11 socket (kernel-mediated
|
||||
# bind, no userspace race). -displayfd reports the chosen number back
|
||||
# to us. A duplicate here would mean we mis-parsed or mis-routed the
|
||||
# displayfd output, or two Xvfbs somehow bound the same socket.
|
||||
vds = [_track(tracked, VirtualDisplay()) for _ in range(N)]
|
||||
|
||||
displays = [vd.get() for vd in vds]
|
||||
|
||||
for d in displays:
|
||||
assert DISPLAY_RE.match(d), d
|
||||
|
||||
unique: Set[str] = set(displays)
|
||||
assert len(unique) == len(displays), (
|
||||
f"duplicate displays: {len(displays) - len(unique)} of {len(displays)}"
|
||||
)
|
||||
|
||||
# Every Xvfb is alive.
|
||||
for vd in vds:
|
||||
assert vd.proc is not None and vd.proc.poll() is None
|
||||
|
||||
# Tear them all down and confirm every Xvfb actually exited -- no
|
||||
# leaked processes.
|
||||
for vd in vds:
|
||||
vd.kill()
|
||||
for vd in vds:
|
||||
_wait_for_exit(vd)
|
||||
for vd in vds:
|
||||
assert vd.proc.poll() is not None
|
||||
|
||||
|
||||
def test_released_display_numbers_can_be_reused_on_the_next_launch(tracked):
|
||||
a = _track(tracked, VirtualDisplay())
|
||||
a_display = a.get()
|
||||
|
||||
a.kill()
|
||||
_wait_for_exit(a)
|
||||
|
||||
# Spawning a new Xvfb after release must succeed. The new display
|
||||
# number may or may not equal a_display -- Xvfb's allocation order is
|
||||
# its concern -- but we must get *some* display.
|
||||
b = _track(tracked, VirtualDisplay())
|
||||
b_display = b.get()
|
||||
assert DISPLAY_RE.match(b_display), b_display
|
||||
|
||||
b.kill()
|
||||
_wait_for_exit(b)
|
||||
|
||||
# Sanity: a_display was a valid form too.
|
||||
assert DISPLAY_RE.match(a_display), a_display
|
||||
Reference in New Issue
Block a user