Windows: fix multi-thread safety (BoboTiG#159)

* Windows: fix multi-thread unsafe and update test

On Windows, the handle of entire window device context is saved to
`srcdc`. But the device context will be released once the thread who
creates it has died, so that `srcdc` is no loner valid. Replace `srcdc`
with `srcdc_dict` to maintain srcdc values created by multiple threads
which ensure the validity of srcdc when it's used.

A threading lock is add to prevent multiple threads from grabbing and
modifying shared class attributes `bmp`/`srcdc`/`memdc` (their windows
object in fact) at same time. Otherwise, unexpected screenshot or
unpredictable error will occur.

Add test_thread_safety in test_windows.py

Author: narumi147

CHANGELOG

@@ -7,6 +7,8 @@ History:
       - MSS: renamed again MSSMixin to MSSBase, now derived from abc.ABCMeta
       - tools: force write of file when saving a PNG file
       - tests: fix tests on macOS with Retina display
+      - Windows: fixed multi-thread safety (fixes #150)
+      - :heart: contributors: @narumishi
 
 5.0.0   2019/12/31
       - removed support for Python 2.7

CHANGES.rst

@@ -8,6 +8,10 @@ base.py
 - ``MSSBase.monitor`` is now an abstract property
 - ``MSSBase.grab()`` is now an abstract method
 
+windows.py
+----------
+ - Replaced ``MSS.srcdc`` with ``MSS.srcdc_dict``
+
 
 5.0.0 (2019-12-31)
 ==================

CONTRIBUTORS

@@ -30,6 +30,9 @@ Jochen 'cycomanic' Schroeder [https://github.com/cycomanic]
 Karan Lyons <[email protected]> [https://karanlyons.com] [https://github.com/karanlyons]
     - MacOS: Proper support for display scaling
 
+narumi [https://github.com/narumishi]
+    - Windows: fix multi-thread unsafe
+
 Oros <[email protected]> [https://ecirtam.net]
     - GNU/Linux tester
 

mss/tests/test_windows.py

@@ -4,6 +4,7 @@
 """
 
 import platform
+import threading
 
 import mss
 import pytest
@@ -48,3 +49,23 @@ def test_region_caching():
         # Grab the area 2 again, the cached BMP is used
         sct.grab(region2)
         assert bmp2 is MSS.bmp
+
+
+def run_child_thread(loops):
+    """Every loop will take about 1 second."""
+    for _ in range(loops):
+        with mss.mss() as sct:
+            sct.grab(sct.monitors[1])
+
+
+def test_thread_safety():
+    """Thread safety test for issue #150.
+    The following code will throw a ScreenShotError exception if thread-safety is not guaranted.
+    """
+    # Let thread 1 finished ahead of thread 2
+    thread1 = threading.Thread(target=run_child_thread, args=(30,))
+    thread2 = threading.Thread(target=run_child_thread, args=(50,))
+    thread1.start()
+    thread2.start()
+    thread1.join()
+    thread2.join()

mss/windows.py

@@ -5,6 +5,7 @@
 
 import sys
 import ctypes
+import threading
 from ctypes.wintypes import (
     BOOL,
     DOUBLE,
@@ -70,10 +71,15 @@ class MSS(MSSBase):
 
     __slots__ = {"_bbox", "_bmi", "_data", "gdi32", "monitorenumproc", "user32"}
 
-    # Class attributes instancied one time to prevent resource leaks.
+    # Class attributes instanced one time to prevent resource leaks.
     bmp = None
     memdc = None
-    srcdc = None
+
+    # A dict to maintain *srcdc* values created by multiple threads.
+    srcdc_dict = {}
+
+    # A threading lock to lock resources(bmp/memdc/srcdc) inside .grab() method.
+    _lock = threading.Lock()
 
     def __init__(self, **_):
         # type: (Any) -> None
@@ -93,9 +99,9 @@ def __init__(self, **_):
         self._bbox = {"height": 0, "width": 0}
         self._data = ctypes.create_string_buffer(0)  # type: ctypes.Array[ctypes.c_char]
 
-        if not MSS.srcdc or not MSS.memdc:
-            MSS.srcdc = self.user32.GetWindowDC(0)
-            MSS.memdc = self.gdi32.CreateCompatibleDC(MSS.srcdc)
+        srcdc = self._get_srcdc()
+        if not MSS.memdc:
+            MSS.memdc = self.gdi32.CreateCompatibleDC(srcdc)
 
         bmi = BITMAPINFO()
         bmi.bmiHeader.biSize = ctypes.sizeof(BITMAPINFOHEADER)
@@ -174,6 +180,20 @@ def _set_dpi_awareness(self):
             # Windows Vista, 7, 8 and Server 2012
             self.user32.SetProcessDPIAware()
 
+    def _get_srcdc(self):
+        """
+        Retrieve a thread-safe HDC from GetWindowDC().
+        In multithreading, if the thread who creates *srcdc* is dead, *srcdc* will
+        no longer be valid to grab the screen. The *srcdc* attribute is replaced
+        with *srcdc_dict* to maintain the *srcdc* values in multithreading.
+        Since the current thread and main thread are always alive, reuse their *srcdc* value first.
+        """
+        cur_thread, main_thread = threading.current_thread(), threading.main_thread()
+        srcdc = MSS.srcdc_dict.get(cur_thread) or MSS.srcdc_dict.get(main_thread)
+        if not srcdc:
+            srcdc = MSS.srcdc_dict[cur_thread] = self.user32.GetWindowDC(0)
+        return srcdc
+
     @property
     def monitors(self):
         # type: () -> Monitors
@@ -251,6 +271,9 @@ def grab(self, monitor):
             Thanks to http://stackoverflow.com/a/3688682
         """
 
+        # Acquire lock to prevent resources from being modified by multiple threads at same time.
+        MSS._lock.acquire()
+
         # Convert PIL bbox style
         if isinstance(monitor, tuple):
             monitor = {
@@ -260,7 +283,7 @@ def grab(self, monitor):
                 "height": monitor[3] - monitor[1],
             }
 
-        srcdc, memdc = MSS.srcdc, MSS.memdc
+        srcdc, memdc = self._get_srcdc(), MSS.memdc
         width, height = monitor["width"], monitor["height"]
 
         if (self._bbox["height"], self._bbox["width"]) != (height, width):
@@ -287,6 +310,7 @@ def grab(self, monitor):
         bits = self.gdi32.GetDIBits(
             memdc, MSS.bmp, 0, height, self._data, self._bmi, DIB_RGB_COLORS
         )
+        MSS._lock.release()
         if bits != height:
             raise ScreenShotError("gdi32.GetDIBits() failed.")