Merged backoff into streaming_bulk

This commit is contained in:
Honza Král
2017-07-18 17:44:43 -04:00
parent 66c380420f
commit b31b5ddae6
2 changed files with 56 additions and 75 deletions
+56 -5
View File
@@ -2,6 +2,7 @@ from __future__ import unicode_literals
import logging
from operator import methodcaller
import time
from ..exceptions import ElasticsearchException, TransportError
from ..compat import map, string_types, Queue
@@ -133,8 +134,10 @@ def _process_bulk_chunk(client, bulk_actions, bulk_data, raise_on_exception=True
raise BulkIndexError('%i document(s) failed to index.' % len(errors), errors)
def streaming_bulk(client, actions, chunk_size=500, max_chunk_bytes=100 * 1024 * 1024,
raise_on_error=True, expand_action_callback=expand_action,
raise_on_exception=True, **kwargs):
raise_on_error=True, expand_action_callback=expand_action,
raise_on_exception=True, max_retries=0, initial_backoff=2,
max_backoff=600, yield_ok=True, **kwargs):
"""
Streaming bulk consumes actions from the iterable passed in and yields
results per action. For non-streaming usecases use
@@ -154,12 +157,58 @@ def streaming_bulk(client, actions, chunk_size=500, max_chunk_bytes=100 * 1024 *
:arg expand_action_callback: callback executed on each action passed in,
should return a tuple containing the action line and the data line
(`None` if data line should be omitted).
:arg max_retries: maximum number of times a document will be retried when
``429`` is received, set to 0 (default) for no retires on ``429``
:arg initial_backoff: number of seconds we should wait before the first
retry. Any subsequent retries will be powers of ``inittial_backoff *
2**retry_number``
:arg max_backoff: maximum number of seconds a retry will wait
"""
actions = map(expand_action_callback, actions)
for bulk_data, bulk_actions in _chunk_actions(actions, chunk_size, max_chunk_bytes, client.transport.serializer):
for result in _process_bulk_chunk(client, bulk_actions, bulk_data, raise_on_exception, raise_on_error, **kwargs):
yield result
for bulk_data, bulk_actions in _chunk_actions(actions, chunk_size,
max_chunk_bytes,
client.transport.serializer):
for attempt in range(max_retries + 1):
to_retry, to_retry_data = [], []
if attempt:
time.sleep(min(max_backoff, initial_backoff * 2**(attempt-1)))
try:
for data, (ok, info) in zip(
bulk_data,
_process_bulk_chunk(client, bulk_actions, bulk_data,
raise_on_exception=raise_on_exception,
raise_on_error=raise_on_error, **kwargs)
):
if not ok:
action, info = info.popitem()
# retry if retries enabled, we get 429, and we are not
# in the last attempt
if max_retries and info['status'] == 429 and (attempt+1) <= max_retries:
to_retry.extend(data)
to_retry_data.append(data)
else:
if action != 'delete':
info['data'] = data[1]
yield ok, {action: info}
elif yield_ok:
yield ok, info
except TransportError as e:
# suppress 429 errors since we will retry them
if not max_retries or e.status_code != 429:
raise
else:
if not to_retry:
break
# retry only subset of documents that didn't succeed
bulk_actions = to_retry
bulk_data = to_retry_data
def bulk(client, actions, stats_only=False, **kwargs):
"""
@@ -190,6 +239,8 @@ def bulk(client, actions, stats_only=False, **kwargs):
# list of errors to be collected is not stats_only
errors = []
# make streaming_bulk yield successful results so we can count them
kwargs['yield_ok'] = True
for ok, item in streaming_bulk(client, actions, **kwargs):
# go through request-reponse pairs and detect failures
if not ok:
-70
View File
@@ -1,70 +0,0 @@
import time
from . import expand_action, _chunk_actions, _process_bulk_chunk
from ..exceptions import TransportError
def backoff_bulk(client, actions, chunk_size=500, max_chunk_bytes=100 * 1024 * 1024,
expand_action_callback=expand_action, max_retries=-1,
initial_backoff=2, max_backoff=600, **kwargs):
"""
Bulk helper that implements proper retry strategy when dealing with
rejecttions - when the target cluster is overloaded, returning ``429``
responses, the bulk actions will be retried after a backoff period to allow
the cluster to recover.
This function is a generator yielding all documents that failed to index.
:arg client: instance of :class:`~elasticsearch.Elasticsearch` to use
:arg actions: iterable containing the actions to be executed
:arg chunk_size: number of docs in one chunk sent to es (default: 500)
:arg max_chunk_bytes: the maximum size of the request in bytes (default: 100MB)
:arg expand_action_callback: callback executed on each action passed in,
should return a tuple containing the action line and the data line
(`None` if data line should be omitted).
:arg max_retries: maximum number of times a document will be retried
:arg initial_backoff: number of seconds we should wait before the first
retry. Any subsequent retries will be powers of ``inittial_backoff *
2**retry_number``
:arg max_backoff: maximum number of seconds a retry will wait
"""
actions = map(expand_action_callback, actions)
for bulk_data, bulk_actions in _chunk_actions(actions, chunk_size,
max_chunk_bytes,
client.transport.serializer):
retry = 0
while max_retries == -1 or retry <= max_retries:
to_retry, to_retry_data = [], []
if retry:
time.sleep(min(max_backoff, initial_backoff * 2**(retry-1)))
try:
for data, (ok, info) in zip(bulk_data,
_process_bulk_chunk(client,
bulk_actions,
bulk_data,
raise_on_exception=True,
raise_on_error=False,
**kwargs)):
if not ok:
action, info = info.popitem()
if info['status'] == 429 and (retry+1) <= max_retries:
to_retry.extend(data)
to_retry_data.append(data)
else:
info['data'] = data
yield {action: info}
except TransportError as e:
if e.status_code != 429:
raise
retry += 1
else:
if not to_retry:
break
bulk_actions = to_retry
bulk_data = to_retry_data
retry += 1