dataTransfer.py

#!/usr/bin/python -u

import os, sys
import time
import subprocess
import threading
import select
from optparse import OptionParser

gThreading = True

class ThreadManager(threading.Thread):

    def __init__(self, cmdline):
        self.pool = []
        self.running = []
        self.parallel = 5
        self.lock = threading.Lock()
        self.cmdline = cmdline
        threading.Thread.__init__(self)

    def clean_pool(self, pool, join=False):
        for th in [x for x in pool if not x.isAlive()]:
            if join: th.join()
            pool.remove(th)
            del th
        return pool

    def spawnThreads(self, threadnum = 5):
        self.pool = []
        self.running = []
        if threadnum: self.parallel = threadnum
        self.addThreads()

    def addThreads(self):
        for r in xrange(0, self.parallel):
            self.pool.append(TrickleClass(self.cmdline))

    def run(self):
        self.started = True
        parallel = self.parallel
        while len(self.running)+len(self.pool) > 0:
            self.clean_pool(self.running, join=True)
            n = min(max(parallel - len(self.running), 0), len(self.pool))
            if n > 0:
                for th in self.pool[0:n]: th.start()
                self.running.extend(self.pool[0:n])
                del self.pool[0:n]
            time.sleep(.01)
        for th in self.running+self.pool: th.join()

    def get_completed_total(self):
        sum = 0
        for thread in self.running:
            sum += thread.get_completed()
        return sum

    def isRunning(self):
        return len(self.running) + len(self.pool)


class TrickleClass(threading.Thread):
    def __init__(self, cmdline):
        self.stdout = None
        self.stderr = None
        self.completed = 0
        self.process = None
        self.pid = 0
        self.data = ""
        self.cmdline = cmdline
        self.lock = threading.RLock()
        threading.Thread.__init__(self)

    def run(self):
        try:
            self.process = subprocess.Popen(self.cmdline, stdout=subprocess.PIPE, bufsize=1, shell=False)
        except OSError:
            print "ERROR"
            return
        self.pid = self.process.pid
        while True:
            if self.update_completed():
                break

    def update_completed(self):
        self.lock.acquire()
        try:
            r = select.select([self.process.stdout.fileno()], [], [], 5)[0]
        except:
            r = False
        if r:
            d = os.read(r[0], 1)
            self.data += d
            lastline = self.data.split("\n")[-1]
            if lastline == "NO DATA":
                gThreading = False
                return True
            elif lastline == "DONE":
                return True
            try:
                self.completed = int(self.data.split("\n")[-2])
            except:
                self.completed = 0
        self.lock.release()
        return False


    def get_completed(self):
        return self.completed


def format_time(seconds):
    m, s = divmod(seconds, 60)
    h, m = divmod(m, 60)
    d, h = divmod(h, 24)
    return "%dd %02d:%02d:%02d" %(d, h, m, s)


if __name__ == "__main__":
    parser = OptionParser()
    parser.add_option('-n', '--threads', action='store', type="int", help='number of threads', default=10)
    parser.add_option('-e', '--exit', action='store_true', default=False)
    parser.add_option('--php', action='store', type="string", help="Path to php executable", default="php")
    parser.add_option('--index', action='store', type="string", help="Path to index.php", default="")
    parser.add_option('-u', '--update', action='store', type="int", help="Number of records completed between updates, default 10", default=10)
    (options, args) = parser.parse_args()
    if not options.index:
        print "Error: --index must be supplied\n"
        parser.print_help()
        sys.exit()
    t1 = time.time()
    base_complete = 0
    total = 0
    cmdline = [options.php, options.index, "--uri=legacy/count"]
    process = subprocess.Popen(cmdline, stdout=subprocess.PIPE, bufsize=0, shell=False)
    pid = process.pid
    records = float(process.stdout.read())

    cmdline = [options.php, options.index, "--uri=/legacy"]
    if options.exit:
        sys.exit()
    while gThreading:
        threadManager = ThreadManager(cmdline)
        threadManager.spawnThreads(options.threads)
        threadManager.start()
        group_total = 0
        while threadManager.isAlive():
            t = threadManager.get_completed_total()
            if t > group_total:
                group_total = t
                total += 1
                if not total % options.update:
                    pct_completed = total / records
                    time_diff = time.time() - t1
                    if time_diff > 0:
                        seconds = (time_diff / pct_completed) - time_diff
                        remaining = "Remaining: %s" %(format_time(seconds))
                        elapsed = "Elasped Time: %s" %(format_time(time_diff))
                        avg = "Avg: %.02f/s"  %(total / time_diff)
                        sys.stdout.write("%s\r" %(" " * 120))
                        sys.stdout.write("%s\t\tTotal Completed: %s\t\t%s\t\t%s\r" % (remaining, total, elapsed, avg))