Roll version of gsutil to 4.15.

third_party/gsutil/gslib/tests/test_parallelism_framework.py

 # -*- coding: utf-8 -*-
 # Copyright 2013 Google Inc. All Rights Reserved.
 #
 # Permission is hereby granted, free of charge, to any person obtaining a
 # copy of this software and associated documentation files (the
 # "Software"), to deal in the Software without restriction, including
 # without limitation the rights to use, copy, modify, merge, publish, dis-
 # tribute, sublicense, and/or sell copies of the Software, and to permit
 # persons to whom the Software is furnished to do so, subject to the fol-
 # lowing conditions:
 #
 # The above copyright notice and this permission notice shall be included
 # in all copies or substantial portions of the Software.
 #
 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 # OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABIL-
 # ITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
 # SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 # WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 # IN THE SOFTWARE.
 """Unit tests for gsutil parallelism framework."""
 
 from __future__ import absolute_import
 
 import functools
+import os
 import signal
+import threading
+import time
 
 from boto.storage_uri import BucketStorageUri
 from gslib import cs_api_map
 from gslib.command import Command
 from gslib.command import CreateGsutilLogger
 from gslib.command import DummyArgChecker
 import gslib.tests.testcase as testcase
 from gslib.tests.testcase.base import RequiresIsolation
 from gslib.tests.util import unittest
+from gslib.util import CheckMultiprocessingAvailableAndInit
 from gslib.util import IS_WINDOWS
-from gslib.util import MultiprocessingIsAvailable
 
 
 # Amount of time for an individual test to run before timing out. We need a
 # reasonably high value since if many tests are running in parallel, an
 # individual test may take a while to complete.
 _TEST_TIMEOUT_SECONDS = 120
 
 
 def Timeout(func):
   """Decorator used to provide a timeout for functions."""
@@ skipping 18 matching lines @@
 class CustomException(Exception):
 
   def __init__(self, exception_str):
     super(CustomException, self).__init__(exception_str)
 
 
 def _ReturnOneValue(cls, args, thread_state=None):
   return 1
 
 
+def _ReturnProcAndThreadId(cls, args, thread_state=None):
+  return (os.getpid(), threading.currentThread().ident)
+
+
+def _SleepThenReturnProcAndThreadId(cls, args, thread_state=None):
+  # This can fail if the total time to spawn new processes and threads takes
+  # longer than 5 seconds, but if that occurs, then we have a performance
+  # problem that needs to be addressed.
+  time.sleep(5)
+  return _ReturnProcAndThreadId(cls, args, thread_state=thread_state)
+
+
 def _FailureFunc(cls, args, thread_state=None):
   raise CustomException('Failing on purpose.')
 
 
 def _FailingExceptionHandler(cls, e):
   cls.failure_count += 1
   raise CustomException('Exception handler failing on purpose.')
 
 
 def _ExceptionHandler(cls, e):
@@ skipping 120 matching lines @@
     }
     default_map = {
         'gs': 'JSON',
         's3': 'XML'
     }
     self.gsutil_api_map = cs_api_map.GsutilApiMapFactory.GetApiMap(
         cs_api_map.GsutilApiClassMapFactory, support_map, default_map)
     self.logger = CreateGsutilLogger('FakeCommand')
     self.parallel_operations = do_parallel
     self.failure_count = 0
-    self.multiprocessing_is_available = MultiprocessingIsAvailable()[0]
+    self.multiprocessing_is_available = (
+        CheckMultiprocessingAvailableAndInit().is_available)
     self.debug = 0
 
 
 class FakeCommandWithoutMultiprocessingModule(FakeCommand):
 
   def __init__(self, do_parallel):
     super(FakeCommandWithoutMultiprocessingModule, self).__init__(do_parallel)
     self.multiprocessing_is_available = False
 
 
@@ skipping 37 matching lines @@
     self._TestBasicApply(3, 3)
 
   @Timeout
   def _TestBasicApply(self, process_count, thread_count):
     args = [()] * (17 * process_count * thread_count + 1)
 
     results = self._RunApply(_ReturnOneValue, args, process_count, thread_count)
     self.assertEqual(len(args), len(results))
 
   @RequiresIsolation
