<commit_before>from futures import ThreadPoolExecutor
import logging
import time
from threading import Lock, RLock, Thread, Event
import Queue
import weakref
from functools import partial

from cassandra import ConsistencyLevel
from cassandra.connection import Connection, ConnectionException
from cassandra.decoder import (QueryMessage, ResultMessage,
                               ErrorMessage, ReadTimeoutErrorMessage,
                               WriteTimeoutErrorMessage,
                               UnavailableErrorMessage,
                               OverloadedErrorMessage,
                               IsBootstrappingErrorMessage)
from cassandra.metadata import Metadata
from cassandra.policies import (RoundRobinPolicy, SimpleConvictionPolicy,
                                ExponentialReconnectionPolicy, HostDistance,
                                RetryPolicy)
from cassandra.query import SimpleStatement
from cassandra.pool import (AuthenticationException, _ReconnectionHandler,
                            _HostReconnectionHandler, HostConnectionPool)


log = logging.getLogger(__name__)


def refresh_schema_and_set_result(keyspace, table, control_conn, response_future):
    try:
        control_conn.refresh_schema(keyspace, table)
    except:
        log.exception("Exception refreshing schema in response to schema change:")
    finally:
        response_future._set_final_result(None)


_NO_RESULT_YET = object()


class ResponseFuture(object):

    def __init__(self, session, message, query):
        self.session = session
        self.message = message
        self.query = query

        # convert the list/generator/etc to an iterator so that subsequent
        # calls to send_request (which retries may do) will resume where
        # they last left off
        self.query_plan = iter(session._load_balancer.make_query_plan(query))

        self._req_id = None
        self._final_result = _NO_RESULT_YET
        self._final_exception = None
        self._current_host = self._current_pool = self._connection = None
        self._event = Event()
        self._errors = {}
        self._query_retries = 0
        self._callback = self._errback = None

    def __del__(self):
        if hasattr(self, 'session'):
            del self.session

    def send_request(self):
        # query_plan is an iterator, so this will resume where we last left
        # off if send_request() is called multiple times
        for host in self.query_plan:
            req_id = self._query(host)
            if req_id is not None:
                self._req_id = req_id
                return

        self._final_exception = NoHostAvailable(self._errors)

    def _query(self, host):
        pool = self.session._pools.get(host)
        if not pool or pool.is_shutdown:
            self._errors[host] = ConnectionException("Pool is shutdown")
            return None

        connection = None
        try:
            # TODO get connectTimeout from cluster settings
            connection = pool.borrow_connection(timeout=2.0)
            request_id = connection.send_msg(self.message, cb=self._set_result)
            self._current_host = host
            self._current_pool = pool
            self._connection = connection
            return request_id
        except Exception, exc:
            self._errors[host] = exc
            if connection:
                pool.return_connection(connection)
            return None

    def _set_result(self, response):
        self._current_pool.return_connection(self._connection)

        if isinstance(response, ResultMessage):
            if response.kind == ResultMessage.KIND_SET_KEYSPACE:
                self.session.set_keyspace(response.results)
                self._set_final_result(None)
            elif response.kind == ResultMessage.KIND_SCHEMA_CHANGE:
                # refresh the schema before responding, but do it in another
                # thread instead of the event loop thread
                self.session.submit(
                    refresh_schema_and_set_result,
                    response.results['keyspace'],
                    response.results['table'],
                    self.session.cluster.control_connection,
                    self)
            else:
                self._set_final_result(getattr(response, 'results', None))
        elif isinstance(response, ErrorMessage):
            retry_policy = self.query.retry_policy
            if not retry_policy:
                retry_policy = self.session.cluster.retry_policy_factory()

            if isinstance(response, ReadTimeoutErrorMessage):
                retry = retry_policy.on_read_timeout(
                    self.query, retry_num=self._query_retries, **response.info)
            elif isinstance(response, WriteTimeoutErrorMessage):
                retry = retry_policy.on_write_timeout(
                    self.query, retry_num=self._query_retries, **response.info)
            elif isinstance(response, UnavailableErrorMessage):
                retry = retry_policy.on_unavailable(
                    self.query, retry_num=self._query_retries, **response.info)
            elif isinstance(response, OverloadedErrorMessage):
                # need to retry against a different host here
                log.warn("Host %s is overloaded, retrying against a different "
                         "host" % (self._current_host))
                self._retry(reuse_connection=False, consistency_level=None)
                return
            elif isinstance(response, IsBootstrappingErrorMessage):
                # need to retry against a different host here
                self._retry(reuse_connection=False, consistency_level=None)
                return
            # TODO need to define the PreparedQueryNotFound class
            # elif isinstance(response, PreparedQueryNotFound):
            #     pass
            else:
                self._set_final_exception(response)
                return

            retry_type, consistency = retry
            if retry_type is RetryPolicy.RETRY:
                self._query_retries += 1
                self._retry(reuse_connection=True, consistency_level=consistency)
            elif retry_type is RetryPolicy.RETHROW:
                self._set_final_exception(response)
            else:  # IGNORE
                self._set_final_result(None)
        elif isinstance(response, Exception):
            self._set_final_exception(response)
        else:
            # we got some other kind of response message
            msg = "Got unexpected message: %r" % (response,)
            exc = ConnectionException(msg, self._current_host)
            self._connection.defunct(exc)
            self._set_final_exception(exc)

    def _set_final_result(self, response):
        del self.session  # clear reference cycles
        self._final_result = response
        self._event.set()
        if self._callback:
            fn, args, kwargs = self._callback
            fn(response, *args, **kwargs)

    def _set_final_exception(self, response):
        del self.session  # clear reference cycles
        self._final_exception = response
        self._event.set()
        if self._errback:
            fn, args, kwargs = self._errback
            fn(response, *args, **kwargs)

    def _retry(self, reuse_connection, consistency_level):
        if consistency_level is not None:
            self.message.consistency_level = consistency_level

        # don't retry on the event loop thread
        self.session.submit(self._retry_task, reuse_connection)

    def _retry_task(self, reuse_connection):
        if reuse_connection and self._query(self._current_host):
            return

        # otherwise, move onto another host
        self.send_request()

    def deliver(self):
        if self._final_result is not _NO_RESULT_YET:
            return self._final_result<commit_msg>Handle responses to orphaned ResponseFutures<commit_after>from futures import ThreadPoolExecutor
