use file locking for _get_next_unused_display

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=[

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)

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)