[gentoo-portage-dev] [PATCH 1/3 v2] Add cached portage.getpid() function

[Zac Medico <zmedico@gentoo.org>]

Since getpid is a syscall, cache results, and update them
via an after fork hook.

Signed-off-by: Zac Medico <zmedico@gentoo.org>
---
 lib/portage/__init__.py                       | 16 +++++++++++++
 .../tests/process/test_AsyncFunction.py       | 24 +++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/lib/portage/__init__.py b/lib/portage/__init__.py
index 916c93510..4d4b590a8 100644
--- a/lib/portage/__init__.py
+++ b/lib/portage/__init__.py
@@ -14,6 +14,7 @@ try:
 	if not hasattr(errno, 'ESTALE'):
 		# ESTALE may not be defined on some systems, such as interix.
 		errno.ESTALE = -1
+	import multiprocessing.util
 	import re
 	import types
 	import platform
@@ -368,6 +369,21 @@ _internal_caller = False
 
 _sync_mode = False
 
+class _ForkWatcher:
+	@staticmethod
+	def hook(_ForkWatcher):
+		_ForkWatcher.current_pid = _os.getpid()
+
+_ForkWatcher.hook(_ForkWatcher)
+
+multiprocessing.util.register_after_fork(_ForkWatcher, _ForkWatcher.hook)
+
+def getpid():
+	"""
+	Cached version of os.getpid(). ForkProcess updates the cache.
+	"""
+	return _ForkWatcher.current_pid
+
 def _get_stdin():
 	"""
 	Buggy code in python's multiprocessing/process.py closes sys.stdin
diff --git a/lib/portage/tests/process/test_AsyncFunction.py b/lib/portage/tests/process/test_AsyncFunction.py
index 55857026d..3b360e02f 100644
--- a/lib/portage/tests/process/test_AsyncFunction.py
+++ b/lib/portage/tests/process/test_AsyncFunction.py
@@ -3,6 +3,7 @@
 
 import sys
 
+import portage
 from portage import os
 from portage.tests import TestCase
 from portage.util._async.AsyncFunction import AsyncFunction
@@ -36,3 +37,26 @@ class AsyncFunctionTestCase(TestCase):
 	def testAsyncFunctionStdin(self):
 		loop = asyncio._wrap_loop()
 		loop.run_until_complete(self._testAsyncFunctionStdin(loop))
+
+	def _test_getpid_fork(self):
+		"""
+		Verify that portage.getpid() cache is updated in a forked child process.
+		"""
+		loop = asyncio._wrap_loop()
+		proc = AsyncFunction(scheduler=loop, target=portage.getpid)
+		proc.start()
+		proc.wait()
+		self.assertEqual(proc.pid, proc.result)
+
+	def test_getpid_fork(self):
+		self._test_getpid_fork()
+
+	def test_getpid_double_fork(self):
+		"""
+		Verify that portage.getpid() cache is updated correctly after
+		two forks.
+		"""
+		loop = asyncio._wrap_loop()
+		proc = AsyncFunction(scheduler=loop, target=self._test_getpid_fork)
+		proc.start()
+		self.assertEqual(proc.wait(), 0)
-- 
2.25.3