# Copyright (c) 2014 Evalf
#
# 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, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following 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 MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS 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.

"""
The parallel module provides tools aimed at parallel computing. At this point
all parallel solutions use the ``fork`` system call and are supported on limited
platforms, notably excluding Windows. On unsupported platforms parallel features
will disable and a warning is printed.
"""

from . import numeric, warnings, util
import os, multiprocessing, mmap, signal, contextlib, builtins, numpy, treelog

_maxprocs = util.settable(int(os.environ.get('NUTILS_NPROCS') or 1))

@util.positional_only
def maxprocs(new: int):
  '''limit number of processes for fork.'''

  if not isinstance(new, int) or new < 1:
    raise ValueError('nprocs requires a positive integer argument')
  return _maxprocs.sets(new)

@contextlib.contextmanager
def fork(nprocs=None):
  '''continue as ``nprocs`` parallel processes by forking ``nprocs-1`` times

  If ``nprocs`` exceeds the configured ``maxprocs`` than it will silently be
  capped. It is up to the user to prepare shared memory and/or locks for
  inter-process communication. As a safety measure nested forks are blocked by
  limiting nprocs to 1; all secondary forks will be silently ignored.
  '''

  if nprocs is None or nprocs > _maxprocs.value:
    nprocs = _maxprocs.value
  if nprocs <= 1:
    yield 0
    return
  if not hasattr(os, 'fork'):
    warnings.warn('fork is unavailable on this platform')
    yield 0
    return
  amchild = False
  try:
    child_pids = []
    for procid in builtins.range(1, nprocs):
      pid = os.fork()
      if not pid: # pragma: no cover
        amchild = True
        signal.signal(signal.SIGINT, signal.SIG_IGN) # disable sigint (ctrl+c) handler
        treelog.current = treelog.NullLog() # silence treelog
        break
      child_pids.append(pid)
    else:
      procid = 0
    with maxprocs(1):
      yield procid
  except BaseException as e:
    if amchild: # pragma: no cover
      try:
        print('[parallel.fork] exception in child process:', e)
      finally:
        os._exit(1) # communicate failure to main process
    for pid in child_pids: # kill all child processes
      os.kill(pid, signal.SIGKILL)
    raise
  else:
    if amchild: # pragma: no cover
      os._exit(0) # communicate success to main process
    with treelog.context('waiting for child processes'):
      nfails = sum(not _wait(pid) for pid in child_pids)
    if nfails: # failure in child process: raise exception
      raise Exception('fork failed in {} out of {} processes'.format(nfails, nprocs))
  finally:
    if amchild: # pragma: no cover
      os._exit(1) # failsafe

def shempty(shape, dtype=float):
  '''create uninitialized array in shared memory'''

  if numeric.isint(shape):
    shape = shape,
  else:
    assert all(numeric.isint(sh) for sh in shape)
  dtype = numpy.dtype(dtype)
  size = util.product(map(int, shape), int(dtype.itemsize))
  if size == 0 or _maxprocs.value == 1:
    return numpy.empty(shape, dtype)
  # `mmap(-1,...)` will allocate *anonymous* memory.  Although linux' man page
  # mmap(2) states that anonymous memory is initialized to zero, we can't rely
  # on this to be true for all platforms (see [SO-mmap]).  [SO-mmap]:
  # https://stackoverflow.com/a/17896084
  return numpy.frombuffer(mmap.mmap(-1, size), dtype).reshape(shape)

def shzeros(shape, dtype=float):
  '''create zero-initialized array in shared memory'''

  array = shempty(shape, dtype=dtype)
  array.fill(0)
  return array

class range:
  '''a shared range-like iterable that yields every index exactly once'''

  def __init__(self, stop):
    self._stop = stop
    self._index = multiprocessing.RawValue('i', 0)
    self._lock = multiprocessing.Lock() # lock to avoid race conditions in incrementing index
  def __iter__(self):
    return self
  def __next__(self):
    with self._lock:
      iiter = self._index.value # claim next value
      if iiter >= self._stop:
        raise StopIteration
      self._index.value = iiter + 1
    return iiter

@contextlib.contextmanager
def ctxrange(name, nitems):
  '''fork and yield shared range-like counter with percentage-style logging'''

  rng = range(nitems) # shared range, must be created pre-fork
  with fork(nitems), treelog.iter.wrap(_pct(name, nitems), rng) as wrprng:
    yield wrprng

def _pct(name, n):
  '''helper function for ctxrange'''

  i = yield name + ' 0%'
  while True:
    i = yield name + ' {:.0f}%'.format(100*(i+1)/n)

def _wait(pid):
  '''wait for process to finish and return True upon success,
  False upon failure while logging the reason.'''
  pid_, status = os.waitpid(pid, 0)
  assert pid_ == pid
  if os.WIFEXITED(status):
    s = os.WEXITSTATUS(status)
    if not s:
      return True
    msg = 'exited with status {}'.format(s)
  elif os.WIFSIGNALED(status):
    s = os.WTERMSIG(status)
    msg = 'was killed with signal {} ({})'.format(s, signal.Signals(s).name)
  elif os.WIFSTOPPED(status):
    s = os.WSTOPSIG(status)
    msg = 'was stopped with signal {} ({})'.format(s, signal.Signals(s).name)
  else:
    msg = 'died of unnatural causes'
  treelog.error('process {} {}'.format(pid, msg))
  return False

# vim:sw=2:sts=2:et

@@ -2720,61 +2720,243 @@
     
-treelog.current = treelog.NullLog() # silence treelog
+setter = treelog.set(treelog.NullLog())%0A        setter.__enter__() # silence treelog%0A        # NOTE for treelog %3E= 2.0 we must hold a reference to setter until%0A        # os.exit_ to save the formerly active logger from being destructed
 %0A   

fix disabling of treelog in forked processes