diff --git a/setup.py b/setup.py index 6c83dad..eb48426 100644 --- a/setup.py +++ b/setup.py @@ -5,13 +5,22 @@ import os -from distutils.core import setup + +try: + from setuptools import setup +except ImportError: + from distutils.core import setup this_dir = os.path.abspath(os.path.dirname(__file__)) with open(os.path.join(this_dir, 'README.rst')) as f: LONG_DESCRIPTION = '\n' + f.read() +tests_require = [] +try: + from unittest import mock # noqa +except ImportError: + tests_require.append('mock') setup( name='xvfbwrapper', @@ -23,6 +32,7 @@ setup( long_description=LONG_DESCRIPTION, url='https://github.com/cgoldberg/xvfbwrapper', download_url='http://pypi.python.org/pypi/xvfbwrapper', + tests_require=tests_require, keywords='xvfb virtual display headless x11'.split(), license='MIT', classifiers=[ diff --git a/test_xvfb.py b/test_xvfb.py index aa0420e..47e79d0 100644 --- a/test_xvfb.py +++ b/test_xvfb.py @@ -1,7 +1,12 @@ #!/usr/bin/env python import os +import sys import unittest +try: + from unittest.mock import patch +except ImportError: + from mock import patch from xvfbwrapper import Xvfb @@ -76,3 +81,21 @@ class TestXvfb(unittest.TestCase): xvfb = Xvfb(foo='bar') with self.assertRaises(RuntimeError): xvfb.start() + + def test_get_next_unused_display_does_not_reuse_lock(self): + xvfb = Xvfb() + xvfb2 = Xvfb() + xvfb3 = Xvfb() + self.addCleanup(xvfb._cleanup_lock_file) + self.addCleanup(xvfb2._cleanup_lock_file) + self.addCleanup(xvfb3._cleanup_lock_file) + with patch('xvfbwrapper.randint', + side_effect=[11, 11, 22, 11, 22, 11, 22, 22, 22, 33]): + self.assertEqual(xvfb._get_next_unused_display(), 11) + if sys.version_info >= (3, 2): + with self.assertWarns(ResourceWarning): + self.assertEqual(xvfb2._get_next_unused_display(), 22) + self.assertEqual(xvfb3._get_next_unused_display(), 33) + else: + self.assertEqual(xvfb2._get_next_unused_display(), 22) + self.assertEqual(xvfb3._get_next_unused_display(), 33) diff --git a/tox.ini b/tox.ini index 1ff8f08..46f7f2d 100644 --- a/tox.ini +++ b/tox.ini @@ -8,6 +8,8 @@ envlist=flake8,py27,py33,py34,py35,pypy [testenv] commands={envpython} -m unittest discover +deps= + py27,pypy: mock [testenv:flake8] deps=flake8 diff --git a/xvfbwrapper.py b/xvfbwrapper.py index a2ad782..b6b6448 100644 --- a/xvfbwrapper.py +++ b/xvfbwrapper.py @@ -9,18 +9,32 @@ import os -import fnmatch import subprocess import tempfile import time +import fcntl +from random import randint + +try: + BlockingIOError +except NameError: + # python 2 + BlockingIOError = IOError -class Xvfb: +class Xvfb(object): - def __init__(self, width=800, height=680, colordepth=24, **kwargs): + # Maximum value to use for a display. 32-bit maxint is the + # highest Xvfb currently supports + MAX_DISPLAY = 2147483647 + SLEEP_TIME_BEFORE_START = 0.1 + + def __init__(self, width=800, height=680, colordepth=24, tempdir=None, + **kwargs): self.width = width self.height = height self.colordepth = colordepth + self._tempdir = tempdir or tempfile.gettempdir() if not self.xvfb_exists(): msg = 'Can not find Xvfb. Please install it and try again.' @@ -55,34 +69,65 @@ class Xvfb: stdout=fnull, stderr=fnull, close_fds=True) - time.sleep(0.1) # give Xvfb time to start + # give Xvfb time to start + time.sleep(self.__class__.SLEEP_TIME_BEFORE_START) ret_code = self.proc.poll() if ret_code is None: self._set_display_var(self.new_display) else: + self._cleanup_lock_file() raise RuntimeError('Xvfb did not start') def stop(self): - if self.orig_display is None: - del os.environ['DISPLAY'] - else: - self._set_display_var(self.orig_display) - if self.proc is not None: - try: - self.proc.terminate() - self.proc.wait() - except OSError: - pass - self.proc = None + try: + if self.orig_display is None: + del os.environ['DISPLAY'] + else: + self._set_display_var(self.orig_display) + if self.proc is not None: + try: + self.proc.terminate() + self.proc.wait() + except OSError: + pass + self.proc = None + finally: + self._cleanup_lock_file() + + def _cleanup_lock_file(self): + ''' + This should always get called if the process exits safely + with Xvfb.stop() (whether called explicitly, or by __exit__). + + If you are ending up with /tmp/X123-lock files when Xvfb is not + running, then Xvfb is not exiting cleanly. Always either call + Xvfb.stop() in a finally block, or use Xvfb as a context manager + to ensure lock files are purged. + + ''' + self._lock_display_file.close() + try: + os.remove(self._lock_display_file.name) + except OSError: + pass def _get_next_unused_display(self): - tmpdir = tempfile.gettempdir() - pattern = '.X*-lock' - lockfile_names = fnmatch.filter(os.listdir(tmpdir), pattern) - existing_displays = [int(name.split('X')[1].split('-')[0]) - for name in lockfile_names] - highest_display = max(existing_displays) if existing_displays else 0 - return highest_display + 1 + ''' + In order to ensure multi-process safety, this method attempts + to acquire an exclusive lock on a temporary file whose name + contains the display number for Xvfb. + ''' + tempfile_path = os.path.join(self._tempdir, '.X{0}-lock') + while True: + rand = randint(1, self.__class__.MAX_DISPLAY) + self._lock_display_file = open(tempfile_path.format(rand), 'w') + try: + fcntl.flock(self._lock_display_file, + fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + continue + else: + return rand def _set_display_var(self, display): os.environ['DISPLAY'] = ':{}'.format(display)