From dc3bb10f2e9e2757fc63d4a971e697ad015e24ca Mon Sep 17 00:00:00 2001
From: Justin Ehlert <justinehlert@flywheel.io>
Date: Wed, 28 Mar 2018 15:59:40 -0500
Subject: [PATCH] Provide an option for soft file locking

---
 xvfbwrapper.py | 41 +++++++++++++++++++++++++++++------------
 1 file changed, 29 insertions(+), 12 deletions(-)

diff --git a/xvfbwrapper.py b/xvfbwrapper.py
index 5fda994..85c747b 100644
--- a/xvfbwrapper.py
+++ b/xvfbwrapper.py
@@ -14,12 +14,7 @@
 
 from random import randint
 
-try:
-    BlockingIOError
-except NameError:
-    # python 2
-    BlockingIOError = IOError
-
+SOFT_FILE_LOCK = os.environ.get('XVFB_WRAPPER_SOFT_FILE_LOCK')
 
 class Xvfb(object):
 
@@ -126,14 +121,36 @@ def _get_next_unused_display(self):
         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
+            if SOFT_FILE_LOCK:
+                self._lock_display_file = self._soft_lock_file(tempfile_path.format(rand))
             else:
+                self._lock_display_file = self._hard_lock_file(tempfile_path.format(rand))
+
+            if self._lock_display_file:
                 return rand
 
     def _set_display_var(self, display):
         os.environ['DISPLAY'] = ':{}'.format(display)
+
+    def _hard_lock_file(self, path):
+        fd = os.open(path, os.O_RDWR | os.O_CREAT | os.O_TRUNC)
+        
+        try:
+            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except (IOError, OSError):
+            os.close(fd)
+        else:
+            return fd
+        
+        return None 
+
+    def _soft_lock_file(self, path):
+        try:
+            fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC | os.O_EXCL)
+        except (IOError, OSError):
+            pass
+        else:
+            return fd
+
+        return None
+