# vim: set ts=8 sts=2 sw=2 tw=99 et:

 import errno

 import shutil

 import os, sys

 import traceback

 import multiprocessing as mp

 from ambuild2 import util

 from ambuild2 import nodetypes

 from ambuild2.ipc import ParentProcessListener, ChildProcessListener

 from ambuild2.ipc import ProcessManager, MessageListener, Error

 from ambuild2.ipc import Channel

 class Task(object):

   def __init__(self, id, entry, outputs):

     self.id = id

     self.type = entry.type

     self.data = entry.blob

     if entry.folder:

       self.folder = entry.folder.path

     else:

       self.folder = None

     self.outputs = outputs

     self.outgoing = []

     self.incoming = set()

   def addOutgoing(self, task):

     self.outgoing.append(task)

     task.incoming.add(self)

   @property

   def folder_name(self):

     if not self.folder:

       return ''

     return self.folder

   def format(self):

     text = ''

     if self.type == nodetypes.Cxx:

       return '[' + self.data['type'] + ']' + ' -> ' + (' '.join([arg for arg in self.data['argv']]))

     if self.type == nodetypes.Symlink:

       return 'ln -s "{0}" "{1}"'.format(self.data[0], os.path.join(self.folder_name, self.data[1]))

     if self.type == nodetypes.Copy:

       return 'cp "{0}" "{1}"'.format(self.data[0], os.path.join(self.folder_name, self.data[1]))

     return (' '.join([arg for arg in self.data]))

 class WorkerChild(ChildProcessListener):

   def __init__(self, pump, channel, buildPath):

     super(WorkerChild, self).__init__(pump, channel)

     self.buildPath = buildPath

     self.pid = os.getpid()

     self.messageMap = {

       'task': lambda channel, message: self.receiveTask(channel, message)

     }

     self.taskMap = {

       'cxx': lambda message: self.doCompile(message),

       'cmd': lambda message: self.doCommand(message),

       'ln': lambda message: self.doSymlink(message),

       'cp': lambda message: self.doCopy(message),

       'rc': lambda message: self.doResource(message),

     }

     self.channel.send({

       'id': 'ready',

       'finished': None

     })

   def receiveTask(self, channel, message):

     task_id = message['task_id']

     task_type = message['task_type']

     task_folder = message['task_folder']

     if not task_folder:

       task_folder = '.'

     # Remove all outputs.

     for output in message['task_outputs']:

       try:

         os.unlink(output)

       except OSError as exn:

         if exn.errno != errno.ENOENT:

           response = {

             'ok': False,

             'cmdline': 'rm {0}'.format(output),

             'stdout': '',

             'stderr': '{0}'.format(exn)

           }

           return self.issueResponse(message, response)

     if not message['task_folder']:

       message['task_folder'] = '.'

     # Do the task.

     response = self.taskMap[task_type](message)

     return self.issueResponse(message, response)

   def issueResponse(self, message, response):

     # Compute new timestamps for all command outputs.

     new_timestamps = []

     if response['ok']:

       for output in message['task_outputs']:

         try:

           new_timestamps.append((output, os.path.getmtime(output)))

         except:

           response['ok'] = False

           response['stderr'] += 'Expected output file, but not found: {0}'.format(output)

           break

     # Send a message back to the master process to update the DAG and spew

     # stdout/stderr if needed.

     response['id'] = 'results'

     response['pid'] = self.pid

     response['task_id'] = message['task_id']

     response['updates'] = new_timestamps

     self.channel.send(response)

   def doCommand(self, message):

     task_folder = message['task_folder']

     argv = message['task_data']

     with util.FolderChanger(task_folder):

       try:

         p, stdout, stderr = util.Execute(argv)

         status = p.returncode == 0

       except Exception as exn:

         status = False

         stdout = ''

         stderr = '{0}'.format(exn)

     reply = {

       'ok': status,

       'cmdline': ' '.join([arg for arg in argv]),

       'stdout': stdout,

       'stderr': stderr,

     }

     return reply

   def doSymlink(self, message):

     task_folder = message['task_folder']

     source_path, output_path = message['task_data']

     with util.FolderChanger(task_folder):

       rcode, stdout, stderr = util.symlink(source_path, output_path)

     reply = {

       'ok': rcode == 0,

       'cmdline': 'ln -s {0} {1}'.format(source_path, output_path),

       'stdout': stdout,

       'stderr': stderr,

     }

     return reply

   def doCopy(self, message):

     task_folder = message['task_folder']

     source_path, output_path = message['task_data']

     with util.FolderChanger(task_folder):

       if os.path.exists(source_path):

         shutil.copy(source_path, output_path)

         ok = True

         stderr = ''

       else:

         ok = False

         stderr = 'File not found: {0}'.format(source_path)

     reply = {

       'ok': ok,

       'cmdline': 'cp "{0}" "{1}"'.format(source_path, os.path.join(task_folder, output_path)),

       'stdout': '',

       'stderr': stderr,

     }

     return reply

   # Adjusts any dependencies relative to the current folder, to be relative to

   # the output folder instead.

   def rewriteDeps(self, deps):

     paths = []

     for inc_path in deps:

       if not os.path.isabs(inc_path):

         inc_path = os.path.abspath(inc_path)

       # Detect whether the include is within the build folder or not.

       build_path = self.buildPath

       if build_path[-1] != os.path.sep:

         build_path += os.path.sep

       prefix = os.path.commonprefix([build_path, inc_path])

       if prefix == build_path:

         # The include is not a system include, i.e. it was generated, so

         # rewrite the path to be relative to the build folder.

         inc_path = os.path.relpath(inc_path, self.buildPath)

       paths.append(inc_path)

     return paths

   def doCompile(self, message):

     task_folder = message['task_folder']

     task_data = message['task_data']

     cc_type = task_data['type']

     argv = task_data['argv']

     with util.FolderChanger(task_folder):

       p, out, err = util.Execute(argv)

       if cc_type == 'gcc':

         err, deps = util.ParseGCCDeps(err)

       elif cc_type == 'msvc':

         out, deps = util.ParseMSVCDeps(out)

       elif cc_type == 'sun':

         err, deps = util.ParseSunDeps(err)

       else:

         raise Exception('unknown compiler type')

       paths = self.rewriteDeps(deps)

     reply = {

       'ok': p.returncode == 0,

       'cmdline': ' '.join([arg for arg in argv]),

       'stdout': out,

       'stderr': err,

       'deps': paths,

     }

     return reply

   def doResource(self, message):

     task_folder = message['task_folder']

     task_data = message['task_data']

     cl_argv = task_data['cl_argv']

     rc_argv = task_data['rc_argv']

     with util.FolderChanger(task_folder):

       # Includes go to stderr when we preprocess to stdout.

       p, out, err = util.Execute(cl_argv)

       out, deps = util.ParseMSVCDeps(err)

       paths = self.rewriteDeps(deps)

       if p.returncode == 0:

         p, out, err = util.Execute(rc_argv)

     reply = {

       'ok': p.returncode == 0,

       'cmdline': (' '.join([arg for arg in cl_argv]) + ' && ' + ' '.join([arg for arg in rc_argv])),

       'stdout': out,

       'stderr': err,

       'deps': paths,

     }

     return reply

 # The WorkerParent is in the same process as the TaskMasterChild.

 class WorkerParent(ParentProcessListener):

   def __init__(self, taskMaster):

     super(WorkerParent, self).__init__('Worker')

     self.taskMaster = taskMaster

     self.messageMap = {

       'ready': lambda child, message: self.receiveReady(child, message),

       'results': lambda child, message: self.receiveResults(child, message),

     }

   def receiveReady(self, child, message):

     self.taskMaster.onWorkerReady(child, message)

   def receiveResults(self, child, message):

     self.taskMaster.onWorkerResults(child, message)

   def receiveError(self, child, error):

     self.taskMaster.onWorkerDied(child, error)

 # The TaskMasterChild is in the same process as the WorkerParent.

 class TaskMasterChild(ChildProcessListener):

   def __init__(self, pump, channel, task_graph, buildPath, num_processes):

     super(TaskMasterChild, self).__init__(pump, channel)

     self.task_graph = task_graph

     self.outstanding = {}

     self.idle = set()

     self.build_failed = False

     self.build_completed = False

     self.messageMap = {

       'stop': lambda channel, message: self.receiveStop(channel, message)

     }

     self.procman = ProcessManager(pump)

     for i in range(num_processes):

       self.procman.spawn(

         WorkerParent(self),

         WorkerChild,

         args=(buildPath,)

       )

     self.channel.send({

       'id': 'spawned',

       'pid': os.getpid(),

       'type': 'taskmaster'

     })

   def receiveStop(self, channel, message):

     if not message['ok']:

       self.terminateBuild()

     self.close_idle()

   def terminateBuild(self):

     if self.build_failed:

       return

     self.build_failed = True

     self.close_idle()

   def receiveClose(self, channel):

     self.procman.shutdown()

     self.pump.cancel()

   def onWorkerResults(self, child, message):

     # Forward the results to the master process.

     self.channel.send(message)

     if not message['ok']:

       self.channel.send({

         'id': 'completed',

         'status': 'failed',

       })

       self.procman.close(child)

       self.terminateBuild()

       return

     task_id = message['task_id']

     task, child = self.outstanding[task_id]

     del self.outstanding[task_id]

     # Enqueue any tasks that can be run if this was their last outstanding

     # dependency.

     for outgoing in task.outgoing:

       outgoing.incoming.remove(task)

       if len(outgoing.incoming) == 0:

         self.task_graph.append(outgoing)

     self.onWorkerReady(child, None)

     # If more stuff was queued, and we have idle processes, use them.

     while len(self.task_graph) and len(self.idle):

       child = self.idle.pop()

       if not self.onWorkerReady(child, None):

         break

   def close_idle(self):

     for child in self.idle:

       self.procman.close(child)

     self.idle = set()

   def onWorkerReady(self, child, message):

     if message and not message['finished']:

       self.channel.send({

         'id': 'spawned',

         'pid': child.pid,

         'type': 'worker'

       })

     # If the build failed, ignore the message, and shutdown the process.

     if self.build_failed:

       self.procman.close(child)

       self.maybe_request_shutdown(child)

       return

     if not len(self.task_graph):

       if len(self.outstanding):

         # There are still tasks left to complete, but they're waiting on

         # others to finish. Mark this process as ready and just ignore the

         # status change for now.

         self.idle.add(child)

       else:

         # There are no tasks remaining, the worker is not needed.

         self.build_completed = True

         self.procman.close(child)

         self.close_idle()

         self.channel.send({

           'id': 'completed',

           'status': 'ok'

         })

       return False

     # Send a task to the worker.

     task = self.task_graph.pop()

     message = {

       'id': 'task',

       'task_id': task.id,

       'task_type': task.type,

       'task_data': task.data,

       'task_folder': task.folder,

       'task_outputs': task.outputs

     }

     child.send(message)

     self.outstanding[task.id] = (task, child)

     return True

   def onWorkerCrashed(self, child, task):

     self.channel.send({

       'id': 'completed',

       'status': 'crashed',

       'task_id': task.id,

     })

     self.terminateBuild()

   def onWorkerDied(self, child, error):

     if error != Error.NormalShutdown:

       for task_id in self.outstanding.keys():

         task, task_child = self.outstanding[task_id]

         if task_child == child:

           # A worker failed, but crashed, so we have to tell the main process.

           self.onWorkerCrashed(child, task)

           break

     self.idle.discard(child)

     self.maybe_request_shutdown(child)

   def maybe_request_shutdown(self, child):

     for other_child in self.procman.children:

       if other_child == child:

         continue

       if other_child.is_alive():

         return

     # If we got here, no other child processes are live, so we can ask for

     # safe shutdown. This is needed to make sure all our messages arrive,

     # since closing one end of the pipe destroys any leftover data.

     self.channel.send({

       'id': 'done'

     })

 class TaskMasterParent(ParentProcessListener):

   def __init__(self, cx, builder, task_graph, max_parallel):

     super(TaskMasterParent, self).__init__('TaskMaster')

     self.cx = cx

     self.builder = builder

     self.build_failed = False

     self.messageMap = {

       'completed': lambda child, message: self.receiveCompleted(child, message),

       'spawned': lambda child, message: self.receiveSpawned(message),

       'results': lambda child, message: self.processResults(message),

       'done': lambda child, message: self.receiveDone(message),

     }

     self.errors_ = []

     # Figure out how many tasks to create.

     if cx.options.jobs == 0:

       # Using 1 process will be strictly worse than an in-process build,

       # since we incur the additional overhead of message passing. Instead,

       # we use two processes as the minimal number. If that turns out to be

       # bad we can create an in-process TaskMaster later.

       if mp.cpu_count() == 1:

         num_processes = 2

       else:

         num_processes = int(mp.cpu_count() * 1.25)

     else:

       num_processes = cx.options.jobs

     # Don't create more processes than we'll need.

     if num_processes > max_parallel:

       num_processes = max_parallel

     # Spawn the task master.

     self.taskMaster = cx.procman.spawn(

       self,

       TaskMasterChild,

       args=(task_graph, cx.buildPath, num_processes)

     )

   def receiveDone(self, message):

     self.cx.procman.shutdown()

   def spewResult(self, message):

     if message['ok']:

       color = util.ConsoleGreen

     else:

       color = util.ConsoleRed

     util.con_out(

       util.ConsoleBlue,

       '[{0}]'.format(message['pid']),

       util.ConsoleNormal,

       ' ',

       color,

       message['cmdline'],

       util.ConsoleNormal

     )

     if len(message['stdout']):

       util.WriteEncodedText(sys.stdout, message['stdout'])

       if message['stdout'][-1] != '\n':

         sys.stdout.write('\n')

     if len(message['stderr']):

       util.WriteEncodedText(sys.stderr, message['stderr'])

       if message['stderr'][-1] != '\n':

         sys.stderr.write('\n')

   def processResults(self, message):

     if not message['ok']:

       self.errors_.append(message)

       self.terminateBuild()

       return

     self.spewResult(message)

     task_id = message['task_id']

     updates = message['updates']

     if not self.builder.updateGraph(task_id, updates, message):

       util.con_out(

         util.ConsoleRed,

         'Failed to update node!',

         util.ConsoleNormal

       )

       self.terminateBuild()

   def terminateBuild(self):

     if self.build_failed:

       return

     self.build_failed = True

     self.taskMaster.send({

       'id': 'stop',

       'ok': False

     })

   def receiveSpawned(self, message):

     util.con_out(

       util.ConsoleHeader,

       'Spawned {0} (pid: {1})'.format(message['type'], message['pid']),

       util.ConsoleNormal

     )

   def receiveError(self, child, error):

     if error != Error.NormalShutdown:

       sys.stderr.write('Received unexpected error from child process {0}: {1}\n'.format(child.pid, error))

       self.terminateBuild()

   def receiveCompleted(self, child, message):

     if message['status'] == 'crashed':

       task = self.builder.commands[message['task_id']]

       sys.stderr.write('Crashed trying to perform update:\n')

       sys.stderr.write('  : {0}\n'.format(task.entry.format()))

       self.terminateBuild()

     elif message['status'] == 'failed':

       self.terminateBuild()

     else:

       self.taskMaster.send({

         'id': 'stop',

         'ok': True

       })

   def run(self):

     try:

       self.cx.messagePump.pump()

     except KeyboardInterrupt:

       self.terminateBuild()

     for message in self.errors_:

       self.spewResult(message)

     return not self.build_failed