def spider_list(self, project, cwd=None):
future = Future()
try:
env = os.environ.copy()
env['SCRAPY_PROJECT'] = project
process = Popen([self.python, '-m', 'scrapyd.runner', 'list'],
env=env,
cwd=cwd,
stdout=PIPE,
stderr=PIPE)
except Exception as e:
logger.error(e)
future.set_exception(e)
return future
def check_process():
logger.debug('poll')
retcode = process.poll()
if retcode is not None:
if retcode == 0:
future.set_result(process.stdout.read().splitlines())
else:
#future.set_exception(ProcessFailed(std_output=process.stdout.read(), err_output=process.stderr.read()))
future.set_exception(
InvalidProjectEgg(detail=process.stderr.read()))
return
IOLoop.current().call_later(1, check_process)
check_process()
return future
def _create_stream(self,
max_buffer_size,
af,
addr,
source_ip=None,
source_port=None):
# Always connect in plaintext; we'll convert to ssl if necessary
# after one connection has completed.
source_port_bind = source_port if isinstance(source_port, int) else 0
source_ip_bind = source_ip
if source_port_bind and not source_ip:
# User required a specific port, but did not specify
# a certain source IP, will bind to the default loopback.
source_ip_bind = '::1' if af == socket.AF_INET6 else '127.0.0.1'
# Trying to use the same address family as the requested af socket:
# - 127.0.0.1 for IPv4
# - ::1 for IPv6
socket_obj = socket.socket(af)
set_close_exec(socket_obj.fileno())
if source_port_bind or source_ip_bind:
# If the user requires binding also to a specific IP/port.
try:
socket_obj.bind((source_ip_bind, source_port_bind))
except socket.error:
socket_obj.close()
# Fail loudly if unable to use the IP/port.
raise
try:
stream = IOStream(socket_obj, max_buffer_size=max_buffer_size)
except socket.error as e:
fu = Future()
fu.set_exception(e)
return fu
else:
return stream, stream.connect(addr)
def _as_future(value):
future = Future()
if isinstance(value, Exception):
future.set_exception(value)
else:
future.set_result(value)
return future
def init(self):
"""
init project isolated workspace,
:return: future
"""
future = Future()
if os.path.exists(self.pip) and os.path.exists(self.python):
future.set_result(self)
return future
logger.debug('workspace dir : %s' % (self.project_workspace_dir,))
logger.info('start creating virtualenv.')
try:
process = Popen([sys.executable, '-m', 'virtualenv',
'--system-site-packages', self.venv_dir])
self.processes.append(process)
except Exception as e:
future.set_exception(e)
return future
def done(process):
self.processes.remove(process)
retcode = process.returncode
if retcode is not None:
if retcode == 0:
logger.info('Create virtualenv done.')
future.set_result(self)
else:
future.set_exception(
ProcessFailed('Error when init workspace virtualenv '))
wait_process(process, done)
return future
def find_project_settings(self, project):
cwd = self.project_workspace_dir
future = Future()
try:
env = os.environ.copy()
env['SCRAPY_PROJECT'] = project
process = Popen([self.python, '-m', 'scrapydd.utils.extract_settings', project], env=env, cwd=cwd,
stdout=PIPE, stderr=PIPE)
self.processes.append(process)
except Exception as e:
logger.error(e)
future.set_exception(e)
return future
def done(process):
retcode = process.poll()
self.processes.remove(process)
if retcode == 0:
proc_output = process.stdout.read()
logger.debug('find_project_settings output:')
logger.debug(proc_output)
return future.set_result(json.loads(proc_output))
else:
# future.set_exception(ProcessFailed(std_output=process.stdout.read(), err_output=process.stderr.read()))
return future.set_exception(InvalidProjectEgg(detail=process.stderr.read()))
wait_process(process, done)
return future
def open(self):
try:
future = self._stream.open()
except Exception as e:
future = Future()
future.set_exception(e)
return future
def pip_install(self, requirements):
logger.debug('installing requirements: %s' % requirements)
future = Future()
try:
process = Popen([self.pip, 'install'] + requirements,
stdout=PIPE,
stderr=PIPE)
except Exception as e:
future.set_exception(e)
return future
def check_process():
logger.debug('poll')
retcode = process.poll()
if retcode is not None:
if retcode == 0:
future.set_result(self)
else:
std_out = process.stdout.read()
err_out = process.stderr.read()
future.set_exception(
ProcessFailed(std_output=std_out, err_output=err_out))
return
IOLoop.current().call_later(1, check_process)
check_process()
return future
def get_event(self, request, tag='', callback=None, timeout=None):
'''
Get an event (async of course) return a future that will get it later
'''
# if the request finished, no reason to allow event fetching, since we
# can't send back to the client
if request._finished:
future = Future()
future.set_exception(TimeoutException())
return future
future = Future()
if callback is not None:
def handle_future(future):
tornado.ioloop.IOLoop.current().add_callback(callback, future)
future.add_done_callback(handle_future)
# add this tag and future to the callbacks
self.tag_map[tag].append(future)
self.request_map[request].append((tag, future))
if timeout:
timeout_future = tornado.ioloop.IOLoop.current().call_later(
timeout, self._timeout_future, tag, future)
self.timeout_map[future] = timeout_future
return future
def test_both(self):
marker = object()
second_marker = object()
result_marker = object()
def _mapper(_):
return marker
def _exception_mapper(_):
return second_marker
first_future = Future()
folded_future = future_fold(first_future,
result_mapper=_mapper,
exception_mapper=_exception_mapper)
folded_future_probe = FutureProbe(folded_future)
second_future = Future()
second_folded_future = future_fold(second_future,
result_mapper=_mapper,
exception_mapper=_exception_mapper)
second_folded_future_probe = FutureProbe(second_folded_future,
stop_cb=self.stop)
first_future.set_result(result_marker)
second_future.set_exception(MyException())
self.wait()
folded_future_probe.assert_single_result_call(self, marker)
second_folded_future_probe.assert_single_result_call(
self, second_marker)
def _start_processing_requests(self):
while True:
data = yield gen.Task(self._stream.read_until, '\r\n')
log.debug('New request: %r', data)
try:
msg = json.loads(data)
key = msg['key']
method = msg['method']
args = msg['args']
kwargs = msg['kwargs']
except (KeyError, ValueError):
log.error('Malformed request data: %s', data)
continue
try:
res = self._handler(method, *args, **kwargs)
if isinstance(res, Future):
future = res
else:
future = Future()
future.set_result(res)
except Exception as e:
log.exception('Failed to handle request: %s', key)
future = concurrent.TracebackFuture()
future.set_exception(e)
future.add_done_callback(partial(self._on_future_finished, key))
def handle_request(server, request, connection_handler):
""" (coroutine) Find request handler function for associated request, run it and return its result.
:param server: a Server object to pass to handler function.
:param request: a request object to pass to handler function.
See diplomacy.communication.requests for possible requests.
:param connection_handler: a ConnectionHandler object to pass to handler function.
:return: (future) either None or a response object.
See module diplomacy.communication.responses for possible responses.
"""
request_handler_fn = MAPPING.get(type(request), None)
if not request_handler_fn:
raise exceptions.RequestException()
game = server.get_game(request.game_id)
# Game not found
if not game or game.is_game_completed or game.is_game_canceled:
future = Future()
future.set_result([responses.REJ(bytes(request))])
return future
if gen.is_coroutine_function(request_handler_fn):
# Throw the future returned by this coroutine.
return request_handler_fn(server, request, connection_handler, game)
# Create and return a future.
future = Future()
try:
result = request_handler_fn(server, request, connection_handler, game)
future.set_result(result)
except exceptions.DiplomacyException as exc:
future.set_exception(exc)
return future
def apply(callback) -> 'promise.Promise':
f = Future()
try:
f.set_result(callback())
except BaseException as e:
f.set_exception(e)
return Promise(f)
def _create_stream(self, max_buffer_size, af, addr, source_ip=None,
source_port=None):
# Always connect in plaintext; we'll convert to ssl if necessary
# after one connection has completed.
source_port_bind = source_port if isinstance(source_port, int) else 0
source_ip_bind = source_ip
if source_port_bind and not source_ip:
# User required a specific port, but did not specify
# a certain source IP, will bind to the default loopback.
source_ip_bind = '::1' if af == socket.AF_INET6 else '127.0.0.1'
# Trying to use the same address family as the requested af socket:
# - 127.0.0.1 for IPv4
# - ::1 for IPv6
socket_obj = socket.socket(af)
set_close_exec(socket_obj.fileno())
if source_port_bind or source_ip_bind:
# If the user requires binding also to a specific IP/port.
try:
socket_obj.bind((source_ip_bind, source_port_bind))
except socket.error:
socket_obj.close()
# Fail loudly if unable to use the IP/port.
raise
try:
stream = IOStream(socket_obj,
max_buffer_size=max_buffer_size)
except socket.error as e:
fu = Future()
fu.set_exception(e)
return fu
else:
return stream, stream.connect(addr)
def settings_module(self):
future = Future()
cwd = self.project_workspace_dir
try:
env = os.environ.copy()
env['SCRAPY_PROJECT'] = self.project_name
env['SCRAPY_EGG'] = 'spider.egg'
process = Popen([self.python, '-m', 'scrapydd.utils.extract_settings_module', 'spider.egg'],
env=env, cwd=cwd, stdout=PIPE,
stderr=PIPE, encoding=PROCESS_ENCODING)
except Exception as e:
logger.error(e)
future.set_exception(e)
return future
def check_process():
logger.debug('poll')
retcode = process.poll()
if retcode is not None:
if retcode == 0:
stdout = process.stdout.read()
stdout = ensure_str(stdout).strip()
future.set_result(stdout)
else:
# future.set_exception(ProcessFailed(std_output=process.stdout.read(), err_output=process.stderr.read()))
future.set_exception(InvalidProjectEgg(detail=process.stderr.read()))
return
IOLoop.current().call_later(1, check_process)
check_process()
return future
class ManualCapClient(BaseCapClient):
def capitalize(self, request_data, callback=None):
logging.debug("capitalize")
self.request_data = request_data
self.stream = IOStream(socket.socket())
self.stream.connect(('127.0.0.1', self.port),
callback=self.handle_connect)
self.future = Future()
if callback is not None:
self.future.add_done_callback(
stack_context.wrap(lambda future: callback(future.result())))
return self.future
def handle_connect(self):
logging.debug("handle_connect")
self.stream.write(utf8(self.request_data + "\n"))
self.stream.read_until(b'\n', callback=self.handle_read)
def handle_read(self, data):
logging.debug("handle_read")
self.stream.close()
try:
self.future.set_result(self.process_response(data))
except CapError as e:
self.future.set_exception(e)
def open(self):
try:
future = self._stream.open()
except Exception as e:
future = Future()
future.set_exception(e)
return future
class ManualCapClient(BaseCapClient):
def capitalize(self, request_data, callback=None):
logging.info("capitalize")
self.request_data = request_data
self.stream = IOStream(socket.socket())
self.stream.connect(('127.0.0.1', self.port),
callback=self.handle_connect)
self.future = Future()
if callback is not None:
self.future.add_done_callback(
stack_context.wrap(lambda future: callback(future.result())))
return self.future
def handle_connect(self):
logging.info("handle_connect")
self.stream.write(utf8(self.request_data + "\n"))
self.stream.read_until(b'\n', callback=self.handle_read)
def handle_read(self, data):
logging.info("handle_read")
self.stream.close()
try:
self.future.set_result(self.process_response(data))
except CapError as e:
self.future.set_exception(e)
def futurized(result):
future = Future()
if isinstance(result, Exception):
future.set_exception(result)
else:
future.set_result(result)
return future
def test_both(self):
marker = object()
second_marker = object()
result_marker = object()
def _mapper(_):
return marker
def _exception_mapper(_):
return second_marker
first_future = Future()
folded_future = future_fold(first_future, result_mapper=_mapper, exception_mapper=_exception_mapper)
folded_future_probe = FutureProbe(folded_future)
second_future = Future()
second_folded_future = future_fold(second_future, result_mapper=_mapper, exception_mapper=_exception_mapper)
second_folded_future_probe = FutureProbe(second_folded_future, stop_cb=self.stop)
first_future.set_result(result_marker)
second_future.set_exception(MyException())
self.wait()
folded_future_probe.assert_single_result_call(self, marker)
second_folded_future_probe.assert_single_result_call(self, second_marker)
def pip_install(self, requirements):
logger.debug('installing requirements: %s' % requirements)
future = Future()
try:
env = os.environ.copy()
stdout_stream = tempfile.NamedTemporaryFile()
stderr_stream = tempfile.NamedTemporaryFile()
process = Popen([self.pip, 'install'] + requirements,
stdout=stdout_stream, stderr=stderr_stream,
env=env,
encoding=PROCESS_ENCODING)
self.processes.append(process)
except Exception as e:
future.set_exception(e)
return future
def done(process):
self.processes.remove(process)
retcode = process.returncode
if retcode is not None:
if retcode == 0:
future.set_result(self)
else:
stdout_stream.seek(0)
stderr_stream.seek(0)
std_out = stdout_stream.read().decode(PROCESS_ENCODING)
err_out = stderr_stream.read().decode(PROCESS_ENCODING)
future.set_exception(ProcessFailed(std_output=std_out,
err_output=err_out))
wait_process(process, done)
return future
def _default_invoke_action_handler(self, parameters):
"""Default handler for onInvokeAction."""
future_invoke = Future()
future_invoke.set_exception(
NotImplementedError("Undefined action handler"))
return future_invoke
def run(child_gr, *args, **kwargs):
try:
result = func(*args, **kwargs)
if not is_future(result):
return child_gr.switch(result)
except Exception as e:
result = Future()
result.set_exception(e)
return ioloop.add_future(result, child_gr.switch)
def run(child_gr, *args, **kwargs):
try:
result = func(*args, **kwargs)
if not is_future(result):
return child_gr.switch(result)
except Exception as e:
result = Future()
result.set_exception(e)
return ioloop.add_future(result, child_gr.switch)
def side_effect(request, **kwargs):
if request is not HTTPRequest:
request = HTTPRequest(request)
buffer = StringIO(body)
response = HTTPResponse(request, status_code, buffer=buffer)
future = Future()
if response.error:
future.set_exception(response.error)
else:
future.set_result(response)
return future
def side_effect(request, **kwargs):
if request is not HTTPRequest:
request = HTTPRequest(request)
buffer = StringIO(body)
response = HTTPResponse(request, status_code, buffer=buffer)
future = Future()
if response.error:
future.set_exception(response.error)
else:
future.set_result(response)
return future
def make_future(v):
'''
TODO: exceptions cases
'''
fut = Future()
if isinstance(v, Exception) or isinstance(v, gen.TimeoutError):
fut.set_exception(v)
else:
fut.set_result(v)
return fut
def test_non_successful_responses_marks_client_as_failed(self):
client = self.get_app().sentry_client
with patch.object(client, "_failed_send") as mock_failed:
with patch.object(client, "_send_remote") as mock_send:
f = Future()
f.set_exception(HTTPError(499, "error"))
mock_send.return_value = f
client.send()
yield gen.sleep(0.01) # we need to run after the async send
assert mock_failed.called
def test_non_successful_responses_marks_client_as_failed(self):
client = self.get_app().sentry_client
with patch.object(client, '_failed_send') as mock_failed:
with patch.object(client, '_send_remote') as mock_send:
f = Future()
f.set_exception(HTTPError(499, 'error'))
mock_send.return_value = f
client.send()
yield gen.sleep(0.01) # we need to run after the async send
assert mock_failed.called
def wrapper(*args, **kwargs):
future = Future()
if replace_callback and 'callback' in kwargs:
callback = kwargs.pop('callback')
IOLoop.current().add_future(
future, lambda future: callback(future.result()))
try:
result = func(*args, **kwargs)
except (Return, StopIteration) as e:
result = _value_from_stopiteration(e)
except Exception:
future_set_exc_info(future, sys.exc_info())
return future
else:
if isinstance(result, GeneratorType):
# Inline the first iteration of Runner.run. This lets us
# avoid the cost of creating a Runner when the coroutine
# never actually yields, which in turn allows us to
# use "optional" coroutines in critical path code without
# performance penalty for the synchronous case.
try:
orig_stack_contexts = stack_context._state.contexts
yielded = next(result)
if stack_context._state.contexts is not orig_stack_contexts:
yielded = Future()
yielded.set_exception(
stack_context.StackContextInconsistentError(
'stack_context inconsistency (probably caused '
'by yield within a "with StackContext" block)'))
except (StopIteration, Return) as e:
future.set_result(_value_from_stopiteration(e))
except Exception:
future_set_exc_info(future, sys.exc_info())
else:
_futures_to_runners[future] = Runner(result, future, yielded)
yielded = None
try:
return future
finally:
# Subtle memory optimization: if next() raised an exception,
# the future's exc_info contains a traceback which
# includes this stack frame. This creates a cycle,
# which will be collected at the next full GC but has
# been shown to greatly increase memory usage of
# benchmarks (relative to the refcount-based scheme
# used in the absence of cycles). We can avoid the
# cycle by clearing the local variable after we return it.
future = None
future.set_result(result)
return future
def _create_stream(self, max_buffer_size, af, addr):
# Always connect in plaintext; we'll convert to ssl if necessary
# after one connection has completed.
try:
stream = IOStream(socket.socket(af),
io_loop=self.io_loop,
max_buffer_size=max_buffer_size)
except socket.error as e:
fu = Future()
fu.set_exception(e)
return fu
else:
return stream.connect(addr)
def _create_stream(self, max_buffer_size, af, addr):
# Always connect in plaintext; we'll convert to ssl if necessary
# after one connection has completed.
try:
stream = IOStream(socket.socket(af),
io_loop=self.io_loop,
max_buffer_size=max_buffer_size)
except socket.error as e:
fu = Future()
fu.set_exception(e)
return fu
else:
return stream.connect(addr)
def _build_response(self, tid):
"""
Prepare for a response, returns a future
:param tid:
:return: Future
"""
f = Future()
if not self._connected:
f.set_exception(ConnectionException("Client is not connected"))
return f
self.transaction.addTransaction(f, tid)
return f
def _build_response(self, tid):
"""
Prepare for a response, returns a future
:param tid:
:return: Future
"""
f = Future()
if not self._connected:
f.set_exception(ConnectionException("Client is not connected"))
return f
self.transaction.addTransaction(f, tid)
return f
def remote_call(self, recipient, method, *args, **kwargs):
future = Future()
future_id = (id(future), time())
self._pendings[future_id] = future
try:
message = (recipient, ('CALL', method, args, kwargs, future_id,
fqc_name(self)))
self._router.put_nowait(message)
# if self._debug:
# logging.info("Endpoint '%s' sent message: %s", fqc_name(self), message)
except queue.Full as exc:
future.set_exception(exc)
logging.error("Cannot remote call to '%s', router query is full",
recipient)
return future
def with_timeout(timeout, future, io_loop=None):
"""Wraps a `.Future` in a timeout.
Raises `TimeoutError` if the input future does not complete before
``timeout``, which may be specified in any form allowed by
`.IOLoop.add_timeout` (i.e. a `datetime.timedelta` or an absolute time
relative to `.IOLoop.time`)
Currently only supports Futures, not other `YieldPoint` classes.
.. versionadded:: 4.0
"""
# TODO: allow yield points in addition to futures?
# Tricky to do with stack_context semantics.
#
# It's tempting to optimize this by cancelling the input future on timeout
# instead of creating a new one, but A) we can't know if we are the only
# one waiting on the input future, so cancelling it might disrupt other
# callers and B) concurrent futures can only be cancelled while they are
# in the queue, so cancellation cannot reliably bound our waiting time.
result = Future()
chain_future(future, result)
if io_loop is None:
io_loop = IOLoop.current()
timeout_handle = io_loop.add_timeout(timeout, lambda: result.set_exception(TimeoutError("Timeout")))
if isinstance(future, Future):
# We know this future will resolve on the IOLoop, so we don't
# need the extra thread-safety of IOLoop.add_future (and we also
# don't care about StackContext here.
future.add_done_callback(lambda future: io_loop.remove_timeout(timeout_handle))
else:
# concurrent.futures.Futures may resolve on any thread, so we
# need to route them back to the IOLoop.
io_loop.add_future(future, lambda future: io_loop.remove_timeout(timeout_handle))
return result
def test_fails_before_timeout(self):
future = Future()
self.io_loop.add_timeout(
datetime.timedelta(seconds=0.1),
lambda: future.set_exception(ZeroDivisionError))
with self.assertRaises(ZeroDivisionError):
yield gen.with_timeout(datetime.timedelta(seconds=3600), future)
def test_unsuccessful_future(self):
class BadFuture(Exception):
pass
future = Future()
def run():
with self.assertRaises(BadFuture):
utils.run_async(future)
self.stop()
gr = greenlet.greenlet(run)
gr.switch()
future.set_exception(BadFuture("you have been exterminated"))
self.wait()
def test_fails_before_timeout(self):
future = Future()
self.io_loop.add_timeout(
datetime.timedelta(seconds=0.1),
lambda: future.set_exception(ZeroDivisionError()))
with self.assertRaises(ZeroDivisionError):
yield gen.with_timeout(datetime.timedelta(seconds=3600), future)
def _create_stream(self,
max_buffer_size,
af,
addr,
source_ip=None,
source_port=None):
# Always connect in plaintext; we'll convert to ssl if necessary
# after one connection has completed.
source_port_bind = source_port if isinstance(source_port, int) else 0
source_ip_bind = source_ip
socket_obj = socket.socket(af)
set_close_exec(socket_obj.fileno())
try:
stream = IOStream(socket_obj,
io_loop=self.io_loop,
max_buffer_size=max_buffer_size)
# connect proxy
if source_port_bind or source_ip_bind:
@gen.coroutine
def _(addr):
proxy_headers = get_proxy_headers(source_ip_bind)
parsed = urlparse(source_ip_bind)
scheme, host, port = parsed.scheme, parsed.hostname, source_port_bind
if 'socks' in scheme:
r = yield self._negotiate_socks(
addr, (source_ip_bind, source_port_bind))
raise gen.Return(r)
elif scheme in ('http', 'https'):
r = yield stream.connect((host, port))
if scheme == 'https':
yield self._connect_tunnel(stream, addr,
proxy_headers)
raise gen.Return(r)
else:
raise AttributeError('Unknown scheme: %s' % scheme)
return _(addr)
else:
return stream.connect(addr)
except socket.error as e:
fu = Future()
fu.set_exception(e)
return fu
def open(self):
if self._unix_socket:
address = self._unix_socket
else:
address = (self._host, self._port)
try:
future = self.connect(address)
if self._timeout > 0:
def timeout():
if self._connecting:
self.close((None, TStreamConnectTimeoutError(), None))
self.io_loop.add_timeout(time.time() + self._timeout, timeout)
except Exception as e:
future = Future()
future.set_exception(e)
return future
def test_exception_to_exception(self):
future = Future()
future_probe = FutureProbe(future)
def _exception_mapper(exception):
if isinstance(exception, MyException):
raise MyOtherException()
else:
return None
res_future = future_fold(future, exception_mapper=_exception_mapper)
res_future_probe = FutureProbe(res_future, stop_cb=self.stop)
future.set_exception(MyException())
self.wait()
future_probe.assert_single_exception_call(self, MyException)
res_future_probe.assert_single_exception_call(self, MyOtherException)
def test_exception_to_exception(self):
future = Future()
future_probe = FutureProbe(future)
def _exception_mapper(exception):
if isinstance(exception, MyException):
raise MyOtherException()
else:
return None
res_future = future_fold(future, exception_mapper=_exception_mapper)
res_future_probe = FutureProbe(res_future, stop_cb=self.stop)
future.set_exception(MyException())
self.wait()
future_probe.assert_single_exception_call(self, MyException)
res_future_probe.assert_single_exception_call(self, MyOtherException)
def open(self):
if self._unix_socket:
address = self._unix_socket
else:
address = (self._host, self._port)
try:
future = self.connect(address)
if self._timeout > 0:
def timeout():
if self._connecting:
self.close((None, TStreamConnectTimeoutError(), None))
self.io_loop.add_timeout(time.time() + self._timeout, timeout)
except Exception as e:
future = Future()
future.set_exception(e)
return future
def acquire(self):
"""Decrement the counter. Returns a Future.
Block if the counter is zero and wait for a `.release`. The Future
raises `.TimeoutError` after the deadline.
"""
waiter = Future()
if self.redis_client.set(self.key, self.value, ex=self.ttl, nx=True):
waiter.set_result(locks._ReleasingContextManager(self))
else:
waiter.set_exception(DDosError("被暴击了"))
# def on_timeout():
# waiter.set_exception(gen.TimeoutError())
# io_loop = ioloop.IOLoop.current()
# timeout_handle = io_loop.add_timeout(timeout, on_timeout)
# waiter.add_done_callback(
# lambda _: io_loop.remove_timeout(timeout_handle))
return waiter
class DBusConnection:
def __init__(self, bus_addr):
self.auth_parser = SASLParser()
self.parser = Parser()
self.router = Router(Future)
self.authentication = Future()
self.unique_name = None
self._sock = socket.socket(family=socket.AF_UNIX)
self.stream = IOStream(self._sock, read_chunk_size=4096)
def connected():
self.stream.write(b'\0' + make_auth_external())
self.stream.connect(bus_addr, connected)
self.stream.read_until_close(streaming_callback=self.data_received)
def _authenticated(self):
self.stream.write(BEGIN)
self.authentication.set_result(True)
self.data_received_post_auth(self.auth_parser.buffer)
def data_received(self, data):
if self.authentication.done():
return self.data_received_post_auth(data)
self.auth_parser.feed(data)
if self.auth_parser.authenticated:
self._authenticated()
elif self.auth_parser.error:
self.authentication.set_exception(AuthenticationError(self.auth_parser.error))
def data_received_post_auth(self, data):
for msg in self.parser.feed(data):
self.router.incoming(msg)
def send_message(self, message):
if not self.authentication.done():
raise RuntimeError("Wait for authentication before sending messages")
future = self.router.outgoing(message)
data = message.serialise()
self.stream.write(data)
return future
class GuarantedHTTPRequest(HTTPRequest):
def __init__(self, *args, **kwargs):
self._first_chunk_recieved = False
self.timeout_handle = None
self._io_loop = ioloop.IOLoop.current()
if "streaming_callback" in kwargs:
self.orig_streaming_callback = kwargs["streaming_callback"]
else:
self.orig_streaming_callback = None
if "inactive_timeout" in kwargs:
self.inactive_timeout = kwargs["inactive_timeout"]
else:
self.inactive_timeout = 1
del kwargs["inactive_timeout"]
self.timeout_future = Future()
kwargs["streaming_callback"] = self.stream_cb
super(GuarantedHTTPRequest, self).__init__(*args, **kwargs)
def stream_cb(self, data):
self._first_chunk_recieved = True
if self.orig_streaming_callback:
self.orig_streaming_callback(data)
io_loop = self._io_loop
if self.timeout_handle:
io_loop.remove_timeout(self.timeout_handle)
self.timeout_handle = io_loop.call_at(io_loop.time()+1,
self.throwStreamingTimeout)
def throwStreamingTimeout(self):
err = HTTPError(
504,
message="No activity from server for {} second(s)".format(
self.inactive_timeout))
self.timeout_future.set_exception(err)
def done(self):
if self._first_chunk_recieved:
io_loop = self._io_loop
io_loop.remove_timeout(self.timeout_handle)
def with_timeout(timeout, future, io_loop=None):
result = Future()
chain_future(future, result)
if io_loop is None:
io_loop = IOLoop.current()
timeout_handle = io_loop.add_timeout(
timeout,
lambda: result.set_exception(TimeoutError("Timeout")))
future.add_done_callback(
lambda future: io_loop.remove_timeout(timeout_handle))
return result
def test_exception_to_value(self):
marker = object()
future = Future()
future_probe = FutureProbe(future)
def _exception_mapper(exception):
# We need to check exception type, but here we can't raise AssertionException.
# So it returns None for failing in assertions bellow.
if isinstance(exception, MyException):
return marker
else:
return None
res_future = future_fold(future, exception_mapper=_exception_mapper)
res_future_probe = FutureProbe(res_future, stop_cb=self.stop)
future.set_exception(MyException())
self.wait()
future_probe.assert_single_exception_call(self, MyException)
res_future_probe.assert_single_result_call(self, marker)
def get_page(self):
future = Future()
if self.get_argument('failed_future', 'false') == 'true':
future.set_exception(Exception('failed future exception'))
else:
future.set_result({'1': 'yay'})
another_future = Future()
another_future.set_result({'2': 'yay'})
def _final_callback(results):
self.json.put({'final_callback_called': True})
self.json.put(
self.group(
{
'1': future,
'2': another_future,
},
_final_callback,
name='test async'
)
)
def issue_command(self, topic, cmd, block=True, timeout=5):
ft = Future()
if topic.find('#') >= 0:
### this is a wildcard topic
ft.set_exception(Exception("not implemented"))
else:
### this is a single topic
if topic in CoCommandBus.topic_pool:
topic_instance = CoCommandBus.topic_pool[topic]
if topic_instance.is_marked_gc():
topic_instance = None
del CoCommandBus.topic_pool[topic]
ft.set_exception(OfflineError())
else:
if not isinstance(cmd, dict):
ft.set_exception(BadCommandError())
else:
topic_instance.issue(ft, cmd, block, timeout)
else:
ft.set_exception(OfflineError())
return ft
def _connect(self):
future = Future()
TornadoConnection(
self._connection_parameters,
on_open_callback=future.set_result,
on_open_error_callback=lambda *a: future.set_exception(ConnectionError(a)),
on_close_callback=self._on_close,
custom_ioloop=self.io_loop
)
log.info(
'PikaClient: Trying to connect to rabbitmq://%s:%s/%s, Object: %r',
self._connection_parameters.host,
self._connection_parameters.port,
self._connection_parameters.virtual_host,
self
)
return future
def with_timeout(timeout, future, io_loop=None):
"""Wraps a `.Future` in a timeout.
"""
result = Future()
chain_future(future, result)
if io_loop is None:
io_loop = IOLoop.current()
timeout_handle = io_loop.add_timeout(
timeout,
lambda: result.set_exception(HostConnectionTimeout("Timeout")))
if isinstance(future, Future):
# We know this future will resolve on the IOLoop, so we don't
# need the extra thread-safety of IOLoop.add_future (and we also
# don't care about StackContext here.
future.add_done_callback(
lambda future: io_loop.remove_timeout(timeout_handle))
else:
# concurrent.futures.Futures may resolve on any thread, so we
# need to route them back to the IOLoop.
io_loop.add_future(
future, lambda future: io_loop.remove_timeout(timeout_handle))
return result
class HTTP1Connection(httputil.HTTPConnection):
"""Implements the HTTP/1.x protocol.
This class can be on its own for clients, or via `HTTP1ServerConnection`
for servers.
"""
def __init__(self, stream, is_client, params=None, context=None):
"""
:arg stream: an `.IOStream`
:arg bool is_client: client or server
:arg params: a `.HTTP1ConnectionParameters` instance or ``None``
:arg context: an opaque application-defined object that can be accessed
as ``connection.context``.
"""
self.is_client = is_client
self.stream = stream
if params is None:
params = HTTP1ConnectionParameters()
self.params = params
self.context = context
self.no_keep_alive = params.no_keep_alive
# The body limits can be altered by the delegate, so save them
# here instead of just referencing self.params later.
self._max_body_size = (self.params.max_body_size or
self.stream.max_buffer_size)
self._body_timeout = self.params.body_timeout
# _write_finished is set to True when finish() has been called,
# i.e. there will be no more data sent. Data may still be in the
# stream's write buffer.
self._write_finished = False
# True when we have read the entire incoming body.
self._read_finished = False
# _finish_future resolves when all data has been written and flushed
# to the IOStream.
self._finish_future = Future()
# If true, the connection should be closed after this request
# (after the response has been written in the server side,
# and after it has been read in the client)
self._disconnect_on_finish = False
self._clear_callbacks()
self.stream.set_close_callback(self._on_connection_close)
# Save the start lines after we read or write them; they
# affect later processing (e.g. 304 responses and HEAD methods
# have content-length but no bodies)
self._request_start_line = None
self._response_start_line = None
self._request_headers = None
# True if we are writing output with chunked encoding.
self._chunking_output = None
# While reading a body with a content-length, this is the
# amount left to read.
self._expected_content_remaining = None
def read_response(self, delegate):
"""Read a single HTTP response.
Typical client-mode usage is to write a request using `write_headers`,
`write`, and `finish`, and then call ``read_response``.
:arg delegate: a `.HTTPMessageDelegate`
Returns a `.Future` that resolves to None after the full response has
been read.
"""
if self.params.use_gzip:
delegate = _GzipMessageDelegate(delegate, self.params.chunk_size)
return self._read_message(delegate)
@gen.coroutine
def _read_message(self, delegate):
need_delegate_close = False
try:
header_future = self.stream.read_until_regex(
b"\r?\n\r?\n",
max_bytes=self.params.max_header_size)
if self.params.header_timeout is None:
header_data = yield header_future
else:
try:
header_data = yield gen.with_timeout(
self.stream.io_loop.time() + self.params.header_timeout,
header_future,
io_loop=self.stream.io_loop)
except gen.TimeoutError:
self.close()
raise gen.Return(False)
start_line, headers = self._parse_headers(header_data)
if self.is_client:
start_line = httputil.parse_response_start_line(start_line)
self._response_start_line = start_line
else:
start_line = httputil.parse_request_start_line(start_line)
self._request_start_line = start_line
self._request_headers = headers
self._disconnect_on_finish = not self._can_keep_alive(
start_line, headers)
need_delegate_close = True
header_future = delegate.headers_received(start_line, headers)
if header_future is not None:
yield header_future
if self.stream is None:
# We've been detached.
need_delegate_close = False
raise gen.Return(False)
skip_body = False
if self.is_client:
if (self._request_start_line is not None and
self._request_start_line.method == 'HEAD'):
skip_body = True
code = start_line.code
if code == 304:
skip_body = True
if code >= 100 and code < 200:
# TODO: client delegates will get headers_received twice
# in the case of a 100-continue. Document or change?
yield self._read_message(delegate)
else:
if (headers.get("Expect") == "100-continue" and
not self._write_finished):
self.stream.write(b"HTTP/1.1 100 (Continue)\r\n\r\n")
if not skip_body:
body_future = self._read_body(headers, delegate)
if body_future is not None:
if self._body_timeout is None:
yield body_future
else:
try:
yield gen.with_timeout(
self.stream.io_loop.time() + self._body_timeout,
body_future, self.stream.io_loop)
except gen.TimeoutError:
gen_log.info("Timeout reading body from %s",
self.context)
self.stream.close()
raise gen.Return(False)
self._read_finished = True
if not self._write_finished or self.is_client:
need_delegate_close = False
delegate.finish()
yield self._finish_future
if self.is_client and self._disconnect_on_finish:
self.close()
if self.stream is None:
raise gen.Return(False)
except httputil.HTTPInputException as e:
gen_log.info("Malformed HTTP message from %s: %s",
self.context, e)
self.close()
raise gen.Return(False)
finally:
if need_delegate_close:
delegate.on_connection_close()
self._clear_callbacks()
raise gen.Return(True)
def _clear_callbacks(self):
"""Clears the callback attributes.
This allows the request handler to be garbage collected more
quickly in CPython by breaking up reference cycles.
"""
self._write_callback = None
self._write_future = None
self._close_callback = None
def set_close_callback(self, callback):
"""Sets a callback that will be run when the connection is closed.
.. deprecated:: 3.3
Use `.HTTPMessageDelegate.on_connection_close` instead.
"""
self._close_callback = stack_context.wrap(callback)
def _on_connection_close(self):
if self._close_callback is not None:
callback = self._close_callback
self._close_callback = None
callback()
if not self._finish_future.done():
self._finish_future.set_result(None)
self._clear_callbacks()
def close(self):
self.stream.close()
self._clear_callbacks()
def detach(self):
"""Take control of the underlying stream.
Returns the underlying `.IOStream` object and stops all further
HTTP processing. May only be called during
`.HTTPMessageDelegate.headers_received`. Intended for implementing
protocols like websockets that tunnel over an HTTP handshake.
"""
stream = self.stream
self.stream = None
return stream
def set_body_timeout(self, timeout):
"""Sets the body timeout for a single request.
Overrides the value from `.HTTP1ConnectionParameters`.
"""
self._body_timeout = timeout
def set_max_body_size(self, max_body_size):
"""Sets the body size limit for a single request.
Overrides the value from `.HTTP1ConnectionParameters`.
"""
self._max_body_size = max_body_size
def write_headers(self, start_line, headers, chunk=None, callback=None):
"""Implements `.HTTPConnection.write_headers`."""
if self.is_client:
self._request_start_line = start_line
# Client requests with a non-empty body must have either a
# Content-Length or a Transfer-Encoding.
self._chunking_output = (
start_line.method in ('POST', 'PUT', 'PATCH') and
'Content-Length' not in headers and
'Transfer-Encoding' not in headers)
else:
self._response_start_line = start_line
self._chunking_output = (
# TODO: should this use
# self._request_start_line.version or
# start_line.version?
self._request_start_line.version == 'HTTP/1.1' and
# 304 responses have no body (not even a zero-length body), and so
# should not have either Content-Length or Transfer-Encoding.
# headers.
start_line.code != 304 and
# No need to chunk the output if a Content-Length is specified.
'Content-Length' not in headers and
# Applications are discouraged from touching Transfer-Encoding,
# but if they do, leave it alone.
'Transfer-Encoding' not in headers)
# If a 1.0 client asked for keep-alive, add the header.
if (self._request_start_line.version == 'HTTP/1.0' and
(self._request_headers.get('Connection', '').lower()
== 'keep-alive')):
headers['Connection'] = 'Keep-Alive'
if self._chunking_output:
headers['Transfer-Encoding'] = 'chunked'
if (not self.is_client and
(self._request_start_line.method == 'HEAD' or
start_line.code == 304)):
self._expected_content_remaining = 0
elif 'Content-Length' in headers:
self._expected_content_remaining = int(headers['Content-Length'])
else:
self._expected_content_remaining = None
lines = [utf8("%s %s %s" % start_line)]
lines.extend([utf8(n) + b": " + utf8(v) for n, v in headers.get_all()])
for line in lines:
if b'\n' in line:
raise ValueError('Newline in header: ' + repr(line))
if self.stream.closed():
self._write_future = Future()
self._write_future.set_exception(iostream.StreamClosedError())
else:
if callback is not None:
self._write_callback = stack_context.wrap(callback)
else:
self._write_future = Future()
data = b"\r\n".join(lines) + b"\r\n\r\n"
if chunk:
data += self._format_chunk(chunk)
self.stream.write(data, self._on_write_complete)
return self._write_future
def _format_chunk(self, chunk):
if self._expected_content_remaining is not None:
self._expected_content_remaining -= len(chunk)
if self._expected_content_remaining < 0:
# Close the stream now to stop further framing errors.
self.stream.close()
raise httputil.HTTPOutputException(
"Tried to write more data than Content-Length")
if self._chunking_output and chunk:
# Don't write out empty chunks because that means END-OF-STREAM
# with chunked encoding
return utf8("%x" % len(chunk)) + b"\r\n" + chunk + b"\r\n"
else:
return chunk
def write(self, chunk, callback=None):
"""Implements `.HTTPConnection.write`.
For backwards compatibility is is allowed but deprecated to
skip `write_headers` and instead call `write()` with a
pre-encoded header block.
"""
if self.stream.closed():
self._write_future = Future()
self._write_future.set_exception(iostream.StreamClosedError())
else:
if callback is not None:
self._write_callback = stack_context.wrap(callback)
else:
self._write_future = Future()
self.stream.write(self._format_chunk(chunk),
self._on_write_complete)
return self._write_future
def finish(self):
"""Implements `.HTTPConnection.finish`."""
if (self._expected_content_remaining is not None and
self._expected_content_remaining != 0 and
not self.stream.closed()):
self.stream.close()
raise httputil.HTTPOutputException(
"Tried to write %d bytes less than Content-Length" %
self._expected_content_remaining)
if self._chunking_output:
if not self.stream.closed():
self.stream.write(b"0\r\n\r\n", self._on_write_complete)
self._write_finished = True
# If the app finished the request while we're still reading,
# divert any remaining data away from the delegate and
# close the connection when we're done sending our response.
# Closing the connection is the only way to avoid reading the
# whole input body.
if not self._read_finished:
self._disconnect_on_finish = True
# No more data is coming, so instruct TCP to send any remaining
# data immediately instead of waiting for a full packet or ack.
self.stream.set_nodelay(True)
if not self.stream.writing():
self._finish_request()
def _on_write_complete(self):
if self._write_callback is not None:
callback = self._write_callback
self._write_callback = None
callback()
if self._write_future is not None:
future = self._write_future
self._write_future = None
future.set_result(None)
# _on_write_complete is enqueued on the IOLoop whenever the
# IOStream's write buffer becomes empty, but it's possible for
# another callback that runs on the IOLoop before it to
# simultaneously write more data and finish the request. If
# there is still data in the IOStream, a future
# _on_write_complete will be responsible for calling
# _finish_request.
if self._write_finished and not self.stream.writing():
self._finish_request()
def _can_keep_alive(self, start_line, headers):
if self.params.no_keep_alive:
return False
connection_header = headers.get("Connection")
if connection_header is not None:
connection_header = connection_header.lower()
if start_line.version == "HTTP/1.1":
return connection_header != "close"
elif ("Content-Length" in headers
or start_line.method in ("HEAD", "GET")):
return connection_header == "keep-alive"
return False
def _finish_request(self):
self._clear_callbacks()
if not self.is_client and self._disconnect_on_finish:
self.close()
return
# Turn Nagle's algorithm back on, leaving the stream in its
# default state for the next request.
self.stream.set_nodelay(False)
if not self._finish_future.done():
self._finish_future.set_result(None)
def _parse_headers(self, data):
data = native_str(data.decode('latin1'))
eol = data.find("\r\n")
start_line = data[:eol]
try:
headers = httputil.HTTPHeaders.parse(data[eol:])
except ValueError:
# probably form split() if there was no ':' in the line
raise httputil.HTTPInputException("Malformed HTTP headers: %r" %
data[eol:100])
return start_line, headers
def _read_body(self, headers, delegate):
content_length = headers.get("Content-Length")
if content_length:
content_length = int(content_length)
if content_length > self._max_body_size:
raise httputil.HTTPInputException("Content-Length too long")
return self._read_fixed_body(content_length, delegate)
if headers.get("Transfer-Encoding") == "chunked":
return self._read_chunked_body(delegate)
if self.is_client:
return self._read_body_until_close(delegate)
return None
@gen.coroutine
def _read_fixed_body(self, content_length, delegate):
while content_length > 0:
body = yield self.stream.read_bytes(
min(self.params.chunk_size, content_length), partial=True)
content_length -= len(body)
if not self._write_finished or self.is_client:
yield gen.maybe_future(delegate.data_received(body))
@gen.coroutine
def _read_chunked_body(self, delegate):
# TODO: "chunk extensions" http://tools.ietf.org/html/rfc2616#section-3.6.1
total_size = 0
while True:
chunk_len = yield self.stream.read_until(b"\r\n", max_bytes=64)
chunk_len = int(chunk_len.strip(), 16)
if chunk_len == 0:
return
total_size += chunk_len
if total_size > self._max_body_size:
raise httputil.HTTPInputException("chunked body too large")
bytes_to_read = chunk_len
while bytes_to_read:
chunk = yield self.stream.read_bytes(
min(bytes_to_read, self.params.chunk_size), partial=True)
bytes_to_read -= len(chunk)
if not self._write_finished or self.is_client:
yield gen.maybe_future(
delegate.data_received(chunk))
# chunk ends with \r\n
crlf = yield self.stream.read_bytes(2)
assert crlf == b"\r\n"
@gen.coroutine
def _read_body_until_close(self, delegate):
body = yield self.stream.read_until_close()
if not self._write_finished or self.is_client:
delegate.data_received(body)
class _Connector(object):
"""A stateless implementation of the "Happy Eyeballs" algorithm.
"Happy Eyeballs" is documented in RFC6555 as the recommended practice
for when both IPv4 and IPv6 addresses are available.
In this implementation, we partition the addresses by family, and
make the first connection attempt to whichever address was
returned first by ``getaddrinfo``. If that connection fails or
times out, we begin a connection in parallel to the first address
of the other family. If there are additional failures we retry
with other addresses, keeping one connection attempt per family
in flight at a time.
http://tools.ietf.org/html/rfc6555
"""
def __init__(self, addrinfo, io_loop, connect):
self.io_loop = io_loop
self.connect = connect
self.future = Future()
self.timeout = None
self.last_error = None
self.remaining = len(addrinfo)
self.primary_addrs, self.secondary_addrs = self.split(addrinfo)
@staticmethod
def split(addrinfo):
"""Partition the ``addrinfo`` list by address family.
Returns two lists. The first list contains the first entry from
``addrinfo`` and all others with the same family, and the
second list contains all other addresses (normally one list will
be AF_INET and the other AF_INET6, although non-standard resolvers
may return additional families).
"""
primary = []
secondary = []
primary_af = addrinfo[0][0]
for af, addr in addrinfo:
if af == primary_af:
primary.append((af, addr))
else:
secondary.append((af, addr))
return primary, secondary
def start(self, timeout=_INITIAL_CONNECT_TIMEOUT):
self.try_connect(iter(self.primary_addrs))
self.set_timout(timeout)
return self.future
def try_connect(self, addrs):
try:
af, addr = next(addrs)
except StopIteration:
# We've reached the end of our queue, but the other queue
# might still be working. Send a final error on the future
# only when both queues are finished.
if self.remaining == 0 and not self.future.done():
self.future.set_exception(self.last_error or
IOError("connection failed"))
return
future = self.connect(af, addr)
future.add_done_callback(functools.partial(self.on_connect_done,
addrs, af, addr))
def on_connect_done(self, addrs, af, addr, future):
self.remaining -= 1
try:
stream = future.result()
except Exception as e:
if self.future.done():
return
# Error: try again (but remember what happened so we have an
# error to raise in the end)
self.last_error = e
self.try_connect(addrs)
if self.timeout is not None:
# If the first attempt failed, don't wait for the
# timeout to try an address from the secondary queue.
self.on_timeout()
return
self.clear_timeout()
if self.future.done():
# This is a late arrival; just drop it.
stream.close()
else:
self.future.set_result((af, addr, stream))
def set_timout(self, timeout):
self.timeout = self.io_loop.add_timeout(self.io_loop.time() + timeout,
self.on_timeout)
def on_timeout(self):
self.timeout = None
self.try_connect(iter(self.secondary_addrs))
def clear_timeout(self):
if self.timeout is not None:
self.io_loop.remove_timeout(self.timeout)
def mkfuture_exception(exception):
future = Future()
future.set_exception(exception)
return future
class WebSocketClientConnection(simple_httpclient._HTTPConnection):
"""WebSocket client connection.
This class should not be instantiated directly; use the
`websocket_connect` function instead.
"""
def __init__(self, request, on_message_callback=None,
compression_options=None, ping_interval=None, ping_timeout=None,
max_message_size=None):
self.compression_options = compression_options
self.connect_future = Future()
self.protocol = None
self.read_future = None
self.read_queue = collections.deque()
self.key = base64.b64encode(os.urandom(16))
self._on_message_callback = on_message_callback
self.close_code = self.close_reason = None
self.ping_interval = ping_interval
self.ping_timeout = ping_timeout
self.max_message_size = max_message_size
scheme, sep, rest = request.url.partition(':')
scheme = {'ws': 'http', 'wss': 'https'}[scheme]
request.url = scheme + sep + rest
request.headers.update({
'Upgrade': 'websocket',
'Connection': 'Upgrade',
'Sec-WebSocket-Key': self.key,
'Sec-WebSocket-Version': '13',
})
if self.compression_options is not None:
# Always offer to let the server set our max_wbits (and even though
# we don't offer it, we will accept a client_no_context_takeover
# from the server).
# TODO: set server parameters for deflate extension
# if requested in self.compression_options.
request.headers['Sec-WebSocket-Extensions'] = (
'permessage-deflate; client_max_window_bits')
self.tcp_client = TCPClient()
super(WebSocketClientConnection, self).__init__(
None, request, lambda: None, self._on_http_response,
104857600, self.tcp_client, 65536, 104857600)
def close(self, code=None, reason=None):
"""Closes the websocket connection.
``code`` and ``reason`` are documented under
`WebSocketHandler.close`.
.. versionadded:: 3.2
.. versionchanged:: 4.0
Added the ``code`` and ``reason`` arguments.
"""
if self.protocol is not None:
self.protocol.close(code, reason)
self.protocol = None
def on_connection_close(self):
if not self.connect_future.done():
self.connect_future.set_exception(StreamClosedError())
self.on_message(None)
self.tcp_client.close()
super(WebSocketClientConnection, self).on_connection_close()
def _on_http_response(self, response):
if not self.connect_future.done():
if response.error:
self.connect_future.set_exception(response.error)
else:
self.connect_future.set_exception(WebSocketError(
"Non-websocket response"))
def headers_received(self, start_line, headers):
if start_line.code != 101:
return super(WebSocketClientConnection, self).headers_received(
start_line, headers)
self.headers = headers
self.protocol = self.get_websocket_protocol()
self.protocol._process_server_headers(self.key, self.headers)
self.protocol.start_pinging()
self.protocol._receive_frame()
if self._timeout is not None:
self.io_loop.remove_timeout(self._timeout)
self._timeout = None
self.stream = self.connection.detach()
self.stream.set_close_callback(self.on_connection_close)
# Once we've taken over the connection, clear the final callback
# we set on the http request. This deactivates the error handling
# in simple_httpclient that would otherwise interfere with our
# ability to see exceptions.
self.final_callback = None
future_set_result_unless_cancelled(self.connect_future, self)
def write_message(self, message, binary=False):
"""Sends a message to the WebSocket server."""
return self.protocol.write_message(message, binary)
def read_message(self, callback=None):
"""Reads a message from the WebSocket server.
If on_message_callback was specified at WebSocket
initialization, this function will never return messages
Returns a future whose result is the message, or None
if the connection is closed. If a callback argument
is given it will be called with the future when it is
ready.
"""
assert self.read_future is None
future = Future()
if self.read_queue:
future_set_result_unless_cancelled(future, self.read_queue.popleft())
else:
self.read_future = future
if callback is not None:
self.io_loop.add_future(future, callback)
return future
def on_message(self, message):
if self._on_message_callback:
self._on_message_callback(message)
elif self.read_future is not None:
future_set_result_unless_cancelled(self.read_future, message)
self.read_future = None
else:
self.read_queue.append(message)
def on_pong(self, data):
pass
def on_ping(self, data):
pass
def get_websocket_protocol(self):
return WebSocketProtocol13(self, mask_outgoing=True,
compression_options=self.compression_options)
class WebSocketClientConnection(simple_httpclient._HTTPConnection):
"""WebSocket client connection."""
def __init__(self, io_loop, request):
self.connect_future = Future()
self.read_future = None
self.read_queue = collections.deque()
self.key = base64.b64encode(os.urandom(16))
scheme, sep, rest = request.url.partition(':')
scheme = {'ws': 'http', 'wss': 'https'}[scheme]
request.url = scheme + sep + rest
request.headers.update({
'Upgrade': 'websocket',
'Connection': 'Upgrade',
'Sec-WebSocket-Key': self.key,
'Sec-WebSocket-Version': '13',
})
self.resolver = Resolver(io_loop=io_loop)
super(WebSocketClientConnection, self).__init__(
io_loop, None, request, lambda: None, self._on_http_response,
104857600, self.resolver)
def _on_close(self):
self.on_message(None)
self.resolver.close()
def _on_http_response(self, response):
if not self.connect_future.done():
if response.error:
self.connect_future.set_exception(response.error)
else:
self.connect_future.set_exception(WebSocketError(
"Non-websocket response"))
def _handle_1xx(self, code):
assert code == 101
assert self.headers['Upgrade'].lower() == 'websocket'
assert self.headers['Connection'].lower() == 'upgrade'
accept = WebSocketProtocol13.compute_accept_value(self.key)
assert self.headers['Sec-Websocket-Accept'] == accept
self.protocol = WebSocketProtocol13(self, mask_outgoing=True)
self.protocol._receive_frame()
if self._timeout is not None:
self.io_loop.remove_timeout(self._timeout)
self._timeout = None
self.connect_future.set_result(self)
def write_message(self, message, binary=False):
"""Sends a message to the WebSocket server."""
self.protocol.write_message(message, binary)
def read_message(self, callback=None):
"""Reads a message from the WebSocket server.
Returns a future whose result is the message, or None
if the connection is closed. If a callback argument
is given it will be called with the future when it is
ready.
"""
assert self.read_future is None
future = Future()
if self.read_queue:
future.set_result(self.read_queue.popleft())
else:
self.read_future = future
if callback is not None:
self.io_loop.add_future(future, callback)
return future
def on_message(self, message):
if self.read_future is not None:
self.read_future.set_result(message)
self.read_future = None
else:
self.read_queue.append(message)
def on_pong(self, data):
pass
def failed(e) -> 'promise.Promise':
f = Future()
f.set_exception(e)
return Promise(f)
class AsyncGroup:
"""
Grouping of several async requests and final callback in such way that final callback is invoked
after the last request is finished.
If any callback throws an exception, all pending callbacks would be aborted and finish_cb
would not be automatically called.
"""
def __init__(self, finish_cb, name=None):
self._counter = 0
self._finish_cb = finish_cb
self._finished = False
self._name = name
self._future = Future()
self._start_time = time.time()
def abort(self):
async_logger.info('aborting %s', self)
self._finished = True
if not self._future.done():
self._future.set_exception(AbortAsyncGroup())
def finish(self):
if self._finished:
async_logger.warning('trying to finish already finished %s', self)
return
self._finished = True
self._future.set_result(None)
try:
self._finish_cb()
finally:
# prevent possible cycle references
self._finish_cb = None
def try_finish(self):
if self._counter == 0:
self.finish()
def try_finish_async(self):
"""Executes finish_cb in next IOLoop iteration"""
if self._counter == 0:
IOLoop.current().add_callback(self.finish)
def _inc(self):
if self._finished:
async_logger.info('ignoring adding callback in %s', self)
raise AbortAsyncGroup()
self._counter += 1
def _dec(self):
self._counter -= 1
def add(self, intermediate_cb):
self._inc()
@wraps(intermediate_cb)
def new_cb(*args, **kwargs):
if self._finished:
async_logger.info('ignoring executing callback in %s', self)
return
try:
self._dec()
intermediate_cb(*args, **kwargs)
except Exception:
self.abort()
raise
self.try_finish()
return new_cb
def add_notification(self):
self._inc()
def new_cb(*args, **kwargs):
self._dec()
self.try_finish()
return new_cb
def add_future(self, future):
IOLoop.current().add_future(future, self.add_notification())
return future
def get_finish_future(self):
return self._future
def __str__(self):
return f'AsyncGroup(name={self._name}, finished={self._finished})'
