Python kazoo.client.KazooState.CONNECTED Examples

def _lock_listener(self, state):
        """
        Listener to handle ZK disconnection/reconnection. Since I don't know of safe way to check if a lock is still in
        ZK after a reconnect, we simply release the lock and try and re-acquire it.

        Args:
            state (kazoo.client.KazooState): The state of the ZK connection

        Returns:
            None
        """

        if state in [KazooState.LOST, KazooState.SUSPENDED]:
            self._logger.warn(u'Disconnected from Zookeeper, waiting to reconnect lock for {}'.format(str(self)))
            self._locked = False
        elif state == KazooState.CONNECTED:
            self._logger.warn(
                    u'Reconnected to Zookeeper, trying to release and re-acquire lock for {}'.format(str(self)))
            self._context.zookeeper_client.handler.spawn(self._release_and_reacquire)
        else:
            self._logger.warn(u'Got unknown state "{}" from Zookeeper'.format(state)) 

def specialized(self, msg):

        assert 'request' in msg, 'bogus message received ?'
        req = msg['request']
        if req == 'state change':

            #
            # - we got a zk state change
            # - when ok is off the scheduling will be interrupted
            #
            state = msg['state']
            self.connected = state == KazooState.CONNECTED

        elif req == 'execute':

            #
            # - request to run some code, append to our FIFO
            #
            self.pending.append(msg)

        else:
            super(ZK, self).specialized(msg) 

def leave_cluster(self):
        self.started_shutdown = True

        # cancel any attempts to acquire master lock which could make us hang
        self.lock.cancel()

        if self.zk.state == KazooState.CONNECTED:
            if self.party and self.party.participating:
                logger.info('Leaving party')
                self.party.leave()
            if self.lock and self.lock.is_acquired:
                logger.info('Releasing lock')
                self.lock.release()

        # Make us not the master
        self.is_master = False

        # Avoid sending metrics that we are still the master when we're not
        metrics.set('is_master_sender', 0) 

def state_listener(self, state):
        '''
        Restarts the session if we get anything besides CONNECTED
        '''
        if state == KazooState.SUSPENDED:
            self.set_valid(False)
            self.call_error(self.BAD_CONNECTION)
        elif state == KazooState.LOST and not self.do_not_restart:
            self.threaded_start()
        elif state == KazooState.CONNECTED:
            # This is going to throw a SUSPENDED kazoo error
            # which will cause the sessions to be wiped and re established.
            # Used b/c of massive connection pool issues
            self.zoo_client.stop() 

def reconnection_listener(self) -> None:
        attempts = 0
        while attempts < 5:
            if self.client.state == KazooState.CONNECTED:
                self.log.warning("Zookeeper connection recovered!")
                self.waiting_for_reconnect = False
                return
            self.log.warning("Waiting for zookeeper connection to recover")
            time.sleep(5)
            attempts += 1
        self.log.error("Connection did not recover, abdicating!")
        self._terminate() 

def test_connection_listener(self):
        with mock.patch(
            "paasta_tools.deployd.leader.PaastaThread", autospec=True
        ) as mock_paasta_thread:
            self.election.connection_listener(KazooState.CONNECTED)
            self.election.connection_listener(KazooState.SUSPENDED)
            mock_paasta_thread.assert_called_with(
                target=self.election.reconnection_listener
            )
            assert self.election.waiting_for_reconnect
            self.election.connection_listener(KazooState.LOST)
            self.mock_control.put.assert_called_with("ABORT") 

def test_reconnection_listener(self):
        self.mock_client.state = KazooState.CONNECTED
        self.election.reconnection_listener()
        assert not self.election.waiting_for_reconnect
        assert not self.mock_control.put.called

        self.mock_client.state = KazooState.SUSPENDED
        self.election.waiting_for_reconnect = True
        with mock.patch("time.sleep", autospec=True):
            self.election.reconnection_listener()
        assert self.election.waiting_for_reconnect
        self.mock_control.put.assert_called_with("ABORT") 

