<commit_before>from __future__ import unicode_literals

import collections
import logging
import random

from tornado import gen, ioloop

from zoonado import protocol, exc
from .connection import Connection
from .states import States, SessionStateMachine
from .retry import RetryPolicy


DEFAULT_ZOOKEEPER_PORT = 2181

MAX_FIND_WAIT = 60  # in seconds

HEARTBEAT_FREQUENCY = 3  # heartbeats per timeout interval


log = logging.getLogger(__name__)


class Session(object):

    def __init__(self, servers, timeout, retry_policy, allow_read_only):
        self.hosts = []
        for server in servers.split(","):
            if ":" in server:
                host, port = server.split(":")
            else:
                host = server
                port = DEFAULT_ZOOKEEPER_PORT

            self.hosts.append((host, port))

        self.conn = None
        self.state = SessionStateMachine()

        self.retry_policy = retry_policy or RetryPolicy.forever()
        self.allow_read_only = allow_read_only

        self.xid = 0
        self.last_zxid = None

        self.session_id = None
        self.timeout = timeout
        self.password = b'\<PASSWORD>'

        self.heartbeat_handle = None

        self.watch_callbacks = collections.defaultdict(set)

        self.closing = False

    @gen.coroutine
    def ensure_safe_state(self, writing=False):
        safe_states = [States.CONNECTED]
        if self.allow_read_only and not writing:
            safe_states.append(States.READ_ONLY)

        if self.state in safe_states:
            return

        yield self.state.wait_for(*safe_states)

    @gen.coroutine
    def start(self):
        io_loop = ioloop.IOLoop.current()
        io_loop.add_callback(self.set_heartbeat)
        io_loop.add_callback(self.repair_loop)

        yield self.ensure_safe_state()

    @gen.coroutine
    def find_server(self, allow_read_only):
        conn = None

        retry_policy = RetryPolicy.exponential_backoff(maximum=MAX_FIND_WAIT)

        while not conn:
            yield retry_policy.enforce()

            servers = random.sample(self.hosts, len(self.hosts))
            for host, port in servers:
                log.info("Connecting to %s:%s", host, port)
                conn = yield self.make_connection(host, port)
                if not conn or (conn.start_read_only and not allow_read_only):
                    continue

            if not conn:
                log.warn("No servers available, will keep trying.")

        old_conn = self.conn
        self.conn = conn

        io_loop = ioloop.IOLoop.current()

        if old_conn:
            io_loop.add_callback(old_conn.close, self.timeout)

        if conn.start_read_only:
            io_loop.add_callback(self.find_server, allow_read_only=False)

    @gen.coroutine
    def make_connection(self, host, port):
        conn = Connection(host, port, watch_handler=self.event_dispatch)
        try:
            yield conn.connect()
        except Exception:
            log.exception("Couldn't connect to %s:%s", host, port)
            return

        raise gen.Return(conn)

    @gen.coroutine
    def establish_session(self):
        log.info("Establising session.")
        zxid, response = yield self.conn.send_connect(
            protocol.ConnectRequest(
                protocol_version=0,
                last_seen_zxid=self.last_zxid or 0,
                timeout=int((self.timeout or 0) * 1000),
                session_id=self.session_id or 0,
                password=<PASSWORD>,
                read_only=self.allow_read_only,
            )
        )
        self.last_zxid = zxid

        if response.session_id == 0:  # invalid session, probably expired
            self.state.transition_to(States.LOST)
            raise exc.SessionLost()

        log.info("Got session id %s", hex(response.session_id))
        log.info("Negotiated timeout: %s seconds", response.timeout / 1000)

        self.session_id = response.session_id
        self.password = response.password
        self.timeout = response.timeout / 1000

        self.last_zxid = zxid

    @gen.coroutine
    def repair_loop(self):
        while not self.closing:
            yield self.state.wait_for(States.SUSPENDED, States.LOST)
            if self.closing:
                break

            yield self.find_server(allow_read_only=self.allow_read_only)

            session_was_lost = self.state == States.LOST

            try:
                yield self.establish_session()
            except exc.SessionLost:
                self.conn.abort(exc.SessionLost)
                yield self.conn.close()
                self.session_id = None
                self.password = b'\<PASSWORD>'
                continue

            if self.conn.start_read_only:
                self.state.transition_to(States.READ_ONLY)
            else:
                self.state.transition_to(States.CONNECTED)<commit_msg>Update session.py<commit_after>from __future__ import unicode_literals

import collections
import logging
import random

from tornado import gen, ioloop

from zoonado import protocol, exc
from .connection import Connection
from .states import States, SessionStateMachine
from .retry import RetryPolicy


