class SyncHTTPClientTest(unittest.TestCase):
def setUp(self):
self.server_ioloop = IOLoop()
event = threading.Event()
@gen.coroutine
def init_server():
sock, self.port = bind_unused_port()
app = Application([("/", HelloWorldHandler)])
self.server = HTTPServer(app)
self.server.add_socket(sock)
event.set()
def start():
self.server_ioloop.run_sync(init_server)
self.server_ioloop.start()
self.server_thread = threading.Thread(target=start)
self.server_thread.start()
event.wait()
self.http_client = HTTPClient()
def tearDown(self):
def stop_server():
self.server.stop()
# Delay the shutdown of the IOLoop by several iterations because
# the server may still have some cleanup work left when
# the client finishes with the response (this is noticeable
# with http/2, which leaves a Future with an unexamined
# StreamClosedError on the loop).
@gen.coroutine
def slow_stop():
# The number of iterations is difficult to predict. Typically,
# one is sufficient, although sometimes it needs more.
for i in range(5):
yield
self.server_ioloop.stop()
self.server_ioloop.add_callback(slow_stop)
self.server_ioloop.add_callback(stop_server)
self.server_thread.join()
self.http_client.close()
self.server_ioloop.close(all_fds=True)
def get_url(self, path):
return "http://127.0.0.1:%d%s" % (self.port, path)
def test_sync_client(self):
response = self.http_client.fetch(self.get_url("/"))
self.assertEqual(b"Hello world!", response.body)
def test_sync_client_error(self):
# Synchronous HTTPClient raises errors directly; no need for
# response.rethrow()
with self.assertRaises(HTTPError) as assertion:
self.http_client.fetch(self.get_url("/notfound"))
self.assertEqual(assertion.exception.code, 404)
def run_worker_fork(q, ip, scheduler_ip, scheduler_port, ncores, nanny_port,
worker_port, local_dir, services, name, memory_limit):
""" Function run by the Nanny when creating the worker """
from distributed import Worker # pragma: no cover
from tornado.ioloop import IOLoop # pragma: no cover
IOLoop.clear_instance() # pragma: no cover
loop = IOLoop() # pragma: no cover
loop.make_current() # pragma: no cover
worker = Worker(scheduler_ip, scheduler_port, ncores=ncores, ip=ip,
service_ports={'nanny': nanny_port}, local_dir=local_dir,
services=services, name=name, memory_limit=memory_limit,
loop=loop) # pragma: no cover
@gen.coroutine # pragma: no cover
def start():
try: # pragma: no cover
yield worker._start(worker_port) # pragma: no cover
except Exception as e: # pragma: no cover
logger.exception(e) # pragma: no cover
q.put(e) # pragma: no cover
else:
assert worker.port # pragma: no cover
q.put({'port': worker.port, 'dir': worker.local_dir}) # pragma: no cover
loop.add_callback(start) # pragma: no cover
try:
loop.start() # pragma: no cover
finally:
loop.stop()
loop.close(all_fds=True)
class SyncHTTPClientTest(unittest.TestCase):
def setUp(self):
if IOLoop.configured_class().__name__ == 'TwistedIOLoop':
# TwistedIOLoop only supports the global reactor, so we can't have
# separate IOLoops for client and server threads.
raise unittest.SkipTest(
'Sync HTTPClient not compatible with TwistedIOLoop')
self.server_ioloop = IOLoop()
@gen.coroutine
def init_server():
sock, self.port = bind_unused_port()
app = Application([('/', HelloWorldHandler)])
self.server = HTTPServer(app)
self.server.add_socket(sock)
self.server_ioloop.run_sync(init_server)
self.server_thread = threading.Thread(target=self.server_ioloop.start)
self.server_thread.start()
self.http_client = HTTPClient()
def tearDown(self):
def stop_server():
self.server.stop()
# Delay the shutdown of the IOLoop by several iterations because
# the server may still have some cleanup work left when
# the client finishes with the response (this is noticeable
# with http/2, which leaves a Future with an unexamined
# StreamClosedError on the loop).
@gen.coroutine
def slow_stop():
# The number of iterations is difficult to predict. Typically,
# one is sufficient, although sometimes it needs more.
for i in range(5):
yield
self.server_ioloop.stop()
self.server_ioloop.add_callback(slow_stop)
self.server_ioloop.add_callback(stop_server)
self.server_thread.join()
self.http_client.close()
self.server_ioloop.close(all_fds=True)
def get_url(self, path):
return 'http://127.0.0.1:%d%s' % (self.port, path)
def test_sync_client(self):
response = self.http_client.fetch(self.get_url('/'))
self.assertEqual(b'Hello world!', response.body)
def test_sync_client_error(self):
# Synchronous HTTPClient raises errors directly; no need for
# response.rethrow()
with self.assertRaises(HTTPError) as assertion:
self.http_client.fetch(self.get_url('/notfound'))
self.assertEqual(assertion.exception.code, 404)
class WebApplication(TornadoApplication):
def __init__(self, app):
super(WebApplication, self).__init__([
(r"/settings/", SettingsHandler),
(r"/status/", StatusHandler),
])
self.app = app
self.loop = None
self.port = None # Actual port that we're running
self.apply_in_progress = False
@staticmethod
def _ensure_event_loop():
try:
_ = asyncio.get_event_loop()
except RuntimeError:
asyncio.set_event_loop(asyncio.new_event_loop())
def routine(self, port: int = None):
self._ensure_event_loop()
if port is None:
sock, port = bind_unused_port()
http_server = HTTPServer(self)
http_server.add_sockets([sock])
else:
self.listen(port)
self.port = port
self.loop = IOLoop().current()
self.loop.start()
def stop(self):
if self.loop:
self.loop.stop()
class HTTPClient(object):
""" 阻塞式的HTTP客户端。使用ioloop+异步HTTP客户端实现，基本只是用于测试。 """
def __init__(self, async_client_class=None, **kwargs):
self._io_loop = IOLoop()
if async_client_class is None:
async_client_class = AsyncHTTPClient
self._async_client = async_client_class(self._io_loop, **kwargs)
self._response = None
self._closed = False
def __del__(self):
self.close()
def close(self):
if not self._closed:
self._async_client.close()
self._io_loop.close()
self._closed = True
def fetch(self, request, **kwargs):
def callback(response):
self._response = response
self._io_loop.stop()
self._async_client.fetch(request, callback, **kwargs)
self._io_loop.start()
response = self._response
self._response = None
response.rethrow()
return response
class Client(object):
def __init__(self, host, port, timeout=None, connect_timeout=-1, unix_socket=None, max_buffer_size=104857600):
self._io_loop = IOLoop()
self._async_client = AsyncClient(host, port, unix_socket, self._io_loop, timeout, connect_timeout, max_buffer_size)
self._response = None
self._closed = False
def __del__(self):
self.close()
@property
def closed(self):
return self._async_client.closed
def close(self):
if not self._closed:
self._async_client.close()
self._io_loop.close()
self._closed = True
def call(self, request):
def callback(response):
self._response = response
self._io_loop.stop()
self._async_client.call(request, callback)
self._io_loop.start()
response = self._response
self._response = None
response.rethrow()
return response
def __str__(self):
return str(self._async_client)
def run_worker_fork(q, ip, scheduler_ip, scheduler_port, ncores, nanny_port,
worker_port, local_dir, services, name, memory_limit, reconnect):
""" Function run by the Nanny when creating the worker """
from distributed import Worker # pragma: no cover
from tornado.ioloop import IOLoop # pragma: no cover
IOLoop.clear_instance() # pragma: no cover
loop = IOLoop() # pragma: no cover
loop.make_current() # pragma: no cover
worker = Worker(scheduler_ip, scheduler_port, ncores=ncores, ip=ip,
service_ports={'nanny': nanny_port}, local_dir=local_dir,
services=services, name=name, memory_limit=memory_limit,
reconnect=reconnect, loop=loop) # pragma: no cover
@gen.coroutine # pragma: no cover
def run():
try: # pragma: no cover
yield worker._start(worker_port) # pragma: no cover
except Exception as e: # pragma: no cover
logger.exception(e) # pragma: no cover
q.put(e) # pragma: no cover
else:
assert worker.port # pragma: no cover
q.put({'port': worker.port, 'dir': worker.local_dir}) # pragma: no cover
while worker.status != 'closed':
yield gen.sleep(0.1)
logger.info("Worker closed")
try:
loop.run_sync(run)
finally:
loop.stop()
loop.close(all_fds=True)
def loop():
IOLoop.clear_instance()
loop = IOLoop()
loop.make_current()
yield loop
loop.stop()
loop.close()
def loop():
IOLoop.clear_instance()
loop = IOLoop()
loop.make_current()
yield loop
loop.stop()
loop.close()
def start(self):
loop = IOLoop()
s = Context().socket(zmq.PULL)
s.set_hwm(ZMQ_HWM)
socket = zmqstream.ZMQStream(s, loop)
socket.on_recv(self._process_message)
logger.debug(f"connecting to: {HUNTER_ADDR}")
socket.connect(HUNTER_ADDR)
# this needs to be done here
self.router = Client(remote=HUNTER_SINK_ADDR,
nowait=True,
autoclose=False)
try:
loop.start()
except KeyboardInterrupt as e:
loop.stop()
for ss in [s, self.router.socket]:
ss.close()
self.stop()
class HTTPClient(object):
""" 阻塞式的HTTP客户端。使用ioloop+异步HTTP客户端实现，基本只是用于测试。 """
def __init__(self, async_client_class=None, **kwargs):
self._io_loop = IOLoop()
if async_client_class is None:
async_client_class = AsyncHTTPClient
self._async_client = async_client_class(self._io_loop, **kwargs)
self._response = None
self._closed = False
def __del__(self):
self.close()
def close(self):
if not self._closed:
self._async_client.close()
self._io_loop.close()
self._closed = True
def fetch(self, request, **kwargs):
def callback(response):
self._response = response
self._io_loop.stop()
self._async_client.fetch(request, callback, **kwargs)
self._io_loop.start()
response = self._response
self._response = None
response.rethrow()
return response
class SocketServerThreadStarter(Thread):
'''
Used to fire up the three services each in its own thread.
'''
def __init__(self, socketServerClassName, port):
'''
Create one thread for one of the services to run in.
@param socketServerClassName: Name of top level server class to run.
@type socketServerClassName: string
@param port: port to listen on
@type port: int
'''
super(SocketServerThreadStarter, self).__init__();
self.socketServerClassName = socketServerClassName;
self.port = port;
self.ioLoop = None;
def stop(self):
self.ioLoop.stop();
def run(self):
'''
Use the service name to instantiate the proper service, passing in the
proper helper class.
'''
super(SocketServerThreadStarter, self).run();
try:
if self.socketServerClassName == 'RootWordSubmissionService':
EchoTreeService.log("Starting EchoTree new tree submissions server %d: accepts word trees submitted from connecting clients." % self.port);
http_server = RootWordSubmissionService(RootWordSubmissionService.handle_request);
http_server.listen(self.port);
self.ioLoop = IOLoop();
self.ioLoop.start();
self.ioLoop.close(all_fds=True);
return;
elif self.socketServerClassName == 'EchoTreeScriptRequestHandler':
EchoTreeService.log("Starting EchoTree script server %d: Returns one script that listens to the new-tree events in the browser." % self.port);
http_server = EchoTreeScriptRequestHandler(EchoTreeScriptRequestHandler.handle_request);
http_server.listen(self.port);
self.ioLoop = IOLoop();
self.ioLoop.start();
self.ioLoop.close(all_fds=True);
return;
else:
raise ValueError("Service class %s is unknown." % self.socketServerClassName);
except Exception:
# Typically an exception is caught here that complains about 'socket in use'
# Should avoid that by sensing busy socket and timing out:
# if e.errno == 98:
# print "Exception: %s. You need to try starting this service again. Socket busy condition will time out within 30 secs or so." % `e`
# else:
# print `e`;
#raise e;
pass
finally:
if self.ioLoop is not None and self.ioLoop.running():
self.ioLoop.stop();
return;
class XDebugServer(TCPServer):
"""Class to listen for xdebug requests"""
def __init__(self):
"""Constructor """
self.ioloop = IOLoop()
super(XDebugServer, self).__init__(io_loop=self.ioloop)
self.listen(9000)
# this is for cross thread communication
self.inport = Queue()
self.outport = Queue()
self._xdebug_connection = None
def listenfunc():
self.ioloop.make_current()
self.ioloop.start()
self.ioloop.close(all_fds=True)
self.listener_thread = threading.Thread(target=listenfunc)
self.listener_thread.daemon = True
self.listener_thread.start()
def handle_stream(self, stream, address):
"""Handle a connection
Only one connection at a time, for now
:stream: @todo
:address: @todo
:returns: @todo
"""
self._xdebug_connection = XDebugConnection(self, stream, address)
def run_command(self, command, data=None):
"""Send status
:returns: @todo
"""
self.inport.put("{} -i 1\0".format(str(command)))
return self.outport.get()
def stop(self):
"""Stop tornado event loop
:returns: @todo
"""
self.ioloop.stop()
self.listener_thread.join()
del self.ioloop
del self.listener_thread
class SocketServerThreadStarter(Thread):
'''
Convenience for firing up various servers. Currently not used.
In its current form it knows to start the service that distributes
a JavaScript script that subscribes to the EchoTree service (the
main class, which inherits from WebSocketHandler. Need to start
the script server (EchoTreeScriptRequestHandler) in main() if
this module is used stand-alone, rather than from some browser-side
script that already knows how to push new root words, and subscribe
to EchoTrees.
'''
def __init__(self, socketServerClassName, port):
'''
Create one thread for one of the services to run in.
@param socketServerClassName: Name of top level server class to run.
@type socketServerClassName: string
@param port: port to listen on
@type port: int
'''
super(SocketServerThreadStarter, self).__init__();
self.socketServerClassName = socketServerClassName;
self.port = port;
self.ioLoop = None;
def stop(self):
self.ioLoop.stop();
def run(self):
'''
Use the service name to instantiate the proper service, passing in the
proper helper class.
'''
super(SocketServerThreadStarter, self).run();
try:
if self.socketServerClassName == 'EchoTreeScriptRequestHandler':
EchoTreeService.log("Starting EchoTree script server %d: Returns one script that listens to the new-tree events in the browser." % self.port);
http_server = EchoTreeScriptRequestHandler(EchoTreeScriptRequestHandler.handle_request);
http_server.listen(self.port);
self.ioLoop = IOLoop();
self.ioLoop.start();
self.ioLoop.close(all_fds=True);
return;
else:
raise ValueError("Service class %s is unknown." % self.socketServerClassName);
except Exception:
# Typically an exception is caught here that complains about 'socket in use'
# Should avoid that by sensing busy socket and timing out:
# if e.errno == 98:
# print "Exception: %s. You need to try starting this service again. Socket busy condition will time out within 30 secs or so." % `e`
# else:
# print `e`;
#raise e;
pass
finally:
if self.ioLoop is not None and self.ioLoop.running():
self.ioLoop.stop();
return;
class TestGitHubParser(TestCase, TimeoutMixin):
def setUp(self):
self._ioloop = IOLoop()
def test_github_parser(self):
result = {}
def callback(commit_result):
for key in commit_result:
result[key] = commit_result[key]
self._ioloop.stop()
github_parser = GitHubParser(self._ioloop, {})
github_parser.parse("https://github.com/julython/"
"julythontweets/commit/25645d2cf6b58d2657cf6eb0fb4ca59d5f2499f4",
callback)
self.add_timeout(3)
self._ioloop.start()
# will block until callback or timeout
self.assertEqual("Josh Marshall",
result["author"]["name"])
self.assertEqual("joshmarshall",
result["author"]["username"])
self.assertEqual(
"25645d2cf6b58d2657cf6eb0fb4ca59d5f2499f4",
result["commit"])
def test_get_commit_atom_link(self):
"""Test extracting commit Atom URL from an HTML page."""
with self.assertRaises(MissingAtomLink):
get_commit_atom_link("WHATEVER")
with self.assertRaises(MissingAtomLink):
get_commit_atom_link("<html></html>")
html = open(_GITHUB_SAMPLE_HTML_PATH).read()
atom_link = get_commit_atom_link(html)
self.assertEqual(
"https://github.com/julython/julythontweets/commits/master.atom",
atom_link)
def test_get_most_recent_commit(self):
"""Test extracting most recent commit from an Atom Feed."""
with self.assertRaises(MissingCommit):
get_most_recent_commit("whatever")
with self.assertRaises(MissingCommit):
get_most_recent_commit("<rss></rss>")
feed = open(_GITHUB_SAMPLE_ATOM_PATH).read()
commit = get_most_recent_commit(feed)
self.assertEqual({
"name": "Josh Marshall",
"url": "https://github.com/joshmarshall",
"username": "******",
"service": "github"
}, commit["author"])
self.assertEqual("25645d2cf6b58d2657cf6eb0fb4ca59d5f2499f4",
commit["commit"])
self.assertEqual("julython/julythontweets", commit["project"]["id"])
self.assertEqual("github", commit["project"]["service"])
class TornadoFlask(Flask):
def run(self, host='0.0.0.0', port=5000):
self.ioloop = IOLoop()
http_server = HTTPServer(WSGIContainer(app), io_loop=self.ioloop)
http_server.listen(port, host)
self.ioloop.start()
def quit(self):
self.ioloop.stop()
class TestGitHubParser(TestCase, TimeoutMixin):
def setUp(self):
self._ioloop = IOLoop()
def test_github_parser(self):
result = {}
def callback(commit_result):
for key in commit_result:
result[key] = commit_result[key]
self._ioloop.stop()
github_parser = GitHubParser(self._ioloop, {})
github_parser.parse(
"https://github.com/julython/"
"julythontweets/commit/25645d2cf6b58d2657cf6eb0fb4ca59d5f2499f4",
callback)
self.add_timeout(3)
self._ioloop.start()
# will block until callback or timeout
self.assertEqual("Josh Marshall", result["author"]["name"])
self.assertEqual("joshmarshall", result["author"]["username"])
self.assertEqual("25645d2cf6b58d2657cf6eb0fb4ca59d5f2499f4",
result["commit"])
def test_get_commit_atom_link(self):
"""Test extracting commit Atom URL from an HTML page."""
with self.assertRaises(MissingAtomLink):
get_commit_atom_link("WHATEVER")
with self.assertRaises(MissingAtomLink):
get_commit_atom_link("<html></html>")
html = open(_GITHUB_SAMPLE_HTML_PATH).read()
atom_link = get_commit_atom_link(html)
self.assertEqual(
"https://github.com/julython/julythontweets/commits/master.atom",
atom_link)
def test_get_most_recent_commit(self):
"""Test extracting most recent commit from an Atom Feed."""
with self.assertRaises(MissingCommit):
get_most_recent_commit("whatever")
with self.assertRaises(MissingCommit):
get_most_recent_commit("<rss></rss>")
feed = open(_GITHUB_SAMPLE_ATOM_PATH).read()
commit = get_most_recent_commit(feed)
self.assertEqual(
{
"name": "Josh Marshall",
"url": "https://github.com/joshmarshall",
"username": "******",
"service": "github"
}, commit["author"])
self.assertEqual("25645d2cf6b58d2657cf6eb0fb4ca59d5f2499f4",
commit["commit"])
self.assertEqual("julython/julythontweets", commit["project"]["id"])
self.assertEqual("github", commit["project"]["service"])
def test_func():
IOLoop.clear_instance()
loop = IOLoop()
loop.make_current()
cor = gen.coroutine(func)
try:
loop.run_sync(cor, timeout=timeout)
finally:
loop.stop()
loop.close(all_fds=True)
class FakeServerContext(object):
def __init__(self, monkeypatch, fail_these, expected_basename):
self._monkeypatch = monkeypatch
self._fail_these = fail_these
self._expected_basename = expected_basename
self._url = None
self._loop = None
self._started = threading.Condition()
self._thread = threading.Thread(target=self._run)
def __exit__(self, type, value, traceback):
if self._loop is not None:
# we can ONLY use add_callback here, since the loop is
# running in a different thread.
self._loop.add_callback(self._stop)
self._thread.join()
def __enter__(self):
self._started.acquire()
self._thread.start()
self._started.wait()
self._started.release()
_monkeypatch_client_config(self._monkeypatch, self._url)
return self._url
def _run(self):
self._loop = IOLoop()
self._server = FakeAnacondaServer(
io_loop=self._loop,
fail_these=self._fail_these,
expected_basename=self._expected_basename)
self._url = self._server.url
def notify_started():
self._started.acquire()
self._started.notify()
self._started.release()
self._loop.add_callback(notify_started)
self._loop.start()
# done
self._server.unlisten()
def _stop(self):
def really_stop():
if self._loop is not None:
self._loop.stop()
self._loop = None
# the delay allows pending next-tick things to go ahead
# and happen, which may avoid some problems with trying to
# output to stdout after pytest closes it
if self._loop is not None:
self._loop.call_later(delay=0.05, callback=really_stop)
def test_func():
IOLoop.clear_instance()
loop = IOLoop()
loop.make_current()
cor = gen.coroutine(func)
try:
loop.run_sync(cor, timeout=timeout)
finally:
loop.stop()
loop.close(all_fds=True)
class HTTPClient(object):
"""A blocking HTTP client.
This interface is provided for convenience and testing; most applications
that are running an IOLoop will want to use `AsyncHTTPClient` instead.
Typical usage looks like this::
http_client = httpclient.HTTPClient()
try:
response = http_client.fetch("http://www.google.com/")
print response.body
except httpclient.HTTPError, e:
print "Error:", e
"""
def __init__(self, async_client_class=None, **kwargs):
self._io_loop = IOLoop()
if async_client_class is None:
async_client_class = AsyncHTTPClient
self._async_client = async_client_class(self._io_loop, **kwargs)
self._response = None
self._closed = False
def __del__(self):
self.close()
def close(self):
"""Closes the HTTPClient, freeing any resources used."""
if not self._closed:
self._async_client.close()
self._io_loop.close()
self._closed = True
def fetch(self, request, **kwargs):
"""Executes a request, returning an `HTTPResponse`.
The request may be either a string URL or an `HTTPRequest` object.
If it is a string, we construct an `HTTPRequest` using any additional
kwargs: ``HTTPRequest(request, **kwargs)``
If an error occurs during the fetch, we raise an `HTTPError`.
"""
def callback(response):
self._response = response
self._io_loop.stop()
self._async_client.fetch(request, callback, **kwargs)
self._io_loop.start()
response = self._response
self._response = None
response.rethrow()
return response
def main(args=None):
if args is None:
args = sys.argv
setup_logger()
router = TornadioRouter(RouterConnection)
bg_loop = IOLoop()
bg_thread = Thread(target=lambda: bg_loop.start())
bg_task = task.Task(bg_loop)
application = Application(
router.apply_routes([
(r"/", IndexHandler),
(r"/start", StartHandler),
(r"/start_mission", StartMissionHandler),
]),
debug=True,
)
application.listen(8000)
StartHandler.triggered.connect(bg_task.setup_api)
StartMissionHandler.triggered.connect(bg_task.start_mission)
event.api_started.connect(RouterConnection.on_api_started)
event.mission_started.connect(RouterConnection.on_mission_started)
event.mission_result.connect(RouterConnection.on_mission_result)
try:
bg_thread.start()
IOLoop.instance().start()
except KeyboardInterrupt:
bg_loop.stop()
IOLoop.instance().stop()
# api_token = args[1]
# client_ = client.Client(api_token)
# event_loop = tornado.ioloop.IOLoop()
# event_loop = task.EventLoop()
# mission = client.Mission(client=client_, event_loop=event_loop)
# mission.start(api_deck_id=2, api_mission_id=5)
# mission.start(api_deck_id=3, api_mission_id=21)
# mission.start(api_deck_id=4, api_mission_id=38)
# nyukyo = client.Nyukyo(client=client_, event_loop=event_loop)
# nyukyo.start()
# event_loop.start()
return 0
def test_func():
IOLoop.clear_instance()
loop = IOLoop()
loop.make_current()
s, workers = loop.run_sync(lambda: start_cluster(ncores))
try:
loop.run_sync(lambda: cor(s, *workers), timeout=timeout)
finally:
loop.run_sync(lambda: end_cluster(s, workers))
loop.stop()
loop.close()
def test_func():
IOLoop.clear_instance()
loop = IOLoop()
loop.make_current()
s, workers = loop.run_sync(lambda: start_cluster(ncores))
try:
loop.run_sync(lambda: cor(s, *workers), timeout=timeout)
finally:
loop.run_sync(lambda: end_cluster(s, workers))
loop.stop()
loop.close()
class _Executor(Thread):
"""
Tasks executor. Task is executed in ioloop for easier stopping it. Subprocess based tasks are killed external
"""
def __init__(self, task, number, *args, **kwargs):
super(_Executor, self).__init__(*args, **kwargs)
self.ioloop = None
self.task = task
self.number = number
def run(self):
if self.task.cancelled:
return
self.task.set_executor(executor=self)
self.ioloop = IOLoop()
self.ioloop.make_current()
self.ioloop.add_callback(self.execute)
self.ioloop.start()
self.task.clear()
self.ioloop.clear_current()
async def execute(self):
"""
Update task and stop ioloop
"""
try:
await self.task()
await self.task.post_run()
except subprocess.CalledProcessError as exception:
log.warning('%s', exception)
except:
log.exception("Exception while executing task on worker %s",
self.number)
finally:
self.ioloop.stop()
self.task.finish_time = int(time.time())
def stop(self):
"""
Stop task. Important especially for Subprocess based tasks
"""
self.task.kill()
# As Subprocess based tasks generate traffic only using external tool, they should exit gracefully
if not isinstance(self.task,
(CommandTask, NmapPortInfoTask, PortScanTask)):
self.ioloop.stop()
def __str__(self):
return str(self.task)
class FakeServerContext(object):
def __init__(self, monkeypatch, fail_these, expected_basename):
self._monkeypatch = monkeypatch
self._fail_these = fail_these
self._expected_basename = expected_basename
self._url = None
self._loop = None
self._started = threading.Condition()
self._thread = threading.Thread(target=self._run)
def __exit__(self, type, value, traceback):
if self._loop is not None:
# we can ONLY use add_callback here, since the loop is
# running in a different thread.
self._loop.add_callback(self._stop)
self._thread.join()
def __enter__(self):
self._started.acquire()
self._thread.start()
self._started.wait()
self._started.release()
_monkeypatch_client_config(self._monkeypatch, self._url)
return self._url
def _run(self):
self._loop = IOLoop()
self._server = FakeAnacondaServer(io_loop=self._loop,
fail_these=self._fail_these,
expected_basename=self._expected_basename)
self._url = self._server.url
def notify_started():
self._started.acquire()
self._started.notify()
self._started.release()
self._loop.add_callback(notify_started)
self._loop.start()
# done
self._server.unlisten()
def _stop(self):
def really_stop():
if self._loop is not None:
self._loop.stop()
self._loop = None
# the delay allows pending next-tick things to go ahead
# and happen, which may avoid some problems with trying to
# output to stdout after pytest closes it
if self._loop is not None:
self._loop.call_later(delay=0.05, callback=really_stop)
def run_worker_fork(q, scheduler_addr, ncores, nanny_port, worker_ip,
worker_port, local_dir, **kwargs):
"""
Create a worker in a forked child.
"""
from distributed import Worker # pragma: no cover
from tornado.ioloop import IOLoop # pragma: no cover
try:
from dask.multiprocessing import initialize_worker_process
except ImportError: # old Dask version
pass
else:
initialize_worker_process()
IOLoop.clear_instance() # pragma: no cover
loop = IOLoop() # pragma: no cover
loop.make_current() # pragma: no cover
worker = Worker(scheduler_addr,
ncores=ncores,
service_ports={'nanny': nanny_port},
local_dir=local_dir,
**kwargs) # pragma: no cover
@gen.coroutine # pragma: no cover
def run():
try: # pragma: no cover
yield worker._start(worker_port) # pragma: no cover
except Exception as e: # pragma: no cover
logger.exception(e) # pragma: no cover
q.put(e) # pragma: no cover
else:
assert worker.port # pragma: no cover
q.put({
'address': worker.address,
'dir': worker.local_dir
}) # pragma: no cover
yield worker.wait_until_closed()
logger.info("Worker closed")
try:
loop.run_sync(run)
except TimeoutError:
logger.info("Worker timed out")
except KeyboardInterrupt:
pass
finally:
loop.stop()
loop.close(all_fds=True)
class HTTPClient(object):
"""A blocking HTTP client.
This interface is provided for convenience and testing; most applications
that are running an IOLoop will want to use `AsyncHTTPClient` instead.
Typical usage looks like this::
http_client = httpclient.HTTPClient()
try:
response = http_client.fetch("http://www.google.com/")
print response.body
except httpclient.HTTPError, e:
print "Error:", e
"""
def __init__(self, async_client_class=None, **kwargs):
self._io_loop = IOLoop()
if async_client_class is None:
async_client_class = AsyncHTTPClient
self._async_client = async_client_class(self._io_loop, **kwargs)
self._response = None
self._closed = False
def __del__(self):
self.close()
def close(self):
"""Closes the HTTPClient, freeing any resources used."""
if not self._closed:
self._async_client.close()
self._io_loop.close()
self._closed = True
def fetch(self, request, **kwargs):
"""Executes a request, returning an `HTTPResponse`.
The request may be either a string URL or an `HTTPRequest` object.
If it is a string, we construct an `HTTPRequest` using any additional
kwargs: ``HTTPRequest(request, **kwargs)``
If an error occurs during the fetch, we raise an `HTTPError`.
"""
def callback(response):
self._response = response
self._io_loop.stop()
self._async_client.fetch(request, callback, **kwargs)
self._io_loop.start()
response = self._response
self._response = None
response.rethrow()
return response
class TestIOLoopCurrent(unittest.TestCase):
def setUp(self):
self.io_loop = IOLoop()
def tearDown(self):
self.io_loop.close()
def test_current(self):
def f():
self.current_io_loop = IOLoop.current()
self.io_loop.stop()
self.io_loop.add_callback(f)
self.io_loop.start()
self.assertIs(self.current_io_loop, self.io_loop)
class TestIOLoopCurrent(unittest.TestCase):
def setUp(self):
self.io_loop = IOLoop()
def tearDown(self):
self.io_loop.close()
def test_current(self):
def f():
self.current_io_loop = IOLoop.current()
self.io_loop.stop()
self.io_loop.add_callback(f)
self.io_loop.start()
self.assertIs(self.current_io_loop, self.io_loop)
class IOLoop(object):
NONE = TornadoIOLoop.NONE
READ = TornadoIOLoop.READ
WRITE = TornadoIOLoop.WRITE
ERROR = TornadoIOLoop.ERROR
def __init__(self):
self._tornado_io_loop = TornadoIOLoop()
def inner(self):
return self._tornado_io_loop
def close(self, all_fds=False):
self._tornado_io_loop.close(all_fds)
def add_handler(self, fd, handler, events):
self._tornado_io_loop.add_handler(fd, handler, events)
def update_handler(self, fd, events):
self._tornado_io_loop.update_handler(fd, events)
def remove_handler(self, fd):
self._tornado_io_loop.remove_handler(fd)
def start(self):
self._tornado_io_loop.start()
def stop(self):
self._tornado_io_loop.stop()
def time(self):
return self._tornado_io_loop.time()
def add_timeout(self, deadline, callback):
return self._tornado_io_loop.add_timeout(deadline, callback)
def remove_timeout(self, timeout):
self._tornado_io_loop.remove_timeout(timeout)
def add_callback(self, callback, *args, **kwargs):
self._tornado_io_loop.add_callback(callback, *args, **kwargs)
def run(self):
try:
self.start()
except KeyboardInterrupt:
print ""
print "Ctrl-C recieved. Exiting."
class Client(object):
def __init__(self, host, port, timeout=None, connect_timeout=-1, unix_socket=None, max_buffer_size=104857600):
def connect():
self._io_loop = IOLoop()
self._async_client = AsyncClient(host, port, timeout, connect_timeout, unix_socket, self._io_loop, max_buffer_size)
self._response = None
self._connect = connect
self._response = None
self._connect()
@property
def closed(self):
return self._async_client is None
def close(self):
if self.closed:
return
if self._async_client:
self._async_client.close()
self._io_loop.close()
self._async_client = None
self._io_loop = None
self._response = None
def call(self, request):
if self.closed:
self._connect()
def callback(response):
self._response = response
self._io_loop.stop()
self._async_client.call(request, callback)
self._io_loop.start()
response = self._response
self._response = None
try:
response.rethrow()
except:
self.close()
raise
return response
def __del__(self):
self.close()
def __str__(self):
return 'STPClient to %s' % str(self._async_client.connection)
class RedisClient(object):
"""A blocking Redis client.
This interface is provided for convenience and testing; most applications
that are running an IOLoop will want to use `AsyncRedisClient` instead.
Typical usage looks like this::
redis_client = redisclient.RedisClient(('127.0.0.1', 6379))
try:
result = redis_client.fetch(('set', 'foo', 'bar'))
print result
except redisclient.RedisError, e:
print "Error:", e
"""
def __init__(self, address):
self.address = address
self._io_loop = IOLoop()
self._async_client = AsyncRedisClient(self.address, self._io_loop)
self._result = None
self._closed = False
def __del__(self):
self.close()
def close(self):
"""Closes the RedisClient, freeing any resources used."""
if not self._closed:
self._async_client.close()
self._io_loop.close()
self._closed = True
def fetch(self, request, **kwargs):
"""Executes a request, returning an `result`.
The request may be a tuple object. like ('set','foo','bar')
If an error occurs during the fetch, we raise an `RedisError`.
"""
def callback(result):
self._result = result
self._io_loop.stop()
self._async_client.fetch(request, callback, **kwargs)
self._io_loop.start()
result = self._result
self._result = None
return result
class SyncHTTPClientTest(unittest.TestCase):
def setUp(self):
if IOLoop.configured_class().__name__ in ('TwistedIOLoop',
'AsyncIOMainLoop'):
# TwistedIOLoop only supports the global reactor, so we can't have
# separate IOLoops for client and server threads.
# AsyncIOMainLoop doesn't work with the default policy
# (although it could with some tweaks to this test and a
# policy that created loops for non-main threads).
raise unittest.SkipTest(
'Sync HTTPClient not compatible with TwistedIOLoop or '
'AsyncIOMainLoop')
self.server_ioloop = IOLoop()
sock, self.port = bind_unused_port()
app = Application([('/', HelloWorldHandler)])
self.server = HTTPServer(app, io_loop=self.server_ioloop)
self.server.add_socket(sock)
self.server_thread = threading.Thread(target=self.server_ioloop.start)
self.server_thread.start()
self.http_client = HTTPClient()
def tearDown(self):
def stop_server():
self.server.stop()
self.server_ioloop.stop()
self.server_ioloop.add_callback(stop_server)
self.server_thread.join()
self.http_client.close()
self.server_ioloop.close(all_fds=True)
def get_url(self, path):
return 'http://localhost:%d%s' % (self.port, path)
def test_sync_client(self):
response = self.http_client.fetch(self.get_url('/'))
self.assertEqual(b'Hello world!', response.body)
def test_sync_client_error(self):
# Synchronous HTTPClient raises errors directly; no need for
# response.rethrow()
with self.assertRaises(HTTPError) as assertion:
self.http_client.fetch(self.get_url('/notfound'))
self.assertEqual(assertion.exception.code, 404)
class RedisClient(object):
"""A blocking Redis client.
This interface is provided for convenience and testing; most applications
that are running an IOLoop will want to use `AsyncRedisClient` instead.
Typical usage looks like this::
redis_client = redisclient.RedisClient(('127.0.0.1', 6379))
try:
result = redis_client.fetch(('set', 'foo', 'bar'))
print result
except redisclient.RedisError, e:
print "Error:", e
"""
def __init__(self, address):
self.address = address
self._io_loop = IOLoop()
self._async_client = AsyncRedisClient(self.address, self._io_loop)
self._result = None
self._closed = False
def __del__(self):
self.close()
def close(self):
"""Closes the RedisClient, freeing any resources used."""
if not self._closed:
self._async_client.close()
self._io_loop.close()
self._closed = True
def fetch(self, request, **kwargs):
"""Executes a request, returning an `result`.
The request may be a tuple object. like ('set','foo','bar')
If an error occurs during the fetch, we raise an `RedisError`.
"""
def callback(result):
self._result = result
self._io_loop.stop()
self._async_client.fetch(request, callback, **kwargs)
self._io_loop.start()
result = self._result
self._result = None
return result
class TestIOLoopCurrent(unittest.TestCase):
def setUp(self):
setup_with_context_manager(self, ignore_deprecation())
self.io_loop = None # type: typing.Optional[IOLoop]
IOLoop.clear_current()
def tearDown(self):
if self.io_loop is not None:
self.io_loop.close()
def test_default_current(self):
self.io_loop = IOLoop()
# The first IOLoop with default arguments is made current.
self.assertIs(self.io_loop, IOLoop.current())
# A second IOLoop can be created but is not made current.
io_loop2 = IOLoop()
self.assertIs(self.io_loop, IOLoop.current())
io_loop2.close()
def test_non_current(self):
self.io_loop = IOLoop(make_current=False)
# The new IOLoop is not initially made current.
self.assertIsNone(IOLoop.current(instance=False))
# Starting the IOLoop makes it current, and stopping the loop
# makes it non-current. This process is repeatable.
for i in range(3):
def f():
self.current_io_loop = IOLoop.current()
assert self.io_loop is not None
self.io_loop.stop()
self.io_loop.add_callback(f)
self.io_loop.start()
self.assertIs(self.current_io_loop, self.io_loop)
# Now that the loop is stopped, it is no longer current.
self.assertIsNone(IOLoop.current(instance=False))
def test_force_current(self):
self.io_loop = IOLoop(make_current=True)
self.assertIs(self.io_loop, IOLoop.current())
with self.assertRaises(RuntimeError):
# A second make_current=True construction cannot succeed.
IOLoop(make_current=True)
# current() was not affected by the failed construction.
self.assertIs(self.io_loop, IOLoop.current())
class StackContextTest(unittest.TestCase):
# Note that this test logs an error even when it passes.
# TODO(bdarnell): better logging setup for unittests
def test_stack_context(self):
self.io_loop = IOLoop()
app = Application([('/', TestRequestHandler, dict(io_loop=self.io_loop))])
server = HTTPServer(app, io_loop=self.io_loop)
server.listen(11000)
client = AsyncHTTPClient(io_loop=self.io_loop)
client.fetch('http://localhost:11000/', self.handle_response)
self.io_loop.start()
self.assertEquals(self.response.code, 500)
self.assertTrue('got expected exception' in self.response.body)
def handle_response(self, response):
self.response = response
self.io_loop.stop()
class SyncHTTPClientTest(unittest.TestCase):
def setUp(self):
if IOLoop.configured_class().__name__ in ('TwistedIOLoop',
'AsyncIOMainLoop'):
# TwistedIOLoop only supports the global reactor, so we can't have
# separate IOLoops for client and server threads.
# AsyncIOMainLoop doesn't work with the default policy
# (although it could with some tweaks to this test and a
# policy that created loops for non-main threads).
raise unittest.SkipTest(
'Sync HTTPClient not compatible with TwistedIOLoop or '
'AsyncIOMainLoop')
self.server_ioloop = IOLoop()
sock, self.port = bind_unused_port()
app = Application([('/', HelloWorldHandler)])
self.server = HTTPServer(app, io_loop=self.server_ioloop)
self.server.add_socket(sock)
self.server_thread = threading.Thread(target=self.server_ioloop.start)
self.server_thread.start()
self.http_client = HTTPClient()
def tearDown(self):
def stop_server():
self.server.stop()
self.server_ioloop.stop()
self.server_ioloop.add_callback(stop_server)
self.server_thread.join()
self.http_client.close()
self.server_ioloop.close(all_fds=True)
def get_url(self, path):
return 'http://localhost:%d%s' % (self.port, path)
def test_sync_client(self):
response = self.http_client.fetch(self.get_url('/'))
self.assertEqual(b'Hello world!', response.body)
def test_sync_client_error(self):
# Synchronous HTTPClient raises errors directly; no need for
# response.rethrow()
with self.assertRaises(HTTPError) as assertion:
self.http_client.fetch(self.get_url('/notfound'))
self.assertEqual(assertion.exception.code, 404)
class ProxyTestServer(object):
def __init__(self,):
self.server_ioloop = IOLoop()
self.access_count = 0
@tornado.gen.coroutine
def init_server():
sock, self.port = bind_unused_port()
app = Application([('/(.*)', ProxyTestHandler, dict(server=self))])
self.server = HTTPServer(app)
self.server.add_socket(sock)
self.server_ioloop.run_sync(init_server)
self.server_thread = threading.Thread(target=self.server_ioloop.start)
def start(self):
self.server_thread.start()
def stop(self):
def stop_server():
self.server.stop()
@tornado.gen.coroutine
def slow_stop():
for i in range(5):
yield
self.server_ioloop.stop()
self.server_ioloop.add_callback(slow_stop)
self.server_ioloop.add_callback(stop_server)
self.server_thread.join()
self.server_ioloop.close(all_fds=True)
def get_access_count(self):
return self.access_count
def clear_access_count(self):
self.access_count = 0
class TestIOLoopCurrent(unittest.TestCase):
def setUp(self):
self.io_loop = None
IOLoop.clear_current()
def tearDown(self):
if self.io_loop is not None:
self.io_loop.close()
def test_default_current(self):
self.io_loop = IOLoop()
# The first IOLoop with default arguments is made current.
self.assertIs(self.io_loop, IOLoop.current())
# A second IOLoop can be created but is not made current.
io_loop2 = IOLoop()
self.assertIs(self.io_loop, IOLoop.current())
io_loop2.close()
def test_non_current(self):
self.io_loop = IOLoop(make_current=False)
# The new IOLoop is not initially made current.
self.assertIsNone(IOLoop.current(instance=False))
# Starting the IOLoop makes it current, and stopping the loop
# makes it non-current. This process is repeatable.
for i in range(3):
def f():
self.current_io_loop = IOLoop.current()
self.io_loop.stop()
self.io_loop.add_callback(f)
self.io_loop.start()
self.assertIs(self.current_io_loop, self.io_loop)
# Now that the loop is stopped, it is no longer current.
self.assertIsNone(IOLoop.current(instance=False))
def test_force_current(self):
self.io_loop = IOLoop(make_current=True)
self.assertIs(self.io_loop, IOLoop.current())
with self.assertRaises(RuntimeError):
# A second make_current=True construction cannot succeed.
IOLoop(make_current=True)
# current() was not affected by the failed construction.
self.assertIs(self.io_loop, IOLoop.current())
class ProxyTestServer(object):
def __init__(self, ):
self.server_ioloop = IOLoop()
self.access_count = 0
@tornado.gen.coroutine
def init_server():
sock, self.port = bind_unused_port()
app = Application([('/(.*)', ProxyTestHandler, dict(server=self))])
self.server = HTTPServer(app)
self.server.add_socket(sock)
self.server_ioloop.run_sync(init_server)
self.server_thread = threading.Thread(target=self.server_ioloop.start)
def start(self):
self.server_thread.start()
def stop(self):
def stop_server():
self.server.stop()
@tornado.gen.coroutine
def slow_stop():
for i in range(5):
yield
self.server_ioloop.stop()
self.server_ioloop.add_callback(slow_stop)
self.server_ioloop.add_callback(stop_server)
self.server_thread.join()
self.server_ioloop.close(all_fds=True)
def get_access_count(self):
return self.access_count
def clear_access_count(self):
self.access_count = 0
class TestIOLoopCurrent(unittest.TestCase):
def setUp(self):
self.io_loop = None
IOLoop.clear_current()
def tearDown(self):
if self.io_loop is not None:
self.io_loop.close()
def test_default_current(self):
self.io_loop = IOLoop()
# The first IOLoop with default arguments is made current.
self.assertIs(self.io_loop, IOLoop.current())
# A second IOLoop can be created but is not made current.
io_loop2 = IOLoop()
self.assertIs(self.io_loop, IOLoop.current())
io_loop2.close()
def test_non_current(self):
self.io_loop = IOLoop(make_current=False)
# The new IOLoop is not initially made current.
self.assertIsNone(IOLoop.current(instance=False))
def f():
# But it is current after it is started.
self.current_io_loop = IOLoop.current()
self.io_loop.stop()
self.io_loop.add_callback(f)
self.io_loop.start()
self.assertIs(self.current_io_loop, self.io_loop)
# Now that the loop is stopped, it is no longer current.
self.assertIsNone(IOLoop.current(instance=False))
def test_force_current(self):
self.io_loop = IOLoop(make_current=True)
self.assertIs(self.io_loop, IOLoop.current())
with self.assertRaises(RuntimeError):
# A second make_current=True construction cannot succeed.
IOLoop(make_current=True)
# current() was not affected by the failed construction.
self.assertIs(self.io_loop, IOLoop.current())
class Stream(threading.Thread):
def __init__(self, stream_ip, stream_name):
super().__init__()
# set ip and name
self._ip = stream_ip
self._name = stream_name
# set logger
self._logger = Logger('stream')
# event loop
self._loop = IOLoop()
def run(self):
# queue main action
self._loop.add_callback(self._once)
# start event loop in this thread
self._loop.start()
# free all resources
self._loop.close()
@coroutine
def stop(self):
ok = yield Task(
lambda callback: self._loop.add_callback(self._stop_, callback))
return ok
def _once(self):
# read and then write rtmp stream (for hls feature)
cmd = 'ffmpeg -i rtmp://{0}/src/{1} -c:a aac -b:a 64k -c:v libx264 -b:v 256k -vf scale=-1:480 -f flv rtmp://{0}/hls/{1}_480p -c:a aac -b:a 128k -c:v libx264 -b:v 512K -vf scale=-1:720 -f flv rtmp://{0}/hls/{1}_720p'.format(
self._ip, self._name)
self._child = Subprocess(shlex.split(cmd),
stdout=Subprocess.STREAM,
stderr=Subprocess.STREAM)
# stop loop
self._loop.stop()
def _stop(self):
try:
pid = self._child.proc.pid
self._logger.debug('_stop')
os.kill(pid, signal.SIGTERM)
except ProcessLookupError as e:
self._logger.exception('error on killing pid {0}'.format(pid))
def run_worker_fork(q, scheduler_addr, ncores, nanny_port, worker_ip,
worker_port, local_dir, **kwargs):
"""
Create a worker by forking.
"""
from distributed import Worker # pragma: no cover
from tornado.ioloop import IOLoop # pragma: no cover
IOLoop.clear_instance() # pragma: no cover
loop = IOLoop() # pragma: no cover
loop.make_current() # pragma: no cover
worker = Worker(scheduler_addr,
ncores=ncores,
service_ports={'nanny': nanny_port},
local_dir=local_dir,
**kwargs) # pragma: no cover
@gen.coroutine # pragma: no cover
def run():
try: # pragma: no cover
yield worker._start((worker_ip, worker_port)) # pragma: no cover
except Exception as e: # pragma: no cover
logger.exception(e) # pragma: no cover
q.put(e) # pragma: no cover
else:
assert worker.port # pragma: no cover
q.put({
'address': worker.address,
'dir': worker.local_dir
}) # pragma: no cover
yield worker.wait_until_closed()
logger.info("Worker closed")
try:
loop.run_sync(run)
except TimeoutError:
logger.info("Worker timed out")
finally:
loop.stop()
loop.close(all_fds=True)
def test_func():
IOLoop.clear_instance()
loop = IOLoop()
loop.make_current()
s, workers = loop.run_sync(
lambda: start_cluster(ncores, loop, Worker=Worker))
args = [s] + workers
if client:
e = Client((s.ip, s.port), loop=loop, start=False)
loop.run_sync(e._start)
args = [e] + args
try:
loop.run_sync(lambda: cor(*args), timeout=timeout)
finally:
if client:
loop.run_sync(e._shutdown)
loop.run_sync(lambda: end_cluster(s, workers))
loop.stop()
loop.close(all_fds=True)
def test_func():
IOLoop.clear_instance()
loop = IOLoop()
loop.make_current()
s, workers = loop.run_sync(lambda: start_cluster(ncores,
Worker=Worker))
args = [s] + workers
if executor:
e = Executor((s.ip, s.port), loop=loop, start=False)
loop.run_sync(e._start)
args = [e] + args
try:
loop.run_sync(lambda: cor(*args), timeout=timeout)
finally:
if executor:
loop.run_sync(e._shutdown)
loop.run_sync(lambda: end_cluster(s, workers))
loop.stop()
loop.close(all_fds=True)
class LoadTest:
def __init__(self, ws_url, concurrency, recycle, duration):
self.ws_url = ws_url
self.concurrency = concurrency
self.recycle = recycle
self.duration = duration
self.loop = IOLoop()
self.result = None
def print_summary(self):
print(f"Total of websockets opened: {self.result.websockets_opened}")
print(f"Total of websockets closed: {self.result.websockets_closed}")
async def websocket_connect(self):
ws = await websocket_connect(self.ws_url)
await gen.moment
self.result.websockets_opened += 1
self.loop.call_later(self.recycle, self.recycle_ws(ws))
def recycle_ws(self, ws):
async def handle():
ws.close()
self.result.websockets_closed += 1
await self.websocket_connect()
return handle
def start(self):
self.loop.make_current()
self.result = Result()
for _ in range(self.concurrency):
self.loop.call_later(0, self.websocket_connect)
self.loop.call_later(self.duration, self.stop)
self.loop.start()
def stop(self):
self.print_summary()
self.loop.stop()
class HTTPClient(object):
"""A blocking HTTP client.
Typical usage looks like this:
http_client = httpclient.HTTPClient()
try:
response = http_client.fetch("http://www.google.com/")
print response.body
except httpclient.HTTPError, e:
print "Error:", e
fetch() can take a string URL or an HTTPRequest instance, which offers
more options, like executing POST/PUT/DELETE requests.
"""
def __init__(self):
self._io_loop = IOLoop()
self._async_client = AsyncHTTPClient(self._io_loop)
self._response = None
def __del__(self):
self._async_client.close()
def fetch(self, request, **kwargs):
"""Executes an HTTPRequest, returning an HTTPResponse.
If an error occurs during the fetch, we raise an HTTPError.
"""
def callback(response):
self._response = response
self._io_loop.stop()
self._async_client.fetch(request, callback, **kwargs)
self._io_loop.start()
response = self._response
self._response = None
response.rethrow()
return response
class LoopAndGroup(object):
def __init__(self, quit_after=None):
self.io_loop = IOLoop()
self.group = _CallbackGroup(self.io_loop)
if quit_after is not None:
self.io_loop.call_later(quit_after / 1000.0, lambda: self.io_loop.stop())
def __exit__(self, type, value, traceback):
run(self.io_loop)
self.io_loop.close()
def __enter__(self):
return self
class HTTPClient(object):
"""A blocking HTTP client.
Typical usage looks like this:
http_client = httpclient.HTTPClient()
try:
response = http_client.fetch("http://www.google.com/")
print response.body
except httpclient.HTTPError, e:
print "Error:", e
fetch() can take a string URL or an HTTPRequest instance, which offers
more options, like executing POST/PUT/DELETE requests.
"""
def __init__(self):
self._io_loop = IOLoop()
self._async_client = AsyncHTTPClient(self._io_loop)
self._response = None
def __del__(self):
self._async_client.close()
def fetch(self, request, **kwargs):
"""Executes an HTTPRequest, returning an HTTPResponse.
If an error occurs during the fetch, we raise an HTTPError.
"""
def callback(response):
self._response = response
self._io_loop.stop()
self._async_client.fetch(request, callback, **kwargs)
self._io_loop.start()
response = self._response
self._response = None
response.rethrow()
return response
class LoopAndGroup(object):
def __init__(self, quit_after=None):
self.io_loop = IOLoop()
self.group = _CallbackGroup(self.io_loop)
if quit_after is not None:
self.io_loop.call_later(quit_after / 1000.0,
lambda: self.io_loop.stop())
def __exit__(self, type, value, traceback):
run(self.io_loop)
self.io_loop.close()
def __enter__(self):
return self
def run_worker_fork(q, ip, scheduler_ip, scheduler_port, ncores, nanny_port,
worker_port, local_dir, services, name, memory_limit, reconnect,
resources, validate):
""" Function run by the Nanny when creating the worker """
from distributed import Worker # pragma: no cover
from tornado.ioloop import IOLoop # pragma: no cover
IOLoop.clear_instance() # pragma: no cover
loop = IOLoop() # pragma: no cover
loop.make_current() # pragma: no cover
worker = Worker(scheduler_ip, scheduler_port, ncores=ncores, ip=ip,
service_ports={'nanny': nanny_port}, local_dir=local_dir,
services=services, name=name, memory_limit=memory_limit,
reconnect=reconnect, validate=validate,
resources=resources, loop=loop) # pragma: no cover
@gen.coroutine # pragma: no cover
def run():
try: # pragma: no cover
yield worker._start(worker_port) # pragma: no cover
except Exception as e: # pragma: no cover
logger.exception(e) # pragma: no cover
q.put(e) # pragma: no cover
else:
assert worker.port # pragma: no cover
q.put({'port': worker.port, 'dir': worker.local_dir}) # pragma: no cover
while worker.status != 'closed':
yield gen.sleep(0.1)
logger.info("Worker closed")
try:
loop.run_sync(run)
finally:
loop.stop()
loop.close(all_fds=True)
def run(self, loop=None):
'''
Start servicing the Tornado event loop.
'''
if not loop:
loop = IOLoop()
loop.make_current()
# bind the socket
self.listen(self._port, self._address)
logger.info('Pensive started on {}:{}'.format(self._address or '*',
self._port))
try:
loop.start()
except KeyboardInterrupt:
pass
loop.stop()
loop.close()
logger.info('Pensive stopped')
class Core(object):
def __init__(self):
self.started = False
self.daemon = None
self.io_loop = IOLoop()
self.pid = os.getpid()
self.showlist = []
self.tz = tz.tzwinlocal() if tz.tzwinlocal else tz.tzlocal()
self.config_file = None
self.data_dir = None
self.cache_dir = None
self.quite = None
self.no_launch = None
self.web_port = None
self.developer = None
self.debug = None
self.newest_version = None
self.newest_version_string = None
self.naming_ep_type = ("%(seasonnumber)dx%(episodenumber)02d",
"s%(seasonnumber)02de%(episodenumber)02d",
"S%(seasonnumber)02dE%(episodenumber)02d",
"%(seasonnumber)02dx%(episodenumber)02d",
"S%(seasonnumber)02d E%(episodenumber)02d")
self.sports_ep_type = ("%(seasonnumber)dx%(episodenumber)02d",
"s%(seasonnumber)02de%(episodenumber)02d",
"S%(seasonnumber)02dE%(episodenumber)02d",
"%(seasonnumber)02dx%(episodenumber)02d",
"S%(seasonnumber)02 dE%(episodenumber)02d")
self.naming_ep_type_text = (
"1x02",
"s01e02",
"S01E02",
"01x02",
"S01 E02",
)
self.naming_multi_ep_type = {
0: ["-%(episodenumber)02d"] * len(self.naming_ep_type),
1: [" - " + x for x in self.naming_ep_type],
2: [x + "%(episodenumber)02d" for x in ("x", "e", "E", "x")]
}
self.naming_multi_ep_type_text = ("extend", "duplicate", "repeat")
self.naming_sep_type = (" - ", " ")
self.naming_sep_type_text = (" - ", "space")
self.user_agent = 'SiCKRAGE.CE.1/({};{};{})'.format(
platform.system(), platform.release(), str(uuid.uuid1()))
self.languages = [
language for language in os.listdir(sickrage.LOCALE_DIR)
if '_' in language
]
self.sys_encoding = get_sys_encoding()
self.client_web_urls = {'torrent': '', 'newznab': ''}
self.adba_connection = None
self.notifier_providers = None
self.metadata_providers = {}
self.search_providers = None
self.log = None
self.config = None
self.alerts = None
self.main_db = None
self.cache_db = None
self.scheduler = None
self.wserver = None
self.google_auth = None
self.name_cache = None
self.show_queue = None
self.search_queue = None
self.postprocessor_queue = None
self.version_updater = None
self.show_updater = None
self.daily_searcher = None
self.backlog_searcher = None
self.proper_searcher = None
self.trakt_searcher = None
self.subtitle_searcher = None
self.auto_postprocessor = None
self.upnp_client = None
self.oidc_client = None
self.quicksearch_cache = None
def start(self):
self.started = True
# thread name
threading.currentThread().setName('CORE')
# patch modules with encoding kludge
patch_modules()
# init core classes
self.notifier_providers = NotifierProviders()
self.metadata_providers = MetadataProviders()
self.search_providers = SearchProviders()
self.log = Logger()
self.config = Config()
self.alerts = Notifications()
self.main_db = MainDB()
self.cache_db = CacheDB()
self.scheduler = TornadoScheduler()
self.wserver = WebServer()
self.name_cache = NameCache()
self.show_queue = ShowQueue()
self.search_queue = SearchQueue()
self.postprocessor_queue = PostProcessorQueue()
self.version_updater = VersionUpdater()
self.show_updater = ShowUpdater()
self.daily_searcher = DailySearcher()
self.failed_snatch_searcher = FailedSnatchSearcher()
self.backlog_searcher = BacklogSearcher()
self.proper_searcher = ProperSearcher()
self.trakt_searcher = TraktSearcher()
self.subtitle_searcher = SubtitleSearcher()
self.auto_postprocessor = AutoPostProcessor()
self.upnp_client = UPNPClient()
self.quicksearch_cache = QuicksearchCache()
# setup oidc client
realm = KeycloakRealm(server_url='https://auth.sickrage.ca',
realm_name='sickrage')
self.oidc_client = realm.open_id_connect(
client_id='sickrage-app',
client_secret='5d4710b2-ca70-4d39-b5a3-0705e2c5e703')
# Check if we need to perform a restore first
if os.path.exists(
os.path.abspath(os.path.join(self.data_dir, 'restore'))):
success = restoreSR(
os.path.abspath(os.path.join(self.data_dir, 'restore')),
self.data_dir)
print("Restoring SiCKRAGE backup: %s!\n" %
("FAILED", "SUCCESSFUL")[success])
if success:
shutil.rmtree(os.path.abspath(
os.path.join(self.data_dir, 'restore')),
ignore_errors=True)
# migrate old database file names to new ones
if os.path.isfile(
os.path.abspath(os.path.join(self.data_dir, 'sickbeard.db'))):
if os.path.isfile(os.path.join(self.data_dir, 'sickrage.db')):
helpers.moveFile(
os.path.join(self.data_dir, 'sickrage.db'),
os.path.join(
self.data_dir, '{}.bak-{}'.format(
'sickrage.db',
datetime.datetime.now().strftime(
'%Y%m%d_%H%M%S'))))
helpers.moveFile(
os.path.abspath(os.path.join(self.data_dir, 'sickbeard.db')),
os.path.abspath(os.path.join(self.data_dir, 'sickrage.db')))
# load config
self.config.load()
# set language
self.config.change_gui_lang(self.config.gui_lang)
# set socket timeout
socket.setdefaulttimeout(self.config.socket_timeout)
# setup logger settings
self.log.logSize = self.config.log_size
self.log.logNr = self.config.log_nr
self.log.logFile = os.path.join(self.data_dir, 'logs', 'sickrage.log')
self.log.debugLogging = self.config.debug
self.log.consoleLogging = not self.quite
# start logger
self.log.start()
# user agent
if self.config.random_user_agent:
self.user_agent = UserAgent().random
urlparse.uses_netloc.append('scgi')
urllib.FancyURLopener.version = self.user_agent
# set torrent client web url
torrent_webui_url(True)
# Check available space
try:
total_space, available_space = getFreeSpace(self.data_dir)
if available_space < 100:
self.log.error(
'Shutting down as SiCKRAGE needs some space to work. You\'ll get corrupted data '
'otherwise. Only %sMB left', available_space)
return
except Exception:
self.log.error('Failed getting disk space: %s',
traceback.format_exc())
# perform database startup actions
for db in [self.main_db, self.cache_db]:
# initialize database
db.initialize()
# check integrity of database
db.check_integrity()
# migrate database
db.migrate()
# misc database cleanups
db.cleanup()
# upgrade database
db.upgrade()
# compact main database
if self.config.last_db_compact < time.time() - 604800: # 7 days
self.main_db.compact()
self.config.last_db_compact = int(time.time())
# load name cache
self.name_cache.load()
# load data for shows from database
self.load_shows()
if self.config.default_page not in ('schedule', 'history', 'IRC'):
self.config.default_page = 'home'
# cleanup cache folder
for folder in ['mako', 'sessions', 'indexers']:
try:
shutil.rmtree(os.path.join(sickrage.app.cache_dir, folder),
ignore_errors=True)
except Exception:
continue
# init anidb connection
if self.config.use_anidb:
def anidb_logger(msg):
return self.log.debug("AniDB: {} ".format(msg))
try:
self.adba_connection = adba.Connection(keepAlive=True,
log=anidb_logger)
self.adba_connection.auth(self.config.anidb_username,
self.config.anidb_password)
except Exception as e:
self.log.warning("AniDB exception msg: %r " % repr(e))
if self.config.web_port < 21 or self.config.web_port > 65535:
self.config.web_port = 8081
if not self.config.web_cookie_secret:
self.config.web_cookie_secret = generate_secret()
# attempt to help prevent users from breaking links by using a bad url
if not self.config.anon_redirect.endswith('?'):
self.config.anon_redirect = ''
if not re.match(r'\d+\|[^|]+(?:\|[^|]+)*', self.config.root_dirs):
self.config.root_dirs = ''
self.config.naming_force_folders = check_force_season_folders()
if self.config.nzb_method not in ('blackhole', 'sabnzbd', 'nzbget'):
self.config.nzb_method = 'blackhole'
if self.config.torrent_method not in ('blackhole', 'utorrent',
'transmission', 'deluge',
'deluged', 'download_station',
'rtorrent', 'qbittorrent',
'mlnet', 'putio'):
self.config.torrent_method = 'blackhole'
if self.config.autopostprocessor_freq < self.config.min_autopostprocessor_freq:
self.config.autopostprocessor_freq = self.config.min_autopostprocessor_freq
if self.config.daily_searcher_freq < self.config.min_daily_searcher_freq:
self.config.daily_searcher_freq = self.config.min_daily_searcher_freq
self.config.min_backlog_searcher_freq = get_backlog_cycle_time()
if self.config.backlog_searcher_freq < self.config.min_backlog_searcher_freq:
self.config.backlog_searcher_freq = self.config.min_backlog_searcher_freq
if self.config.version_updater_freq < self.config.min_version_updater_freq:
self.config.version_updater_freq = self.config.min_version_updater_freq
if self.config.subtitle_searcher_freq < self.config.min_subtitle_searcher_freq:
self.config.subtitle_searcher_freq = self.config.min_subtitle_searcher_freq
if self.config.failed_snatch_age < self.config.min_failed_snatch_age:
self.config.failed_snatch_age = self.config.min_failed_snatch_age
if self.config.proper_searcher_interval not in ('15m', '45m', '90m',
'4h', 'daily'):
self.config.proper_searcher_interval = 'daily'
if self.config.showupdate_hour < 0 or self.config.showupdate_hour > 23:
self.config.showupdate_hour = 0
if self.config.subtitles_languages[0] == '':
self.config.subtitles_languages = []
# add version checker job
self.scheduler.add_job(
self.version_updater.run,
IntervalTrigger(hours=self.config.version_updater_freq),
name=self.version_updater.name,
id=self.version_updater.name)
# add network timezones updater job
self.scheduler.add_job(update_network_dict,
IntervalTrigger(days=1),
name="TZUPDATER",
id="TZUPDATER")
# add show updater job
self.scheduler.add_job(self.show_updater.run,
IntervalTrigger(
days=1,
start_date=datetime.datetime.now().replace(
hour=self.config.showupdate_hour)),
name=self.show_updater.name,
id=self.show_updater.name)
# add daily search job
self.scheduler.add_job(
self.daily_searcher.run,
IntervalTrigger(minutes=self.config.daily_searcher_freq,
start_date=datetime.datetime.now() +
datetime.timedelta(minutes=4)),
name=self.daily_searcher.name,
id=self.daily_searcher.name)
# add failed snatch search job
self.scheduler.add_job(
self.failed_snatch_searcher.run,
IntervalTrigger(hours=1,
start_date=datetime.datetime.now() +
datetime.timedelta(minutes=4)),
name=self.failed_snatch_searcher.name,
id=self.failed_snatch_searcher.name)
# add backlog search job
self.scheduler.add_job(
self.backlog_searcher.run,
IntervalTrigger(minutes=self.config.backlog_searcher_freq,
start_date=datetime.datetime.now() +
datetime.timedelta(minutes=30)),
name=self.backlog_searcher.name,
id=self.backlog_searcher.name)
# add auto-postprocessing job
self.scheduler.add_job(
self.auto_postprocessor.run,
IntervalTrigger(minutes=self.config.autopostprocessor_freq),
name=self.auto_postprocessor.name,
id=self.auto_postprocessor.name)
# add find proper job
self.scheduler.add_job(
self.proper_searcher.run,
IntervalTrigger(minutes={
'15m': 15,
'45m': 45,
'90m': 90,
'4h': 4 * 60,
'daily': 24 * 60
}[self.config.proper_searcher_interval]),
name=self.proper_searcher.name,
id=self.proper_searcher.name)
# add trakt.tv checker job
self.scheduler.add_job(self.trakt_searcher.run,
IntervalTrigger(hours=1),
name=self.trakt_searcher.name,
id=self.trakt_searcher.name)
# add subtitles finder job
self.scheduler.add_job(
self.subtitle_searcher.run,
IntervalTrigger(hours=self.config.subtitle_searcher_freq),
name=self.subtitle_searcher.name,
id=self.subtitle_searcher.name)
# add upnp client job
self.scheduler.add_job(
self.upnp_client.run,
IntervalTrigger(seconds=self.upnp_client._nat_portmap_lifetime),
name=self.upnp_client.name,
id=self.upnp_client.name)
# start scheduler service
self.scheduler.start()
# start queue's
self.search_queue.start()
self.show_queue.start()
self.postprocessor_queue.start()
# start webserver
self.wserver.start()
# start ioloop
self.io_loop.start()
def shutdown(self, restart=False):
if self.started:
self.log.info('SiCKRAGE IS SHUTTING DOWN!!!')
# shutdown webserver
if self.wserver:
self.wserver.shutdown()
# shutdown show queue
if self.show_queue:
self.log.debug("Shutting down show queue")
self.show_queue.shutdown()
del self.show_queue
# shutdown search queue
if self.search_queue:
self.log.debug("Shutting down search queue")
self.search_queue.shutdown()
del self.search_queue
# shutdown post-processor queue
if self.postprocessor_queue:
self.log.debug("Shutting down post-processor queue")
self.postprocessor_queue.shutdown()
del self.postprocessor_queue
# log out of ADBA
if self.adba_connection:
self.log.debug("Shutting down ANIDB connection")
self.adba_connection.stop()
# save all show and config settings
self.save_all()
# close databases
for db in [self.main_db, self.cache_db]:
if db.opened:
self.log.debug(
"Shutting down {} database connection".format(db.name))
db.close()
# shutdown logging
if self.log:
self.log.close()
if restart:
os.execl(sys.executable, sys.executable, *sys.argv)
if sickrage.app.daemon:
sickrage.app.daemon.stop()
self.started = False
self.io_loop.stop()
def save_all(self):
# write all shows
self.log.info("Saving all shows to the database")
for show in self.showlist:
try:
show.saveToDB()
except Exception:
continue
# save config
self.config.save()
def load_shows(self):
"""
Populates the showlist and quicksearch cache with shows and episodes from the database
"""
self.quicksearch_cache.load()
for dbData in self.main_db.all('tv_shows'):
try:
self.log.debug("Loading data for show: [{}]".format(
dbData['show_name']))
self.showlist.append(
TVShow(int(dbData['indexer']), int(dbData['indexer_id'])))
self.quicksearch_cache.add_show(dbData['indexer_id'])
except Exception as e:
self.log.debug("Show error in [%s]: %s" %
(dbData['location'], str(e)))
class TestTwitterWatcher(TestCase, TimeoutMixin):
def setUp(self):
TweetStream.fetch = self._catch_tweetstream_fetch
self._tweetstream_fetch_called = False
self._ioloop = IOLoop()
def tearDown(self):
TweetStream.fetch = _ORIGINAL_TWEETSTREAM_FETCH
def _catch_tweetstream_fetch(self, *args, **kwargs):
"""Just assert that this is called."""
self._tweetstream_fetch_called = True
self._ioloop.stop()
def test_twitter_watcher(self):
def callback(result):
self.fail("Callback should not have been called.")
config = {
"twitter_consumer_key": "key",
"twitter_consumer_secret": "secret",
"twitter_access_token": "token",
"twitter_access_token_secret": "token_secret",
"twitter_search_term": "searching",
"parsers": {},
"callbacks": []
}
for key in config.keys():
bad_config = config.copy()
del bad_config[key]
with self.assertRaises(MissingTwitterConfiguration):
TwitterWatcher(self._ioloop, callback, bad_config)
self.add_timeout(3)
twitter_watcher = TwitterWatcher(self._ioloop, callback, config)
twitter_watcher.start()
self._ioloop.start()
# blocks until the fetch callback or timeout
self.assertTrue(self._tweetstream_fetch_called)
def test_twitter_watcher_parse(self):
class FakeParser(object):
def __init__(self, *args, **kwargs):
pass
def parse(self, message, callback):
callback({
"message": message
})
result = {}
def callback(parsed_result):
"""When all is said and parsed, this gets called."""
for key in parsed_result:
result[key] = parsed_result[key]
self.stop()
parsers = {
"www.google.com": FakeParser()
}
config = {
"twitter_consumer_key": "key",
"twitter_consumer_secret": "secret",
"twitter_access_token": "token",
"twitter_access_token_secret": "token_secret",
"twitter_search_term": "ireallywanttoknow",
"parsers": parsers,
"callbacks": []
}
# should do a full redirect on this
original_url = "http://bit.ly/julythongoogle"
final_url = "http://www.google.com/search?q=julython"
twitter_watcher = TwitterWatcher(self._ioloop, callback, config)
twitter_watcher.extract_from_tweet({
"text": "This is an awesome tweet! %s" % original_url,
"user": {
"name": "Josh Marshall",
"screen_name": "joshmarshall",
"profile_image_url": "http://stuff.com/profile.jpg"
}
})
self.add_timeout(5)
self._ioloop.start()
self.assertEqual(final_url, result["message"])
def test_twitter_callbacks(self):
result = {}
def fake_callback(tweet):
result["tweet"] = tweet
config = {
"twitter_consumer_key": "key",
"twitter_consumer_secret": "secret",
"twitter_access_token": "token",
"twitter_access_token_secret": "token_secret",
"twitter_search_term": "whateverworks",
"parsers": {},
"callbacks": [fake_callback]
}
tweet = {
"text": "This is a great tweet.",
"user": {
"screen_name": "joshmarshall",
"profile_image_url": "http://stuff.com/profile.jpg"
}
}
twitter_watcher = TwitterWatcher(self._ioloop, None, config)
twitter_watcher.extract_from_tweet(tweet)
self.assertEqual(result["tweet"], {
"username": "******",
"message": "This is a great tweet.",
"picture_url": "http://stuff.com/profile.jpg"
})
def run_game_data(nb_daide_clients, rules, csv_file):
""" Start a server and a client to test DAIDE communications
:param port: The port of the DAIDE server
:param csv_file: the csv file containing the list of DAIDE communications
"""
server = Server()
io_loop = IOLoop()
io_loop.make_current()
common.Tornado.stop_loop_on_callback_error(io_loop)
@gen.coroutine
def coroutine_func():
""" Concrete call to main function. """
port = random.randint(9000, 9999)
while is_port_opened(port, HOSTNAME):
port = random.randint(9000, 9999)
nb_human_players = 1 if nb_daide_clients < 7 else 0
server.start(port=port)
server_game = ServerGame(map_name='standard',
n_controls=nb_daide_clients +
nb_human_players,
rules=rules,
server=server)
# Register game on server.
game_id = server_game.game_id
server.add_new_game(server_game)
server.start_new_daide_server(game_id)
# Creating human player
human_username = '******'
human_password = '******'
# Creating bot player to play for dummy powers
bot_username = constants.PRIVATE_BOT_USERNAME
bot_password = constants.PRIVATE_BOT_PASSWORD
# Connecting
connection = yield connect(HOSTNAME, port)
human_channel = yield connection.authenticate(human_username,
human_password)
bot_channel = yield connection.authenticate(bot_username, bot_password)
# Joining human to game
channels = {BOT_KEYWORD: bot_channel}
if nb_human_players:
yield human_channel.join_game(game_id=game_id,
power_name='AUSTRIA')
channels['AUSTRIA'] = human_channel
comms_simulator = ClientsCommsSimulator(nb_daide_clients, csv_file,
game_id, channels)
yield comms_simulator.retrieve_game_port(HOSTNAME, port)
# done_future is only used to prevent pylint E1101 errors on daide_future
done_future = Future()
daide_future = comms_simulator.execute()
chain_future(daide_future, done_future)
for _ in range(3 + nb_daide_clients):
if done_future.done() or server_game.count_controlled_powers() >= (
nb_daide_clients + nb_human_players):
break
yield gen.sleep(2.5)
else:
raise TimeoutError()
# Waiting for process to finish
while not done_future.done() and server_game.status == strings.ACTIVE:
yield gen.sleep(2.5)
yield daide_future
try:
io_loop.run_sync(coroutine_func)
finally:
server.stop_daide_server(None)
if server.backend.http_server:
server.backend.http_server.stop()
io_loop.stop()
io_loop.clear_current()
io_loop.close()
server = None
Server.__cache__.clear()
def wrap_up(self, resp):
self.print_stats()
IOLoop.stop()
class Executor(object):
""" Distributed executor with data dependencies
This executor resembles executors in concurrent.futures but also allows
Futures within submit/map calls.
Provide center address on initialization
>>> executor = Executor(('127.0.0.1', 8787)) # doctest: +SKIP
Use ``submit`` method like normal
>>> a = executor.submit(add, 1, 2) # doctest: +SKIP
>>> b = executor.submit(add, 10, 20) # doctest: +SKIP
Additionally, provide results of submit calls (futures) to further submit
calls:
>>> c = executor.submit(add, a, b) # doctest: +SKIP
This allows for the dynamic creation of complex dependencies.
"""
def __init__(self, center, start=True, delete_batch_time=1):
self.center = coerce_to_rpc(center)
self.futures = dict()
self.refcount = defaultdict(lambda: 0)
self.dask = dict()
self.restrictions = dict()
self.loop = IOLoop()
self.report_queue = Queue()
self.scheduler_queue = Queue()
self._shutdown_event = Event()
self._delete_batch_time = delete_batch_time
if start:
self.start()
def start(self):
""" Start scheduler running in separate thread """
from threading import Thread
self.loop.add_callback(self._go)
self._loop_thread = Thread(target=self.loop.start)
self._loop_thread.start()
def __enter__(self):
if not self.loop._running:
self.start()
return self
def __exit__(self, type, value, traceback):
self.shutdown()
def _inc_ref(self, key):
self.refcount[key] += 1
def _dec_ref(self, key):
self.refcount[key] -= 1
if self.refcount[key] == 0:
del self.refcount[key]
self._release_key(key)
def _release_key(self, key):
""" Release key from distributed memory """
self.futures[key]['event'].clear()
logger.debug("Release key %s", key)
del self.futures[key]
self.scheduler_queue.put_nowait({
'op': 'release-held-data',
'key': key
})
@gen.coroutine
def report(self):
""" Listen to scheduler """
while True:
msg = yield self.report_queue.get()
if msg['op'] == 'close':
break
if msg['op'] == 'task-finished':
if msg['key'] in self.futures:
self.futures[msg['key']]['status'] = 'finished'
self.futures[msg['key']]['event'].set()
if msg['op'] == 'lost-data':
if msg['key'] in self.futures:
self.futures[msg['key']]['status'] = 'lost'
self.futures[msg['key']]['event'].clear()
if msg['op'] == 'task-erred':
if msg['key'] in self.futures:
self.futures[msg['key']]['status'] = 'error'
self.futures[msg['key']]['event'].set()
@gen.coroutine
def _shutdown(self):
""" Send shutdown signal and wait until _go completes """
self.report_queue.put_nowait({'op': 'close'})
self.scheduler_queue.put_nowait({'op': 'close'})
yield self._shutdown_event.wait()
def shutdown(self):
""" Send shutdown signal and wait until scheduler terminates """
self.report_queue.put_nowait({'op': 'close'})
self.scheduler_queue.put_nowait({'op': 'close'})
self.loop.stop()
self._loop_thread.join()
@gen.coroutine
def _go(self):
""" Setup and run all other coroutines. Block until finished. """
self.who_has, self.has_what, self.ncores = yield [
self.center.who_has(),
self.center.has_what(),
self.center.ncores()
]
self.waiting = {}
self.processing = {}
self.stacks = {}
worker_queues = {worker: Queue() for worker in self.ncores}
delete_queue = Queue()
coroutines = ([
self.report(),
scheduler(self.scheduler_queue, self.report_queue, worker_queues,
delete_queue, self.who_has, self.has_what, self.ncores,
self.dask, self.restrictions, self.waiting, self.stacks,
self.processing),
delete(self.scheduler_queue, delete_queue, self.center.ip,
self.center.port, self._delete_batch_time)
] + [
worker(self.scheduler_queue, worker_queues[w], w, n)
for w, n in self.ncores.items()
])
results = yield All(coroutines)
self._shutdown_event.set()
def submit(self, func, *args, **kwargs):
""" Submit a function application to the scheduler
Parameters
----------
func: callable
*args:
**kwargs:
pure: bool (defaults to True)
Whether or not the function is pure. Set ``pure=False`` for
impure functions like ``np.random.random``.
workers: set, iterable of sets
A set of worker hostnames on which computations may be performed.
Leave empty to default to all workers (common case)
Examples
--------
>>> c = executor.submit(add, a, b) # doctest: +SKIP
Returns
-------
Future
See Also
--------
distributed.executor.Executor.submit:
"""
if not callable(func):
raise TypeError(
"First input to submit must be a callable function")
key = kwargs.pop('key', None)
pure = kwargs.pop('pure', True)
workers = kwargs.pop('workers', None)
if key is None:
if pure:
key = funcname(func) + '-' + tokenize(func, kwargs, *args)
else:
key = funcname(func) + '-' + next(tokens)
if key in self.futures:
return Future(key, self)
if kwargs:
task = (apply, func, args, kwargs)
else:
task = (func, ) + args
if workers is not None:
restrictions = {key: workers}
else:
restrictions = {}
if key not in self.futures:
self.futures[key] = {'event': Event(), 'status': 'waiting'}
logger.debug("Submit %s(...), %s", funcname(func), key)
self.scheduler_queue.put_nowait({
'op': 'update-graph',
'dsk': {
key: task
},
'keys': [key],
'restrictions': restrictions
})
return Future(key, self)
def map(self, func, *iterables, **kwargs):
""" Map a function on a sequence of arguments
Arguments can be normal objects or Futures
Parameters
----------
func: callable
iterables: Iterables
pure: bool (defaults to True)
Whether or not the function is pure. Set ``pure=False`` for
impure functions like ``np.random.random``.
workers: set, iterable of sets
A set of worker hostnames on which computations may be performed.
Leave empty to default to all workers (common case)
Examples
--------
>>> L = executor.map(func, sequence) # doctest: +SKIP
Returns
-------
list of futures
See also
--------
distributed.executor.Executor.submit
"""
pure = kwargs.pop('pure', True)
workers = kwargs.pop('workers', None)
if not callable(func):
raise TypeError("First input to map must be a callable function")
iterables = [list(it) for it in iterables]
if pure:
keys = [
funcname(func) + '-' + tokenize(func, kwargs, *args)
for args in zip(*iterables)
]
else:
uid = str(uuid.uuid4())
keys = [
funcname(func) + '-' + uid + '-' + next(tokens)
for i in range(min(map(len, iterables)))
]
if not kwargs:
dsk = {
key: (func, ) + args
for key, args in zip(keys, zip(*iterables))
}
else:
dsk = {
key: (apply, func, args, kwargs)
for key, args in zip(keys, zip(*iterables))
}
for key in dsk:
if key not in self.futures:
self.futures[key] = {'event': Event(), 'status': 'waiting'}
if isinstance(workers, (list, set)):
if workers and isinstance(first(workers), (list, set)):
if len(workers) != len(keys):
raise ValueError("You only provided %d worker restrictions"
" for a sequence of length %d" %
(len(workers), len(keys)))
restrictions = dict(zip(keys, workers))
else:
restrictions = {key: workers for key in keys}
elif workers is None:
restrictions = {}
else:
raise TypeError("Workers must be a list or set of workers or None")
logger.debug("map(%s, ...)", funcname(func))
self.scheduler_queue.put_nowait({
'op': 'update-graph',
'dsk': dsk,
'keys': keys,
'restrictions': restrictions
})
return [Future(key, self) for key in keys]
@gen.coroutine
def _gather(self, futures):
futures2, keys = unpack_remotedata(futures)
keys = list(keys)
while True:
yield All([self.futures[key]['event'].wait() for key in keys])
try:
data = yield _gather(self.center, keys)
except KeyError as e:
self.scheduler_queue.put_nowait({
'op': 'missing-data',
'missing': e.args
})
for key in e.args:
self.futures[key]['event'].clear()
else:
break
data = dict(zip(keys, data))
result = pack_data(futures2, data)
raise gen.Return(result)
def gather(self, futures):
""" Gather futures from distributed memory
Accepts a future or any nested core container of futures
Examples
--------
>>> from operator import add # doctest: +SKIP
>>> e = Executor('127.0.0.1:8787') # doctest: +SKIP
>>> x = e.submit(add, 1, 2) # doctest: +SKIP
>>> e.gather(x) # doctest: +SKIP
3
>>> e.gather([x, [x], x]) # doctest: +SKIP
[3, [3], 3]
"""
return sync(self.loop, self._gather, futures)
@gen.coroutine
def _get(self, dsk, keys, restrictions=None):
flatkeys = list(flatten(keys))
for key in flatkeys:
if key not in self.futures:
self.futures[key] = {'event': Event(), 'status': None}
futures = {key: Future(key, self) for key in flatkeys}
self.scheduler_queue.put_nowait({
'op': 'update-graph',
'dsk': dsk,
'keys': flatkeys,
'restrictions': restrictions or {}
})
packed = pack_data(keys, futures)
result = yield self._gather(packed)
raise gen.Return(result)
def get(self, dsk, keys, **kwargs):
""" Gather futures from distributed memory
Parameters
----------
dsk: dict
keys: object, or nested lists of objects
restrictions: dict (optional)
A mapping of {key: {set of worker hostnames}} that restricts where
jobs can take place
Examples
--------
>>> from operator import add # doctest: +SKIP
>>> e = Executor('127.0.0.1:8787') # doctest: +SKIP
>>> e.get({'x': (add, 1, 2)}, 'x') # doctest: +SKIP
3
"""
return sync(self.loop, self._get, dsk, keys, **kwargs)
class SyncHTTPClientTest(unittest.TestCase):
def setUp(self):
self.server_ioloop = IOLoop()
event = threading.Event()
@gen.coroutine
def init_server():
sock, self.port = bind_unused_port()
app = Application([("/", HelloWorldHandler)])
self.server = HTTPServer(app)
self.server.add_socket(sock)
event.set()
def start():
self.server_ioloop.run_sync(init_server)
self.server_ioloop.start()
self.server_thread = threading.Thread(target=start)
self.server_thread.start()
event.wait()
self.http_client = HTTPClient()
def tearDown(self):
def stop_server():
self.server.stop()
# Delay the shutdown of the IOLoop by several iterations because
# the server may still have some cleanup work left when
# the client finishes with the response (this is noticeable
# with http/2, which leaves a Future with an unexamined
# StreamClosedError on the loop).
@gen.coroutine
def slow_stop():
yield self.server.close_all_connections()
# The number of iterations is difficult to predict. Typically,
# one is sufficient, although sometimes it needs more.
for i in range(5):
yield
self.server_ioloop.stop()
self.server_ioloop.add_callback(slow_stop)
self.server_ioloop.add_callback(stop_server)
self.server_thread.join()
self.http_client.close()
self.server_ioloop.close(all_fds=True)
def get_url(self, path):
return "http://127.0.0.1:%d%s" % (self.port, path)
def test_sync_client(self):
response = self.http_client.fetch(self.get_url("/"))
self.assertEqual(b"Hello world!", response.body)
def test_sync_client_error(self):
# Synchronous HTTPClient raises errors directly; no need for
# response.rethrow()
with self.assertRaises(HTTPError) as assertion:
self.http_client.fetch(self.get_url("/notfound"))
self.assertEqual(assertion.exception.code, 404)
class MockHTTP(object):
"""A Mock HTTP Server for unit testing web services calls.
Basic Usage::
mock_server = MockHTTP(42424)
mock_server.expects(GET, '/index.html').will(body='A HTML body.')
mock_server.expects(GET, '/asdf').will(http_code=404)
urlopen('http://localhost:42424/index.html').read() == 'A HTML body.'
urlopen('http://localhost:42424/asdf') # HTTPError: 404
mock_server.verify()"""
def __init__(self, shutdown_on_verify=True):
"""Create a MockHTTP server listening on localhost at the given port."""
self.shutdown_on_verify = shutdown_on_verify
self._server_io_loop = IOLoop()
self.reset()
self.sock, self.port = bind_unused_port()
self.server_address = "http://localhost:%s" % self.port
self._server_is_down = False
self.thread = threading.Thread(target=_http_server_thread,
kwargs=dict(io_loop=self._server_io_loop, mock=self))
self.thread.start()
def get_url(self, path):
if not path.startswith("/"):
return self.get_url("/" + path)
return self.server_address + path
def expects(self, method, path, *args, **kwargs):
"""Declares an HTTP Request that this MockHTTP expects.
:param method: The HTTP method expected to use to access this URL.
:param path: The expected path segment of this URL.
:param body: The expected contents of the request body, as a string. If\
you expect to send unicode, encode it as utf-8 first. *Default:* The\
contents of the request body are irrelevant.
:param params: Expected query parameters as a dictionary mapping query\
parameter name to expected value. Checks to make sure that all expected\
query parameters are present and have specified values. *Default:* No\
query parameters are expected.
:param headers: Expected headers as a dictionary mapping header name to\
expected value. Checks to make sure that all expected headers are\
present and have the specified values. *Default:* No headers are\
required.
:param times: The number of times this URL expects to be requested. Use\
mock_http.never, mock_http.once, or mock_http.at_least_once for this.\
*Default:* It does not matter how many times the URL is accessed.
:param name: A name that can be used to refer to this expectation later.\
*Default:* The expectation has no name.
:param after: This URL must be accessed after a previously-named URL.\
*Default:* The URL can be accessed at any time.
:returns: The :class:`Expectation` object describing how this URL is\
expected. You'll probably want to call :meth:`Expectation.will` on it\
to describe how the URL should be responded to.
"""
expectation = Expectation(self, method, path, *args, **kwargs)
self.expected[method][path] = expectation
log.debug("Added expectation: %s at %s", expectation, self)
return expectation
def reset(self):
"""
Clears expectations
"""
self.last_failure = None
self.expected = defaultdict(dict)
self.expected_by_name = {}
log.debug("Reset expectations MockHTTP %s", self)
def shutdown(self):
"""Close down the server"""
if not self._server_is_down:
log.debug("Shutting down %s", self)
self._server_io_loop.stop()
log.debug("Shutdown %s: Ok", self)
self._server_is_down = True
else:
log.debug("Attempt to shut down server when its already down")
def verify(self):
"""Close down the server and verify that this MockHTTP has met all its
expectations.
:returns: True, if all went as expected.
:raises MockHTTPExpectationFailure: Or a subclass, describing the last\
unexpected thing that happened."""
if self.shutdown_on_verify:
log.debug("Shutting down before verify %s", self)
self.shutdown()
if self.last_failure is not None:
raise self.last_failure
for method, expected in self.expected.iteritems():
for path, expectation in expected.iteritems():
if (isinstance(expectation.times, int) and expectation.times != expectation.invoked_times) or \
(isinstance(expectation.times, ExpectedTime)
and not expectation.times.passes_test(expectation.invoked_times)):
raise UnretrievedURLException("%s not %s" % (path, method))
return True
def is_expected(self, method, path, params, headers, body):
"""Test to see whether a request is expected.
.. todo::
Gracefully handle multiple expectations for the same URL and method.
:raises MockHTTPExpectationFailure: Or a subclass, describing why this\
request is unexpected.
:returns: The :class:`Expectation` object that expects this request."""
try:
log.debug("Looking for [%s %s] at expectations", method, path)
expectation = self.expected[method][path]
log.debug("Found [%s %s] matching expectation %s", method, path, expectation)
if expectation.check(method, path, params, headers, body):
return expectation
except KeyError:
failure = UnexpectedURLException('Unexpected URL: %s' % path)
self.last_failure = failure
log.debug("Didnt found [%s %s] at expectations", method, path)
raise failure
except MockHTTPExpectationFailure, failure:
self.last_failure = failure
log.debug("Expectation %s check failed at [%s %s] with %s", expectation, method, path, failure)
raise
