ltsp-cluster-lbserver 2.0.0-0ubuntu5 / usr / lib / python2.7 / dist-packages / lbserver / main.py

#!/usr/bin/env python

"""
LBServer is a service to be executed on the load balancer server.
"""

import SocketServer
import BaseHTTPServer
import SimpleXMLRPCServer
import traceback
from threading import Thread
from lbserver import *
from lbserverfactory import LBServerFactory

def xmlEscape(txt):
    "Escape XML entities"
    txt = txt.replace("&", "&")
    txt = txt.replace("<", "&lt;")
    txt = txt.replace(">", "&gt;")
    return txt


def xmlUnescape(txt):
    "Unescape XML entities"
    txt = txt.replace("&amp;", "&")
    txt = txt.replace("&lt;", "<")
    txt = txt.replace("&gt;", ">")
    return txt

def mergeConfig(domOrig,domToMerge):
    "Merge xml config files"
    oldRoot = domOrig.childNodes[0]
    for n in oldRoot.childNodes:
        if n.nodeName == u'nodes':
            oldNodes = n
        elif n.nodeName == u'rules':
            oldRules = n
    for n in domToMerge.childNodes[0].childNodes:
        if n.nodeName == u'nodes':
            newNodes = n
        elif n.nodeName == u'rules':
            newRules = n
    oldRoot.replaceChild(newNodes, oldNodes)
    oldRoot.replaceChild(newRules, oldRules)
    return domOrig

class LogFile:
    """
    Log file wrapper
    """
    def __init__(self,file):
        self.file=file
    def close(self):
        self.file.close()
    def flush(self):
        self.file.flush()
    def write(self,str):
        self.file.write(str)
        self.flush()
    def writelines(self,sequence):
        self.file.write(sequence)
        self.flush()
    def isatty(self):
        return False

class XMLRPCServer(SimpleXMLRPCServer.SimpleXMLRPCServer):
    """
    XMLRPCServer
    XMLRPC Server overrided to support custom options
    """
    # Override of TIME_WAIT
    allow_reuse_address = True


class HTTPServer(SocketServer.ThreadingMixIn,BaseHTTPServer.HTTPServer):
    """
    HTTPServer
    HTTP Server overrided to support custom options
    Inherit ThreadingMixIn to be MultiThreaded
    """
    # Override of TIME_WAIT
    allow_reuse_address = True


class LBServerContainer:
    """
    LBserverContainer
    Manage LBServer instance and configuration
    """
    def __init__(self,configFile):
        "Constructor, take config file name"
        self.loading=False
        self._configFile = configFile
        self._loadConfigFile(self._configFile)
    def _loadConfigFile(self,fileName):
        "Create new lbAgent for config file"
        import os
        print "Config change, reload LBServer..."
        self.loading=True
        try:
            factory = LBServerFactory(fileName)
            lbServer = factory.newLBServer()
            lbServer.refreshNodesValues()
            self._lbServer = lbServer
            # Save last modification time
            self._lastm = os.stat(fileName).st_mtime
        finally:
            self.loading=False
    def _checkConfigChange(self):
        "Check for changes in the config file"
        import os
        if not(self.loading) and self._lastm != os.stat(self._configFile).st_mtime:
            self._loadConfigFile(self._configFile)
    def getLastm(self):
        "Return last timestamp of the config file"
        return str(self._lastm)
    def updateConfig(self,configString):
        "Update configuration from serialized XML string"
        import os
        from xml.dom.minidom import parseString
        try:
            # Start new config
            oldDom = parseString(open(self._configFile).read())
            newDom = parseString(xmlUnescape(configString))
            newDom = mergeConfig(oldDom,newDom)
            tmpFileName = self._configFile+"~"
            file = open(tmpFileName, "w")
            newDom.writexml(file)
            file.close()
            # Test new config
            factory = LBServerFactory(tmpFileName)
            lbServer = factory.newLBServer()
            # Replace config file
            os.remove(self._configFile)
            os.rename(tmpFileName,self._configFile)
            self._checkConfigChange()
        except:
            traceback.print_exception(sys.exc_info()[0],sys.exc_info()[1],sys.exc_info()[2])
    def getConfigString(self):
        "Return configuration in a xml string"
        xml = open(self._configFile).read()
        return xmlEscape(xml)
    def getLBServer(self):
        "Return contained instance of LBServer"
        self._checkConfigChange()
        return self._lbServer


class NodesMonitor(Thread):
    """
    NodesMonitor
    Maintain nodes status
    """
    global container
    def __init__(self):
        Thread.__init__(self)
    def run(self):
        import time
        while True:
            time.sleep(container.getLBServer().refresh)
            try:
                container.getLBServer().refreshNodesValues()
            except:
                traceback.print_exception(sys.exc_info()[0],sys.exc_info()[1],sys.exc_info()[2])


