forked from opensearch-project/opensearch-benchmark
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add an asyncio-based load generator (#935)
With this commit we add a new experimental subcommand `race-aync` to Rally. It allows to specify significantly more clients than the current `race` subcommand. The reason for this is that under the hood, `race-async` uses `asyncio` and runs all clients in a single event loop. Contrary to that, `race` uses an actor system under the hood and maps each client to one process. As the new subcommand is very experimental and not yet meant to be used broadly, there is no accompanying user documentation in this PR. Instead, we plan to build on top of this PR and expand the load generator to take advantage of multiple cores before we consider this usable in production (it will likely keep its experimental status though). In this PR we also implement a compatibility layer into the current load generator so both work internally now with `asyncio`. Consequently, we have already adapted all Rally tracks with a backwards-compatibility layer (see elastic/rally-tracks#97 and elastic/rally-eventdata-track#80). Closes #852 Relates #916
- Loading branch information
1 parent
b33da19
commit 3b5eee2
Showing
25 changed files
with
2,656 additions
and
1,099 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,134 @@ | ||
import asyncio | ||
import ssl | ||
import warnings | ||
|
||
import aiohttp | ||
from aiohttp.client_exceptions import ServerFingerprintMismatch | ||
import async_timeout | ||
|
||
from elasticsearch.exceptions import ConnectionError, ConnectionTimeout, ImproperlyConfigured, SSLError | ||
from elasticsearch.connection import Connection | ||
from elasticsearch.compat import urlencode | ||
from elasticsearch.connection.http_urllib3 import create_ssl_context | ||
|
||
|
||
# This is only needed because https://github.com/elastic/elasticsearch-py-async/pull/68 is not merged yet | ||
# In addition we have raised the connection limit in TCPConnector from 100 to 10000. | ||
|
||
# We want to keep the diff as small as possible thus suppressing pylint warnings that we would not allow in Rally | ||
# pylint: disable=W0706 | ||
class AIOHttpConnection(Connection): | ||
def __init__(self, host='localhost', port=9200, http_auth=None, | ||
use_ssl=False, verify_certs=False, ca_certs=None, client_cert=None, | ||
client_key=None, loop=None, use_dns_cache=True, headers=None, | ||
ssl_context=None, trace_config=None, **kwargs): | ||
super().__init__(host=host, port=port, **kwargs) | ||
|
||
self.loop = asyncio.get_event_loop() if loop is None else loop | ||
|
||
if http_auth is not None: | ||
if isinstance(http_auth, str): | ||
http_auth = tuple(http_auth.split(':', 1)) | ||
|
||
if isinstance(http_auth, (tuple, list)): | ||
http_auth = aiohttp.BasicAuth(*http_auth) | ||
|
||
headers = headers or {} | ||
headers.setdefault('content-type', 'application/json') | ||
|
||
# if providing an SSL context, raise error if any other SSL related flag is used | ||
if ssl_context and (verify_certs or ca_certs): | ||
raise ImproperlyConfigured("When using `ssl_context`, `use_ssl`, `verify_certs`, `ca_certs` are not permitted") | ||
|
||
if use_ssl or ssl_context: | ||
cafile = ca_certs | ||
if not cafile and not ssl_context and verify_certs: | ||
# If no ca_certs and no sslcontext passed and asking to verify certs | ||
# raise error | ||
raise ImproperlyConfigured("Root certificates are missing for certificate " | ||
"validation. Either pass them in using the ca_certs parameter or " | ||
"install certifi to use it automatically.") | ||
if verify_certs or ca_certs: | ||
warnings.warn('Use of `verify_certs`, `ca_certs` have been deprecated in favor of using SSLContext`', DeprecationWarning) | ||
|
||
if not ssl_context: | ||
# if SSLContext hasn't been passed in, create one. | ||
# need to skip if sslContext isn't avail | ||
try: | ||
ssl_context = create_ssl_context(cafile=cafile) | ||
except AttributeError: | ||
ssl_context = None | ||
|
||
if not verify_certs and ssl_context is not None: | ||
ssl_context.check_hostname = False | ||
ssl_context.verify_mode = ssl.CERT_NONE | ||
warnings.warn( | ||
'Connecting to %s using SSL with verify_certs=False is insecure.' % host) | ||
if ssl_context: | ||
verify_certs = True | ||
use_ssl = True | ||
|
||
trace_configs = [trace_config] if trace_config else None | ||
|
||
self.session = aiohttp.ClientSession( | ||
auth=http_auth, | ||
timeout=self.timeout, | ||
connector=aiohttp.TCPConnector( | ||
loop=self.loop, | ||
verify_ssl=verify_certs, | ||
use_dns_cache=use_dns_cache, | ||
ssl_context=ssl_context, | ||
# this has been changed from the default (100) | ||
limit=100000 | ||
), | ||
headers=headers, | ||
trace_configs=trace_configs | ||
) | ||
|
||
self.base_url = 'http%s://%s:%d%s' % ( | ||
's' if use_ssl else '', | ||
host, port, self.url_prefix | ||
) | ||
|
||
@asyncio.coroutine | ||
def close(self): | ||
yield from self.session.close() | ||
|
||
@asyncio.coroutine | ||
def perform_request(self, method, url, params=None, body=None, timeout=None, ignore=(), headers=None): | ||
url_path = url | ||
if params: | ||
url_path = '%s?%s' % (url, urlencode(params or {})) | ||
url = self.base_url + url_path | ||
|
||
start = self.loop.time() | ||
response = None | ||
try: | ||
with async_timeout.timeout(timeout or self.timeout.total, loop=self.loop): | ||
response = yield from self.session.request(method, url, data=body, headers=headers) | ||
raw_data = yield from response.text() | ||
duration = self.loop.time() - start | ||
|
||
except asyncio.CancelledError: | ||
raise | ||
|
||
except Exception as e: | ||
self.log_request_fail(method, url, url_path, body, self.loop.time() - start, exception=e) | ||
if isinstance(e, ServerFingerprintMismatch): | ||
raise SSLError('N/A', str(e), e) | ||
if isinstance(e, asyncio.TimeoutError): | ||
raise ConnectionTimeout('TIMEOUT', str(e), e) | ||
raise ConnectionError('N/A', str(e), e) | ||
|
||
finally: | ||
if response is not None: | ||
yield from response.release() | ||
|
||
# raise errors based on http status codes, let the client handle those if needed | ||
if not (200 <= response.status < 300) and response.status not in ignore: | ||
self.log_request_fail(method, url, url_path, body, duration, status_code=response.status, response=raw_data) | ||
self._raise_error(response.status, raw_data) | ||
|
||
self.log_request_success(method, url, url_path, body, response.status, raw_data, duration) | ||
|
||
return response.status, response.headers, raw_data |
Oops, something went wrong.