import eventlet
eventlet.monkey_patch()

import json
import time
import zookeeper

from eventlet import tpool

TIMES = []

class Session(object):
    """Session interface that makes use of the python-zookeeper bindings.
    """
    def __init__(self):
        self.connected = False
        self.handle = None

    def handle_session_event(self, handle, event_type, conn_state, path):
        """Fires when there is a change in the state of the zk session.

        :param handle: ZK session handle.
        :param event_type: The type of zk event.
        :param conn_state: The new state of the connection.
        :param path: The path that resulted in this callback (if any).
        """
        if conn_state == zookeeper.CONNECTED_STATE:
            TIMES.append(time.time() - self.time_started)
            self.connected = True
            return
        if conn_state == zookeeper.CONNECTING_STATE:
            self.connected = False
            return

        raise NotImplementedError('Unhandled Event')

    def connect(self):
        session_id = -1
        zk_pass = ''
        recv_timeout = 10000
        zk_endpoint = 'localhost:2181'

        self.time_started = time.time()
        self.handle = tpool.execute(zookeeper.init,
                                    zk_endpoint,
                                    self.handle_session_event,
                                    recv_timeout,
                                    (session_id, zk_pass))


        for x in range(15):
            if self.connected:
                return
            eventlet.sleep(0.01)
        else:
            print "Session did not establish within 150ms"
            while 1:
                if self.connected:
                    return
                eventlet.sleep(0.01)
            return

    def disconnect(self):
        """Close the active zookeeper connection.
        """
        if self.handle is not None:
            tpool.execute(zookeeper.close, self.handle)
            self.connected = False
            self.handle = None

if __name__ == '__main__':
    zklog_fd = open('/tmp/zk_log.log', 'a+')
    zookeeper.set_log_stream(zklog_fd)

    count = 0
    for x in range(3000):
        session = Session()
        session.connect()
        session.disconnect()
        count += 1
        if count % 100 == 0:
            print count

    with open('output', 'w') as f:
        json.dump(TIMES, f)