class TestPartitionOwnerRegistry(KazooTestCase): """Test the methods of :class:`samsa.consumer.PartitionOwnerRegistry`. """ def setUp(self): super(TestPartitionOwnerRegistry, self).setUp() self.c = Cluster(self.client) self.c.brokers = mock.MagicMock() broker = mock.Mock() broker.id = 1 self.c.brokers.__getitem__.return_value = broker self.consumer = mock.Mock() self.consumer.id = 1234 self.topic = mock.Mock() self.topic.name = 'topic' self.por = PartitionOwnerRegistry( self.consumer, self.c, self.topic, 'group' ) # Create 5 partitions with on the same topic and broker self.partitions = [] for i in xrange(5): self.partitions.append( Partition(self.c, self.topic, broker, i) ) def test_crd(self): """Test partition *c*reate, *r*ead, and *d*elete. """ # Add the first 3 partitions to the registry and see that they're set self.por.add(self.partitions[:3]) self.assertEquals( self.por.get(), set(self.partitions[:3]) ) # Remove the first partition and see that only [1, 2] exist self.por.remove([self.partitions[0]]) self.assertEquals( self.por.get(), set(self.partitions[1:3]) ) def test_grows(self): """Test that the reference returned by :func:`samsa.consumer.partitions.PartitionOwnerRegistry.get` reflects the latest state. """ partitions = self.por.get() self.assertEquals(len(partitions), 0) self.por.add(self.partitions) self.assertEquals(len(partitions), len(self.partitions))
class TestPartitionOwnerRegistry(KazooTestCase): """Test the methods of :class:`samsa.consumer.PartitionOwnerRegistry`. """ @mock.patch('samsa.cluster.BrokerMap') def setUp(self, bm, *args): super(TestPartitionOwnerRegistry, self).setUp() self.c = Cluster(self.client) broker = mock.Mock() broker.id = 1 self.c.brokers.__getitem__.return_value = broker self.consumer = mock.Mock() self.consumer.id = '1234' self.topic = mock.Mock() self.topic.name = 'topic' self.por = PartitionOwnerRegistry(self.consumer, self.c, self.topic, 'group') # Create 5 partitions with on the same topic and broker self.partitions = [] self.message_set_queue = Queue.Queue() for i in xrange(5): self.partitions.append( OwnedPartition(Partition(self.c, self.topic, broker, i), 'group', self.message_set_queue)) @mock.patch.object(OwnedPartition, 'start') def test_crd(self, *args): """Test partition *c*reate, *r*ead, and *d*elete. """ # Add the first 3 partitions to the registry and see that they're set self.por.add(self.partitions[:3]) self.assertEquals(self.por.get(), set(self.partitions[:3])) # Remove the first partition and see that only [1, 2] exist self.por.remove([self.partitions[0]]) self.assertEquals(self.por.get(), set(self.partitions[1:3])) @mock.patch.object(OwnedPartition, 'start') def test_grows(self, *args): """Test that the reference returned by :func:`samsa.consumer.partitions.PartitionOwnerRegistry.get` reflects the latest state. """ partitions = self.por.get() self.assertEquals(len(partitions), 0) self.por.add(self.partitions) self.assertEquals(len(partitions), len(self.partitions))
class Consumer(object): """Primary API for consuming kazoo messages as a group. """ def __init__(self, cluster, topic, group): """ :param cluster: :type cluster: :class:`samsa.cluster.Cluster`. :param topic: The topic to consume messages from. :type topic: :class:`samsa.topics.Topic`. :param group: The consumer group to join. :type group: str. """ self.config = ConsumerConfig().build() self.cluster = cluster self.topic = topic self.group = group self.id = "%s:%s" % (socket.gethostname(), uuid4()) self.id_path = '/consumers/%s/ids' % self.group self.partition_owner_registry = PartitionOwnerRegistry( self, cluster, topic, group) self.partitions = self.partition_owner_registry.get() path = '%s/%s' % (self.id_path, self.id) self.cluster.zookeeper.create( path, self.topic.name, ephemeral=True, makepath=True ) self._rebalance() def _rebalance(self, event=None): """Joins a consumer group and claims partitions. """ logger.info('Rebalancing consumer %s for topic %s.' % ( self.id, self.topic.name) ) zk = self.cluster.zookeeper broker_path = '/brokers/ids' try: zk.get_children(broker_path, watch=self._rebalance) except NoNodeException: raise ImproperlyConfiguredError( 'The broker_path "%s" does not exist in your ' 'ZooKeeper cluster -- is your Kafka cluster running?' % broker_path) # 3. all consumers in the same group as Ci that consume topic T consumer_ids = zk.get_children(self.id_path, watch=self._rebalance) participants = [] for id_ in consumer_ids: topic, stat = zk.get("%s/%s" % (self.id_path, id_)) if topic == self.topic.name: participants.append(id_) # 5. participants.sort() # 6. i = participants.index(self.id) parts_per_consumer = len(self.topic.partitions) / len(participants) remainder_ppc = len(self.topic.partitions) % len(participants) start = parts_per_consumer * i + min(i, remainder_ppc) num_parts = parts_per_consumer + (0 if (i + 1 > remainder_ppc) else 1) # 7. assign partitions from i*N to (i+1)*N - 1 to consumer Ci new_partitions = itertools.islice( self.topic.partitions, start, start + num_parts ) new_partitions = set(new_partitions) self.stop_partitions() # 8. remove current entries from the partition owner registry self.partition_owner_registry.remove( self.partitions - new_partitions ) # 9. add newly assigned partitions to the partition owner registry for i in xrange(self.config['rebalance_retries_max']): try: # N.B. self.partitions will always reflect the most current # view of owned partitions. Therefor retrying this method # will progress. self.partition_owner_registry.add( new_partitions - self.partitions ) break except PartitionOwnedError, e: logger.debug("Someone still owns partition %s. Retrying" % e) time.sleep(i ** 2) continue else:
class Consumer(object): """Primary API for consuming kazoo messages as a group. """ def __init__(self, cluster, topic, group, backoff_increment=1, connect_retries=4, fetch_size=307200, offset_reset='nearest', rebalance_retries=4, ): """ For more info see: samsa.topics.Topic.subscribe :param cluster: :type cluster: :class:`samsa.cluster.Cluster`. :param topic: The topic to consume messages from. :type topic: :class:`samsa.topics.Topic`. :param group: The consumer group to join. :param backoff_increment: How fast to incrementally backoff when a partition has no messages to read. :param connect_retries: Retries before giving up on connecting :param fetch_size: Default fetch size (in bytes) to get from Kafka :param offset_reset: Where to reset when an OffsetOutOfRange happens :param rebalance_retries: Retries before giving up on rebalance :rtype: :class:`samsa.consumer.consumer.Consumer` """ self.connect_retries = connect_retries self.rebalance_retries = rebalance_retries self.cluster = cluster self.topic = topic self.group = group self.id = "%s:%s" % (socket.gethostname(), uuid4()) self.id_path = '/consumers/%s/ids' % self.group self.partition_owner_registry = PartitionOwnerRegistry( self, cluster, topic, group, backoff_increment=backoff_increment, fetch_size=fetch_size, offset_reset=offset_reset) self.partitions = self.partition_owner_registry.get() # Keep track of the partition being read and how much has been read self._current_partition = None self._current_read_ct = 0 # Watches # TODO: This is a *ton* of watches, some of which are duplicated # elsewhere. This should be cleaned up and all watches put # in a single zookeeper connector, like in the Scala driver. self._broker_watcher = None self._consumer_watcher = None self._topic_watcher = None self._topics_watcher = None self._rebalancing = True # To stop rebalance while setting watches self._add_self() def _add_self(self): """Add this consumer to the zookeeper participants. Ensures we don't add more participants than partitions """ for i in xrange(self.connect_retries): time.sleep(i**2) # first run is 0, ensures we sleep before retry participants = self._get_participants() if len(self.topic.partitions) > len(participants): break # some room to spare else: logger.debug("More consumers than partitions. " "Waiting %is to retry" % (i+1) ** 2) else: raise NoAvailablePartitionsError("Couldn't acquire partition. " "More consumers than partitions.") path = '%s/%s' % (self.id_path, self.id) self.cluster.zookeeper.create( path, self.topic.name, ephemeral=True, makepath=True) # Set all our watches and then rebalance self._rebalancing = False broker_path = '/brokers/ids' try: self._broker_watcher = ChildrenWatch( self.cluster.zookeeper, broker_path, self._brokers_changed ) except NoNodeException: raise ImproperlyConfiguredError( 'The broker_path "%s" does not exist in your ' 'ZooKeeper cluster -- is your Kafka cluster running?' % broker_path) topics_path = '/brokers/topics' self._topics_watcher = ChildrenWatch( self.cluster.zookeeper, '/brokers/topics', self._topics_changed ) self._rebalancing = True # Final watch will trigger rebalance self._consumer_watcher = ChildrenWatch( self.cluster.zookeeper, self.id_path, self._consumers_changed ) def _brokers_changed(self, brokers): """Watcher for consumer group changes""" if not self._rebalancing: return logger.info("Rebalance triggered by /brokers/ids change") self._rebalance(self.cluster.zookeeper.get_children(self.id_path)) def _consumers_changed(self, consumer_ids): """Watcher for consumer group changes """ if not self._rebalancing: return logger.info("Rebalance triggered by %s change" % self.id_path) self._rebalance(consumer_ids) def _topic_changed(self, broker_ids): """Watcher for brokers/partition count for a topic """ if not self._rebalancing: return topic_path = '/brokers/topics/%s' % self.topic.name logger.info("Rebalance triggered by %s change" % topic_path) self._rebalance(self.cluster.zookeeper.get_children(self.id_path)) def _topics_changed(self, topics): """Watch for the topic we want to show up, then stop watch """ if self.topic.name in topics: self._topic_watcher = ChildrenWatch( self.cluster.zookeeper, '/brokers/topics/%s' % self.topic.name, self._topic_changed ) return False # stop watch def _get_participants(self, consumer_ids=None): """Get a the other consumers of this topic :param consumer_ids: List of consumer_ids (from ChildrenWatch) """ zk = self.cluster.zookeeper if not consumer_ids: try: consumer_ids = zk.get_children(self.id_path) except NoNodeException: logger.debug("Consumer group doesn't exist. " "No participants to find") return [] participants = [] for id_ in consumer_ids: try: topic, stat = zk.get("%s/%s" % (self.id_path, id_)) if topic == self.topic.name: participants.append(id_) except NoNodeException: pass # disappeared between ``get_children`` and ``get`` participants.sort() return participants def _decide_partitions(self, participants): """Use consumers and partitions to determined owned partitions Give a set of subscribed consumers, every individual consumer should be able to figure out the same distribution of available partitions. It's very, very important this gives the same result on all machines, so things like participant and partition lists are always sorted. The algorithm is to distribute blocks of partitions based on how many participants there are. If there are partitions remaining, the last R participants get one extra, where R is the remainder. """ # Freeze and sort partitions so we always have the same results p_to_str = lambda p: '-'.join( [p.topic.name, str(p.broker.id), str(p.number)] ) all_partitions = list(self.topic.partitions) all_partitions.sort(key=p_to_str) # get start point, # of partitions, and remainder idx = participants.index(self.id) parts_per_consumer = len(all_partitions) / len(participants) remainder_ppc = len(all_partitions) % len(participants) start = parts_per_consumer * idx + min(idx, remainder_ppc) num_parts = parts_per_consumer + (0 if (idx + 1 > remainder_ppc) else 1) # assign partitions from i*N to (i+1)*N - 1 to consumer Ci new_partitions = itertools.islice( all_partitions, start, start + num_parts ) new_partitions = set(new_partitions) logger.info( 'Rebalancing %i participants for %i partitions. ' 'My Partitions: %s -- Consumers: %s --- All Partitions: %s', len(participants), len(all_partitions), [p_to_str(p) for p in new_partitions], str(participants), [p_to_str(p) for p in all_partitions] ) return new_partitions def _rebalance(self, consumer_ids): """Joins a consumer group and claims partitions. """ logger.info('Rebalancing consumer %s for topic %s.' % ( self.id, self.topic.name) ) participants = self._get_participants(consumer_ids=consumer_ids) new_partitions = self._decide_partitions(participants) for i in xrange(self.rebalance_retries): if i > 0: logger.debug("Retrying in %is" % ((i+1) ** 2)) time.sleep(i ** 2) # Make sure nothing's changed while we waited participants = self._get_participants() new_partitions = self._decide_partitions(participants) # Remove old partitions and acquire new ones. old_partitions = self.partitions - new_partitions self.stop_partitions(partitions=old_partitions) self.partition_owner_registry.remove(old_partitions) try: self.partition_owner_registry.add( new_partitions - self.partitions ) break except PartitionOwnedError, e: logger.debug("Someone still owns partition %s.", e) continue else: