use file locking for _get_next_unused_display

This commit is contained in:
Stephen J. Fuhry 2016-06-30 18:15:18 -04:00
parent 2bdee33b63
commit b7c6daad71
4 changed files with 103 additions and 23 deletions

View File

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

View File

@ -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)

View File

@ -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

View File

@ -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)