source: mergebot/trunk/mergebot/WorkQueue.py @ 16

#!/usr/bin/env python
import os
import sys
import time
import traceback
import trac.env

from TrackerTools import Task, GetWorkDir
from CreateDaemon import createDaemon

class WorkQueue(object):
    # TODO: add locking to file accesses.
    delim = " "
    def __init__(self, workdir, name):
        self.name = name
        if not os.path.exists(workdir):
            os.makedirs(workdir)
        self.queuefile = os.path.join(workdir, "queue.%s" % name)
        self.inprocessfile = os.path.join(workdir, "queue.%s.current" % name)
        self.pidfile = os.path.join(workdir, "queue.%s.pid" % name)
    def drivequeue(self):
        "return true if this actor is already running"
        if os.path.exists(self.pidfile):
            # Check that it's running...
            pid = open(self.pidfile).read()
            process = os.popen("ps h --pid %s" % pid).read()
            if process:
                return 1
            # otherwise, just overwrite the existing one.
        open(self.pidfile, "w").write("%s\n" % os.getpid())
        return 0
    def releasequeue(self):
        os.remove(self.pidfile)
    def readqueue(self):
        if os.path.exists(self.queuefile):
            lines = open( self.queuefile ).readlines()
        else:
            lines = []
        # The -1 is to strip the trailing newline.
        queued = map(lambda x: x[:-1].split(self.delim), lines)
        return queued
    def writequeue(self, tasks):
        if tasks:
            lines = map( lambda task:"%s\n" % (self.delim.join(task)), tasks )
        else:
            lines = []
        open(self.queuefile, "w").writelines(lines)
    def readcurrent(self):
        if os.path.exists(self.inprocessfile):
            current = open( self.inprocessfile ).readlines()
        else:
            current = []
        if current:
            task = current[0].split(self.delim)
        else:
            task = None
        return task
    def writecurrent(self, task):
        if task:
            line = "%s\n" % (self.delim.join(task))
        else:
            line = ""
        open(self.inprocessfile, "w").writelines([line])
    def enqueue(self, task):
        queued = self.readqueue()
        if task in queued:
            return 1
        queued.append(task)
        self.writequeue(queued)
        return 0
    def dequeue(self):
        # Check that we don't have something half-way dequeued.
        # The problem is that if we have something half-way dequeued, there was
        # a problem completing that task.
        # TODO: So instead of returning it, we need to ring alarm bells, but
        # dequeue the next task.
        queued = self.readqueue()
        if not queued:
            return None
        self.writequeue(queued[1:])
        self.writecurrent(queued[0])
        return queued[0]
    def completetask(self):
        self.writecurrent(None)

class MergeBotActor(object):
    def __init__(self, tracEnv, name, action):
        self.tracEnv = tracEnv
        self.name = name
        self.action = action
        workdir = GetWorkDir(tracEnv)
        self.queue = WorkQueue(workdir, name)
    def AddTask(self, task):
        if self.queue.enqueue(task):
            # must be a dup
            return
        task_obj = Task(self.tracEnv, task[0])
        task_obj.SetMergeBotStatus("to%s"%self.name)
        task_obj.Save()
    def Run(self):
        self.tracEnv.log.debug("Starting %s action" % (self.name))
        # The child will need the absolute path to the environment.
        trac_env_path = os.path.abspath(self.tracEnv.path)
        pid = os.fork()
        if pid:
            # Parent
            self.tracEnv.log.debug("Running %s action in child %s" % (self.name, pid))
        else:
            # Child.  We daemonize and do the work.
            createDaemon()  # This closes all open files, and sets the working directory to /
            # And because of that, we must open our own copy of the Trac environment.
            tracEnv = trac.env.open_environment(trac_env_path)
            if self.queue.drivequeue():
                # something is already working
                tracEnv.log.debug("The %s queue is already running." % (self.name, ))
                sys.exit(0)
            try:
                tracEnv.log.debug("Child running %s action" % (self.name, ))
                while 1:
                    task = self.queue.dequeue()
                    if not task:
                        tracEnv.log.debug("Queue %s emptied." % (self.name, ))
                        # we're done.
                        break
                    tracEnv.log.debug("%s working on %s..." % (self.name, task))
                    # FIXME: Need to rethink the status update logic.
                    args = [tracEnv] + task
                    tracEnv.log.debug("Running %s action with arguments %s" % (self.name, repr(args)))
                    try:
                        result, task_obj = self.action(*args)
                        if result:
                            tracEnv.log.debug("Action %s completed with result %s" % (self.name, result))
                            task_obj.SetMergeBotStatus(result)
                        task_obj.Save()
                    except Exception, e:
                        tracEnv.log.error("BUG!!!  Task %r failed while running"
                            " the %s queue.  Continuing." % (args, self.name))
                        tracEnv.log.exception(e)
                        
                    self.queue.completetask()
                self.queue.releasequeue()
            except Exception, e:
                tracEnv.log.error("BUG!!!  Failed while running the %s queue" % (self.name, ))
                tracEnv.log.exception(e)
                sys.exit(1)
            
            tracEnv.log.debug("Child completed action %s, exiting" % (self.name, ))
            sys.exit(0)
    def GetStatus(self):
        """Returns a tuple (current, (queue)), where each element is a list of
        arguments"""
        return (self.queue.readcurrent(), self.queue.readqueue() )

def VersionToDir(version):
    """Given the version from the Trac version list, determine what the path
    should be under that component.  trunk -> trunk, but the rest will be
    branches/something."""
    if version == "trunk":
        versiondir = "trunk"
    elif version.startswith("#"):
        ticketnum = version[1:] # Strip the "#"
        versiondir = "branches/ticket-%s" % (ticketnum)
    else:
        versiondir = "branches/release-%s" % (version)
    return versiondir

def testcase():
    wc = WorkQueue(os.path.join(os.getcwd(), "test"), "noop")
    print "Initial state."
    print wc.readcurrent()
    print wc.readqueue()

    while wc.dequeue():
        wc.completetask()

    print wc.enqueue(["task1"])
    print wc.enqueue(["task2"])

    task = wc.dequeue()
    wc.completetask()

    print "Final state."
    print wc.readcurrent()
    print wc.readqueue()

if __name__ == "__main__":
    testcase()

# vim:foldcolumn=4 foldmethod=indent
# vim:tabstop=4 shiftwidth=4 softtabstop=4 expandtab