<s> import logging <EOL> import os <EOL> import os . path <EOL> import socket <EOL> import time <EOL> import threading <EOL> import traceback <EOL> import shutil <EOL> import subprocess <EOL> import ufora <EOL> import ufora . core . SubprocessRunner as SubprocessRunner <EOL> import ufora . distributed . ServerUtils . SimpleServer as SimpleServer <EOL> import ufora . distributed . util . common as common <EOL> from ufora . test . TestScriptRunner import TestScriptRunner <EOL> def createTestRunner ( testDir = '<STR_LIT>' , timeout = None ) : <EOL> assert '<STR_LIT>' in os . environ <EOL> assert '<STR_LIT>' in os . environ <EOL> assert '<STR_LIT>' in os . environ <EOL> assert '<STR_LIT>' in os . environ <EOL> machinesInCluster = sorted ( os . getenv ( '<STR_LIT>' ) . split ( '<STR_LIT:U+002C>' ) ) <EOL> ownAddress = os . getenv ( '<STR_LIT>' ) <EOL> assert ownAddress in machinesInCluster <EOL> testId = os . getenv ( '<STR_LIT>' ) <EOL> testRunnerClass = MultiMachineMasterRunner if ownAddress == machinesInCluster [ <NUM_LIT:0> ] else MultiMachineWorkerRunner <EOL> return testRunnerClass ( machinesInCluster , ownAddress , testId , testDir , timeout = timeout ) <EOL> def extractErrorMessagesFromLogfile ( logfile , targetDir ) : <EOL> try : <EOL> lineNumber = <NUM_LIT:0> <EOL> curBlock = [ ] <EOL> output = [ ] <EOL> def contains ( msg ) : <EOL> return "<STR_LIT>" in msg or "<STR_LIT>" in msg <EOL> for line in open ( logfile , "<STR_LIT:r>" ) : <EOL> lineNumber += <NUM_LIT:1> <EOL> if line [ <NUM_LIT:0> ] in ( '<STR_LIT:U+0020>' , '<STR_LIT:\t>' ) or not curBlock : <EOL> curBlock . append ( line . rstrip ( ) ) <EOL> else : <EOL> if contains ( "<STR_LIT:\n>" . join ( curBlock ) ) : <EOL> output . append ( "<STR_LIT:\n>" . join ( curBlock ) ) <EOL> curBlock = [ ] <EOL> curBlock . append ( line . rstrip ( ) ) <EOL> if contains ( "<STR_LIT:\n>" . join ( curBlock ) ) : <EOL> output . append ( "<STR_LIT:\n>" . join ( curBlock ) ) <EOL> if output : <EOL> fname = os . path . split ( logfile ) [ <NUM_LIT:1> ] <EOL> finalFname = os . path . join ( targetDir , fname + "<STR_LIT>" ) <EOL> with open ( finalFname , "<STR_LIT:w>" ) as resultFile : <EOL> for o in output : <EOL> print >> resultFile , o <EOL> logging . info ( <EOL> "<STR_LIT>" , <EOL> lineNumber , <EOL> logfile , <EOL> len ( output ) , <EOL> finalFname <EOL> ) <EOL> else : <EOL> logging . info ( "<STR_LIT>" , lineNumber , logfile ) <EOL> except : <EOL> logging . error ( "<STR_LIT>" , <EOL> logfile , <EOL> traceback . format_exc ( ) ) <EOL> class MultiMachineTestRunner ( object ) : <EOL> ProtoclVersion = '<STR_LIT>' <EOL> TestControlPort = <NUM_LIT> <EOL> def __init__ ( self , machinesInCluster , ownAddress , testId , testDir , timeout = None ) : <EOL> self . machinesInCluster = machinesInCluster <EOL> self . ownAddress = ownAddress <EOL> assert self . ownAddress in self . machinesInCluster <EOL> self . testId = testId <EOL> self . testDir = testDir <EOL> self . timeout = timeout <EOL> assert os . path . exists ( self . testDir ) , "<STR_LIT>" % self . testDir <EOL> @ property <EOL> def clusterMaster ( self ) : <EOL> return self . machinesInCluster [ <NUM_LIT:0> ] <EOL> @ property <EOL> def isClusterMaster ( self ) : <EOL> return self . clusterMaster == self . ownAddress <EOL> @ property <EOL> def serviceName ( self ) : <EOL> raise NotImplementedError ( "<STR_LIT>" ) <EOL> def run_ ( self ) : <EOL> raise NotImplementedError ( "<STR_LIT>" ) <EOL> def getWorkerEnvironment ( self ) : <EOL> env = dict ( os . environ ) <EOL> env . update ( { <EOL> "<STR_LIT>" : self . clusterMaster , <EOL> "<STR_LIT>" : self . ownAddress <EOL> } ) <EOL> return env <EOL> def run ( self ) : <EOL> logging . info ( "<STR_LIT>" , <EOL> self . machinesInCluster , self . ownAddress , self . clusterMaster , self . isClusterMaster ) <EOL> try : <EOL> outputDir = os . path . join ( os . getenv ( '<STR_LIT>' ) , '<STR_LIT>' ) <EOL> os . makedirs ( outputDir ) <EOL> return self . run_ ( ) <EOL> except Exception as e : <EOL> logging . error ( "<STR_LIT>" , e , traceback . format_exc ( ) ) <EOL> return False <EOL> finally : <EOL> self . postExecutionCleanup ( outputDir ) <EOL> def postExecutionCleanup ( self , outputDir ) : <EOL> self . removePath ( os . path . join ( outputDir , "<STR_LIT>" ) ) <EOL> logDir = os . path . join ( outputDir , "<STR_LIT>" ) <EOL> logging . info ( "<STR_LIT>" , logDir ) <EOL> if not os . path . exists ( logDir ) : <EOL> logging . info ( "<STR_LIT>" ) <EOL> return <EOL> for logfile in os . listdir ( logDir ) : <EOL> fullPath = os . path . join ( logDir , logfile ) <EOL> if os . path . isfile ( fullPath ) : <EOL> extractErrorMessagesFromLogfile ( fullPath , outputDir ) <EOL> try : <EOL> targetLogPath = os . path . join ( outputDir , "<STR_LIT>" ) <EOL> subprocess . call ( [ '<STR_LIT>' % ( targetLogPath , logDir ) ] , shell = True ) <EOL> except : <EOL> logging . error ( "<STR_LIT>" , <EOL> logDir , <EOL> targetLogPath , <EOL> traceback . format_exc ( ) ) <EOL> @ staticmethod <EOL> def removePath ( path ) : <EOL> if not os . path . exists ( path ) : <EOL> return <EOL> for tries in range ( <NUM_LIT:3> ) : <EOL> try : <EOL> shutil . rmtree ( path ) <EOL> return True <EOL> except : <EOL> if tries == <NUM_LIT:2> : <EOL> logging . warn ( "<STR_LIT>" , <EOL> path , <EOL> traceback . format_exc ( ) ) <EOL> return False <EOL> def controlService ( self , serviceName , controlArg = None , env = None ) : <EOL> logging . info ( "<STR_LIT>" , controlArg , serviceName ) <EOL> binDir = os . path . join ( ufora . rootPythonPath , '<STR_LIT>' ) <EOL> exitCode = self . runCommandAndLogOutput ( binDir , serviceName , controlArg , env = env ) <EOL> if exitCode != <NUM_LIT:0> : <EOL> logging . error ( "<STR_LIT>" , controlArg , serviceName , exitCode ) <EOL> raise Exception ( "<STR_LIT>" % ( controlArg , serviceName , exitCode ) ) <EOL> logging . info ( "<STR_LIT>" , serviceName , controlArg ) <EOL> @ staticmethod <EOL> def runCommandAndLogOutput ( path , command , argument = None , env = None ) : <EOL> toRun = [ os . path . join ( path , command ) ] <EOL> if argument : <EOL> toRun . append ( argument ) <EOL> logging . info ( "<STR_LIT>" , toRun ) <EOL> exitCode , stdOut , stdErr = SubprocessRunner . callAndReturnResultAndOutput ( <EOL> toRun , <EOL> env = env <EOL> ) <EOL> logging . info ( "<STR_LIT>" , exitCode ) <EOL> logging . info ( "<STR_LIT>" , command , stdOut ) <EOL> logging . info ( "<STR_LIT>" , command , stdErr ) <EOL> return exitCode <EOL> class MultiMachineWorkerRunner ( MultiMachineTestRunner ) : <EOL> def __init__ ( self , machinesInCluster , ownAddress , testId , testDir , timeout = None ) : <EOL> super ( MultiMachineWorkerRunner , self ) . __init__ ( <EOL> machinesInCluster , ownAddress , testDir , testDir , timeout <EOL> ) <EOL> @ property <EOL> def serviceName ( self ) : <EOL> return "<STR_LIT>" <EOL> def run_ ( self ) : <EOL> '''<STR_LIT>''' <EOL> sock = None <EOL> t0 = time . time ( ) <EOL> while not sock : <EOL> sock = self . connectToMaster ( ) <EOL> if not sock : <EOL> if time . time ( ) - t0 > <NUM_LIT> : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> time . sleep ( <NUM_LIT:1.0> ) <EOL> try : <EOL> self . startWorker ( ) <EOL> self . waitForStopCommand ( sock ) <EOL> finally : <EOL> logging . info ( "<STR_LIT>" ) <EOL> self . stopWorker ( ) <EOL> def connectToMaster ( self ) : <EOL> logging . info ( "<STR_LIT>" , <EOL> self . clusterMaster , <EOL> MultiMachineWorkerRunner . TestControlPort ) <EOL> while True : <EOL> socketError = None <EOL> sock = socket . socket ( socket . AF_INET , socket . SOCK_STREAM ) <EOL> sock . setsockopt ( socket . SOL_SOCKET , socket . SO_REUSEADDR , <NUM_LIT:1> ) <EOL> try : <EOL> sock . connect ( ( self . clusterMaster , MultiMachineWorkerRunner . TestControlPort ) ) <EOL> common . writeString ( sock , MultiMachineTestRunner . ProtoclVersion ) <EOL> common . writeString ( sock , self . testId ) <EOL> response = common . readString ( sock ) <EOL> if response != '<STR_LIT>' : <EOL> logging . error ( "<STR_LIT>" , response ) <EOL> sock . close ( ) <EOL> sock = None <EOL> logging . info ( "<STR_LIT>" , response ) <EOL> return sock <EOL> except common . SocketException : <EOL> socketError = traceback . format_exc ( ) <EOL> except socket . error : <EOL> socketError = traceback . format_exc ( ) <EOL> if socketError : <EOL> sock . close ( ) <EOL> logging . info ( "<STR_LIT>" , socketError ) <EOL> time . sleep ( <NUM_LIT:0.5> ) <EOL> @ staticmethod <EOL> def waitForStopCommand ( sock ) : <EOL> message = common . readString ( sock ) <EOL> if message != '<STR_LIT>' : <EOL> logging . error ( "<STR_LIT>" , message ) <EOL> else : <EOL> logging . info ( "<STR_LIT>" ) <EOL> def startWorker ( self ) : <EOL> self . controlService ( '<STR_LIT>' , '<STR_LIT:start>' , env = self . getWorkerEnvironment ( ) ) <EOL> def stopWorker ( self ) : <EOL> self . controlService ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> class MultiMachineMasterRunner ( MultiMachineTestRunner ) : <EOL> def __init__ ( self , machinesInCluster , ownAddress , testId , testDir , timeout = None ) : <EOL> super ( MultiMachineMasterRunner , self ) . __init__ ( <EOL> machinesInCluster , ownAddress , testDir , testDir , timeout <EOL> ) <EOL> self . testControlServer = SimpleServer . SimpleServer ( MultiMachineMasterRunner . TestControlPort ) <EOL> self . testControlServer . _onConnect = self . onConnect <EOL> self . serverThread = None <EOL> self . lock = threading . Lock ( ) <EOL> self . allWorkersConnectedEvent = threading . Event ( ) <EOL> if len ( self . machinesInCluster ) == <NUM_LIT:1> : <EOL> self . allWorkersConnectedEvent . set ( ) <EOL> self . connectedWorkers = [ ] <EOL> @ property <EOL> def serviceName ( self ) : <EOL> return "<STR_LIT>" <EOL> def run_ ( self ) : <EOL> try : <EOL> self . startTestControlServer ( ) <EOL> self . waitForAllWorkers ( ) <EOL> self . startMaster ( ) <EOL> self . startWorkers ( ) <EOL> return self . runTests ( ) <EOL> finally : <EOL> self . stopWorkers ( ) <EOL> self . stopMaster ( ) <EOL> self . stopTestControlServer ( ) <EOL> def startTestControlServer ( self ) : <EOL> logging . info ( "<STR_LIT>" , <EOL> MultiMachineMasterRunner . TestControlPort <EOL> ) <EOL> self . serverThread = threading . Thread ( target = self . serverListenLoop ) <EOL> self . serverThread . start ( ) <EOL> def serverListenLoop ( self ) : <EOL> while not self . testControlServer . shouldStop ( ) : <EOL> try : <EOL> self . testControlServer . runListenLoop ( ) <EOL> except socket . error as e : <EOL> if e . errno == <NUM_LIT> : <EOL> time . sleep ( <NUM_LIT:1.0> ) <EOL> def stopTestControlServer ( self ) : <EOL> logging . info ( "<STR_LIT>" ) <EOL> self . testControlServer . stop ( ) <EOL> logging . info ( "<STR_LIT>" ) <EOL> self . serverThread . join ( ) <EOL> logging . info ( "<STR_LIT>" ) <EOL> def waitForAllWorkers ( self ) : <EOL> logging . info ( "<STR_LIT>" ) <EOL> self . allWorkersConnectedEvent . wait ( ) <EOL> def startMaster ( self ) : <EOL> self . controlService ( '<STR_LIT:start>' ) <EOL> self . controlService ( '<STR_LIT>' , '<STR_LIT:start>' , env = self . getWorkerEnvironment ( ) ) <EOL> def stopMaster ( self ) : <EOL> logging . info ( "<STR_LIT>" ) <EOL> self . controlService ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> self . controlService ( '<STR_LIT>' ) <EOL> def startWorkers ( self ) : <EOL> self . forEachWorker ( <EOL> lambda workerSocket : common . writeString ( workerSocket , "<STR_LIT>" ) <EOL> ) <EOL> def stopWorkers ( self ) : <EOL> logging . info ( "<STR_LIT>" ) <EOL> self . forEachWorker ( <EOL> lambda workerSocket : common . writeString ( workerSocket , "<STR_LIT>" ) <EOL> ) <EOL> def forEachWorker ( self , func ) : <EOL> for workerSocket , address in self . connectedWorkers : <EOL> try : <EOL> func ( workerSocket ) <EOL> except socket . error : <EOL> logging . error ( "<STR_LIT>" , address , traceback . format_exc ( ) ) <EOL> workerSocket . close ( ) <EOL> def runTests ( self ) : <EOL> scriptRunner = TestScriptRunner ( testRoot = self . testDir , defaultTimeout = self . timeout ) <EOL> return scriptRunner . run ( ) <EOL> def onConnect ( self , workerSocket , address ) : <EOL> host , port = address <EOL> try : <EOL> if not self . protocolHandshake ( workerSocket ) : <EOL> logging . error ( "<STR_LIT>" , host ) <EOL> self . endSession ( workerSocket , '<STR_LIT>' ) <EOL> elif not host in self . machinesInCluster : <EOL> logging . error ( "<STR_LIT>" , host ) <EOL> self . endSession ( workerSocket , '<STR_LIT>' ) <EOL> elif not self . verifyTestId ( workerSocket ) : <EOL> logging . error ( "<STR_LIT>" , host ) <EOL> self . endSession ( workerSocket , '<STR_LIT>' ) <EOL> else : <EOL> self . addConnectedWorker ( ( workerSocket , host ) ) <EOL> except socket . error : <EOL> logging . error ( "<STR_LIT>" , host , traceback . format_exc ( ) ) <EOL> def protocolHandshake ( self , workerSocket ) : <EOL> workerVersion = common . readString ( workerSocket ) <EOL> return workerVersion == MultiMachineTestRunner . ProtoclVersion <EOL> def verifyTestId ( self , workerSocket ) : <EOL> workerTestId = common . readString ( workerSocket ) <EOL> return workerTestId == self . testId <EOL> def endSession ( self , workerSocket , message = None ) : <EOL> if message : <EOL> common . writeString ( workerSocket , message ) <EOL> workerSocket . close ( ) <EOL> def addConnectedWorker ( self , worker ) : <EOL> with self . lock : <EOL> logging . info ( "<STR_LIT>" , worker ) <EOL> self . connectedWorkers . append ( worker ) <EOL> if len ( self . connectedWorkers ) == len ( self . machinesInCluster ) - <NUM_LIT:1> : <EOL> logging . info ( "<STR_LIT>" ) <EOL> self . allWorkersConnectedEvent . set ( ) </s>
