From 620afc667f1cac6e32604a7226129466358c1851 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Honza=20Kr=C3=A1l?= Date: Wed, 12 Jul 2017 23:48:58 +0200 Subject: [PATCH] Set a limit on the internal queue of the parallel_bulk helper Fixes #547 Solution taken from https://bugs.python.org/msg90067 --- elasticsearch/compat.py | 2 ++ elasticsearch/helpers/__init__.py | 18 ++++++++++++++---- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/elasticsearch/compat.py b/elasticsearch/compat.py index a5b615d2..c3d5fa28 100644 --- a/elasticsearch/compat.py +++ b/elasticsearch/compat.py @@ -7,7 +7,9 @@ if PY2: from urllib import quote_plus, urlencode, unquote from urlparse import urlparse from itertools import imap as map + from Queue import Queue else: string_types = str, bytes from urllib.parse import quote_plus, urlencode, urlparse, unquote map = map + from queue import Queue diff --git a/elasticsearch/helpers/__init__.py b/elasticsearch/helpers/__init__.py index d7ab0de2..5f0f240f 100644 --- a/elasticsearch/helpers/__init__.py +++ b/elasticsearch/helpers/__init__.py @@ -4,7 +4,8 @@ import logging from operator import methodcaller from ..exceptions import ElasticsearchException, TransportError -from ..compat import map, string_types +from ..compat import map, string_types, Queue + logger = logging.getLogger('elasticsearch.helpers') @@ -204,7 +205,7 @@ def bulk(client, actions, stats_only=False, **kwargs): return success, failed if stats_only else errors def parallel_bulk(client, actions, thread_count=4, chunk_size=500, - max_chunk_bytes=100 * 1024 * 1024, + max_chunk_bytes=100 * 1024 * 1024, queue_size=4, expand_action_callback=expand_action, **kwargs): """ Parallel version of the bulk helper run in multiple threads at once. @@ -221,13 +222,22 @@ def parallel_bulk(client, actions, thread_count=4, chunk_size=500, :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 queue_size: size of the task queue between the main thread (producing + chunks to send) and the processing threads. """ # Avoid importing multiprocessing unless parallel_bulk is used # to avoid exceptions on restricted environments like App Engine - from multiprocessing.dummy import Pool + from multiprocessing.pool import ThreadPool + actions = map(expand_action_callback, actions) - pool = Pool(thread_count) + class BlockingPool(ThreadPool): + def _setup_queues(self): + super(BlockingPool, self)._setup_queues() + self._inqueue = Queue(queue_size) + self._quick_put = self._inqueue.put + + pool = BlockingPool(thread_count) try: for result in pool.imap(