+  def testNoTasksSingleProcessSingleThread(self):
+    self._TestApplyWithNoTasks(1, 1)
+
+  @RequiresIsolation
+  def testNoTasksSingleProcessMultiThread(self):
+    self._TestApplyWithNoTasks(1, 3)
+
+  @RequiresIsolation
+  @unittest.skipIf(IS_WINDOWS, 'Multiprocessing is not supported on Windows')
+  def testNoTasksMultiProcessSingleThread(self):
+    self._TestApplyWithNoTasks(3, 1)
+
+  @RequiresIsolation
+  @unittest.skipIf(IS_WINDOWS, 'Multiprocessing is not supported on Windows')
+  def testNoTasksMultiProcessMultiThread(self):
+    self._TestApplyWithNoTasks(3, 3)
+
+  @Timeout
+  def _TestApplyWithNoTasks(self, process_count, thread_count):
+    """Tests that calling Apply with no tasks releases locks/semaphores."""
+    empty_args = [()]
+
+    for _ in range(process_count * thread_count + 1):
+      self._RunApply(_ReturnOneValue, empty_args, process_count, thread_count)
+
+    # Ensure that work can still be performed.
+    self._TestBasicApply(process_count, thread_count)
+
+  @RequiresIsolation
+  @unittest.skipIf(IS_WINDOWS, 'Multiprocessing is not supported on Windows')
+  def testApplySaturatesMultiProcessSingleThread(self):
+    self._TestApplySaturatesAvailableProcessesAndThreads(3, 1)
+
+  @RequiresIsolation
+  def testApplySaturatesSingleProcessMultiThread(self):
+    self._TestApplySaturatesAvailableProcessesAndThreads(1, 3)
+
+  @RequiresIsolation
+  @unittest.skipIf(IS_WINDOWS, 'Multiprocessing is not supported on Windows')
+  def testApplySaturatesMultiProcessMultiThread(self):
+    self._TestApplySaturatesAvailableProcessesAndThreads(3, 3)
+
+  @RequiresIsolation
+  def _TestApplySaturatesAvailableProcessesAndThreads(self, process_count,
+                                                      thread_count):
+    """Tests that created processes and threads evenly share tasks."""
+    calls_per_thread = 2
+    args = [()] * (process_count * thread_count * calls_per_thread)
+    expected_calls_per_thread = calls_per_thread
+
+    if not self.command_class(True).multiprocessing_is_available:
+      # When multiprocessing is unavailable, only a single process is used.
+      # Calls should be evenly distributed across threads.
+      expected_calls_per_thread = calls_per_thread * process_count
+
+    results = self._RunApply(_SleepThenReturnProcAndThreadId, args,
+                             process_count, thread_count)
+    usage_dict = {}  # (process_id, thread_id): number of tasks performed
+    for (process_id, thread_id) in results:
+      usage_dict[(process_id, thread_id)] = (
+          usage_dict.get((process_id, thread_id), 0) + 1)
+
+    for (id_tuple, num_tasks_completed) in usage_dict.iteritems():
+      self.assertEqual(num_tasks_completed, expected_calls_per_thread,
+                       'Process %s thread %s completed %s tasks. Expected: %s' %
+                       (id_tuple[0], id_tuple[1], num_tasks_completed,
+                        expected_calls_per_thread))
+
+  @RequiresIsolation
   def testIteratorFailureSingleProcessSingleThread(self):
     self._TestIteratorFailure(1, 1)
 
   @RequiresIsolation
   def testIteratorFailureSingleProcessMultiThread(self):
     self._TestIteratorFailure(1, 3)
 
   @RequiresIsolation
   @unittest.skipIf(IS_WINDOWS, 'Multiprocessing is not supported on Windows')
   def testIteratorFailureMultiProcessSingleThread(self):
@@ skipping 61 matching lines @@
     command_inst.arg_length_sum = 19
     args = ['foo', ['bar', 'baz'], [], ['x', 'y'], [], 'abcd']
     self._RunApply(_IncrementByLength, args, process_count,
                    thread_count, command_inst=command_inst,
                    shared_attrs=['arg_length_sum'])
     expected_sum = 19
     for arg in args:
       expected_sum += len(arg)
     self.assertEqual(expected_sum, command_inst.arg_length_sum)
 
-    # Test that shared variables work when the iterator fails.
-    command_inst = self.command_class(True)
-    args = FailingIterator(10, [1, 3, 5])
-    self._RunApply(_ReturnOneValue, args, process_count, thread_count,
-                   command_inst=command_inst, shared_attrs=['failure_count'])
-    self.assertEqual(3, command_inst.failure_count)
+    # Test that shared variables work when the iterator fails at the beginning,
+    # middle, and end.
+    for (failing_iterator, expected_failure_count) in (
+        (FailingIterator(5, [0]), 1),
+        (FailingIterator(10, [1, 3, 5]), 3),
+        (FailingIterator(5, [4]), 1)):
+      command_inst = self.command_class(True)
+      args = failing_iterator
+      self._RunApply(_ReturnOneValue, args, process_count, thread_count,
+                     command_inst=command_inst, shared_attrs=['failure_count'])
+      self.assertEqual(
+          expected_failure_count, command_inst.failure_count,
+          msg='Failure count did not match. Expected: %s, actual: %s '
+          'for failing iterator of size %s, failing indices %s' %
+          (expected_failure_count, command_inst.failure_count,
+           failing_iterator.size, failing_iterator.failure_indices))
 
   @RequiresIsolation
   def testThreadsSurviveExceptionsInFuncSingleProcessSingleThread(self):
     self._TestThreadsSurviveExceptionsInFunc(1, 1)
 
   @RequiresIsolation
   def testThreadsSurviveExceptionsInFuncSingleProcessMultiThread(self):
     self._TestThreadsSurviveExceptionsInFunc(1, 3)
 
   @RequiresIsolation
@@ skipping 185 matching lines @@
     args = [2 * x for x in args]
     results = self._RunApply(_ReturnOneValue, args, process_count, thread_count,
                              arg_checker=_SkipEvenNumbersArgChecker)
     self.assertEqual(0, len(results))
 
 
 class TestParallelismFrameworkWithoutMultiprocessing(TestParallelismFramework):
   """Tests parallelism framework works with multiprocessing module unavailable.
 
   Notably, this test has no way to override previous calls
-  to gslib.util.MultiprocessingIsAvailable to prevent the initialization of
-  all of the global variables in command.py, so this still behaves slightly
-  differently than the behavior one would see on a machine where the
-  multiprocessing functionality is actually not available (in particular, it
-  will not catch the case where a global variable that is not available for
-  the sequential path is referenced before initialization).
+  to gslib.util.CheckMultiprocessingAvailableAndInit to prevent the
+  initialization of all of the global variables in command.py, so this still
+  behaves slightly differently than the behavior one would see on a machine
+  where the multiprocessing functionality is actually not available (in
+  particular, it will not catch the case where a global variable that is not
+  available for the sequential path is referenced before initialization).
   """
   command_class = FakeCommandWithoutMultiprocessingModule