Merge pull request #17 from MontmereLimited/file_locking_for_display

Fix multiprocessing X display collisions by utilizing `flock`

.gitignore

@@ -1,11 +1,89 @@
-MANIFEST
+# Byte-compiled / optimized / DLL files
+__pycache__/
 *.py[cod]
-__pycache__
-_build
-build
-dist
-sdist
-.tox
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
 *.egg
-*.egg-info
-eggs
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*,cover
+.hypothesis/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# IPython Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# dotenv
+.env
+
+# virtualenv
+venv/
+ENV/
+
+# Spyder project settings
+.spyderproject
+
+# Rope project settings
+.ropeproject

.travis.yml

@@ -1,10 +1,10 @@
 language: python
 python:
   - "2.7"
-  - "3.2"
   - "3.3"
   - "3.4"
   - "3.5"
+  - "pypy"
 before_install:
   - "sudo apt-get update -qq"
   - "sudo apt-get install -y xvfb"

setup.py

@@ -5,17 +5,26 @@
 
 
 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',
-    version='0.2.9dev',
+    version='0.2.9.dev0',
     py_modules=['xvfbwrapper'],
     author='Corey Goldberg',
     author_email='cgoldberg _at_ gmail.com',
@@ -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=[

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
 
@@ -14,6 +19,12 @@ class TestXvfb(unittest.TestCase):
     def setUp(self):
         self.reset_display()
 
+    def test_xvfb_binary_not_exists(self):
+        with patch('xvfbwrapper.Xvfb.xvfb_exists') as xvfb_exists:
+            xvfb_exists.return_value = False
+            with self.assertRaises(EnvironmentError):
+                Xvfb()
+
     def test_start(self):
         xvfb = Xvfb()
         self.addCleanup(xvfb.stop)
@@ -76,3 +87,27 @@ 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)
+        side_effect = [11, 11, 22, 11, 22, 11, 22, 22, 22, 33]
+        with patch('xvfbwrapper.randint',
+                   side_effect=side_effect) as mockrandint:
+            self.assertEqual(xvfb._get_next_unused_display(), 11)
+            self.assertEqual(mockrandint.call_count, 1)
+            if sys.version_info >= (3, 2):
+                with self.assertWarns(ResourceWarning):
+                    self.assertEqual(xvfb2._get_next_unused_display(), 22)
+                    self.assertEqual(mockrandint.call_count, 3)
+                    self.assertEqual(xvfb3._get_next_unused_display(), 33)
+                    self.assertEqual(mockrandint.call_count, 10)
+            else:
+                self.assertEqual(xvfb2._get_next_unused_display(), 22)
+                self.assertEqual(mockrandint.call_count, 3)
+                self.assertEqual(xvfb3._get_next_unused_display(), 33)
+                self.assertEqual(mockrandint.call_count, 10)

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

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)