Spaces:
Running
Running
| # Copyright DataStax, Inc. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| from collections import namedtuple | |
| from heapq import heappush, heappop | |
| from itertools import cycle | |
| from threading import Condition | |
| import sys | |
| from cassandra.cluster import ResultSet, EXEC_PROFILE_DEFAULT | |
| import logging | |
| log = logging.getLogger(__name__) | |
| ExecutionResult = namedtuple('ExecutionResult', ['success', 'result_or_exc']) | |
| def execute_concurrent(session, statements_and_parameters, concurrency=100, raise_on_first_error=True, results_generator=False, execution_profile=EXEC_PROFILE_DEFAULT): | |
| """ | |
| Executes a sequence of (statement, parameters) tuples concurrently. Each | |
| ``parameters`` item must be a sequence or :const:`None`. | |
| The `concurrency` parameter controls how many statements will be executed | |
| concurrently. When :attr:`.Cluster.protocol_version` is set to 1 or 2, | |
| it is recommended that this be kept below 100 times the number of | |
| core connections per host times the number of connected hosts (see | |
| :meth:`.Cluster.set_core_connections_per_host`). If that amount is exceeded, | |
| the event loop thread may attempt to block on new connection creation, | |
| substantially impacting throughput. If :attr:`~.Cluster.protocol_version` | |
| is 3 or higher, you can safely experiment with higher levels of concurrency. | |
| If `raise_on_first_error` is left as :const:`True`, execution will stop | |
| after the first failed statement and the corresponding exception will be | |
| raised. | |
| `results_generator` controls how the results are returned. | |
| * If :const:`False`, the results are returned only after all requests have completed. | |
| * If :const:`True`, a generator expression is returned. Using a generator results in a constrained | |
| memory footprint when the results set will be large -- results are yielded | |
| as they return instead of materializing the entire list at once. The trade for lower memory | |
| footprint is marginal CPU overhead (more thread coordination and sorting out-of-order results | |
| on-the-fly). | |
| `execution_profile` argument is the execution profile to use for this | |
| request, it is passed directly to :meth:`Session.execute_async`. | |
| A sequence of ``ExecutionResult(success, result_or_exc)`` namedtuples is returned | |
| in the same order that the statements were passed in. If ``success`` is :const:`False`, | |
| there was an error executing the statement, and ``result_or_exc`` will be | |
| an :class:`Exception`. If ``success`` is :const:`True`, ``result_or_exc`` | |
| will be the query result. | |
| Example usage:: | |
| select_statement = session.prepare("SELECT * FROM users WHERE id=?") | |
| statements_and_params = [] | |
| for user_id in user_ids: | |
| params = (user_id, ) | |
| statements_and_params.append((select_statement, params)) | |
| results = execute_concurrent( | |
| session, statements_and_params, raise_on_first_error=False) | |
| for (success, result) in results: | |
| if not success: | |
| handle_error(result) # result will be an Exception | |
| else: | |
| process_user(result[0]) # result will be a list of rows | |
| Note: in the case that `generators` are used, it is important to ensure the consumers do not | |
| block or attempt further synchronous requests, because no further IO will be processed until | |
| the consumer returns. This may also produce a deadlock in the IO event thread. | |
| """ | |
| if concurrency <= 0: | |
| raise ValueError("concurrency must be greater than 0") | |
| if not statements_and_parameters: | |
| return [] | |
| executor = ConcurrentExecutorGenResults(session, statements_and_parameters, execution_profile) \ | |
| if results_generator else ConcurrentExecutorListResults(session, statements_and_parameters, execution_profile) | |
| return executor.execute(concurrency, raise_on_first_error) | |
| class _ConcurrentExecutor(object): | |
| max_error_recursion = 100 | |
| def __init__(self, session, statements_and_params, execution_profile): | |
| self.session = session | |
| self._enum_statements = enumerate(iter(statements_and_params)) | |
| self._execution_profile = execution_profile | |
| self._condition = Condition() | |
| self._fail_fast = False | |
| self._results_queue = [] | |
| self._current = 0 | |
| self._exec_count = 0 | |
| self._exec_depth = 0 | |
| def execute(self, concurrency, fail_fast): | |
| self._fail_fast = fail_fast | |
| self._results_queue = [] | |
| self._current = 0 | |
| self._exec_count = 0 | |
| with self._condition: | |
| for n in range(concurrency): | |
| if not self._execute_next(): | |
| break | |
| return self._results() | |
| def _execute_next(self): | |
| # lock must be held | |
| try: | |
| (idx, (statement, params)) = next(self._enum_statements) | |
| self._exec_count += 1 | |
| self._execute(idx, statement, params) | |
| return True | |
| except StopIteration: | |
| pass | |
| def _execute(self, idx, statement, params): | |
| self._exec_depth += 1 | |
| try: | |
| future = self.session.execute_async(statement, params, timeout=None, execution_profile=self._execution_profile) | |
| args = (future, idx) | |
| future.add_callbacks( | |
| callback=self._on_success, callback_args=args, | |
| errback=self._on_error, errback_args=args) | |
| except Exception as exc: | |
| # If we're not failing fast and all executions are raising, there is a chance of recursing | |
| # here as subsequent requests are attempted. If we hit this threshold, schedule this result/retry | |
| # and let the event loop thread return. | |
| if self._exec_depth < self.max_error_recursion: | |
| self._put_result(exc, idx, False) | |
| else: | |
| self.session.submit(self._put_result, exc, idx, False) | |
| self._exec_depth -= 1 | |
| def _on_success(self, result, future, idx): | |
| future.clear_callbacks() | |
| self._put_result(ResultSet(future, result), idx, True) | |
| def _on_error(self, result, future, idx): | |
| self._put_result(result, idx, False) | |
| class ConcurrentExecutorGenResults(_ConcurrentExecutor): | |
| def _put_result(self, result, idx, success): | |
| with self._condition: | |
| heappush(self._results_queue, (idx, ExecutionResult(success, result))) | |
| self._execute_next() | |
| self._condition.notify() | |
| def _results(self): | |
| with self._condition: | |
| while self._current < self._exec_count: | |
| while not self._results_queue or self._results_queue[0][0] != self._current: | |
| self._condition.wait() | |
| while self._results_queue and self._results_queue[0][0] == self._current: | |
| _, res = heappop(self._results_queue) | |
| try: | |
| self._condition.release() | |
| if self._fail_fast and not res[0]: | |
| raise res[1] | |
| yield res | |
| finally: | |
| self._condition.acquire() | |
| self._current += 1 | |
| class ConcurrentExecutorListResults(_ConcurrentExecutor): | |
| _exception = None | |
| def execute(self, concurrency, fail_fast): | |
| self._exception = None | |
| return super(ConcurrentExecutorListResults, self).execute(concurrency, fail_fast) | |
| def _put_result(self, result, idx, success): | |
| self._results_queue.append((idx, ExecutionResult(success, result))) | |
| with self._condition: | |
| self._current += 1 | |
| if not success and self._fail_fast: | |
| if not self._exception: | |
| self._exception = result | |
| self._condition.notify() | |
| elif not self._execute_next() and self._current == self._exec_count: | |
| self._condition.notify() | |
| def _results(self): | |
| with self._condition: | |
| while self._current < self._exec_count: | |
| self._condition.wait() | |
| if self._exception and self._fail_fast: | |
| raise self._exception | |
| if self._exception and self._fail_fast: # raise the exception even if there was no wait | |
| raise self._exception | |
| return [r[1] for r in sorted(self._results_queue)] | |
| def execute_concurrent_with_args(session, statement, parameters, *args, **kwargs): | |
| """ | |
| Like :meth:`~cassandra.concurrent.execute_concurrent()`, but takes a single | |
| statement and a sequence of parameters. Each item in ``parameters`` | |
| should be a sequence or :const:`None`. | |
| Example usage:: | |
| statement = session.prepare("INSERT INTO mytable (a, b) VALUES (1, ?)") | |
| parameters = [(x,) for x in range(1000)] | |
| execute_concurrent_with_args(session, statement, parameters, concurrency=50) | |
| """ | |
| return execute_concurrent(session, zip(cycle((statement,)), parameters), *args, **kwargs) | |