Linux: the implementation is now thread-safe (fixes BoboTiG#169)

Author: BoboTiG

CHANGELOG

@@ -5,6 +5,7 @@ History:
 5.1.1   2020/xx/xx
       - removed usage of deprecated "license_file" option for "license_files"
       - fixed flake8 usage in pre-commit
+      - Linux: the implementation is now thread-safe (fixes #169)
       - :heart: contributors: @
 
 5.1.0   2020/04/30

mss/linux.py

@@ -6,6 +6,7 @@
 import ctypes
 import ctypes.util
 import os
+import threading
 from types import SimpleNamespace
 from typing import TYPE_CHECKING
 
@@ -187,6 +188,12 @@ class MSS(MSSBase):
     # Instancied one time to prevent resource leak.
     display = None
 
+    # A dict to maintain *display* values created by multiple threads.
+    _display_dict = {}  # type: Dict[threading.Thread, int]
+
+    # A threading lock to lock resources.
+    _lock = threading.Lock()
+
     def __init__(self, display=None):
         # type: (Optional[Union[bytes, str]]) -> None
         """ GNU/Linux initialisations. """
@@ -221,14 +228,29 @@ def __init__(self, display=None):
 
         self._set_cfunctions()
 
-        if not MSS.display:
-            MSS.display = self.xlib.XOpenDisplay(display)
-        self.root = self.xlib.XDefaultRootWindow(MSS.display)
+        self.root = self.xlib.XDefaultRootWindow(self._get_display(display))
 
         # Fix for XRRGetScreenResources and XGetImage:
         #     expected LP_Display instance instead of LP_XWindowAttributes
         self.drawable = ctypes.cast(self.root, ctypes.POINTER(Display))
 
+    def _get_display(self, disp=None):
+        """
+        Retrieve a thread-safe display from XOpenDisplay().
+        In multithreading, if the thread who creates *display* is dead, *display* will
+        no longer be valid to grab the screen. The *display* attribute is replaced
+        with *_display_dict* to maintain the *display* values in multithreading.
+        Since the current thread and main thread are always alive, reuse their
+        *display* value first.
+        """
+        cur_thread, main_thread = threading.current_thread(), threading.main_thread()
+        display = MSS._display_dict.get(cur_thread) or MSS._display_dict.get(
+            main_thread
+        )
+        if not display:
+            display = MSS._display_dict[cur_thread] = self.xlib.XOpenDisplay(disp)
+        return display
+
     def _set_cfunctions(self):
         """
         Set all ctypes functions and attach them to attributes.
@@ -324,7 +346,7 @@ def get_error_details(self):
             ERROR.details = None
             xserver_error = ctypes.create_string_buffer(1024)
             self.xlib.XGetErrorText(
-                MSS.display,
+                self._get_display(),
                 details.get("xerror_details", {}).get("error_code", 0),
                 xserver_error,
                 len(xserver_error),
@@ -335,53 +357,66 @@ def get_error_details(self):
 
         return details
 
-    @property
-    def monitors(self):
-        # type: () -> Monitors
-        """ Get positions of monitors (see parent class property). """
+    def _monitors_impl(self):
+        # type: () -> None
+        """
+        Get positions of monitors (has to be run using a threading lock).
+        It will populate self._monitors.
+        """
 
-        if not self._monitors:
-            display = MSS.display
-            int_ = int
-            xrandr = self.xrandr
+        display = self._get_display()
+        int_ = int
+        xrandr = self.xrandr
+
+        # All monitors
+        gwa = XWindowAttributes()
+        self.xlib.XGetWindowAttributes(display, self.root, ctypes.byref(gwa))
+        self._monitors.append(
+            {
+                "left": int_(gwa.x),
+                "top": int_(gwa.y),
+                "width": int_(gwa.width),
+                "height": int_(gwa.height),
+            }
+        )
+
+        # Each monitors
+        mon = xrandr.XRRGetScreenResourcesCurrent(display, self.drawable).contents
+        crtcs = mon.crtcs
+        for idx in range(mon.ncrtc):
+            crtc = xrandr.XRRGetCrtcInfo(display, mon, crtcs[idx]).contents
+            if crtc.noutput == 0:
+                xrandr.XRRFreeCrtcInfo(crtc)
+                continue
 
-            # All monitors
-            gwa = XWindowAttributes()
-            self.xlib.XGetWindowAttributes(display, self.root, ctypes.byref(gwa))
             self._monitors.append(
                 {
-                    "left": int_(gwa.x),
-                    "top": int_(gwa.y),
-                    "width": int_(gwa.width),
-                    "height": int_(gwa.height),
+                    "left": int_(crtc.x),
+                    "top": int_(crtc.y),
+                    "width": int_(crtc.width),
+                    "height": int_(crtc.height),
                 }
             )
+            xrandr.XRRFreeCrtcInfo(crtc)
+        xrandr.XRRFreeScreenResources(mon)
 
-            # Each monitors
-            mon = xrandr.XRRGetScreenResourcesCurrent(display, self.drawable).contents
-            crtcs = mon.crtcs
-            for idx in range(mon.ncrtc):
-                crtc = xrandr.XRRGetCrtcInfo(display, mon, crtcs[idx]).contents
-                if crtc.noutput == 0:
-                    xrandr.XRRFreeCrtcInfo(crtc)
-                    continue
-
-                self._monitors.append(
-                    {
-                        "left": int_(crtc.x),
-                        "top": int_(crtc.y),
-                        "width": int_(crtc.width),
-                        "height": int_(crtc.height),
-                    }
-                )
-                xrandr.XRRFreeCrtcInfo(crtc)
-            xrandr.XRRFreeScreenResources(mon)
+    @property
+    def monitors(self):
+        # type: () -> Monitors
+        """ Get positions of monitors (see parent class property). """
+
+        if not self._monitors:
+            with MSS._lock:
+                self._monitors_impl()
 
         return self._monitors
 
-    def grab(self, monitor):
+    def _grab_impl(self, monitor):
         # type: (Monitor) -> ScreenShot
-        """ Retrieve all pixels from a monitor. Pixels have to be RGB. """
+        """
+        Retrieve all pixels from a monitor. Pixels have to be RGB.
+        That method has to be run using a threading lock.
+        """
 
         # Convert PIL bbox style
         if isinstance(monitor, tuple):
@@ -393,7 +428,7 @@ def grab(self, monitor):
             }
 
         ximage = self.xlib.XGetImage(
-            MSS.display,
+            self._get_display(),
             self.drawable,
             monitor["left"],
             monitor["top"],
@@ -424,3 +459,10 @@ def grab(self, monitor):
             self.xlib.XDestroyImage(ximage)
 
         return self.cls_image(data, monitor)
+
+    def grab(self, monitor):
+        # type: (Monitor) -> ScreenShot
+        """ Retrieve all pixels from a monitor. Pixels have to be RGB. """
+
+        with MSS._lock:
+            return self._grab_impl(monitor)

mss/tests/test_implementation.py

@@ -168,3 +168,32 @@ def test_grab_with_tuple_percents(sct, pixel_ratio):
     assert im.size == im2.size
     assert im.pos == im2.pos
     assert im.rgb == im2.rgb
+
+
+def test_thread_safety():
+    """Regression test for issue #169."""
+    import threading
+    import time
+
+    def record(check):
+        """Record for one second."""
+
+        start_time = time.time()
+        while time.time() - start_time < 1:
+            with mss.mss() as sct:
+                sct.grab(sct.monitors[1])
+
+        check[threading.current_thread()] = True
+
+    checkpoint = {}
+    t1 = threading.Thread(target=record, args=(checkpoint,))
+    t2 = threading.Thread(target=record, args=(checkpoint,))
+
+    t1.start()
+    time.sleep(0.5)
+    t2.start()
+
+    t1.join()
+    t2.join()
+
+    assert len(checkpoint) == 2