root/spreader/spreader.py

import select
import socket
import md5

from twisted.internet import task
from twisted.application import service
from twisted.python import usage

from logging import getLogger


log = getLogger()

import spread

class SpreaderOptions(usage.Options):
    optParameters = [
            ['spread_poll_interval', None, '0.05'],
            ['spread_host', None, 'localhost'],
            ['spread_port', None, str(spread.DEFAULT_SPREAD_PORT)],
            ['spread_private', None, 'dummy'],
            ['spread_membership', None, '1'],
            ['spread_timeout', None, '0.01'],
        ]
    
class ParallelSpreaderOptions(usage.Options):
    optParameters = [
            ['spread_groups', None, ''],
        ] + SpreaderOptions.optParameters

class Spreader(service.Service):
    mailbox = None
    privateGroup = None

    def __init__(self, config, service_):
        self.pollInterval = float(config['spread_poll_interval'])
        self.service = service_
        self.poller = task.LoopingCall(self.poll)
        self.spreadName = "%s@%s"%(config['spread_port'], config['spread_host'])
        self.privateName = "%s-%s" % (config['spread_private'], 
                                      socket.gethostname())
        self.priority = 0
        self.membership = int(config['spread_membership'])
        self.timeout = float(config['spread_timeout'])

    def startService(self):
        service.Service.startService(self)
        self.connect()
        self.poller.start(self.pollInterval)

    def stopService(self):
        service.Service.stopService(self)
        if self.poller.running:
            self.poller.stop()
        self.disconnect()

    def poll(self):
        if not self.mailbox:
            return
        ready = select.select([self.mailbox.fileno()], [], [], self.timeout)
        if not len(ready[0]):
           return
        m = self.mailbox.receive()
        if isinstance(m, spread.RegularMsgType):
            self.regularMessageReceived(m)
        elif isinstance(m, spread.MembershipMsgType):
            self.membershipMessageReceived(m)

    def regularMessageReceived(self, message):
        self.messageReceived(message.sender, message.groups, message.message,
                             message.msg_type, message.endian)

    def membershipMessageReceived(self, message):
        if message.reason == 0:
            self.groupTransitioned(message.group)
        elif message.reason == spread.CAUSED_BY_JOIN:
            self.groupJoined(message.group, message.extra)
            self.groupChanged(message.group, message.members)
        elif (message.reason == spread.CAUSED_BY_LEAVE 
              or message.reason == spread.CAUSED_BY_DISCONNECT):
            self.groupLeft(message.group, message.extra)
            self.groupChanged(message.group, message.members)
        else:
            raise Exception(message)

    def messageReceived(self, sender, groups, message, msg_type, extra):
        log.debug("messageReceived: %s %s %s", sender, groups, message)
        pass

    def groupTransitioned(self, group):
        log.debug("groupTransitioned: %s", group)
        pass

    def groupJoined(self, group, who):
        log.debug("groupJoined: %s %s", group, who)
        pass

    def groupLeft(self, group, who):
        log.debug("groupLeft: %s %s", group, who)
        pass

    def groupChanged(self, group, current):
        log.debug("groupChanged: %s %s", group, current)
        pass

    def connect(self):
        self.mailbox = spread.connect(self.spreadName, self.privateName, 
                                      self.priority, self.membership)
        self.privateGroup = self.mailbox.private_group
    
    def disconnect(self):
        if self.mailbox:
            self.mailbox.disconnect()
            self.mailbox = None
            self.privateGroup = None

    def join(self, group):
        if not self.mailbox:
            raise Exception("no mailbox")
        self.mailbox.join(group)

    def leave(self, group):
        if not self.mailbox:
            raise Exception("no mailbox")
        self.mailbox.leave(group)

    def multicast(self, service_type, group, message, message_type=0):
        if not self.mailbox:
            raise Exception("no mailbox")
        return self.mailbox.multicast(service_type, group, message, message_type)

    def multigroup_multicast(self, *a, **kw):
        if not self.mailbox:
            raise Exception("no mailbox")
        return self.mailbox.multigroup_multicast(*a, **kw)


class ParallelSpreader(Spreader):
    groups = None

    def __init__(self, config, service_):
        Spreader.__init__(self, config, service_)
        self.groups = {}
        for k in config['spread_groups'].split(","):
            self.groups[k] = {"index": None,
                              "members": None,
                              "handle_all": None}
        
    def connect(self):
        Spreader.connect(self)
        # join the groups we want to listen to
        for k in self.groups.keys():
            self.join(k)

    def groupChanged(self, group, current):
        Spreader.groupChanged(self, group, current)
        current = list(current)
        if self.privateGroup not in current:
            self.groups[group]['index'] = None
            return
       
        # update the group entry to have the new members,
        # disable handle_all, and find our new index
        self.groups[group]['members'] = current
        self.groups[group]['handle_all'] = False
        self.groups[group]['index'] = current.index(self.privateGroup)

    def groupTransitioned(self, group):
        Spreader.groupTransitioned(self, group)
        self.groups[group]['handle_all'] = True
    
    def regularMessageReceived(self, message):
        if not self.shouldHandle(message.groups, message.msg_type):
            return
        Spreader.regularMessageReceived(self, message)
    
    def safecast(self, group, message, message_type=None):
        if message_type is None:
            message_type = self.hash(message)
        service_type = spread.SAFE_MESS
        return self.multicast(service_type, group, message, message_type)

    def hash(self, s):
        return int(md5.new(s).hexdigest()[:2], 16)

    def shouldHandle(self, groups, id):
        # XXX doesn't really handle multiple groups the best way probably
        for g in groups:
            try:
               index = self.groups[g]['index']
            except ValueError:
                # if we are not in the group, ignore (this shouldn't happen)
                continue

            # if we are transitioning, handle anything you see
            if self.groups[g]['handle_all']:
                return True

            # if i hash right, do it too
            who = id % len(self.groups[g]['members'])
            if who == index:
                return True
        return False