class TCPClientTest(AsyncTestCase):
def setUp(self):
super(TCPClientTest, self).setUp()
self.server = None
self.client = TCPClient()
def start_server(self, family):
if family == socket.AF_UNSPEC and "TRAVIS" in os.environ:
self.skipTest("dual-stack servers often have port conflicts on travis")
self.server = TestTCPServer(family)
return self.server.port
def stop_server(self):
if self.server is not None:
self.server.stop()
self.server = None
def tearDown(self):
self.client.close()
self.stop_server()
super(TCPClientTest, self).tearDown()
def skipIf10.0.0.7V4(self):
# The port used here doesn't matter, but some systems require it
# to be non-zero if we do not also pass AI_PASSIVE.
addrinfo = self.io_loop.run_sync(lambda: Resolver().resolve("10.0.0.7", 80))
families = set(addr[0] for addr in addrinfo)
if socket.AF_INET6 not in families:
self.skipTest("10.0.0.7 does not resolve to ipv6")
class CircleClient(object):
def __init__(self, endpoint=DEFAULT_ENDPOINT_DEALER, timeout=5.0):
self.endpoint = endpoint
self._id = cast_bytes(uuid.uuid4().hex)
self.timeout = timeout
self.stream = None
self.endpoint = endpoint
self.client = TCPClient()
def stop(self):
self.client.close()
def send_message(self, command, **props):
return self.call(make_message(command, **props))
def call(self, cmd):
result = IOLoop.instance().run_sync(lambda: self._call(cmd))
return result
@gen.coroutine
def _call(self, cmd):
if isinstance(cmd, basestring):
raise DeprecationWarning('call() takes a mapping')
call_id = uuid.uuid4().hex
cmd['id'] = call_id
host, port = self.endpoint.split(':')
try:
cmd = json.dumps(cmd)
self.stream = yield gen_timeout(self.timeout,
self.client.connect(host, port))
yield gen_timeout(self.timeout, self.stream.write(cmd + MSG_END))
except StreamClosedError:
raise CallError("Can't connect circled. Maybe it is closed.")
except gen.TimeoutError:
raise CallError('Connect timed out ({} seconds).'.format(
self.timeout))
except ValueError as e:
raise CallError(str(e))
while True:
try:
msg = yield gen_timeout(self.timeout,
self.stream.read_until(MSG_END))
msg = rstrip(msg, MSG_END)
res = json.loads(msg)
if 'id' in res and res['id'] not in (call_id, None):
# we got the wrong message
continue
raise gen.Return(res)
except gen.TimeoutError:
raise CallError('Run timed out ({} seconds).'.format(
self.timeout))
except ValueError as e:
raise CallError(str(e))
class TCPClientTest(AsyncTestCase):
def setUp(self):
super(TCPClientTest, self).setUp()
self.server = None
self.client = TCPClient()
def start_server(self, family):
if family == socket.AF_UNSPEC and 'TRAVIS' in os.environ:
self.skipTest("dual-stack servers often have port conflicts on travis")
self.server = TestTCPServer(family)
return self.server.port
def stop_server(self):
if self.server is not None:
self.server.stop()
self.server = None
def tearDown(self):
self.client.close()
self.stop_server()
super(TCPClientTest, self).tearDown()
def skipIfLocalhostV4(self):
# The port used here doesn't matter, but some systems require it
# to be non-zero if we do not also pass AI_PASSIVE.
addrinfo = self.io_loop.run_sync(lambda: Resolver().resolve('localhost', 80))
families = set(addr[0] for addr in addrinfo)
if socket.AF_INET6 not in families:
self.skipTest("localhost does not resolve to ipv6")
@gen_test
def do_test_connect(self, family, host, source_ip=None, source_port=None):
port = self.start_server(family)
stream = yield self.client.connect(host, port,
source_ip=source_ip,
source_port=source_port)
server_stream = yield self.server.queue.get()
with closing(stream):
stream.write(b"hello")
data = yield server_stream.read_bytes(5)
self.assertEqual(data, b"hello")
def test_connect_ipv4_ipv4(self):
self.do_test_connect(socket.AF_INET, '127.0.0.1')
def test_connect_ipv4_dual(self):
self.do_test_connect(socket.AF_INET, 'localhost')
@skipIfNoIPv6
def test_connect_ipv6_ipv6(self):
self.skipIfLocalhostV4()
self.do_test_connect(socket.AF_INET6, '::1')
@skipIfNoIPv6
def test_connect_ipv6_dual(self):
self.skipIfLocalhostV4()
if Resolver.configured_class().__name__.endswith('TwistedResolver'):
self.skipTest('TwistedResolver does not support multiple addresses')
self.do_test_connect(socket.AF_INET6, 'localhost')
def test_connect_unspec_ipv4(self):
self.do_test_connect(socket.AF_UNSPEC, '127.0.0.1')
@skipIfNoIPv6
def test_connect_unspec_ipv6(self):
self.skipIfLocalhostV4()
self.do_test_connect(socket.AF_UNSPEC, '::1')
def test_connect_unspec_dual(self):
self.do_test_connect(socket.AF_UNSPEC, 'localhost')
@gen_test
def test_refused_ipv4(self):
cleanup_func, port = refusing_port()
self.addCleanup(cleanup_func)
with self.assertRaises(IOError):
yield self.client.connect('127.0.0.1', port)
def test_source_ip_fail(self):
'''
Fail when trying to use the source IP Address '8.8.8.8'.
'''
self.assertRaises(socket.error,
self.do_test_connect,
socket.AF_INET,
'127.0.0.1',
source_ip='8.8.8.8')
def test_source_ip_success(self):
'''
Success when trying to use the source IP Address '127.0.0.1'
'''
self.do_test_connect(socket.AF_INET, '127.0.0.1', source_ip='127.0.0.1')
@skipIfNonUnix
def test_source_port_fail(self):
'''
Fail when trying to use source port 1.
'''
self.assertRaises(socket.error,
self.do_test_connect,
socket.AF_INET,
'127.0.0.1',
source_port=1)
@gen_test
def test_connect_timeout(self):
timeout = 0.05
class TimeoutResolver(Resolver):
def resolve(self, *args, **kwargs):
return Future() # never completes
with self.assertRaises(TimeoutError):
yield TCPClient(resolver=TimeoutResolver()).connect(
'1.2.3.4', 12345, timeout=timeout)
class SimpleAsyncHTTPClient(AsyncHTTPClient):
"""Non-blocking HTTP client with no external dependencies.
This class implements an HTTP 1.1 client on top of Tornado's IOStreams.
It does not currently implement all applicable parts of the HTTP
specification, but it does enough to work with major web service APIs.
Some features found in the curl-based AsyncHTTPClient are not yet
supported. In particular, proxies are not supported, connections
are not reused, and callers cannot select the network interface to be
used.
"""
def initialize(self, io_loop, max_clients=10,
hostname_mapping=None, max_buffer_size=104857600,
resolver=None, defaults=None, max_header_size=None,
max_body_size=None):
"""Creates a AsyncHTTPClient.
Only a single AsyncHTTPClient instance exists per IOLoop
in order to provide limitations on the number of pending connections.
force_instance=True may be used to suppress this behavior.
max_clients is the number of concurrent requests that can be
in progress. Note that this arguments are only used when the
client is first created, and will be ignored when an existing
client is reused.
hostname_mapping is a dictionary mapping hostnames to IP addresses.
It can be used to make local DNS changes when modifying system-wide
settings like /etc/hosts is not possible or desirable (e.g. in
unittests).
max_buffer_size is the number of bytes that can be read by IOStream. It
defaults to 100mb.
"""
super(SimpleAsyncHTTPClient, self).initialize(io_loop,
defaults=defaults)
self.max_clients = max_clients
self.queue = collections.deque()
self.active = {}
self.waiting = {}
self.max_buffer_size = max_buffer_size
self.max_header_size = max_header_size
self.max_body_size = max_body_size
# TCPClient could create a Resolver for us, but we have to do it
# ourselves to support hostname_mapping.
if resolver:
self.resolver = resolver
self.own_resolver = False
else:
self.resolver = Resolver(io_loop=io_loop)
self.own_resolver = True
if hostname_mapping is not None:
self.resolver = OverrideResolver(resolver=self.resolver,
mapping=hostname_mapping)
self.tcp_client = TCPClient(resolver=self.resolver, io_loop=io_loop)
def close(self):
super(SimpleAsyncHTTPClient, self).close()
if self.own_resolver:
self.resolver.close()
self.tcp_client.close()
def fetch_impl(self, request, callback):
key = object()
self.queue.append((key, request, callback))
if not len(self.active) < self.max_clients:
timeout_handle = self.io_loop.add_timeout(
self.io_loop.time() + min(request.connect_timeout,
request.request_timeout),
functools.partial(self._on_timeout, key))
else:
timeout_handle = None
self.waiting[key] = (request, callback, timeout_handle)
self._process_queue()
if self.queue:
gen_log.debug("max_clients limit reached, request queued. "
"%d active, %d queued requests." % (
len(self.active), len(self.queue)))
def _process_queue(self):
with stack_context.NullContext():
while self.queue and len(self.active) < self.max_clients:
key, request, callback = self.queue.popleft()
if key not in self.waiting:
continue
self._remove_timeout(key)
self.active[key] = (request, callback)
release_callback = functools.partial(self._release_fetch, key)
self._handle_request(request, release_callback, callback)
def _connection_class(self):
return _HTTPConnection
def _handle_request(self, request, release_callback, final_callback):
self._connection_class()(
self.io_loop, self, request, release_callback,
final_callback, self.max_buffer_size, self.tcp_client,
self.max_header_size, self.max_body_size)
def _release_fetch(self, key):
del self.active[key]
self._process_queue()
def _remove_timeout(self, key):
if key in self.waiting:
request, callback, timeout_handle = self.waiting[key]
if timeout_handle is not None:
self.io_loop.remove_timeout(timeout_handle)
del self.waiting[key]
def _on_timeout(self, key):
request, callback, timeout_handle = self.waiting[key]
self.queue.remove((key, request, callback))
timeout_response = HTTPResponse(
request, 599, error=HTTPError(599, "Timeout"),
request_time=self.io_loop.time() - request.start_time)
self.io_loop.add_callback(callback, timeout_response)
del self.waiting[key]
class SimpleAsyncHTTPClient(AsyncHTTPClient):
"""Non-blocking HTTP client with no external dependencies.
This class implements an HTTP 1.1 client on top of Tornado's IOStreams.
Some features found in the curl-based AsyncHTTPClient are not yet
supported. In particular, proxies are not supported, connections
are not reused, and callers cannot select the network interface to be
used.
"""
def initialize(self, max_clients=10,
hostname_mapping=None, max_buffer_size=104857600,
resolver=None, defaults=None, max_header_size=None,
max_body_size=None):
"""Creates a AsyncHTTPClient.
Only a single AsyncHTTPClient instance exists per IOLoop
in order to provide limitations on the number of pending connections.
``force_instance=True`` may be used to suppress this behavior.
Note that because of this implicit reuse, unless ``force_instance``
is used, only the first call to the constructor actually uses
its arguments. It is recommended to use the ``configure`` method
instead of the constructor to ensure that arguments take effect.
``max_clients`` is the number of concurrent requests that can be
in progress; when this limit is reached additional requests will be
queued. Note that time spent waiting in this queue still counts
against the ``request_timeout``.
``hostname_mapping`` is a dictionary mapping hostnames to IP addresses.
It can be used to make local DNS changes when modifying system-wide
settings like ``/etc/hosts`` is not possible or desirable (e.g. in
unittests).
``max_buffer_size`` (default 100MB) is the number of bytes
that can be read into memory at once. ``max_body_size``
(defaults to ``max_buffer_size``) is the largest response body
that the client will accept. Without a
``streaming_callback``, the smaller of these two limits
applies; with a ``streaming_callback`` only ``max_body_size``
does.
.. versionchanged:: 4.2
Added the ``max_body_size`` argument.
"""
super(SimpleAsyncHTTPClient, self).initialize(defaults=defaults)
self.max_clients = max_clients
self.queue = collections.deque()
self.active = {}
self.waiting = {}
self.max_buffer_size = max_buffer_size
self.max_header_size = max_header_size
self.max_body_size = max_body_size
# TCPClient could create a Resolver for us, but we have to do it
# ourselves to support hostname_mapping.
if resolver:
self.resolver = resolver
self.own_resolver = False
else:
self.resolver = Resolver()
self.own_resolver = True
if hostname_mapping is not None:
self.resolver = OverrideResolver(resolver=self.resolver,
mapping=hostname_mapping)
self.tcp_client = TCPClient(resolver=self.resolver)
def close(self):
super(SimpleAsyncHTTPClient, self).close()
if self.own_resolver:
self.resolver.close()
self.tcp_client.close()
def fetch_impl(self, request, callback):
key = object()
self.queue.append((key, request, callback))
if not len(self.active) < self.max_clients:
timeout_handle = self.io_loop.add_timeout(
self.io_loop.time() + min(request.connect_timeout,
request.request_timeout),
functools.partial(self._on_timeout, key, "in request queue"))
else:
timeout_handle = None
self.waiting[key] = (request, callback, timeout_handle)
self._process_queue()
if self.queue:
gen_log.debug("max_clients limit reached, request queued. "
"%d active, %d queued requests." % (
len(self.active), len(self.queue)))
def _process_queue(self):
with stack_context.NullContext():
while self.queue and len(self.active) < self.max_clients:
key, request, callback = self.queue.popleft()
if key not in self.waiting:
continue
self._remove_timeout(key)
self.active[key] = (request, callback)
release_callback = functools.partial(self._release_fetch, key)
self._handle_request(request, release_callback, callback)
def _connection_class(self):
return _HTTPConnection
def _handle_request(self, request, release_callback, final_callback):
self._connection_class()(
self, request, release_callback,
final_callback, self.max_buffer_size, self.tcp_client,
self.max_header_size, self.max_body_size)
def _release_fetch(self, key):
del self.active[key]
self._process_queue()
def _remove_timeout(self, key):
if key in self.waiting:
request, callback, timeout_handle = self.waiting[key]
if timeout_handle is not None:
self.io_loop.remove_timeout(timeout_handle)
del self.waiting[key]
def _on_timeout(self, key, info=None):
"""Timeout callback of request.
Construct a timeout HTTPResponse when a timeout occurs.
:arg object key: A simple object to mark the request.
:info string key: More detailed timeout information.
"""
request, callback, timeout_handle = self.waiting[key]
self.queue.remove((key, request, callback))
error_message = "Timeout {0}".format(info) if info else "Timeout"
timeout_response = HTTPResponse(
request, 599, error=HTTPTimeoutError(error_message),
request_time=self.io_loop.time() - request.start_time)
self.io_loop.add_callback(callback, timeout_response)
del self.waiting[key]
class WebSocketClientConnection(simple_httpclient._HTTPConnection):
"""WebSocket client connection.
This class should not be instantiated directly; use the
`websocket_connect` function instead.
"""
def __init__(self, io_loop, request):
self.connect_future = TracebackFuture()
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.tcp_client = TCPClient(io_loop=io_loop)
super(WebSocketClientConnection, self).__init__(
io_loop, None, request, lambda: None, self._on_http_response,
104857600, self.tcp_client, 65536)
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
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.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
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 = TracebackFuture()
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
class WebSocketClientConnection(simple_httpclient._HTTPConnection):
"""WebSocket client connection.
This class should not be instantiated directly; use the
`websocket_connect` function instead.
"""
def __init__(self, io_loop, request, on_message_callback=None,
compression_options=None):
self.compression_options = compression_options
self.connect_future = TracebackFuture()
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
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(io_loop=io_loop)
super(WebSocketClientConnection, self).__init__(
io_loop, 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._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
self.connect_future.set_result(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 = TracebackFuture()
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._on_message_callback:
self._on_message_callback(message)
elif 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 get_websocket_protocol(self):
return WebSocketProtocol13(self, mask_outgoing=True,
compression_options=self.compression_options)
class SimpleAsyncHTTPClient(AsyncHTTPClient):
"""Non-blocking HTTP client with no external dependencies.
This class implements an HTTP 1.1 client on top of Tornado's IOStreams.
Some features found in the curl-based AsyncHTTPClient are not yet
supported. In particular, proxies are not supported, connections
are not reused, and callers cannot select the network interface to be
used.
"""
def initialize(self,
io_loop,
max_clients=10,
hostname_mapping=None,
max_buffer_size=104857600,
resolver=None,
defaults=None,
max_header_size=None,
max_body_size=None):
"""Creates a AsyncHTTPClient.
Only a single AsyncHTTPClient instance exists per IOLoop
in order to provide limitations on the number of pending connections.
``force_instance=True`` may be used to suppress this behavior.
Note that because of this implicit reuse, unless ``force_instance``
is used, only the first call to the constructor actually uses
its arguments. It is recommended to use the ``configure`` method
instead of the constructor to ensure that arguments take effect.
``max_clients`` is the number of concurrent requests that can be
in progress; when this limit is reached additional requests will be
queued. Note that time spent waiting in this queue still counts
against the ``request_timeout``.
``hostname_mapping`` is a dictionary mapping hostnames to IP addresses.
It can be used to make local DNS changes when modifying system-wide
settings like ``/etc/hosts`` is not possible or desirable (e.g. in
unittests).
``max_buffer_size`` (default 100MB) is the number of bytes
that can be read into memory at once. ``max_body_size``
(defaults to ``max_buffer_size``) is the largest response body
that the client will accept. Without a
``streaming_callback``, the smaller of these two limits
applies; with a ``streaming_callback`` only ``max_body_size``
does.
.. versionchanged:: 4.2
Added the ``max_body_size`` argument.
"""
super(SimpleAsyncHTTPClient, self).initialize(io_loop,
defaults=defaults)
self.max_clients = max_clients
self.queue = collections.deque()
self.active = {}
self.waiting = {}
self.max_buffer_size = max_buffer_size
self.max_header_size = max_header_size
self.max_body_size = max_body_size
# TCPClient could create a Resolver for us, but we have to do it
# ourselves to support hostname_mapping.
if resolver:
self.resolver = resolver
self.own_resolver = False
else:
self.resolver = Resolver(io_loop=io_loop)
self.own_resolver = True
if hostname_mapping is not None:
self.resolver = OverrideResolver(resolver=self.resolver,
mapping=hostname_mapping)
self.tcp_client = TCPClient(resolver=self.resolver, io_loop=io_loop)
def close(self):
super(SimpleAsyncHTTPClient, self).close()
if self.own_resolver:
self.resolver.close()
self.tcp_client.close()
def fetch_impl(self, request, callback):
key = object()
self.queue.append((key, request, callback))
if not len(self.active) < self.max_clients:
timeout_handle = self.io_loop.add_timeout(
self.io_loop.time() +
min(request.connect_timeout, request.request_timeout),
functools.partial(self._on_timeout, key, "in request queue"))
else:
timeout_handle = None
self.waiting[key] = (request, callback, timeout_handle)
self._process_queue()
if self.queue:
gen_log.debug("max_clients limit reached, request queued. "
"%d active, %d queued requests." %
(len(self.active), len(self.queue)))
def _process_queue(self):
with stack_context.NullContext():
while self.queue and len(self.active) < self.max_clients:
key, request, callback = self.queue.popleft()
if key not in self.waiting:
continue
self._remove_timeout(key)
self.active[key] = (request, callback)
release_callback = functools.partial(self._release_fetch, key)
self._handle_request(request, release_callback, callback)
def _connection_class(self):
return _HTTPConnection
def _handle_request(self, request, release_callback, final_callback):
self._connection_class()(self.io_loop, self, request, release_callback,
final_callback, self.max_buffer_size,
self.tcp_client, self.max_header_size,
self.max_body_size)
def _release_fetch(self, key):
del self.active[key]
self._process_queue()
def _remove_timeout(self, key):
if key in self.waiting:
request, callback, timeout_handle = self.waiting[key]
if timeout_handle is not None:
self.io_loop.remove_timeout(timeout_handle)
del self.waiting[key]
def _on_timeout(self, key, info=None):
"""Timeout callback of request.
Construct a timeout HTTPResponse when a timeout occurs.
:arg object key: A simple object to mark the request.
:info string key: More detailed timeout information.
"""
request, callback, timeout_handle = self.waiting[key]
self.queue.remove((key, request, callback))
error_message = "Timeout {0}".format(info) if info else "Timeout"
timeout_response = HTTPResponse(request,
599,
error=HTTPError(599, error_message),
request_time=self.io_loop.time() -
request.start_time)
self.io_loop.add_callback(callback, timeout_response)
del self.waiting[key]
class TCPClientTest(AsyncTestCase):
def setUp(self):
super(TCPClientTest, self).setUp()
self.server = None
self.client = TCPClient()
def start_server(self, family):
self.server = TestTCPServer(family)
return self.server.port
def stop_server(self):
if self.server is not None:
self.server.stop()
self.server = None
def tearDown(self):
self.client.close()
self.stop_server()
super(TCPClientTest, self).tearDown()
def skipIfLocalhostV4(self):
Resolver().resolve('localhost', 0, callback=self.stop)
addrinfo = self.wait()
families = set(addr[0] for addr in addrinfo)
if socket.AF_INET6 not in families:
self.skipTest("localhost does not resolve to ipv6")
@gen_test
def do_test_connect(self, family, host):
port = self.start_server(family)
stream = yield self.client.connect(host, port)
with closing(stream):
stream.write(b"hello")
data = yield self.server.streams[0].read_bytes(5)
self.assertEqual(data, b"hello")
def test_connect_ipv4_ipv4(self):
self.do_test_connect(socket.AF_INET, '127.0.0.1')
def test_connect_ipv4_dual(self):
self.do_test_connect(socket.AF_INET, 'localhost')
@skipIfNoIPv6
def test_connect_ipv6_ipv6(self):
self.skipIfLocalhostV4()
self.do_test_connect(socket.AF_INET6, '::1')
@skipIfNoIPv6
def test_connect_ipv6_dual(self):
self.skipIfLocalhostV4()
if Resolver.configured_class().__name__.endswith('TwistedResolver'):
self.skipTest('TwistedResolver does not support multiple addresses')
self.do_test_connect(socket.AF_INET6, 'localhost')
def test_connect_unspec_ipv4(self):
self.do_test_connect(socket.AF_UNSPEC, '127.0.0.1')
@skipIfNoIPv6
def test_connect_unspec_ipv6(self):
self.skipIfLocalhostV4()
self.do_test_connect(socket.AF_UNSPEC, '::1')
def test_connect_unspec_dual(self):
self.do_test_connect(socket.AF_UNSPEC, 'localhost')
@gen_test
def test_refused_ipv4(self):
sock, port = bind_unused_port()
sock.close()
with self.assertRaises(IOError):
yield self.client.connect('127.0.0.1', port)
class TCPClientTest(AsyncTestCase):
def setUp(self):
super().setUp()
self.server = None
self.client = TCPClient()
def start_server(self, family):
if family == socket.AF_UNSPEC and "TRAVIS" in os.environ:
self.skipTest("dual-stack servers often have port conflicts on travis")
self.server = TestTCPServer(family)
return self.server.port
def stop_server(self):
if self.server is not None:
self.server.stop()
self.server = None
def tearDown(self):
self.client.close()
self.stop_server()
super().tearDown()
def skipIfLocalhostV4(self):
# The port used here doesn't matter, but some systems require it
# to be non-zero if we do not also pass AI_PASSIVE.
addrinfo = self.io_loop.run_sync(lambda: Resolver().resolve("localhost", 80))
families = set(addr[0] for addr in addrinfo)
if socket.AF_INET6 not in families:
self.skipTest("localhost does not resolve to ipv6")
@gen_test
def do_test_connect(self, family, host, source_ip=None, source_port=None):
port = self.start_server(family)
stream = yield self.client.connect(
host, port, source_ip=source_ip, source_port=source_port
)
assert self.server is not None
server_stream = yield self.server.queue.get()
with closing(stream):
stream.write(b"hello")
data = yield server_stream.read_bytes(5)
self.assertEqual(data, b"hello")
def test_connect_ipv4_ipv4(self):
self.do_test_connect(socket.AF_INET, "127.0.0.1")
def test_connect_ipv4_dual(self):
self.do_test_connect(socket.AF_INET, "localhost")
@skipIfNoIPv6
def test_connect_ipv6_ipv6(self):
self.skipIfLocalhostV4()
self.do_test_connect(socket.AF_INET6, "::1")
@skipIfNoIPv6
def test_connect_ipv6_dual(self):
self.skipIfLocalhostV4()
if Resolver.configured_class().__name__.endswith("TwistedResolver"):
self.skipTest("TwistedResolver does not support multiple addresses")
self.do_test_connect(socket.AF_INET6, "localhost")
def test_connect_unspec_ipv4(self):
self.do_test_connect(socket.AF_UNSPEC, "127.0.0.1")
@skipIfNoIPv6
def test_connect_unspec_ipv6(self):
self.skipIfLocalhostV4()
self.do_test_connect(socket.AF_UNSPEC, "::1")
def test_connect_unspec_dual(self):
self.do_test_connect(socket.AF_UNSPEC, "localhost")
@gen_test
def test_refused_ipv4(self):
cleanup_func, port = refusing_port()
self.addCleanup(cleanup_func)
with self.assertRaises(IOError):
yield self.client.connect("127.0.0.1", port)
def test_source_ip_fail(self):
"""Fail when trying to use the source IP Address '8.8.8.8'."""
self.assertRaises(
socket.error,
self.do_test_connect,
socket.AF_INET,
"127.0.0.1",
source_ip="8.8.8.8",
)
def test_source_ip_success(self):
"""Success when trying to use the source IP Address '127.0.0.1'."""
self.do_test_connect(socket.AF_INET, "127.0.0.1", source_ip="127.0.0.1")
@skipIfNonUnix
def test_source_port_fail(self):
"""Fail when trying to use source port 1."""
if getpass.getuser() == "root":
# Root can use any port so we can't easily force this to fail.
# This is mainly relevant for docker.
self.skipTest("running as root")
self.assertRaises(
socket.error,
self.do_test_connect,
socket.AF_INET,
"127.0.0.1",
source_port=1,
)
@gen_test
def test_connect_timeout(self):
timeout = 0.05
class TimeoutResolver(Resolver):
def resolve(self, *args, **kwargs):
return Future() # never completes
with self.assertRaises(TimeoutError):
yield TCPClient(resolver=TimeoutResolver()).connect(
"1.2.3.4", 12345, timeout=timeout
)
class TCPClientTest(AsyncTestCase):
def setUp(self):
super(TCPClientTest, self).setUp()
self.server = None
self.client = TCPClient()
def start_server(self, family):
if family == socket.AF_UNSPEC and 'TRAVIS' in os.environ:
self.skipTest(
"dual-stack servers often have port conflicts on travis")
self.server = TestTCPServer(family)
return self.server.port
def stop_server(self):
if self.server is not None:
self.server.stop()
self.server = None
def tearDown(self):
self.client.close()
self.stop_server()
super(TCPClientTest, self).tearDown()
def skipIfLocalhostV4(self):
# The port used here doesn't matter, but some systems require it
# to be non-zero if we do not also pass AI_PASSIVE.
Resolver().resolve('localhost', 80, callback=self.stop)
addrinfo = self.wait()
families = set(addr[0] for addr in addrinfo)
if socket.AF_INET6 not in families:
self.skipTest("localhost does not resolve to ipv6")
@gen_test
def do_test_connect(self, family, host, source_ip=None, source_port=None):
port = self.start_server(family)
stream = yield self.client.connect(host,
port,
source_ip=source_ip,
source_port=source_port)
server_stream = yield self.server.queue.get()
with closing(stream):
stream.write(b"hello")
data = yield server_stream.read_bytes(5)
self.assertEqual(data, b"hello")
def test_connect_ipv4_ipv4(self):
self.do_test_connect(socket.AF_INET, '127.0.0.1')
def test_connect_ipv4_dual(self):
self.do_test_connect(socket.AF_INET, 'localhost')
@skipIfNoIPv6
def test_connect_ipv6_ipv6(self):
self.skipIfLocalhostV4()
self.do_test_connect(socket.AF_INET6, '::1')
@skipIfNoIPv6
def test_connect_ipv6_dual(self):
self.skipIfLocalhostV4()
if Resolver.configured_class().__name__.endswith('TwistedResolver'):
self.skipTest(
'TwistedResolver does not support multiple addresses')
self.do_test_connect(socket.AF_INET6, 'localhost')
def test_connect_unspec_ipv4(self):
self.do_test_connect(socket.AF_UNSPEC, '127.0.0.1')
@skipIfNoIPv6
def test_connect_unspec_ipv6(self):
self.skipIfLocalhostV4()
self.do_test_connect(socket.AF_UNSPEC, '::1')
def test_connect_unspec_dual(self):
self.do_test_connect(socket.AF_UNSPEC, 'localhost')
@gen_test
def test_refused_ipv4(self):
cleanup_func, port = refusing_port()
self.addCleanup(cleanup_func)
with self.assertRaises(IOError):
yield self.client.connect('127.0.0.1', port)
def test_source_ip_fail(self):
'''
Fail when trying to use the source IP Address '8.8.8.8'.
'''
self.assertRaises(socket.error,
self.do_test_connect,
socket.AF_INET,
'127.0.0.1',
source_ip='8.8.8.8')
def test_source_ip_success(self):
'''
Success when trying to use the source IP Address '127.0.0.1'
'''
self.do_test_connect(socket.AF_INET,
'127.0.0.1',
source_ip='127.0.0.1')
@skipIfNonUnix
def test_source_port_fail(self):
'''
Fail when trying to use source port 1.
'''
self.assertRaises(socket.error,
self.do_test_connect,
socket.AF_INET,
'127.0.0.1',
source_port=1)
@gen_test
def test_connect_timeout(self):
timeout = 0.05
class TimeoutResolver(Resolver):
def resolve(self, *args, **kwargs):
return Future() # never completes
with self.assertRaises(TimeoutError):
yield TCPClient(resolver=TimeoutResolver()).connect(
'1.2.3.4', 12345, timeout=timeout)
class Emitter(object):
def __init__(self, port, n=1000, values=1, duration=3.0):
self.port = port
self.n = n
self.values = values
self.duration = duration
self.message = self.hello
self.i = 0
self.pcb = None
self.client = None
def start(self):
self.client = TCPClient()
self.pcb = PeriodicCallback(self.send, 1000.0 / self.n)
self.pcb.start()
IOLoop.current().call_later(self.duration + 0.5, self.stop)
IOLoop.current().start()
IOLoop.clear_current()
def stop(self):
if self.pcb is not None:
self.pcb.stop()
if self.client is not None:
self.client.close()
IOLoop.current().stop()
@gen.coroutine
def send(self):
if self.i >= self.duration * self.n * self.values:
self.pcb.stop()
return
try:
stream = yield self.client.connect('127.0.0.1', self.port)
with closing(stream):
messages = b''.join(self.message() for _ in range(self.values))
stream.write(messages)
self.i += self.values
except StreamClosedError:
return
def hello(self):
return b'hello\n'
def r(self):
s = random.randint(1, 10)
v = s / 10.0 + (1.5 - s / 10.0) * random.random()
return (s, v)
def text(self):
return 'sensor{}|{}\n'.format(*self.r()).encode('utf8')
def json(self):
s, v = self.r()
return (json.dumps({
'sensor{}'.format(s): v,
}) + '\n').encode('utf8')
def bello(self):
# 5 bytes
return b'bello'
def struct(self):
# 8 bytes
return struct.pack('If', *self.r())
class WebSocketClientConnection(simple_httpclient._HTTPConnection):
"""WebSocket 客户端连接
这个类不应当直接被实例化, 请使用 `websocket_connect`
"""
def __init__(self, io_loop, request, on_message_callback=None,
compression_options=None):
self.compression_options = compression_options
self.connect_future = TracebackFuture()
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
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(io_loop=io_loop)
super(WebSocketClientConnection, self).__init__(
io_loop, None, request, lambda: None, self._on_http_response,
104857600, self.tcp_client, 65536, 104857600)
def close(self, code=None, reason=None):
"""关闭 websocket 连接
``code`` 和 ``reason`` 的文档在 `WebSocketHandler.close` 下已给出.
.. versionadded:: 3.2
.. versionchanged:: 4.0
添加 ``code`` 和 ``reason`` 这两个参数
"""
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._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
self.connect_future.set_result(self)
def write_message(self, message, binary=False):
"""发送消息到 websocket 服务器."""
return self.protocol.write_message(message, binary)
def read_message(self, callback=None):
"""读取来自 WebSocket 服务器的消息.
如果在 WebSocket 初始化时指定了 on_message_callback ,那么这个方法永远不会返回消息
如果连接已经关闭,返回结果会是一个结果是 message 的 future 对象或者是 None.
如果 future 给出了回调参数, 这个参数将会在 future 完成时调用.
"""
assert self.read_future is None
future = TracebackFuture()
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._on_message_callback:
self._on_message_callback(message)
elif 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 get_websocket_protocol(self):
return WebSocketProtocol13(self, mask_outgoing=True,
compression_options=self.compression_options)
class XTCPClient(object):
def __init__(self,
io_loop=None,
max_clients=10,
max_buffer_size=None,
max_response_size=None):
self._io_loop = io_loop or IOLoop.instance()
self.max_clients = max_clients
self.max_buffer_size = max_buffer_size or 104857600 # 100M
self.max_response_size = max_response_size or 10 * 1024 * 1024 # 10M
self.queue = collections.deque()
self.active = {}
self.waiting = {}
self._client_closed = False
self.tcp_client = TCPClient(io_loop=self._io_loop)
def __del__(self):
self.close()
def close(self):
if not self._client_closed:
self._client_closed = True
self.tcp_client.close()
self._io_loop.close()
def acquire(self, request):
response = self._io_loop.run_sync(
functools.partial(self._acquire_by_request, request))
return response
def _acquire_by_request(self, request):
future = Future()
def _handle_response(success, response=None):
if success is True:
response_message = response
if request.user_request_callback is not None:
response_message = request.user_request_callback(response)
future.set_result(response_message)
else:
future.set_exc_info(response)
request.request_callback = _handle_response
self._acquire_loop_by_request(request)
return future
def _acquire_loop_by_request(self, request):
key = object()
self.queue.append((key, request))
if not len(self.active) < self.max_clients:
waiting_timeout_handle = self._io_loop.add_timeout(
self._io_loop.time() + self.request.waiting_timeout,
functools.partial(self._on_waiting_timeout, key))
else:
waiting_timeout_handle = None
self.waiting[key] = (request, waiting_timeout_handle)
self._process_queue()
if self.queue:
xtcp_logger.debug(
"max_clients limits reached. {} active, {} queued requests".
format(len(self.active), len(self.queue)))
def _on_waiting_timeout(self, key):
xtcp_logger.debug("_on_waiting_timeout : {}".format(self.waiting[key]))
request, callback, waiting_timeout_handle = self.waiting[key]
self.queue.remove((key, request, callback))
del self.waiting[key]
def _process_queue(self):
while self.queue and len(self.active) < self.max_clients:
key, request = self.queue.popleft()
if key not in self.waiting:
continue
self._remove_waiting_timeout_request(key)
self.active[key] = request
self._handle_request(request,
functools.partial(self._release_request, key))
def _handle_request(self, request, release_callback):
connection = _ClientConnection(self,
io_loop=self._io_loop,
request=request,
release_callback=release_callback)
connection.connect()
def _release_request(self, key):
del self.active[key]
self._process_queue()
def _remove_waiting_timeout_request(self, key):
if key in self.waiting:
_, timeout_handle = self.waiting[key]
if timeout_handle is not None:
self._io_loop.remove_timeout(timeout_handle)
del self.waiting[key]
class TCPClientTest(AsyncTestCase):
def setUp(self):
super(TCPClientTest, self).setUp()
self.server = None
self.client = TCPClient()
def start_server(self, family):
if family == socket.AF_UNSPEC and 'TRAVIS' in os.environ:
self.skipTest("dual-stack servers often have port conflicts on travis")
self.server = TestTCPServer(family)
return self.server.port
def stop_server(self):
if self.server is not None:
self.server.stop()
self.server = None
def tearDown(self):
self.client.close()
self.stop_server()
super(TCPClientTest, self).tearDown()
def skipIfLocalhostV4(self):
# The port used here doesn't matter, but some systems require it
# to be non-zero if we do not also pass AI_PASSIVE.
Resolver().resolve('localhost', 80, callback=self.stop)
addrinfo = self.wait()
families = set(addr[0] for addr in addrinfo)
if socket.AF_INET6 not in families:
self.skipTest("localhost does not resolve to ipv6")
@gen_test
def do_test_connect(self, family, host):
port = self.start_server(family)
stream = yield self.client.connect(host, port)
with closing(stream):
stream.write(b"hello")
data = yield self.server.streams[0].read_bytes(5)
self.assertEqual(data, b"hello")
def test_connect_ipv4_ipv4(self):
self.do_test_connect(socket.AF_INET, '127.0.0.1')
def test_connect_ipv4_dual(self):
self.do_test_connect(socket.AF_INET, 'localhost')
@skipIfNoIPv6
def test_connect_ipv6_ipv6(self):
self.skipIfLocalhostV4()
self.do_test_connect(socket.AF_INET6, '::1')
@skipIfNoIPv6
def test_connect_ipv6_dual(self):
self.skipIfLocalhostV4()
if Resolver.configured_class().__name__.endswith('TwistedResolver'):
self.skipTest('TwistedResolver does not support multiple addresses')
self.do_test_connect(socket.AF_INET6, 'localhost')
def test_connect_unspec_ipv4(self):
self.do_test_connect(socket.AF_UNSPEC, '127.0.0.1')
@skipIfNoIPv6
def test_connect_unspec_ipv6(self):
self.skipIfLocalhostV4()
self.do_test_connect(socket.AF_UNSPEC, '::1')
def test_connect_unspec_dual(self):
self.do_test_connect(socket.AF_UNSPEC, 'localhost')
@gen_test
def test_refused_ipv4(self):
sock, port = bind_unused_port()
sock.close()
with self.assertRaises(IOError):
yield self.client.connect('127.0.0.1', port)
class TCPClientTest(AsyncTestCase):
def setUp(self):
super(TCPClientTest, self).setUp()
self.server = None
self.client = TCPClient()
def start_server(self, family):
if family == socket.AF_UNSPEC and 'TRAVIS' in os.environ:
self.skipTest(
"dual-stack servers often have port conflicts on travis")
self.server = TestTCPServer(family)
return self.server.port
def stop_server(self):
if self.server is not None:
self.server.stop()
self.server = None
def tearDown(self):
self.client.close()
self.stop_server()
super(TCPClientTest, self).tearDown()
def skipIfLocalhostV4(self):
Resolver().resolve('localhost', 0, callback=self.stop)
addrinfo = self.wait()
families = set(addr[0] for addr in addrinfo)
if socket.AF_INET6 not in families:
self.skipTest("localhost does not resolve to ipv6")
@gen_test
def do_test_connect(self, family, host):
port = self.start_server(family)
stream = yield self.client.connect(host, port)
with closing(stream):
stream.write(b"hello")
data = yield self.server.streams[0].read_bytes(5)
self.assertEqual(data, b"hello")
def test_connect_ipv4_ipv4(self):
self.do_test_connect(socket.AF_INET, '127.0.0.1')
def test_connect_ipv4_dual(self):
self.do_test_connect(socket.AF_INET, 'localhost')
@skipIfNoIPv6
def test_connect_ipv6_ipv6(self):
self.skipIfLocalhostV4()
self.do_test_connect(socket.AF_INET6, '::1')
@skipIfNoIPv6
def test_connect_ipv6_dual(self):
self.skipIfLocalhostV4()
if Resolver.configured_class().__name__.endswith('TwistedResolver'):
self.skipTest(
'TwistedResolver does not support multiple addresses')
self.do_test_connect(socket.AF_INET6, 'localhost')
def test_connect_unspec_ipv4(self):
self.do_test_connect(socket.AF_UNSPEC, '127.0.0.1')
@skipIfNoIPv6
def test_connect_unspec_ipv6(self):
self.skipIfLocalhostV4()
self.do_test_connect(socket.AF_UNSPEC, '::1')
def test_connect_unspec_dual(self):
self.do_test_connect(socket.AF_UNSPEC, 'localhost')
@gen_test
def test_refused_ipv4(self):
sock, port = bind_unused_port()
sock.close()
with self.assertRaises(IOError):
yield self.client.connect('127.0.0.1', port)
class WebSocketClientConnection(simple_httpclient._HTTPConnection):
"""WebSocket 客户端连接
这个类不应当直接被实例化, 请使用 `websocket_connect`
"""
def __init__(self,
io_loop,
request,
on_message_callback=None,
compression_options=None):
self.compression_options = compression_options
self.connect_future = TracebackFuture()
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
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(io_loop=io_loop)
super(WebSocketClientConnection,
self).__init__(io_loop, None, request, lambda: None,
self._on_http_response, 104857600,
self.tcp_client, 65536, 104857600)
def close(self, code=None, reason=None):
"""关闭 websocket 连接
``code`` 和 ``reason`` 的文档在 `WebSocketHandler.close` 下已给出.
.. versionadded:: 3.2
.. versionchanged:: 4.0
添加 ``code`` 和 ``reason`` 这两个参数
"""
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._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
self.connect_future.set_result(self)
def write_message(self, message, binary=False):
"""发送消息到 websocket 服务器."""
return self.protocol.write_message(message, binary)
def read_message(self, callback=None):
"""读取来自 WebSocket 服务器的消息.
如果在 WebSocket 初始化时指定了 on_message_callback ,那么这个方法永远不会返回消息
如果连接已经关闭,返回结果会是一个结果是 message 的 future 对象或者是 None.
如果 future 给出了回调参数, 这个参数将会在 future 完成时调用.
"""
assert self.read_future is None
future = TracebackFuture()
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._on_message_callback:
self._on_message_callback(message)
elif 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 get_websocket_protocol(self):
return WebSocketProtocol13(
self,
mask_outgoing=True,
compression_options=self.compression_options)
class WebSocketClientConnection(simple_httpclient._HTTPConnection):
"""WebSocket client connection.
This class should not be instantiated directly; use the
`websocket_connect` function instead.
"""
def __init__(self, io_loop, request):
self.connect_future = TracebackFuture()
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.tcp_client = TCPClient(io_loop=io_loop)
super(WebSocketClientConnection,
self).__init__(io_loop, None, request, lambda: None,
self._on_http_response, 104857600,
self.tcp_client, 65536)
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
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.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
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 = TracebackFuture()
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
class Emitter(object):
def __init__(self, port, n=1000, values=1, duration=3.0):
self.port = port
self.n = n
self.values = values
self.duration = duration
self.message = self.hello
self.i = 0
self.pcb = None
self.client = None
def start(self):
self.client = TCPClient()
self.pcb = PeriodicCallback(self.send, 1000.0 / self.n)
self.pcb.start()
IOLoop.current().call_later(self.duration + 0.5, self.stop)
IOLoop.current().start()
IOLoop.clear_current()
def stop(self):
if self.pcb is not None:
self.pcb.stop()
if self.client is not None:
self.client.close()
IOLoop.current().stop()
@gen.coroutine
def send(self):
if self.i >= self.duration * self.n * self.values:
self.pcb.stop()
return
try:
stream = yield self.client.connect('127.0.0.1', self.port)
with closing(stream):
messages = b''.join(self.message() for _ in range(self.values))
stream.write(messages)
self.i += self.values
except StreamClosedError:
return
def hello(self):
return b'hello\n'
def r(self):
s = random.randint(1, 10)
v = s / 10.0 + (1.5 - s / 10.0) * random.random()
return (s, v)
def text(self):
return 'sensor{}|{}\n'.format(*self.r()).encode('utf8')
def json(self):
s, v = self.r()
return (json.dumps({
'sensor{}'.format(s): v,
}) + '\n').encode('utf8')
def bello(self):
# 5 bytes
return b'bello'
def struct(self):
# 8 bytes
return struct.pack('If', *self.r())
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.
If the stream is closed, raises `WebSocketClosedError`.
Returns a `.Future` which can be used for flow control.
.. versionchanged:: 5.0
Exception raised on a closed stream changed from `.StreamClosedError`
to `WebSocketClosedError`.
"""
return self.protocol.write_message(message, binary=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 PropertyEstimatorClient:
"""The PropertyEstimatorClient is the main object that users of the
property estimator will interface with. It is responsible for requesting
that a PropertyEstimatorServer estimates a set of physical properties,
as well as querying for when those properties have been estimated.
The PropertyEstimatorClient supports two main workflows: one where
a PropertyEstimatorServer lives on a remote supercomputing cluster
where all of the expensive calculations will be run, and one where
the users local machine acts as both the server and the client, and
all calculations will be performed locally.
Warnings
--------
While the API of this class in now close to being final, the internals and implementation
are still heavily under development and is subject to rapid changes.
Examples
--------
Setting up the client instance:
>>> from propertyestimator.client import PropertyEstimatorClient
>>> property_estimator = PropertyEstimatorClient()
If the PropertyEstimatorServer is not running on the local machine, you will
need to specify its address and the port that it is listening on:
>>> from propertyestimator.client import ConnectionOptions
>>>
>>> connection_options = ConnectionOptions(server_address='server_address',
>>> server_port=8000)
>>> property_estimator = PropertyEstimatorClient(connection_options)
To asynchronously submit a request to the running server using the default estimator
options:
>>> # Load in the data set of properties which will be used for comparisons
>>> from propertyestimator.datasets import ThermoMLDataSet
>>> data_set = ThermoMLDataSet.from_doi('10.1016/j.jct.2016.10.001')
>>> # Filter the dataset to only include densities measured between 130-260 K
>>> from propertyestimator.properties import Density
>>>
>>> data_set.filter_by_properties(types=[Density])
>>> data_set.filter_by_temperature(min_temperature=130*unit.kelvin, max_temperature=260*unit.kelvin)
>>>
>>> # Load initial parameters
>>> from openforcefield.typing.engines.smirnoff import ForceField
>>> parameters = ForceField(['smirnoff99Frosst.offxml'])
>>>
>>> request = property_estimator.request_estimate(data_set, parameters)
The status of the request can be asynchronously queried by calling
>>> results = request.results()
or the main thread can be blocked until the results are
available by calling
>>> results = request.results(synchronous=True)
How the property set will be estimated can easily be controlled by passing a
PropertyEstimatorOptions object to the estimate commands.
The calculations layers which will be used to estimate the properties can be
controlled for example like so:
>>> from propertyestimator.layers import ReweightingLayer, SimulationLayer
>>>
>>> options = PropertyEstimatorOptions(allowed_calculation_layers = [ReweightingLayer,
>>> SimulationLayer])
>>>
>>> request = property_estimator.request_estimate(data_set, parameters, options)
Options for how properties should be estimated can be set on a per property basis. For example
the relative uncertainty that properties should estimated to within can be set as:
>>> from propertyestimator.properties.properties import PropertyWorkflowOptions
>>>
>>> workflow_options = PropertyWorkflowOptions(PropertyWorkflowOptions.ConvergenceMode.RelativeUncertainty,
>>> relative_uncertainty_fraction=0.1)
>>> options.workflow_options = {
>>> 'Density': workflow_options,
>>> 'Dielectric': workflow_options
>>> }
Or alternatively, as absolute uncertainty tolerance can be set as:
>>> density_options = PropertyWorkflowOptions(PropertyWorkflowOptions.ConvergenceMode.AbsoluteUncertainty,
>>> absolute_uncertainty=0.0002 * unit.gram / unit.milliliter)
>>> dielectric_options = PropertyWorkflowOptions(PropertyWorkflowOptions.ConvergenceMode.AbsoluteUncertainty,
>>> absolute_uncertainty=0.02 * unit.dimensionless)
>>>
>>> options.workflow_options = {
>>> 'Density': density_options,
>>> 'Dielectric': dielectric_options
>>> }
"""
@property
def server_address(self):
return self._connection_options.server_address
@property
def server_port(self):
return self._connection_options.server_port
class Request:
"""An object representation of a estimation request which has
been sent to a `PropertyEstimatorServer` instance. This object
can be used to query and retrieve the results of the request, or
be stored to retrieve the request at some point in the future."""
@property
def id(self):
"""str: The id of the submitted request."""
return self._id
@property
def server_address(self):
"""str: The address of the server that the request was sent to."""
return self._server_address
@property
def server_port(self):
"""The port that the server is listening on."""
return self._server_port
def __init__(self, request_id, connection_options, client=None):
"""Constructs a new Request object.
Parameters
----------
request_id: str
The id of the submitted request.
connection_options: ConnectionOptions
The options that were used to connect to the server that the request was sent to.
client: PropertyEstimatorClient, optional
The client that was used to submit the request.
"""
self._id = request_id
self._server_address = connection_options.server_address
self._server_port = connection_options.server_port
self._client = client
if client is None:
connection_options = ConnectionOptions(
server_address=connection_options.server_address,
server_port=connection_options.server_port)
self._client = PropertyEstimatorClient(connection_options)
def __str__(self):
return 'EstimateRequest id: {} server_address: {} server_port: {}'.format(
self._id, self._server_address, self._server_port)
def __repr__(self):
return '<EstimateRequest id: {} server_address: {} server_port: {}>'.format(
self._id, self._server_address, self._server_port)
def json(self):
"""Returns a JSON representation of the `Request` object.
Returns
-------
str:
The JSON representation of the `Request` object.
"""
return json.dumps({
'id': self._id,
'server_address': self._id,
'server_port': self._id
})
@classmethod
def from_json(cls, json_string):
"""Creates a new `Request` object from a JSON representation.
Parameters
----------
json_string: str
The JSON representation of the `Request` object.
Returns
-------
str:
The created `Request` object.
"""
json_dict = json.loads(json_string)
return cls(json_dict['id'], json_dict['server_address'],
json_dict['server_port'])
def results(self, synchronous=False, polling_interval=5):
"""Retrieve the results of an estimate request.
Parameters
----------
synchronous: bool
If true, this method will block the main thread until the server
either returns a result or an error.
polling_interval: int
If running synchronously, this is the time interval (seconds) between
checking if the calculation has finished.
Returns
-------
PropertyEstimatorResult or PropertyEstimatorException:
Returns either the results of the requested estimate, or any
exceptions which were raised.
If the method is run synchronously then this method will block the main
thread until all of the requested properties have been estimated, or
an exception is returned.
"""
return self._client._retrieve_estimate(self._id, synchronous,
polling_interval)
def __init__(self, connection_options=ConnectionOptions()):
"""Constructs a new PropertyEstimatorClient object.
Parameters
----------
connection_options: ConnectionOptions
The options used when connecting to the calculation server.
"""
self._connection_options = connection_options
if connection_options.server_address is None:
raise ValueError('The address of the server which will run'
'these calculations must be given.')
self._tcp_client = TCPClient()
def request_estimate(self, property_set, force_field, options=None):
"""Requests that a PropertyEstimatorServer attempt to estimate the
provided property set using the supplied force field and estimator options.
Parameters
----------
property_set : PhysicalPropertyDataSet
The set of properties to attempt to estimate.
force_field : ForceField
The OpenFF force field to use for the calculations.
options : PropertyEstimatorOptions, optional
A set of estimator options. If None, default options
will be used.
Returns
-------
PropertyEstimatorClient.Request
An object which will provide access the the results of the request.
"""
if property_set is None or force_field is None:
raise ValueError('Both a data set and parameter set must be '
'present to compute physical properties.')
if options is None:
options = PropertyEstimatorOptions()
if len(options.allowed_calculation_layers) == 0:
raise ValueError(
'A submission contains no allowed calculation layers.')
properties_list = []
property_types = set()
# Refactor the properties into a list, and extract the types
# of properties to be estimated (e.g 'Denisty', 'DielectricConstant').
for substance_tag in property_set.properties:
for physical_property in property_set.properties[substance_tag]:
properties_list.append(physical_property)
type_name = type(physical_property).__name__
if type_name not in registered_properties:
raise ValueError(
f'The property estimator does not support {type_name} properties.'
)
if type_name in property_types:
continue
property_types.add(type_name)
if options.workflow_options is None:
options.workflow_options = {}
# Assign default workflows in the cases where the user hasn't
# provided one, and validate all of the workflows to be used
# in the estimation.
for type_name in property_types:
if type_name not in options.workflow_schemas:
options.workflow_schemas[type_name] = {}
if type_name not in options.workflow_options:
options.workflow_options[type_name] = PropertyWorkflowOptions()
for calculation_layer in options.allowed_calculation_layers:
if (calculation_layer
not in options.workflow_schemas[type_name] or
options.workflow_schemas[type_name][calculation_layer]
is None):
property_type = registered_properties[type_name]()
options.workflow_schemas[type_name][calculation_layer] = \
property_type.get_default_workflow_schema(calculation_layer,
options.workflow_options[type_name])
workflow = options.workflow_schemas[type_name][
calculation_layer]
if workflow is None:
# Not all properties may support every calculation layer.
continue
# Will raise the correct exception for non-valid interfaces.
workflow.validate_interfaces()
# Enforce the global option of whether to allow merging or not.
for protocol_schema_name in workflow.protocols:
protocol_schema = workflow.protocols[protocol_schema_name]
if not options.allow_protocol_merging:
protocol_schema.inputs['.allow_merging'] = False
submission = PropertyEstimatorSubmission(properties=properties_list,
force_field=force_field,
options=options)
request_id = IOLoop.current().run_sync(
lambda: self._send_calculations_to_server(submission))
request_object = PropertyEstimatorClient.Request(
request_id, self._connection_options, self)
return request_object
def _retrieve_estimate(self,
request_id,
synchronous=False,
polling_interval=5):
"""A method to retrieve the status of a requested estimate from the server.
Parameters
----------
request_id: str
The id of the estimate request which was returned by the server
upon making the request.
synchronous: bool
If true, this method will block the main thread until the server
either returns a result or an error.
polling_interval: int
If running synchronously, this is the time interval (seconds) between
checking if the calculation has finished.
Returns
-------
PropertyEstimatorResult or PropertyEstimatorException:
Returns either the results of the requested estimate, or any
exceptions which were raised.
If the method is run synchronously then this method will block the main
thread until all of the requested properties have been estimated, or
an exception is returned.
"""
# If running asynchronously, just return whatever the server
# sends back.
if synchronous is False:
return IOLoop.current().run_sync(
lambda: self._send_query_server(request_id))
assert polling_interval >= 0
response = None
should_run = True
while should_run:
if polling_interval > 0:
sleep(polling_interval)
response = IOLoop.current().run_sync(
lambda: self._send_query_server(request_id))
if isinstance(response, PropertyEstimatorResult) and len(
response.queued_properties) > 0:
continue
logging.info(f'The server has completed request {request_id}.')
should_run = False
return response
async def _send_calculations_to_server(self, submission):
"""Attempts to connect to the calculation server, and
submit the requested calculations.
Notes
-----
This method is based on the StackOverflow response from
A. Jesse Jiryu Davis: https://stackoverflow.com/a/40257248
Parameters
----------
submission: PropertyEstimatorSubmission
The jobs to submit.
Returns
-------
str, optional:
The id which the server has assigned the submitted calculations.
This can be used to query the server for when the calculation
has completed.
Returns None if the calculation could not be submitted.
"""
request_id = None
try:
# Attempt to establish a connection to the server.
logging.info("Attempting Connection to {}:{}".format(
self._connection_options.server_address,
self._connection_options.server_port))
stream = await self._tcp_client.connect(
self._connection_options.server_address,
self._connection_options.server_port)
logging.info("Connected to {}:{}".format(
self._connection_options.server_address,
self._connection_options.server_port))
stream.set_nodelay(True)
# Encode the submission json into an encoded
# packet ready to submit to the server. The
# Length of the packet is encoded in the first
# four bytes.
message_type = pack_int(PropertyEstimatorMessageTypes.Submission)
encoded_json = submission.json().encode()
length = pack_int(len(encoded_json))
await stream.write(message_type + length + encoded_json)
logging.info(
"Sent calculations to {}:{}. Waiting for a response from"
" the server...".format(
self._connection_options.server_address,
self._connection_options.server_port))
# Wait for confirmation that the server has submitted
# the jobs. The first four bytes of the response should
# be the length of the message being sent.
header = await stream.read_bytes(4)
length = unpack_int(header)[0]
# Decode the response from the server. If everything
# went well, this should be a list of ids of the submitted
# calculations.
encoded_json = await stream.read_bytes(length)
request_id = json.loads(encoded_json.decode())
logging.info('Received job id from server: {}'.format(request_id))
stream.close()
self._tcp_client.close()
except StreamClosedError as e:
# Handle no connections to the server gracefully.
logging.info(
"Error connecting to {}:{} : {}. Please ensure the server is running and"
"that the server address / port is correct.".format(
self._connection_options.server_address,
self._connection_options.server_port, e))
# Return the ids of the submitted jobs.
return request_id
async def _send_query_server(self, request_id):
"""Attempts to connect to the calculation server, and
submit the requested calculations.
Notes
-----
This method is based on the StackOverflow response from
A. Jesse Jiryu Davis: https://stackoverflow.com/a/40257248
Parameters
----------
request_id: str
The id of the job to query.
Returns
-------
str, optional:
The status of the submitted job.
Returns None if the calculation has not yet completed.
"""
server_response = None
try:
# Attempt to establish a connection to the server.
stream = await self._tcp_client.connect(
self._connection_options.server_address,
self._connection_options.server_port)
stream.set_nodelay(True)
# Encode the request id into the message.
message_type = pack_int(PropertyEstimatorMessageTypes.Query)
encoded_request_id = request_id.encode()
length = pack_int(len(encoded_request_id))
await stream.write(message_type + length + encoded_request_id)
# Wait for the server response.
header = await stream.read_bytes(4)
length = unpack_int(header)[0]
# Decode the response from the server. If everything
# went well, this should be the finished calculation.
if length > 0:
encoded_json = await stream.read_bytes(length)
server_response = encoded_json.decode()
stream.close()
self._tcp_client.close()
except StreamClosedError as e:
# Handle no connections to the server gracefully.
logging.info(
"Error connecting to {}:{} : {}. Please ensure the server is running and"
"that the server address / port is correct.".format(
self._connection_options.server_address,
self._connection_options.server_port, e))
if server_response is not None:
server_response = TypedBaseModel.parse_json(server_response)
# Return the ids of the submitted jobs.
return server_response
class SimpleAsyncHTTPClient(AsyncHTTPClient):
"""Non-blocking HTTP client with no external dependencies.
This class implements an HTTP 1.1 client on top of Tornado's IOStreams.
It does not currently implement all applicable parts of the HTTP
specification, but it does enough to work with major web service APIs.
Some features found in the curl-based AsyncHTTPClient are not yet
supported. In particular, proxies are not supported, connections
are not reused, and callers cannot select the network interface to be
used.
"""
def initialize(self,
io_loop,
max_clients=10,
hostname_mapping=None,
max_buffer_size=104857600,
resolver=None,
defaults=None,
max_header_size=None):
"""Creates a AsyncHTTPClient.
Only a single AsyncHTTPClient instance exists per IOLoop
in order to provide limitations on the number of pending connections.
force_instance=True may be used to suppress this behavior.
max_clients is the number of concurrent requests that can be
in progress. Note that this arguments are only used when the
client is first created, and will be ignored when an existing
client is reused.
hostname_mapping is a dictionary mapping hostnames to IP addresses.
It can be used to make local DNS changes when modifying system-wide
settings like /etc/hosts is not possible or desirable (e.g. in
unittests).
max_buffer_size is the number of bytes that can be read by IOStream. It
defaults to 100mb.
"""
super(SimpleAsyncHTTPClient, self).initialize(io_loop,
defaults=defaults)
self.max_clients = max_clients
self.queue = collections.deque()
self.active = {}
self.waiting = {}
self.max_buffer_size = max_buffer_size
self.max_header_size = max_header_size
# TCPClient could create a Resolver for us, but we have to do it
# ourselves to support hostname_mapping.
if resolver:
self.resolver = resolver
self.own_resolver = False
else:
self.resolver = Resolver(io_loop=io_loop)
self.own_resolver = True
if hostname_mapping is not None:
self.resolver = OverrideResolver(resolver=self.resolver,
mapping=hostname_mapping)
self.tcp_client = TCPClient(resolver=self.resolver, io_loop=io_loop)
def close(self):
super(SimpleAsyncHTTPClient, self).close()
if self.own_resolver:
self.resolver.close()
self.tcp_client.close()
def fetch_impl(self, request, callback):
key = object()
self.queue.append((key, request, callback))
if not len(self.active) < self.max_clients:
timeout_handle = self.io_loop.add_timeout(
self.io_loop.time() +
min(request.connect_timeout, request.request_timeout),
functools.partial(self._on_timeout, key))
else:
timeout_handle = None
self.waiting[key] = (request, callback, timeout_handle)
self._process_queue()
if self.queue:
gen_log.debug("max_clients limit reached, request queued. "
"%d active, %d queued requests." %
(len(self.active), len(self.queue)))
def _process_queue(self):
with stack_context.NullContext():
while self.queue and len(self.active) < self.max_clients:
key, request, callback = self.queue.popleft()
if key not in self.waiting:
continue
self._remove_timeout(key)
self.active[key] = (request, callback)
release_callback = functools.partial(self._release_fetch, key)
self._handle_request(request, release_callback, callback)
def _handle_request(self, request, release_callback, final_callback):
_HTTPConnection(self.io_loop, self, request, release_callback,
final_callback, self.max_buffer_size, self.tcp_client,
self.max_header_size)
def _release_fetch(self, key):
del self.active[key]
self._process_queue()
def _remove_timeout(self, key):
if key in self.waiting:
request, callback, timeout_handle = self.waiting[key]
if timeout_handle is not None:
self.io_loop.remove_timeout(timeout_handle)
del self.waiting[key]
def _on_timeout(self, key):
request, callback, timeout_handle = self.waiting[key]
self.queue.remove((key, request, callback))
timeout_response = HTTPResponse(request,
599,
error=HTTPError(599, "Timeout"),
request_time=self.io_loop.time() -
request.start_time)
self.io_loop.add_callback(callback, timeout_response)
del self.waiting[key]
class XTCPClient(object):
def __init__(self, io_loop=None, max_clients=10, max_buffer_size=None, max_response_size=None):
self._io_loop = io_loop or IOLoop.instance()
self.max_clients = max_clients
self.max_buffer_size = max_buffer_size or 104857600 # 100M
self.max_response_size = max_response_size or 10 * 1024 * 1024 # 10M
self.queue = collections.deque()
self.active = {}
self.waiting = {}
self._client_closed = False
self.tcp_client = TCPClient(io_loop=self._io_loop)
def __del__(self):
self.close()
def close(self):
if not self._client_closed:
self._client_closed = True
self.tcp_client.close()
self._io_loop.close()
def acquire(self, request):
response = self._io_loop.run_sync(functools.partial(self._acquire_by_request, request))
return response
def _acquire_by_request(self, request):
future = Future()
def _handle_response(success, response=None):
if success is True:
response_message = response
if request.user_request_callback is not None:
response_message = request.user_request_callback(response)
future.set_result(response_message)
else:
future.set_exc_info(response)
request.request_callback = _handle_response
self._acquire_loop_by_request(request)
return future
def _acquire_loop_by_request(self, request):
key = object()
self.queue.append((key, request))
if not len(self.active) < self.max_clients:
waiting_timeout_handle = self._io_loop.add_timeout(
self._io_loop.time() + self.request.waiting_timeout, functools.partial(self._on_waiting_timeout, key))
else:
waiting_timeout_handle = None
self.waiting[key] = (request, waiting_timeout_handle)
self._process_queue()
if self.queue:
xtcp_logger.debug("max_clients limits reached. {} active, {} queued requests".format(len(self.active), len(self.queue)))
def _on_waiting_timeout(self, key):
xtcp_logger.debug("_on_waiting_timeout : {}".format(self.waiting[key]))
request, callback, waiting_timeout_handle = self.waiting[key]
self.queue.remove((key, request, callback))
del self.waiting[key]
def _process_queue(self):
while self.queue and len(self.active) < self.max_clients:
key, request = self.queue.popleft()
if key not in self.waiting:
continue
self._remove_waiting_timeout_request(key)
self.active[key] = request
self._handle_request(request, functools.partial(self._release_request, key))
def _handle_request(self, request, release_callback):
connection = _ClientConnection(
self, io_loop=self._io_loop, request=request, release_callback=release_callback)
connection.connect()
def _release_request(self, key):
del self.active[key]
self._process_queue()
def _remove_waiting_timeout_request(self, key):
if key in self.waiting:
_, timeout_handle = self.waiting[key]
if timeout_handle is not None:
self._io_loop.remove_timeout(timeout_handle)
del self.waiting[key]
