From 055f822b656e209bdec3dc773ef10d3fa871c9c6 Mon Sep 17 00:00:00 2001
From: Michael Hanselmann <hansmi@google.com>
Date: Wed, 17 Feb 2010 16:31:28 +0100
Subject: [PATCH] Add function to reset tempfile module after fork

On fork, the tempfile module's pseudo random generator is
not reset. If several processes (e.g. two children or parent
and child) try to create a temporary file, they'll conflict.
This function can be used to reset the name generator which
contains the pseudo random generator.

A unittest is included. It is in a separate script because
it changes a variable in the tempfile module to speed up the
test.

Signed-off-by: Michael Hanselmann <hansmi@google.com>
Reviewed-by: Iustin Pop <iustin@google.com>
---
 Makefile.am                    |   3 +-
 lib/utils.py                   |  23 ++++++
 test/tempfile_fork_unittest.py | 125 +++++++++++++++++++++++++++++++++
 3 files changed, 150 insertions(+), 1 deletion(-)
 create mode 100755 test/tempfile_fork_unittest.py

diff --git a/Makefile.am b/Makefile.am
index feb2dc517..a0d87f7dc 100644
--- a/Makefile.am
+++ b/Makefile.am
@@ -332,7 +332,8 @@ dist_TESTS = \
 	test/ganeti.ssh_unittest.py \
 	test/ganeti.utils_unittest.py \
 	test/ganeti.workerpool_unittest.py \
-	test/docs_unittest.py
+	test/docs_unittest.py \
+	test/tempfile_fork_unittest.py
 
 nodist_TESTS =
 
diff --git a/lib/utils.py b/lib/utils.py
index 77ab93895..0d8659953 100644
--- a/lib/utils.py
+++ b/lib/utils.py
@@ -334,6 +334,29 @@ def RenameFile(old, new, mkdir=False, mkdir_mode=0750):
     raise
 
 
+def ResetTempfileModule():
+  """Resets the random name generator of the tempfile module.
+
+  This function should be called after C{os.fork} in the child process to
+  ensure it creates a newly seeded random generator. Otherwise it would
+  generate the same random parts as the parent process. If several processes
+  race for the creation of a temporary file, this could lead to one not getting
+  a temporary name.
+
+  """
+  # pylint: disable-msg=W0212
+  if hasattr(tempfile, "_once_lock") and hasattr(tempfile, "_name_sequence"):
+    tempfile._once_lock.acquire()
+    try:
+      # Reset random name generator
+      tempfile._name_sequence = None
+    finally:
+      tempfile._once_lock.release()
+  else:
+    logging.critical("The tempfile module misses at least one of the"
+                     " '_once_lock' and '_name_sequence' attributes")
+
+
 def _FingerprintFile(filename):
   """Compute the fingerprint of a file.
 
diff --git a/test/tempfile_fork_unittest.py b/test/tempfile_fork_unittest.py
new file mode 100755
index 000000000..396c393a7
--- /dev/null
+++ b/test/tempfile_fork_unittest.py
@@ -0,0 +1,125 @@
+#!/usr/bin/python
+#
+
+# Copyright (C) 2010 Google Inc.
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; either version 2 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program; if not, write to the Free Software
+# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
+# 02110-1301, USA.
+
+
+"""Script for testing utils.ResetTempfileModule"""
+
+import os
+import sys
+import errno
+import shutil
+import tempfile
+import unittest
+import logging
+
+from ganeti import utils
+
+import testutils
+
+
+# This constant is usually at a much higher value. Setting it lower for test
+# purposes.
+tempfile.TMP_MAX = 3
+
+
+class TestResetTempfileModule(unittest.TestCase):
+  def setUp(self):
+    self.tmpdir = tempfile.mkdtemp()
+
+  def tearDown(self):
+    shutil.rmtree(self.tmpdir)
+
+  def testNoReset(self):
+    self._Test(False)
+
+  def testReset(self):
+    self._Test(True)
+
+  def _Test(self, reset):
+    self.failIf(tempfile.TMP_MAX > 10)
+
+    # Initialize tempfile module
+    (fd, _) = tempfile.mkstemp(dir=self.tmpdir, prefix="init.", suffix="")
+    os.close(fd)
+
+    (notify_read, notify_write) = os.pipe()
+
+    pid = os.fork()
+    if pid == 0:
+      # Child
+      try:
+        try:
+          if reset:
+            utils.ResetTempfileModule()
+
+          os.close(notify_write)
+
+          # Wait for parent to close pipe
+          os.read(notify_read, 1)
+
+          try:
+            # This is a short-lived process, not caring about closing file
+            # descriptors
+            (_, path) = tempfile.mkstemp(dir=self.tmpdir,
+                                         prefix="test.", suffix="")
+          except EnvironmentError, err:
+            if err.errno == errno.EEXIST:
+              # Couldnt' create temporary file (e.g. because we run out of
+              # retries)
+              os._exit(2)
+            raise
+
+          logging.debug("Child created %s", path)
+
+          os._exit(0)
+        except Exception:
+          logging.exception("Unhandled error")
+      finally:
+        os._exit(1)
+
+    # Parent
+    os.close(notify_read)
+
+    # Create parent's temporary files
+    for _ in range(tempfile.TMP_MAX):
+      (fd, path) = tempfile.mkstemp(dir=self.tmpdir,
+                                    prefix="test.", suffix="")
+      os.close(fd)
+      logging.debug("Parent created %s", path)
+
+    # Notify child by closing pipe
+    os.close(notify_write)
+
+    (_, status) = os.waitpid(pid, 0)
+
+    self.failIf(os.WIFSIGNALED(status))
+
+    if reset:
+      # If the tempfile module was reset, it should not fail to create
+      # temporary files
+      expected = 0
+    else:
+      expected = 2
+
+    self.assertEqual(os.WEXITSTATUS(status), expected)
+
+
+if __name__ == "__main__":
+  testutils.GanetiTestProgram()
-- 
GitLab