class ConfigMonitor(Thread):
    """
    ConfigMonitor
    Load config from master LB
    """
    global container
    def __init__(self):
        Thread.__init__(self)
        self._lastcfgm=''
    def run(self):
        import time
        while True:
            time.sleep(container.getLBServer().cfgrefresh)
            try:
                rpcserver = xmlrpclib.ServerProxy(container.getLBServer().masterurl)
                cfg = rpcserver.get_config(self._lastcfgm)
                if cfg.has_key("xml"):
                    # Prevent call from the same server
                    if cfg["lastm"]!=container.getLastm():
                        container.updateConfig(cfg["xml"])
                        self._lastcfgm=cfg["lastm"]
            except:
                traceback.print_exception(sys.exc_info()[0],sys.exc_info()[1],sys.exc_info()[2])

class MgmtService(Thread):
    """
    MgmtService
    Thread for management with xmlrpc
    """
    global container
    def __init__(self,address):
        Thread.__init__(self)
        self._address = address
    def getStatus(self):
        return container.getLBServer().getStatusStruct()
    def getConfig(self,lastm):
        if container.getLastm() != lastm:
            return {"lastm":container.getLastm(), "xml":container.getConfigString()}
        else:
            return {"lastm":container.getLastm()}
    def run(self):
        # Init server
        xmlrpcServer = XMLRPCServer(self._address)
        xmlrpcServer.register_function(self.getStatus,'get_status')
        xmlrpcServer.register_function(self.getConfig,'get_config')
        # Launch xmlrpc server
        try:
            xmlrpcServer.serve_forever()
        finally:
            xmlrpcServer.server_close()


class LBRequestHandler(BaseHTTPServer.BaseHTTPRequestHandler):
    """
    LBHTTPHandler
    Handler for load balancing http requests
    """
    global container
    def do_GET(self):
        import cgi
        import string
        # Send headers
        self.send_response(200)
        self.send_header("Content-type", "text/html") 
        self.end_headers()
        # Server selection by group
        params = {}
        if '?' in self.path:
            params = cgi.parse_qs(string.split(self.path,'?')[1])

        # Get Server IP
        serverIp=""
        if params.has_key("group"):
            serverIp = container.getLBServer().getNextServerForGroup(params["group"][0])
        else:
            serverIp = container.getLBServer().getNextServer()

        # Send server
        if not serverIp:
            serverIp=""
        self.wfile.write(serverIp+'\n')

def main():
    import sys
    import os
    import socket
    import random
    global container

    # Parse parameters
    logFile = None
    configFile = 'lbsconfig.xml'
    for arg in sys.argv[1:]:
        if arg.startswith("--logfile="):
            logFile=arg.split('=')[1]
        elif len(arg)>2:
            configFile=arg
    if not os.path.exists(configFile):
        print "Config file not found. Please specify config file as last argument."
        sys.exit(1)

    # Init logs
    if (logFile != None):
        try:
            logObject =  LogFile(open(logFile,'a',0))
            sys.stderr = logObject
            sys.stdout = logObject
        except:
            pass

    # Randomize
    random.seed()

    # Create LBServer Container
    container = LBServerContainer(configFile)

    # Set default timeout for all request
    socket.setdefaulttimeout(container.getLBServer().timeout)

    # Launch nodes monitor
    nodesMonitor = NodesMonitor()
    nodesMonitor.setDaemon(True)
    print "Start NodesMonitor..."
    nodesMonitor.start()

    # Launch config monitor
    if container.getLBServer().isslave and container.getLBServer().masterurl != None:
        configMonitor = ConfigMonitor()
        configMonitor.setDaemon(True)
        print "Start ConfigMonitor Slave..."
        configMonitor.start()

    # Launch xmlrpc management service
    if container.getLBServer().mgmtenabled:
        mgmtService = MgmtService(container.getLBServer().mgmtlisten)
        mgmtService.setDaemon(True)
        print "Start MGMT Service..."
        mgmtService.start()

    # Init http server
    print "Start HTTP Server..."
    httpServer = HTTPServer(container.getLBServer().lblisten, LBRequestHandler)

    # Launch http server
    try:
        httpServer.serve_forever()
    finally:
        httpServer.server_close()

if __name__ == "__main__":

    import utils

    utils.get_pidfile_lock('/var/run/ltsp-lbserver.pid')
    if "--daemon" in sys.argv[1:]:
        utils.daemonize()

    main()