fix #1111: use a lock to make Process.oneshot() thread safe

giampaolo · Dec 13, 2018 · b3b5d42 · b3b5d42
1 parent 495bb45
commit b3b5d42
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 41 deletions.
diff --git a/HISTORY.rst b/HISTORY.rst
@@ -13,6 +13,7 @@ XXXX-XX-XX
 
 **Bug fixes**
 
+- 1111_: Process.oneshot() is now thread safe.
 - 1354_: [Linux] disk_io_counters() fails on Linux kernel 4.18+.
 - 1357_: [Linux] Process' memory_maps() and io_counters() method are no longer
   exposed if not supported by the kernel.

diff --git a/psutil/__init__.py b/psutil/__init__.py
@@ -31,6 +31,7 @@
 import signal
 import subprocess
 import sys
+import threading
 import time
 try:
     import pwd
@@ -352,7 +353,7 @@ def _init(self, pid, _ignore_nsp=False):
         self._create_time = None
         self._gone = False
         self._hash = None
-        self._oneshot_inctx = False
+        self._lock = threading.RLock()
         # used for caching on Windows only (on POSIX ppid may change)
         self._ppid = None
         # platform-specific modules define an _psplatform.Process
@@ -456,40 +457,45 @@ def oneshot(self):
         ...
         >>>
         """
-        if self._oneshot_inctx:
-            # NOOP: this covers the use case where the user enters the
-            # context twice. Since as_dict() internally uses oneshot()
-            # I expect that the code below will be a pretty common
-            # "mistake" that the user will make, so let's guard
-            # against that:
-            #
-            # >>> with p.oneshot():
-            # ...    p.as_dict()
-            # ...
-            yield
-        else:
-            self._oneshot_inctx = True
-            try:
-                # cached in case cpu_percent() is used
-                self.cpu_times.cache_activate(self)
-                # cached in case memory_percent() is used
-                self.memory_info.cache_activate(self)
-                # cached in case parent() is used
-                self.ppid.cache_activate(self)
-                # cached in case username() is used
-                if POSIX:
-                    self.uids.cache_activate(self)
-                # specific implementation cache
-                self._proc.oneshot_enter()
+        with self._lock:
+            if hasattr(self, "_cache"):
+                # NOOP: this covers the use case where the user enters the
+                # context twice:
+                #
+                # >>> with p.oneshot():
+                # ...    with p.oneshot():
+                # ...
+                #
+                # Also, since as_dict() internally uses oneshot()
+                # I expect that the code below will be a pretty common
+                # "mistake" that the user will make, so let's guard
+                # against that:
+                #
+                # >>> with p.oneshot():
+                # ...    p.as_dict()
+                # ...
                 yield
-            finally:
-                self.cpu_times.cache_deactivate(self)
-                self.memory_info.cache_deactivate(self)
-                self.ppid.cache_deactivate(self)
-                if POSIX:
-                    self.uids.cache_deactivate(self)
-                self._proc.oneshot_exit()
-                self._oneshot_inctx = False
+            else:
+                try:
+                    # cached in case cpu_percent() is used
+                    self.cpu_times.cache_activate(self)
+                    # cached in case memory_percent() is used
+                    self.memory_info.cache_activate(self)
+                    # cached in case parent() is used
+                    self.ppid.cache_activate(self)
+                    # cached in case username() is used
+                    if POSIX:
+                        self.uids.cache_activate(self)
+                    # specific implementation cache
+                    self._proc.oneshot_enter()
+                    yield
+                finally:
+                    self.cpu_times.cache_deactivate(self)
+                    self.memory_info.cache_deactivate(self)
+                    self.ppid.cache_deactivate(self)
+                    if POSIX:
+                        self.uids.cache_deactivate(self)
+                    self._proc.oneshot_exit()
 
     def as_dict(self, attrs=None, ad_value=None):
         """Utility method returning process information as a

diff --git a/psutil/_common.py b/psutil/_common.py
@@ -336,14 +336,17 @@ def memoize_when_activated(fun):
     """
     @functools.wraps(fun)
     def wrapper(self):
-        if not hasattr(self, "_cache"):
+        try:
+            # case 1: we previously entered oneshot() ctx
+            ret = self._cache[fun]
+        except AttributeError:
+            # case 2: we never entered oneshot() ctx
             return fun(self)
-        else:
-            try:
-                ret = self._cache[fun]
-            except KeyError:
-                ret = self._cache[fun] = fun(self)
-            return ret
+        except KeyError:
+            # case 3: we entered oneshot() ctx but there's no cache
+            # for this entry yet
+            ret = self._cache[fun] = fun(self)
+        return ret
 
     def cache_activate(proc):
         """Activate cache. Expects a Process instance. Cache will be