supersize: Track symbol aliases and shared symbols

tools/binary_size/libsupersize/concurrent.py

+# Copyright 2017 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Helpers related to multiprocessing."""
+
+import atexit
+import logging
+import multiprocessing
+import multiprocessing.dummy
+import os
+import sys
+import threading
+import traceback
+
+
+DISABLE_ASYNC = os.environ.get('SUPERSIZE_DISABLE_ASYNC') == '1'
+if DISABLE_ASYNC:
+  logging.debug('Running in synchronous mode.')
+
+_all_pools = None
+_is_child_process = False
+_silence_exceptions = False
+
+
+class _ImmediateResult(object):
+  def __init__(self, value):
+    self._value = value
+
+  def get(self):
+    return self._value
+
+  def wait(self):
+    pass
+
+  def ready(self):
+    return True
+
+  def successful(self):
+    return True
+
+
+class _ExceptionWrapper(object):
+  """Used to marshal exception messages back to main process."""
+  def __init__(self, msg):
+    self.msg = msg
+
+
+class _FuncWrapper(object):
+  """Runs on the fork()'ed side to catch exceptions and spread *args."""
+  def __init__(self, func):
+    global _is_child_process
+    _is_child_process = True
+    self._func = func
+
+  def __call__(self, args, _=None):
+    try:
+      return self._func(*args)
+    except:  # pylint: disable=bare-except
+      # multiprocessing is supposed to catch and return exceptions automatically
+      # but it doesn't seem to work properly :(.
+      logging.warning('CAUGHT EXCEPTION')
+      return _ExceptionWrapper(traceback.format_exc())
+
+
+class _WrappedResult(object):
+  """Allows for host-side logic to be run after child process has terminated.
+
+  * Unregisters associated pool _all_pools.
+  * Raises exception caught by _FuncWrapper.
+  * Allows for custom unmarshalling of return value.
+  """
+  def __init__(self, result, pool=None, decode_func=None):
+    self._result = result
+    self._pool = pool
+    self._decode_func = decode_func
+
+  def get(self):
+    self.wait()
+    value = self._result.get()
+    _CheckForException(value)
+    if not self._decode_func or not self._result.successful():
+      return value
+    return self._decode_func(value)
+
+  def wait(self):
+    self._result.wait()
+    if self._pool:
+      _all_pools.remove(self._pool)
+      self._pool = None
+
+  def ready(self):
+    return self._result.ready()
+
+  def successful(self):
+    return self._result.successful()
+
+
+def _TerminatePools():
+  """Calls .terminate() on all active process pools.
+
+  Not supposed to be necessary according to the docs, but seems to be required
+  when child process throws an exception or Ctrl-C is hit.
+  """
+  global _silence_exceptions
+  _silence_exceptions = True
+  # Child processes cannot have pools, but atexit runs this function because
+  # it was registered before fork()ing.
+  if _is_child_process:
+    return
+  def close_pool(pool):
+    try:
+      pool.terminate()
+    except:  # pylint: disable=bare-except
+      pass
+
+  for i, pool in enumerate(_all_pools):
+    # Without calling terminate() on a separate thread, the call can block
+    # forever.
+    thread = threading.Thread(name='Pool-Terminate-{}'.format(i),
+                              target=close_pool, args=(pool,))
+    thread.daemon = True
+    thread.start()
+
+
+def _CheckForException(value):
+  if isinstance(value, _ExceptionWrapper):
+    global _silence_exceptions
+    if not _silence_exceptions:
+      _silence_exceptions = True
+      logging.error('Subprocess raised an exception:\n%s', value.msg)
+    sys.exit(1)
+
+
+def _MakeProcessPool(*args):
+  global _all_pools
+  ret = multiprocessing.Pool(*args)
+  if _all_pools is None:
+    _all_pools = []
+    atexit.register(_TerminatePools)
+  _all_pools.append(ret)
+  return ret
+
+
+def ForkAndCall(func, args, decode_func=None):
+  """Runs |func| in a fork'ed process.
+
+  Returns:
+    A Result object (call .get() to get the return value)
+  """
+  if DISABLE_ASYNC:
+    pool = None
+    result = _ImmediateResult(func(*args))
+  else:
+    pool = _MakeProcessPool(1)
+    result = pool.apply_async(_FuncWrapper(func), (args,))
+    pool.close()
+  return _WrappedResult(result, pool=pool, decode_func=decode_func)
+
+
+def BulkForkAndCall(func, arg_tuples):
+  """Calls |func| in a fork'ed process for each set of args within |arg_tuples|.
+
+  Yields the return values as they come in.
+  """
+  pool_size = min(len(arg_tuples), multiprocessing.cpu_count())
+  if DISABLE_ASYNC:
+    for args in arg_tuples:
+      yield func(*args)
+    return
+  pool = _MakeProcessPool(pool_size)
+  wrapped_func = _FuncWrapper(func)
+  for result in pool.imap_unordered(wrapped_func, arg_tuples):
+    _CheckForException(result)
+    yield result
+  pool.close()
+  pool.join()
+  _all_pools.remove(pool)
+
+
+def CallOnThread(func, *args, **kwargs):
+  """Calls |func| on a new thread and returns a promise for its return value."""
+  if DISABLE_ASYNC:
+    return _ImmediateResult(func(*args, **kwargs))
+  pool = multiprocessing.dummy.Pool(1)
+  result = pool.apply_async(func, args=args, kwds=kwargs)
+  pool.close()
+  return result
+
+
+def EncodeDictOfLists(d, key_transform=None):
+  """Serializes a dict where values are lists of strings."""
+  keys = iter(d)
+  if key_transform:
+    keys = (key_transform(k) for k in keys)
+  keys = '\x01'.join(keys)
+  values = '\x01'.join('\x02'.join(x) for x in d.itervalues())
+  return keys, values
+
+
+def DecodeDictOfLists(encoded_keys, encoded_values, key_transform=None):
+  """Deserializes a dict where values are lists of strings."""
+  keys = encoded_keys.split('\x01')
+  if key_transform:
+    keys = (key_transform(k) for k in keys)
+  values = encoded_values.split('\x01')
+  ret = {}
+  for i, key in enumerate(keys):
+    ret[key] = values[i].split('\x02')
+  return ret