use file locking for _get_next_unused_display
This commit is contained in:
parent
2bdee33b63
commit
b7c6daad71
12
setup.py
12
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=[
|
||||
|
23
test_xvfb.py
23
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)
|
||||
|
2
tox.ini
2
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
|
||||
|
@ -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)
|
||||
|
Loading…
Reference in New Issue
Block a user