# Multiprocess a command.
#
# dmandelin@mozilla.com
#
# Python version of pork-barrel, with some differences:
#   Collects output into one file with delimiters
#   Uses format strings to control command line creation
#   Not as much timing info
#
# It is also intended that this can be extended in other ways by users,
# such as customizing the way output is assembled.

from __future__ import division

import os, sys, time
from threading import Thread, Lock
from popen2 import Popen4
from Queue import Queue, Empty

WORKER_COUNT = 2
COMMAND = None

tasks = Queue()
output_lock = Lock()

class Worker(object):
  def __init__(self):
    self.thread = None

  def start(self):
    self.thread = Thread(target=self.run)
    self.thread.setDaemon(True)
    self.thread.start()

  def run(self):
    try:
      while True:
        filename = tasks.get(False)
        print >> sys.stderr, tasks.qsize(), 'tasks left'
        self.do_task(filename)
        # only in Python 2.5
        #tasks.task_done()
    except Empty:
      pass
      
  def do_task(self, filename):
    cmd = COMMAND%filename

    print >> sys.stderr, '>>> %s'%cmd
    pipe = Popen4(cmd)
    pipe.tochild.close()
    output = pipe.fromchild.read()
    pipe.fromchild.close()
    rc = pipe.wait()
    print >> sys.stderr, '<<< %s'%cmd
    
    output_lock.acquire()
    print '# %s'%cmd
    if os.WIFEXITED(rc):
      ec = os.WEXITSTATUS(rc)
      if ec == 0:
        print '#     PORKBAR_OK: rc=0'
      else:
        print '#     PORKBAR_FAIL_EXIT: rc=%d'%ec
    else:
      print '#     PORKBAR_FAIL_ABNORMAL_TERMINATION'
    sys.stdout.write(output)
    output_lock.release()
    
def run(listfile):
  task_count = 0
  for filename in open(listfile):
    filename = filename.strip('\r\n')
    tasks.put_nowait(filename)
    task_count += 1

  t0 = time.clock()

  workers = [ Worker() for _ in range(WORKER_COUNT) ]
  for w in workers:
    w.start()
  for w in workers:
    w.thread.join()

  t1 = time.clock()
  dt = t1-t0
  # This time isn't very useful on Unix for this program
  #print "time taken = %.3f (%.3f/task)"%(dt, dt/task_count)

if __name__ == '__main__':
  if len(sys.argv) != 4:
    print 'usage: python pork-barrel.py <cpu-count> <filename-file> <cmd>'
    print '    by default filename is appended to <cmd>, can also use %s in cmd'
    sys.exit(1)

  _, cpu_count, listfile, COMMAND = sys.argv
  WORKER_COUNT = int(cpu_count)
  if COMMAND.find('%s') == -1: COMMAND += ' %s'
  run(listfile)
