evennia/contrib/procpools/ampoule/test/test_process.py

867 lines
28 KiB
Python

from signal import SIGHUP
import math
import os
import os.path
from cStringIO import StringIO as sio
import tempfile
from twisted.internet import error, defer, reactor
from twisted.python import failure, reflect
from twisted.trial import unittest
from twisted.protocols import amp
from contrib.procpools.ampoule import main, child, commands, pool
class ShouldntHaveBeenCalled(Exception):
pass
def _raise(_):
raise ShouldntHaveBeenCalled(_)
class _FakeT(object):
closeStdinCalled = False
def __init__(self, s):
self.s = s
def closeStdin(self):
self.closeStdinCalled = True
def write(self, data):
self.s.write(data)
class FakeAMP(object):
connector = None
reason = None
def __init__(self, s):
self.s = s
def makeConnection(self, connector):
if self.connector is not None:
raise Exception("makeConnection called twice")
self.connector = connector
def connectionLost(self, reason):
if self.reason is not None:
raise Exception("connectionLost called twice")
self.reason = reason
def dataReceived(self, data):
self.s.write(data)
class Ping(amp.Command):
arguments = [('data', amp.String())]
response = [('response', amp.String())]
class Pong(amp.Command):
arguments = [('data', amp.String())]
response = [('response', amp.String())]
class Pid(amp.Command):
response = [('pid', amp.Integer())]
class Reactor(amp.Command):
response = [('classname', amp.String())]
class NoResponse(amp.Command):
arguments = [('arg', amp.String())]
requiresAnswer = False
class GetResponse(amp.Command):
response = [("response", amp.String())]
class Child(child.AMPChild):
def ping(self, data):
return self.callRemote(Pong, data=data)
Ping.responder(ping)
class PidChild(child.AMPChild):
def pid(self):
import os
return {'pid': os.getpid()}
Pid.responder(pid)
class NoResponseChild(child.AMPChild):
_set = False
def noresponse(self, arg):
self._set = arg
return {}
NoResponse.responder(noresponse)
def getresponse(self):
return {"response": self._set}
GetResponse.responder(getresponse)
class ReactorChild(child.AMPChild):
def reactor(self):
from twisted.internet import reactor
return {'classname': reactor.__class__.__name__}
Reactor.responder(reactor)
class First(amp.Command):
arguments = [('data', amp.String())]
response = [('response', amp.String())]
class Second(amp.Command):
pass
class WaitingChild(child.AMPChild):
deferred = None
def first(self, data):
self.deferred = defer.Deferred()
return self.deferred.addCallback(lambda _: {'response': data})
First.responder(first)
def second(self):
self.deferred.callback('')
return {}
Second.responder(second)
class Die(amp.Command):
pass
class BadChild(child.AMPChild):
def die(self):
self.shutdown = False
self.transport.loseConnection()
return {}
Die.responder(die)
class Write(amp.Command):
response = [("response", amp.String())]
pass
class Writer(child.AMPChild):
def __init__(self, data='hello'):
child.AMPChild.__init__(self)
self.data = data
def write(self):
return {'response': self.data}
Write.responder(write)
class GetCWD(amp.Command):
response = [("cwd", amp.String())]
class TempDirChild(child.AMPChild):
def __init__(self, directory=None):
child.AMPChild.__init__(self)
self.directory = directory
def __enter__(self):
directory = tempfile.mkdtemp()
os.chdir(directory)
if self.directory is not None:
os.mkdir(self.directory)
os.chdir(self.directory)
def __exit__(self, exc_type, exc_val, exc_tb):
cwd = os.getcwd()
os.chdir('..')
os.rmdir(cwd)
def getcwd(self):
return {'cwd': os.getcwd()}
GetCWD.responder(getcwd)
class TestAMPConnector(unittest.TestCase):
def setUp(self):
"""
The only reason why this method exists is to let 'trial ampoule'
to install the signal handlers (#3178 for reference).
"""
super(TestAMPConnector, self).setUp()
d = defer.Deferred()
reactor.callLater(0, d.callback, None)
return d
def _makeConnector(self, s, sa):
a = FakeAMP(sa)
ac = main.AMPConnector(a)
assert ac.name is not None
ac.transport = _FakeT(s)
return ac
def test_protocol(self):
"""
Test that outReceived writes to AMP and that it triggers the
finished deferred once the process ended.
"""
s = sio()
sa = sio()
ac = self._makeConnector(s, sa)
for x in xrange(99):
ac.childDataReceived(4, str(x))
ac.processEnded(failure.Failure(error.ProcessDone(0)))
return ac.finished.addCallback(
lambda _: self.assertEqual(sa.getvalue(), ''.join(str(x) for x in xrange(99)))
)
def test_protocol_failing(self):
"""
Test that a failure in the process termination is correctly
propagated to the finished deferred.
"""
s = sio()
sa = sio()
ac = self._makeConnector(s, sa)
ac.finished.addCallback(_raise)
fail = failure.Failure(error.ProcessTerminated())
self.assertFailure(ac.finished, error.ProcessTerminated)
ac.processEnded(fail)
def test_startProcess(self):
"""
Test that startProcess actually starts a subprocess and that
it receives data back from the process through AMP.
"""
s = sio()
a = FakeAMP(s)
STRING = "ciao"
BOOT = """\
import sys, os
def main(arg):
os.write(4, arg)
main(sys.argv[1])
"""
starter = main.ProcessStarter(bootstrap=BOOT,
args=(STRING,),
packages=("twisted", "ampoule"))
amp, finished = starter.startPythonProcess(main.AMPConnector(a))
def _eb(reason):
print reason
finished.addErrback(_eb)
return finished.addCallback(lambda _: self.assertEquals(s.getvalue(), STRING))
def test_failing_deferToProcess(self):
"""
Test failing subprocesses and the way they terminate and preserve
failing information.
"""
s = sio()
a = FakeAMP(s)
STRING = "ciao"
BOOT = """\
import sys
def main(arg):
raise Exception(arg)
main(sys.argv[1])
"""
starter = main.ProcessStarter(bootstrap=BOOT, args=(STRING,), packages=("twisted", "ampoule"))
ready, finished = starter.startPythonProcess(main.AMPConnector(a), "I'll be ignored")
self.assertFailure(finished, error.ProcessTerminated)
finished.addErrback(lambda reason: self.assertEquals(reason.getMessage(), STRING))
return finished
def test_env_setting(self):
"""
Test that and environment variable passed to the process starter
is correctly passed to the child process.
"""
s = sio()
a = FakeAMP(s)
STRING = "ciao"
BOOT = """\
import sys, os
def main():
os.write(4, os.getenv("FOOBAR"))
main()
"""
starter = main.ProcessStarter(bootstrap=BOOT,
packages=("twisted", "ampoule"),
env={"FOOBAR": STRING})
amp, finished = starter.startPythonProcess(main.AMPConnector(a), "I'll be ignored")
def _eb(reason):
print reason
finished.addErrback(_eb)
return finished.addCallback(lambda _: self.assertEquals(s.getvalue(), STRING))
def test_startAMPProcess(self):
"""
Test that you can start an AMP subprocess and that it correctly
accepts commands and correctly answers them.
"""
STRING = "ciao"
starter = main.ProcessStarter(packages=("twisted", "ampoule"))
c, finished = starter.startAMPProcess(child.AMPChild)
c.callRemote(commands.Echo, data=STRING
).addCallback(lambda response:
self.assertEquals(response['response'], STRING)
).addCallback(lambda _: c.callRemote(commands.Shutdown))
return finished
def test_BootstrapContext(self):
starter = main.ProcessStarter(packages=('twisted', 'ampoule'))
c, finished = starter.startAMPProcess(TempDirChild)
cwd = []
def checkBootstrap(response):
cwd.append(response['cwd'])
self.assertNotEquals(cwd, os.getcwd())
d = c.callRemote(GetCWD)
d.addCallback(checkBootstrap)
d.addCallback(lambda _: c.callRemote(commands.Shutdown))
finished.addCallback(lambda _: self.assertFalse(os.path.exists(cwd[0])))
return finished
def test_BootstrapContextInstance(self):
starter = main.ProcessStarter(packages=('twisted', 'ampoule'))
c, finished = starter.startAMPProcess(TempDirChild,
ampChildArgs=('foo',))
cwd = []
def checkBootstrap(response):
cwd.append(response['cwd'])
self.assertTrue(cwd[0].endswith('/foo'))
d = c.callRemote(GetCWD)
d.addCallback(checkBootstrap)
d.addCallback(lambda _: c.callRemote(commands.Shutdown))
finished.addCallback(lambda _: self.assertFalse(os.path.exists(cwd[0])))
return finished
def test_startAMPAndParentProtocol(self):
"""
Test that you can start an AMP subprocess and the children can
call methods on their parent.
"""
DATA = "CIAO"
APPEND = "123"
class Parent(amp.AMP):
def pong(self, data):
return {'response': DATA+APPEND}
Pong.responder(pong)
starter = main.ProcessStarter(packages=("twisted", "ampoule"))
subp, finished = starter.startAMPProcess(ampChild=Child, ampParent=Parent)
subp.callRemote(Ping, data=DATA
).addCallback(lambda response:
self.assertEquals(response['response'], DATA+APPEND)
).addCallback(lambda _: subp.callRemote(commands.Shutdown))
return finished
def test_roundtripError(self):
"""
Test that invoking a child using an unreachable class raises
a L{RunTimeError} .
"""
class Child(child.AMPChild):
pass
starter = main.ProcessStarter(packages=("twisted", "ampoule"))
self.assertRaises(RuntimeError, starter.startAMPProcess, ampChild=Child)
class TestProcessPool(unittest.TestCase):
def test_startStopWorker(self):
"""
Test that starting and stopping a worker keeps the state of
the process pool consistent.
"""
pp = pool.ProcessPool()
self.assertEquals(pp.started, False)
self.assertEquals(pp.finished, False)
self.assertEquals(pp.processes, set())
self.assertEquals(pp._finishCallbacks, {})
def _checks():
self.assertEquals(pp.started, False)
self.assertEquals(pp.finished, False)
self.assertEquals(len(pp.processes), 1)
self.assertEquals(len(pp._finishCallbacks), 1)
return pp.stopAWorker()
def _closingUp(_):
self.assertEquals(pp.started, False)
self.assertEquals(pp.finished, False)
self.assertEquals(len(pp.processes), 0)
self.assertEquals(pp._finishCallbacks, {})
pp.startAWorker()
return _checks().addCallback(_closingUp).addCallback(lambda _: pp.stop())
def test_startAndStop(self):
"""
Test that a process pool's start and stop method create the
expected number of workers and keep state consistent in the
process pool.
"""
pp = pool.ProcessPool()
self.assertEquals(pp.started, False)
self.assertEquals(pp.finished, False)
self.assertEquals(pp.processes, set())
self.assertEquals(pp._finishCallbacks, {})
def _checks(_):
self.assertEquals(pp.started, True)
self.assertEquals(pp.finished, False)
self.assertEquals(len(pp.processes), pp.min)
self.assertEquals(len(pp._finishCallbacks), pp.min)
return pp.stop()
def _closingUp(_):
self.assertEquals(pp.started, True)
self.assertEquals(pp.finished, True)
self.assertEquals(len(pp.processes), 0)
self.assertEquals(pp._finishCallbacks, {})
return pp.start().addCallback(_checks).addCallback(_closingUp)
def test_adjustPoolSize(self):
"""
Test that calls to pool.adjustPoolSize are correctly handled.
"""
pp = pool.ProcessPool(min=10)
self.assertEquals(pp.started, False)
self.assertEquals(pp.finished, False)
self.assertEquals(pp.processes, set())
self.assertEquals(pp._finishCallbacks, {})
def _resize1(_):
self.assertEquals(pp.started, True)
self.assertEquals(pp.finished, False)
self.assertEquals(len(pp.processes), pp.min)
self.assertEquals(len(pp._finishCallbacks), pp.min)
return pp.adjustPoolSize(min=2, max=3)
def _resize2(_):
self.assertEquals(pp.started, True)
self.assertEquals(pp.finished, False)
self.assertEquals(pp.max, 3)
self.assertEquals(pp.min, 2)
self.assertEquals(len(pp.processes), pp.max)
self.assertEquals(len(pp._finishCallbacks), pp.max)
def _resize3(_):
self.assertRaises(AssertionError, pp.adjustPoolSize, min=-1, max=5)
self.assertRaises(AssertionError, pp.adjustPoolSize, min=5, max=1)
return pp.stop()
return pp.start(
).addCallback(_resize1
).addCallback(_resize2
).addCallback(_resize3)
def test_childRestart(self):
"""
Test that a failing child process is immediately restarted.
"""
pp = pool.ProcessPool(ampChild=BadChild, min=1)
STRING = "DATA"
def _checks(_):
d = pp._finishCallbacks.values()[0]
pp.doWork(Die).addErrback(lambda _: None)
return d.addBoth(_checksAgain)
def _checksAgain(_):
return pp.doWork(commands.Echo, data=STRING
).addCallback(lambda result: self.assertEquals(result['response'], STRING))
return pp.start(
).addCallback(_checks
).addCallback(lambda _: pp.stop())
def test_parentProtocolChange(self):
"""
Test that the father can use an AMP protocol too.
"""
DATA = "CIAO"
APPEND = "123"
class Parent(amp.AMP):
def pong(self, data):
return {'response': DATA+APPEND}
Pong.responder(pong)
pp = pool.ProcessPool(ampChild=Child, ampParent=Parent)
def _checks(_):
return pp.doWork(Ping, data=DATA
).addCallback(lambda response:
self.assertEquals(response['response'], DATA+APPEND)
)
return pp.start().addCallback(_checks).addCallback(lambda _: pp.stop())
def test_deferToAMPProcess(self):
"""
Test that deferToAMPProcess works as expected.
"""
def cleanupGlobalPool():
d = pool.pp.stop()
pool.pp = None
return d
self.addCleanup(cleanupGlobalPool)
STRING = "CIAOOOO"
d = pool.deferToAMPProcess(commands.Echo, data=STRING)
d.addCallback(self.assertEquals, {"response": STRING})
return d
def test_checkStateInPool(self):
"""
Test that busy and ready lists are correctly maintained.
"""
pp = pool.ProcessPool(ampChild=WaitingChild)
DATA = "foobar"
def _checks(_):
d = pp.callRemote(First, data=DATA)
self.assertEquals(pp.started, True)
self.assertEquals(pp.finished, False)
self.assertEquals(len(pp.processes), pp.min)
self.assertEquals(len(pp._finishCallbacks), pp.min)
self.assertEquals(len(pp.ready), pp.min-1)
self.assertEquals(len(pp.busy), 1)
child = pp.busy.pop()
pp.busy.add(child)
child.callRemote(Second)
return d
return pp.start(
).addCallback(_checks
).addCallback(lambda _: pp.stop())
def test_growingToMax(self):
"""
Test that the pool grows over time until it reaches max processes.
"""
MAX = 5
pp = pool.ProcessPool(ampChild=WaitingChild, min=1, max=MAX)
def _checks(_):
self.assertEquals(pp.started, True)
self.assertEquals(pp.finished, False)
self.assertEquals(len(pp.processes), pp.min)
self.assertEquals(len(pp._finishCallbacks), pp.min)
D = "DATA"
d = [pp.doWork(First, data=D) for x in xrange(MAX)]
self.assertEquals(pp.started, True)
self.assertEquals(pp.finished, False)
self.assertEquals(len(pp.processes), pp.max)
self.assertEquals(len(pp._finishCallbacks), pp.max)
[child.callRemote(Second) for child in pp.processes]
return defer.DeferredList(d)
return pp.start(
).addCallback(_checks
).addCallback(lambda _: pp.stop())
def test_growingToMaxAndShrinking(self):
"""
Test that the pool grows but after 'idle' time the number of
processes goes back to the minimum.
"""
MAX = 5
MIN = 1
IDLE = 1
pp = pool.ProcessPool(ampChild=WaitingChild, min=MIN, max=MAX, maxIdle=IDLE)
def _checks(_):
self.assertEquals(pp.started, True)
self.assertEquals(pp.finished, False)
self.assertEquals(len(pp.processes), pp.min)
self.assertEquals(len(pp._finishCallbacks), pp.min)
D = "DATA"
d = [pp.doWork(First, data=D) for x in xrange(MAX)]
self.assertEquals(pp.started, True)
self.assertEquals(pp.finished, False)
self.assertEquals(len(pp.processes), pp.max)
self.assertEquals(len(pp._finishCallbacks), pp.max)
[child.callRemote(Second) for child in pp.processes]
return defer.DeferredList(d).addCallback(_realChecks)
def _realChecks(_):
from twisted.internet import reactor
d = defer.Deferred()
def _cb():
def __(_):
try:
self.assertEquals(pp.started, True)
self.assertEquals(pp.finished, False)
self.assertEquals(len(pp.processes), pp.min)
self.assertEquals(len(pp._finishCallbacks), pp.min)
d.callback(None)
except Exception, e:
d.errback(e)
return pp._pruneProcesses().addCallback(__)
# just to be shure we are called after the pruner
pp.looping.stop() # stop the looping, we don't want it to
# this right here
reactor.callLater(IDLE, _cb)
return d
return pp.start(
).addCallback(_checks
).addCallback(lambda _: pp.stop())
def test_recycling(self):
"""
Test that after a given number of calls subprocesses are
recycled.
"""
MAX = 1
MIN = 1
RECYCLE_AFTER = 1
pp = pool.ProcessPool(ampChild=PidChild, min=MIN, max=MAX, recycleAfter=RECYCLE_AFTER)
self.addCleanup(pp.stop)
def _checks(_):
self.assertEquals(pp.started, True)
self.assertEquals(pp.finished, False)
self.assertEquals(len(pp.processes), pp.min)
self.assertEquals(len(pp._finishCallbacks), pp.min)
return pp.doWork(Pid
).addCallback(lambda response: response['pid'])
def _checks2(pid):
return pp.doWork(Pid
).addCallback(lambda response: response['pid']
).addCallback(self.assertNotEquals, pid)
d = pp.start()
d.addCallback(_checks)
d.addCallback(_checks2)
return d
def test_recyclingWithQueueOverload(self):
"""
Test that we get the correct number of different results when
we overload the pool of calls.
"""
MAX = 5
MIN = 1
RECYCLE_AFTER = 10
CALLS = 60
pp = pool.ProcessPool(ampChild=PidChild, min=MIN, max=MAX, recycleAfter=RECYCLE_AFTER)
self.addCleanup(pp.stop)
def _check(results):
s = set()
for succeed, response in results:
s.add(response['pid'])
# For the first C{MAX} calls, each is basically guaranteed to go to
# a different child. After that, though, there are no guarantees.
# All the rest might go to a single child, since the child to
# perform a job is selected arbitrarily from the "ready" set. Fair
# distribution of jobs needs to be implemented; right now it's "set
# ordering" distribution of jobs.
self.assertTrue(len(s) > MAX)
def _work(_):
l = [pp.doWork(Pid) for x in xrange(CALLS)]
d = defer.DeferredList(l)
return d.addCallback(_check)
d = pp.start()
d.addCallback(_work)
return d
def test_disableProcessRecycling(self):
"""
Test that by setting 0 to recycleAfter we actually disable process recycling.
"""
MAX = 1
MIN = 1
RECYCLE_AFTER = 0
pp = pool.ProcessPool(ampChild=PidChild, min=MIN, max=MAX, recycleAfter=RECYCLE_AFTER)
def _checks(_):
self.assertEquals(pp.started, True)
self.assertEquals(pp.finished, False)
self.assertEquals(len(pp.processes), pp.min)
self.assertEquals(len(pp._finishCallbacks), pp.min)
return pp.doWork(Pid
).addCallback(lambda response: response['pid'])
def _checks2(pid):
return pp.doWork(Pid
).addCallback(lambda response: response['pid']
).addCallback(self.assertEquals, pid
).addCallback(lambda _: pid)
def finish(reason):
return pp.stop().addCallback(lambda _: reason)
return pp.start(
).addCallback(_checks
).addCallback(_checks2
).addCallback(_checks2
).addCallback(finish)
def test_changeChildrenReactor(self):
"""
Test that by passing the correct argument children change their
reactor type.
"""
MAX = 1
MIN = 1
FIRST = "select"
SECOND = "poll"
def checkDefault():
pp = pool.ProcessPool(
starter=main.ProcessStarter(
childReactor=FIRST,
packages=("twisted", "ampoule")),
ampChild=ReactorChild, min=MIN, max=MAX)
pp.start()
return pp.doWork(Reactor
).addCallback(self.assertEquals, {'classname': "SelectReactor"}
).addCallback(lambda _: pp.stop())
def checkPool(_):
pp = pool.ProcessPool(
starter=main.ProcessStarter(
childReactor=SECOND,
packages=("twisted", "ampoule")),
ampChild=ReactorChild, min=MIN, max=MAX)
pp.start()
return pp.doWork(Reactor
).addCallback(self.assertEquals, {'classname': "PollReactor"}
).addCallback(lambda _: pp.stop())
return checkDefault(
).addCallback(checkPool)
try:
from select import poll
except ImportError:
test_changeChildrenReactor.skip = "This architecture doesn't support select.poll, I can't run this test"
def test_commandsWithoutResponse(self):
"""
Test that if we send a command without a required answer we
actually don't have any problems.
"""
DATA = "hello"
pp = pool.ProcessPool(ampChild=NoResponseChild, min=1, max=1)
def _check(_):
return pp.doWork(GetResponse
).addCallback(self.assertEquals, {"response": DATA})
def _work(_):
return pp.doWork(NoResponse, arg=DATA)
return pp.start(
).addCallback(_work
).addCallback(_check
).addCallback(lambda _: pp.stop())
def test_SupplyChildArgs(self):
"""Ensure that arguments for the child constructor are passed in."""
pp = pool.ProcessPool(Writer, ampChildArgs=['body'], min=0)
def _check(result):
return pp.doWork(Write).addCallback(
self.assertEquals, {'response': 'body'})
return pp.start(
).addCallback(_check
).addCallback(lambda _: pp.stop())
def processTimeoutTest(self, timeout):
pp = pool.ProcessPool(WaitingChild, min=1, max=1)
def _work(_):
d = pp.callRemote(First, data="ciao", _timeout=timeout)
self.assertFailure(d, error.ProcessTerminated)
return d
return pp.start(
).addCallback(_work
).addCallback(lambda _: pp.stop())
def test_processTimeout(self):
"""
Test that a call that doesn't finish within the given timeout
time is correctly handled.
"""
return self.processTimeoutTest(1)
def test_processTimeoutZero(self):
"""
Test that the process is correctly handled when the timeout is zero.
"""
return self.processTimeoutTest(0)
def test_processDeadline(self):
pp = pool.ProcessPool(WaitingChild, min=1, max=1)
def _work(_):
d = pp.callRemote(First, data="ciao", _deadline=reactor.seconds())
self.assertFailure(d, error.ProcessTerminated)
return d
return pp.start(
).addCallback(_work
).addCallback(lambda _: pp.stop())
def test_processBeforeDeadline(self):
pp = pool.ProcessPool(PidChild, min=1, max=1)
def _work(_):
d = pp.callRemote(Pid, _deadline=reactor.seconds() + 10)
d.addCallback(lambda result: self.assertNotEqual(result['pid'], 0))
return d
return pp.start(
).addCallback(_work
).addCallback(lambda _: pp.stop())
def test_processTimeoutSignal(self):
"""
Test that a call that doesn't finish within the given timeout
time is correctly handled.
"""
pp = pool.ProcessPool(WaitingChild, min=1, max=1,
timeout_signal=SIGHUP)
def _work(_):
d = pp.callRemote(First, data="ciao", _timeout=1)
d.addCallback(lambda d: self.fail())
text = 'signal %d' % SIGHUP
d.addErrback(
lambda f: self.assertTrue(text in f.value[0],
'"%s" not in "%s"' % (text, f.value[0])))
return d
return pp.start(
).addCallback(_work
).addCallback(lambda _: pp.stop())
def test_processGlobalTimeout(self):
"""
Test that a call that doesn't finish within the given global
timeout time is correctly handled.
"""
pp = pool.ProcessPool(WaitingChild, min=1, max=1, timeout=1)
def _work(_):
d = pp.callRemote(First, data="ciao")
self.assertFailure(d, error.ProcessTerminated)
return d
return pp.start(
).addCallback(_work
).addCallback(lambda _: pp.stop())