DEFAULT_ZOOKEEPER_PORT = 2181

MAX_FIND_WAIT = 60  # in seconds

HEARTBEAT_FREQUENCY = 3  # heartbeats per timeout interval


log = logging.getLogger(__name__)


class Session(object):

    def __init__(self, servers, timeout, retry_policy, allow_read_only):
        self.hosts = []
        for server in servers.split(","):
            if ":" in server:
                host, port = server.split(":")
            else:
                host = server
                port = DEFAULT_ZOOKEEPER_PORT

            self.hosts.append((host, port))

        self.conn = None
        self.state = SessionStateMachine()

        self.retry_policy = retry_policy or RetryPolicy.forever()
        self.allow_read_only = allow_read_only

        self.xid = 0
        self.last_zxid = None

        self.session_id = None
        self.timeout = timeout
        self.password = b'\x00'

        self.heartbeat_handle = None

        self.watch_callbacks = collections.defaultdict(set)

        self.closing = False

    @gen.coroutine
    def ensure_safe_state(self, writing=False):
        safe_states = [States.CONNECTED]
        if self.allow_read_only and not writing:
            safe_states.append(States.READ_ONLY)

        if self.state in safe_states:
            return

        yield self.state.wait_for(*safe_states)

    @gen.coroutine
    def start(self):
        io_loop = ioloop.IOLoop.current()
        io_loop.add_callback(self.set_heartbeat)
        io_loop.add_callback(self.repair_loop)

        yield self.ensure_safe_state()

    @gen.coroutine
    def find_server(self, allow_read_only):
        conn = None

        retry_policy = RetryPolicy.exponential_backoff(maximum=MAX_FIND_WAIT)

        while not conn:
            yield retry_policy.enforce()

            servers = random.sample(self.hosts, len(self.hosts))
            for host, port in servers:
                log.info("Connecting to %s:%s", host, port)
                conn = yield self.make_connection(host, port)
                if not conn or (conn.start_read_only and not allow_read_only):
                    continue

            if not conn:
                log.warn("No servers available, will keep trying.")

        old_conn = self.conn
        self.conn = conn

        io_loop = ioloop.IOLoop.current()

        if old_conn:
            io_loop.add_callback(old_conn.close, self.timeout)

        if conn.start_read_only:
            io_loop.add_callback(self.find_server, allow_read_only=False)

    @gen.coroutine
    def make_connection(self, host, port):
        conn = Connection(host, port, watch_handler=self.event_dispatch)
        try:
            yield conn.connect()
        except Exception:
            log.exception("Couldn't connect to %s:%s", host, port)
            return

        raise gen.Return(conn)

    @gen.coroutine
    def establish_session(self):
        log.info("Establising session.")
        zxid, response = yield self.conn.send_connect(
            protocol.ConnectRequest(
                protocol_version=0,
                last_seen_zxid=self.last_zxid or 0,
                timeout=int((self.timeout or 0) * 1000),
                session_id=self.session_id or 0,
                password=self.password,
                read_only=self.allow_read_only,
            )
        )
        self.last_zxid = zxid

        if response.session_id == 0:  # invalid session, probably expired
            self.state.transition_to(States.LOST)
            raise exc.SessionLost()

        log.info("Got session id %s", hex(response.session_id))
        log.info("Negotiated timeout: %s seconds", response.timeout / 1000)

        self.session_id = response.session_id
        self.password = <PASSWORD>
        self.timeout = response.timeout / 1000

        self.last_zxid = zxid

    @gen.coroutine
    def repair_loop(self):
        while not self.closing:
            yield self.state.wait_for(States.SUSPENDED, States.LOST)
            if self.closing:
                break

            yield self.find_server(allow_read_only=self.allow_read_only)

            session_was_lost = self.state == States.LOST

            try:
                yield self.establish_session()
            except exc.SessionLost:
                self.conn.abort(exc.SessionLost)
                yield self.conn.close(self.timeout)
                self.session_id = None
                self.password = b'\<PASSWORD>'
                continue

            if self.conn.start_read_only:
                self.state.transition_to(States.READ_ONLY)
            else:
                self.state.transition_to(States.CONNECTED)
{'fraction_non_alphanumeric': 0.08320618451836029, 'fraction_numerical': 0.005289390702878649, 'mean_word_length': 2.014101778050276, 'pattern_counts': {'":': 4, '<': 8, '<?xml version=': 0, '>': 8, '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': '7177290', 'n_tokens_mistral': 2959, 'n_tokens_neox': 2774, 'n_words': 604}