import logging
import time
from threading import Lock, RLock, Thread, Event
import Queue
import weakref
from functools import partial

from cassandra import ConsistencyLevel
from cassandra.connection import Connection, ConnectionException
from cassandra.decoder import (QueryMessage, ResultMessage,
                               ErrorMessage, ReadTimeoutErrorMessage,
                               WriteTimeoutErrorMessage,
                               UnavailableErrorMessage,
                               OverloadedErrorMessage,
                               IsBootstrappingErrorMessage)
from cassandra.metadata import Metadata
from cassandra.policies import (RoundRobinPolicy, SimpleConvictionPolicy,
                                ExponentialReconnectionPolicy, HostDistance,
                                RetryPolicy)
from cassandra.query import SimpleStatement
from cassandra.pool import (AuthenticationException, _ReconnectionHandler,
                            _HostReconnectionHandler, HostConnectionPool)


log = logging.getLogger(__name__)


def refresh_schema_and_set_result(keyspace, table, control_conn, response_future):
    try:
        control_conn.refresh_schema(keyspace, table)
    except:
        log.exception("Exception refreshing schema in response to schema change:")
    finally:
        response_future._set_final_result(None)


_NO_RESULT_YET = object()


class ResponseFuture(object):

    def __init__(self, session, message, query):
        self.session = session
        self.message = message
        self.query = query

        # convert the list/generator/etc to an iterator so that subsequent
        # calls to send_request (which retries may do) will resume where
        # they last left off
        self.query_plan = iter(session._load_balancer.make_query_plan(query))

        self._req_id = None
        self._final_result = _NO_RESULT_YET
        self._final_exception = None
        self._current_host = self._current_pool = self._connection = None
        self._event = Event()
        self._errors = {}
        self._query_retries = 0
        self._callback = self._errback = None

    def __del__(self):
        if hasattr(self, 'session'):
            del self.session

    def send_request(self):
        # query_plan is an iterator, so this will resume where we last left
        # off if send_request() is called multiple times
        for host in self.query_plan:
            req_id = self._query(host)
            if req_id is not None:
                self._req_id = req_id
                return

        self._final_exception = NoHostAvailable(self._errors)

    def _query(self, host):
        pool = self.session._pools.get(host)
        if not pool or pool.is_shutdown:
            self._errors[host] = ConnectionException("Pool is shutdown")
            return None

        connection = None
        try:
            # TODO get connectTimeout from cluster settings
            connection = pool.borrow_connection(timeout=2.0)
            request_id = connection.send_msg(self.message, cb=self._set_result)
            self._current_host = host
            self._current_pool = pool
            self._connection = connection
            return request_id
        except Exception, exc:
            self._errors[host] = exc
            if connection:
                pool.return_connection(connection)
            return None

    def _set_result(self, response):
        if self._current_pool:
            self._current_pool.return_connection(self._connection)

        if isinstance(response, ResultMessage):
            if response.kind == ResultMessage.KIND_SET_KEYSPACE:
                session = getattr(self, 'session', None)
                if session:
                    session.keyspace = response.results
                self._set_final_result(None)
            elif response.kind == ResultMessage.KIND_SCHEMA_CHANGE:
                # refresh the schema before responding, but do it in another
                # thread instead of the event loop thread
                self.session.submit(
                    refresh_schema_and_set_result,
                    response.results['keyspace'],
                    response.results['table'],
                    self.session.cluster.control_connection,
                    self)
            else:
                self._set_final_result(getattr(response, 'results', None))
        elif isinstance(response, ErrorMessage):
            retry_policy = self.query.retry_policy
            if not retry_policy:
                retry_policy = self.session.cluster.retry_policy_factory()

            if isinstance(response, ReadTimeoutErrorMessage):
                retry = retry_policy.on_read_timeout(
                    self.query, retry_num=self._query_retries, **response.info)
            elif isinstance(response, WriteTimeoutErrorMessage):
                retry = retry_policy.on_write_timeout(
                    self.query, retry_num=self._query_retries, **response.info)
            elif isinstance(response, UnavailableErrorMessage):
                retry = retry_policy.on_unavailable(
                    self.query, retry_num=self._query_retries, **response.info)
            elif isinstance(response, OverloadedErrorMessage):
                # need to retry against a different host here
                log.warn("Host %s is overloaded, retrying against a different "
                         "host" % (self._current_host))
                self._retry(reuse_connection=False, consistency_level=None)
                return
            elif isinstance(response, IsBootstrappingErrorMessage):
                # need to retry against a different host here
                self._retry(reuse_connection=False, consistency_level=None)
                return
            # TODO need to define the PreparedQueryNotFound class
            # elif isinstance(response, PreparedQueryNotFound):
            #     pass
            else:
                self._set_final_exception(response)
                return

            retry_type, consistency = retry
            if retry_type is RetryPolicy.RETRY:
                self._query_retries += 1
                self._retry(reuse_connection=True, consistency_level=consistency)
            elif retry_type is RetryPolicy.RETHROW:
                self._set_final_exception(response)
            else:  # IGNORE
                self._set_final_result(None)
        elif isinstance(response, Exception):
            self._set_final_exception(response)
        else:
            # we got some other kind of response message
            msg = "Got unexpected message: %r" % (response,)
            exc = ConnectionException(msg, self._current_host)
            self._connection.defunct(exc)
            self._set_final_exception(exc)

    def _set_final_result(self, response):
        if hasattr(self, 'session'):
            del self.session  # clear reference cycles
        self._final_result = response
        self._event.set()
        if self._callback:
            fn, args, kwargs = self._callback
            fn(response, *args, **kwargs)

    def _set_final_exception(self, response):
        if hasattr(self, 'session'):
            del self.session  # clear reference cycles
        self._final_exception = response
        self._event.set()
        if self._errback:
            fn, args, kwargs = self._errback
            fn(response, *args, **kwargs)

    def _retry(self, reuse_connection, consistency_level):
        if consistency_level is not None:
            self.message.consistency_level = consistency_level

        # don't retry on the event loop thread
        self.session.submit(self._retry_task, reuse_connection)

    def _retry_task(self, reuse_connection):
        if reuse_connection and self._query(self._current_host):
            return

        # otherwise, move onto another host
        self.send_request()

    def deliver(self):
        if self._final_result is not _NO_RESULT_YET:
            return self._final_result
{'fraction_non_alphanumeric': 0.06563104373323352, 'fraction_numerical': 0.0004990953896063385, 'mean_word_length': 2.082692307692308, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2929988', 'n_tokens_mistral': 4001, 'n_tokens_neox': 3869, 'n_words': 1000}