def specialized(self, msg):

        assert 'request' in msg, 'bogus message received ?'
        req = msg['request']
        if req == 'state change':

            #
            # - we got a zk state change
            # - we only use the switch to CONNECTED to go from wait_for_cnx() to spin()
            # - ZK disconnects (LOST or SUSPENDED) are simply flagged when exceptions are raised
            #
            state = msg['state']
            current = 'connected' if self.connected else 'disconnected'
            logger.debug('%s : zk state change -> "%s" (%s)' % (self.path, str(state), current))
            if self.connected and state != KazooState.CONNECTED:
                logger.warning('%s : lost connection (%s) / forcing a reset' % (self.path, str(state)))
                self.force_reset = 1
                self.connected = 0

            elif state == KazooState.CONNECTED:
                self.connected = 1

        elif req == 'reset':

            #
            # - we got a request to explicitly force a reset
            # - this is typically invoked from the CLI
            #
            self.force_reset = 1

        else:
            super(ZK, self).specialized(msg) 

def update_status(self):
        if self.started_shutdown:
            return

        if self.zk.state == KazooState.CONNECTED:
            if self.lock.is_acquired:
                self.is_master = True
            else:
                try:
                    self.is_master = self.lock.acquire(blocking=False, timeout=2)

                # This one is expected when we're recovering from ZK being down
                except kazoo.exceptions.CancelledError:
                    self.is_master = False

                except kazoo.exceptions.LockTimeout:
                    self.is_master = False
                    logger.exception('Failed trying to acquire lock (shouldn\'t happen as we\'re using nonblocking locks)')

                except kazoo.exceptions.KazooException:
                    self.is_master = False
                    logger.exception('ZK problem while Failed trying to acquire lock')
        else:
            logger.error('ZK connection is in %s state', self.zk.state)
            self.is_master = False

        if self.zk.state == KazooState.CONNECTED:

            if self.is_master:
                slaves = [self.address_to_tuple(host) for host in self.party if host != self.me]
                self.slave_count = len(slaves)
                self.slaves = cycle(slaves)
            else:
                self.slaves = cycle([])
                self.slave_count = 0

            # Keep us as part of the party, so the current master sees us as a slave
            if not self.party.participating:
                try:
                    self.party.join()
                except kazoo.exceptions.KazooException:
                    logger.exception('ZK problem while trying to join party')
        else:
            self.slaves = cycle([])
            self.slave_count = 0 

def wait_for_cnx(self, data):

        if self.force_reset or self.terminate:
            raise Aborted('resetting')

        #
        # - loop back if we haven't received a CONNECTED event from the driver
        #
        if not self.connected:
            return 'wait_for_cnx', data, SAMPLING

        #
        # - the /pods node holds all our ephemeral per-container data (one container == one child node)
        # - the /hash node stores the last recorded md5 hash (local pods + dependencies), which we use to
        #   flag any change amongst the pods or their dependencies
        #
        data.zk.ensure_path('%s/pods' % self.prefix)
        data.zk.ensure_path('%s/hash' % self.prefix)
        try:

            #
            # - register ourselves by creating an ephemeral
            # - this is where we can store arbitrary information (e.g our breadcrumbs)
            # - we ask for a sequence counter as well which we then keep (e.g in case of connection loss or reset
            #   we guarantee the pod won't get assigned a new index)
            # - this is *critical* for some use-cases (e.g Kafka where the broker index must remain the same)
            #
            path = data.zk.create('%s/pods/%s.' % (self.prefix, self.id), ephemeral=True, sequence=True)
            tokens = path.split('.')
            if self.seq is None:
                self.seq = int(tokens[-1])
            self.breadcrumbs['seq'] = self.seq
            js = json.dumps(self.breadcrumbs)
            data.zk.set(path, js)

        except NodeExistsError:

            #
            # - if the node is already there we just recovered from a zookeeper connection loss
            #   and /snapshot has not been phased out yet .. this is not an issue, simply pause a bit
            #   to re-attempt later
            #
            logger.debug('%s : pod %s is already there (probably a zk reconnect)' % (self.path, self.id))
            return 'wait_for_cnx', data, 5.0 * SAMPLING

        logger.debug('%s : registered as %s (#%d)' % (self.path, self.id, self.seq))
        data.connected_at = time.time()
        return 'spin', data, 0