67 lines
1.6 KiB
Python
67 lines
1.6 KiB
Python
from multiprocessing.dummy import Pool
|
|
from queue import Empty, Queue
|
|
|
|
from threading import Event
|
|
|
|
from . import streaming_bulk
|
|
|
|
def consume(queue, done):
|
|
"""
|
|
Create an iterator on top of a Queue.
|
|
"""
|
|
while True:
|
|
try:
|
|
yield queue.get(True, .01)
|
|
except Empty:
|
|
if done.is_set():
|
|
break
|
|
|
|
def wrapped_bulk(client, input, output, done, **kwargs):
|
|
"""
|
|
Wrap a call to streaming_bulk by feeding it data frm a queue and writing
|
|
the outputs to another queue.
|
|
"""
|
|
try:
|
|
for result in streaming_bulk(client, consume(input, done), **kwargs):
|
|
output.put(result)
|
|
except:
|
|
done.set()
|
|
raise
|
|
|
|
def feed_data(actions, input, done):
|
|
"""
|
|
Feed data from an iterator into a queue.
|
|
"""
|
|
for a in actions:
|
|
input.put(a, True)
|
|
|
|
# error short-circuit
|
|
if done.is_set():
|
|
break
|
|
done.set()
|
|
|
|
|
|
def parallel_bulk(client, actions, thread_count=5, **kwargs):
|
|
"""
|
|
Paralel version of the bulk helper. It runs a thread pool with a thread for
|
|
a producer and ``thread_count`` threads for.
|
|
"""
|
|
done = Event()
|
|
input, output = Queue(), Queue()
|
|
pool = Pool(thread_count + 1)
|
|
|
|
results = [
|
|
pool.apply_async(wrapped_bulk, (client, input, output, done), kwargs)
|
|
for _ in range(thread_count)]
|
|
pool.apply_async(feed_data, (actions, input, done))
|
|
|
|
while True:
|
|
try:
|
|
yield output.get(True, .01)
|
|
except Empty:
|
|
if done.is_set() and all(r.ready() for r in results):
|
|
break
|
|
|
|
pool.close()
|
|
pool.join()
|