private IDictionary <string, IList <KafkaMessageStream <TData> > > Consume <TData>(IDictionary <string, int> topicCountDict, IDecoder <TData> decoder)
{
Logger.Debug("entering consume");
if (topicCountDict == null)
{
throw new ArgumentNullException(nameof(topicCountDict));
}
var dirs = new ZKGroupDirs(this.config.GroupId);
var result = new Dictionary <string, IList <KafkaMessageStream <TData> > >();
string consumerIdString = GetConsumerIdString();
var topicCount = new TopicCount(consumerIdString, topicCountDict);
//// create a queue per topic per consumer thread
var consumerThreadIdsPerTopicMap = topicCount.GetConsumerThreadIdsPerTopic();
foreach (var topic in consumerThreadIdsPerTopicMap.Keys)
{
var streamList = new List <KafkaMessageStream <TData> >();
foreach (string threadId in consumerThreadIdsPerTopicMap[topic])
{
var stream = new BlockingCollection <FetchedDataChunk>(new ConcurrentQueue <FetchedDataChunk>());
this.queues.Add(new Tuple <string, string>(topic, threadId), stream);
streamList.Add(new KafkaMessageStream <TData>(topic, stream, this.config.Timeout, decoder));
}
result.Add(topic, streamList);
Logger.InfoFormat("adding topic {0} and stream to map...", topic);
}
// listener to consumer and partition changes
var loadBalancerListener = new ZKRebalancerListener <TData>(
this.config,
consumerIdString,
this.topicRegistry,
this.GetZkClient(),
this,
queues,
this.fetcher,
result,
topicCount);
if (this.consumerRebalanceHandler != null)
{
loadBalancerListener.ConsumerRebalance += this.consumerRebalanceHandler;
}
stopAsyncRebalancing.Add(loadBalancerListener.StopRebalance);
this.RegisterConsumerInZk(dirs, consumerIdString, topicCount);
//// register listener for session expired event
var zkSessionExpireListener = new ZKSessionExpireListener <TData>(dirs, consumerIdString, topicCount, loadBalancerListener, this);
if (this.zkSessionDisconnectedHandler != null)
{
zkSessionExpireListener.ZKSessionDisconnected += this.zkSessionDisconnectedHandler;
}
if (this.zkSessionExpiredHandler != null)
{
zkSessionExpireListener.ZKSessionExpired += this.zkSessionExpiredHandler;
}
this.GetZkClient().Subscribe(zkSessionExpireListener);
this.subscribedZookeeperStateCollection.Add(zkSessionExpireListener);
this.GetZkClient().Subscribe(dirs.ConsumerRegistryDir, loadBalancerListener);
this.subscribedChildCollection.Add(new Tuple <string, IZooKeeperChildListener>(dirs.ConsumerRegistryDir, loadBalancerListener));
result.ForEach(topicAndStreams =>
{
// register on broker partition path changes
string partitionPath = ZooKeeperClient.DefaultBrokerTopicsPath + "/" + topicAndStreams.Key;
if (this.GetZkClient().Exists(partitionPath))
{
this.GetZkClient().Subscribe(partitionPath, loadBalancerListener);
this.subscribedChildCollection.Add(new Tuple <string, IZooKeeperChildListener>(partitionPath, loadBalancerListener));
// Create a mapping of all topic partitions and their current leaders
var topicsAndPartitions = ZkUtils.GetPartitionsForTopics(this.GetZkClient(), new[] { topicAndStreams.Key });
Dictionary <string, int> partitionLeaderMap = new Dictionary <string, int>();
foreach (var partitionId in topicsAndPartitions[topicAndStreams.Key])
{
// Find/parse current partition leader for this partition and add it
// to the mapping object
var partitionStatePath = partitionPath + "/partitions/" + partitionId + "/state";
this.GetZkClient().MakeSurePersistentPathExists(partitionStatePath);
int?partitionLeader = ZkUtils.GetLeaderForPartition(this.GetZkClient(), topicAndStreams.Key, int.Parse(partitionId));
partitionLeaderMap.Add(partitionStatePath, partitionLeader.GetValueOrDefault(-1));
}
// listen for changes on the state nodes for the partitions
// this will indicate when a leader switches, or the in sync replicas change
var leaderListener = new ZkPartitionLeaderListener <TData>(loadBalancerListener, partitionLeaderMap);
foreach (var partitionId in topicsAndPartitions[topicAndStreams.Key])
{
var partitionStatePath = partitionPath + "/partitions/" + partitionId + "/state";
this.GetZkClient().Subscribe(partitionStatePath, leaderListener);
this.subscribedZookeeperDataCollection.Add(new Tuple <string, IZooKeeperDataListener>(partitionStatePath, leaderListener));
}
}
else
{
Logger.WarnFormat("The topic path at {0}, does not exist.", partitionPath);
}
});
//// explicitly trigger load balancing for this consumer
Logger.Info("Performing rebalancing. A new consumer has been added to consumer group: " + dirs.ConsumerRegistryDir + ", consumer: " + consumerIdString);
Logger.InfoFormat("Subscribe count: subscribedChildCollection:{0} , subscribedZookeeperStateCollection:{1} subscribedZookeeperDataCollection:{2} "
, subscribedChildCollection.Count, subscribedZookeeperStateCollection.Count, subscribedZookeeperDataCollection.Count);
//// When a new consumer join, need wait for rebalance finish to make sure Fetcher thread started.
loadBalancerListener.AsyncRebalance(DefaultWaitTimeForInitialRebalanceInSeconds * 1000);
return(result);
}
private bool Rebalance()
{
var myTopicThresdIdsMap = this.GetTopicCount(this.consumerIdString).GetConsumerThreadIdsPerTopic();
var cluster = new Cluster(zkClient);
var consumersPerTopicMap = this.GetConsumersPerTopic(this.config.GroupId);
var partitionsPerTopicMap = ZkUtils.GetPartitionsForTopics(this.zkClient, myTopicThresdIdsMap.Keys);
var relevantTopicThreadIdsMap = GetRelevantTopicMap(
myTopicThresdIdsMap,
partitionsPerTopicMap,
this.oldPartitionsPerTopicMap,
consumersPerTopicMap,
this.oldConsumersPerTopicMap);
if (relevantTopicThreadIdsMap.Count <= 0)
{
Logger.InfoFormat(CultureInfo.CurrentCulture, "Consumer {0} with {1} doesn't need to rebalance.", this.consumerIdString, consumersPerTopicMap);
return(true);
}
Logger.Info("Committing all offsets");
this.zkConsumerConnector.CommitOffsets();
Logger.Info("Releasing parittion ownership");
this.ReleasePartitionOwnership();
var queuesToBeCleared = new List <BlockingCollection <FetchedDataChunk> >();
foreach (var item in relevantTopicThreadIdsMap)
{
this.topicRegistry.Remove(item.Key);
this.topicRegistry.Add(item.Key, new Dictionary <Partition, PartitionTopicInfo>());
var topicDirs = new ZKGroupTopicDirs(config.GroupId, item.Key);
var curConsumers = consumersPerTopicMap[item.Key];
var curPartitions = new List <string>(partitionsPerTopicMap[item.Key]);
var numberOfPartsPerConsumer = curPartitions.Count / curConsumers.Count;
var numberOfConsumersWithExtraPart = curPartitions.Count % curConsumers.Count;
Logger.InfoFormat(
CultureInfo.CurrentCulture,
"Consumer {0} rebalancing the following partitions: {1} for topic {2} with consumers: {3}",
this.consumerIdString,
string.Join(",", curPartitions),
item.Key,
string.Join(",", curConsumers));
foreach (string consumerThreadId in item.Value)
{
var myConsumerPosition = curConsumers.IndexOf(consumerThreadId);
if (myConsumerPosition < 0)
{
continue;
}
var startPart = (numberOfPartsPerConsumer * myConsumerPosition) +
Math.Min(myConsumerPosition, numberOfConsumersWithExtraPart);
var numberOfParts = numberOfPartsPerConsumer + (myConsumerPosition + 1 > numberOfConsumersWithExtraPart ? 0 : 1);
if (numberOfParts <= 0)
{
Logger.WarnFormat(CultureInfo.CurrentCulture, "No broker partitions consumed by consumer thread {0} for topic {1}", consumerThreadId, item.Key);
}
else
{
for (int i = startPart; i < startPart + numberOfParts; i++)
{
var partition = curPartitions[i];
Logger.InfoFormat(CultureInfo.CurrentCulture, "{0} attempting to claim partition {1}", consumerThreadId, partition);
if (!this.ProcessPartition(topicDirs, partition, item.Key, consumerThreadId))
{
return(false);
}
}
queuesToBeCleared.Add(queues[new Tuple <string, string>(item.Key, consumerThreadId)]);
}
}
}
this.UpdateFetcher(cluster, queuesToBeCleared);
this.oldPartitionsPerTopicMap = partitionsPerTopicMap;
this.oldConsumersPerTopicMap = consumersPerTopicMap;
return(true);
}
private bool Rebalance(Cluster.Cluster cluster, CancellationTokenSource cancellationTokenSource)
{
var topicCount = GetTopicCount(consumerIdString);
var topicThreadIdsMap = topicCount.GetConsumerThreadIdsPerTopic();
if (!topicThreadIdsMap.Any())
{
Logger.ErrorFormat("Consumer ID is not registered to any topics in ZK. Exiting rebalance");
return(false);
}
var consumersPerTopicMap = GetConsumersPerTopic(config.GroupId);
var brokers = ZkUtils.GetAllBrokersInCluster(zkClient);
if (!brokers.Any())
{
Logger.Warn("No brokers found when trying to rebalance.");
zkClient.Subscribe(ZooKeeperClient.DefaultBrokerIdsPath, this);
zkConsumerConnector.subscribedChildCollection.Add(
new Tuple <string, IZooKeeperChildListener>(ZooKeeperClient.DefaultBrokerIdsPath, this));
Logger.ErrorFormat(
"Subscribe count: subscribedChildCollection:{0} , subscribedZookeeperStateCollection:{1} subscribedZookeeperDataCollection:{2} "
, zkConsumerConnector.subscribedChildCollection.Count,
zkConsumerConnector.subscribedZookeeperStateCollection.Count,
zkConsumerConnector.subscribedZookeeperDataCollection.Count);
return(false);
}
var partitionsPerTopicMap = ZkUtils.GetPartitionsForTopics(zkClient, topicThreadIdsMap.Keys);
// Check if we've been canceled externally before we dive into the rebalance
if (cancellationTokenSource.IsCancellationRequested)
{
Logger.ErrorFormat(
"Rebalance operation has been canceled externally by a future rebalance event. Exiting immediately");
return(false);
}
CloseFetchers(cluster, topicThreadIdsMap, zkConsumerConnector);
ReleasePartitionOwnership(topicThreadIdsMap);
try
{
foreach (var item in topicThreadIdsMap)
{
var topic = item.Key;
var consumerThreadIdSet = item.Value;
topicRegistry.Add(topic, new ConcurrentDictionary <int, PartitionTopicInfo>());
var topicDirs = new ZKGroupTopicDirs(config.GroupId, topic);
var curConsumers = new List <string>(consumersPerTopicMap[topic]);
curConsumers.Sort();
var curPartitions = partitionsPerTopicMap[topic].OrderBy(p => int.Parse(p)).ToList();
Logger.InfoFormat(
"{4} Partitions. {5} ConsumerClients. Consumer {0} rebalancing the following partitions: {1} for topic {2} with consumers: {3}",
consumerIdString,
string.Join(",", curPartitions),
topic,
string.Join(",", curConsumers),
curPartitions.Count,
curConsumers.Count);
var numberOfPartsPerConsumer = curPartitions.Count / curConsumers.Count;
Logger.Info("Number of partitions per consumer is: " + numberOfPartsPerConsumer);
var numberOfConsumersWithExtraPart = curPartitions.Count % curConsumers.Count;
Logger.Info("Number of consumers with an extra partition are: " + numberOfConsumersWithExtraPart);
foreach (var consumerThreadId in consumerThreadIdSet)
{
var myConsumerPosition = curConsumers.IndexOf(consumerThreadId);
Logger.Info("Consumer position for consumer " + consumerThreadId + " is: " +
myConsumerPosition);
if (myConsumerPosition < 0)
{
continue;
}
var startPart = numberOfPartsPerConsumer * myConsumerPosition +
Math.Min(myConsumerPosition, numberOfConsumersWithExtraPart);
Logger.Info("Starting partition is: " + startPart);
var numberOfParts = numberOfPartsPerConsumer +
(myConsumerPosition + 1 > numberOfConsumersWithExtraPart ? 0 : 1);
Logger.Info("Number of partitions to work on is: " + numberOfParts);
if (numberOfParts <= 0)
{
Logger.InfoFormat("No broker partitions consumed by consumer thread {0} for topic {1}",
consumerThreadId, item.Key);
}
else
{
for (var i = startPart; i < startPart + numberOfParts; i++)
{
var partition = curPartitions[i];
Logger.InfoFormat("{0} attempting to claim partition {1}", consumerThreadId, partition);
var ownPartition = ProcessPartition(topicDirs, partition, topic, consumerThreadId,
curConsumers, curPartitions, cancellationTokenSource);
if (!ownPartition)
{
Logger.InfoFormat(
"{0} failed to claim partition {1} for topic {2}. Exiting rebalance",
consumerThreadId, partition, topic);
return(false);
}
}
}
}
}
}
catch (Exception ex)
{
Logger.ErrorFormat("error when rebalance: {0}", ex.FormatException());
return(false);
}
// If we get here, we know that we have owned all partitions successfully,
// therefore it is safe to update fetcher threads and begin dequeuing
Logger.Info("All partitions were successfully owned. Updating fetchers");
UpdateFetcher(cluster);
return(true);
}
private bool Rebalance()
{
var myTopicThresdIdsMap = this.GetTopicCount(this.consumerIdString).GetConsumerThreadIdsPerTopic();
var cluster = new Cluster(zkClient);
var consumersPerTopicMap = this.GetConsumersPerTopic(this.config.GroupId);
var partitionsPerTopicMap = ZkUtils.GetPartitionsForTopics(this.zkClient, myTopicThresdIdsMap.Keys);
var relevantTopicThreadIdsMap = GetRelevantTopicMap(
myTopicThresdIdsMap,
partitionsPerTopicMap,
this.oldPartitionsPerTopicMap,
consumersPerTopicMap,
this.oldConsumersPerTopicMap);
if (relevantTopicThreadIdsMap.Count <= 0)
{
Logger.InfoFormat(CultureInfo.CurrentCulture, "Consumer {0} with {1} doesn't need to rebalance.", this.consumerIdString, consumersPerTopicMap);
return(true);
}
this.CloseFetchers(cluster, myTopicThresdIdsMap, this.zkConsumerConnector);
Logger.Info("Releasing parittion ownership");
this.ReleasePartitionOwnership();
var currentTopicRegistry = new ConcurrentDictionary <string, IDictionary <Partition, PartitionTopicInfo> >();
var partitionOwnershipDecision = new Dictionary <Tuple <string, string>, string>();
foreach (var item in myTopicThresdIdsMap)
{
currentTopicRegistry.GetOrAdd(item.Key, new ConcurrentDictionary <Partition, PartitionTopicInfo>());
var topicDirs = new ZKGroupTopicDirs(config.GroupId, item.Key);
var curConsumers = consumersPerTopicMap[item.Key];
var curPartitions = new List <string>(partitionsPerTopicMap[item.Key]);
var numberOfPartsPerConsumer = curPartitions.Count / curConsumers.Count;
var numberOfConsumersWithExtraPart = curPartitions.Count % curConsumers.Count;
Logger.InfoFormat(
CultureInfo.CurrentCulture,
"Consumer {0} rebalancing the following partitions: {1} for topic {2} with consumers: {3}",
this.consumerIdString,
string.Join(",", curPartitions),
item.Key,
string.Join(",", curConsumers));
foreach (string consumerThreadId in item.Value)
{
var myConsumerPosition = curConsumers.IndexOf(consumerThreadId);
if (myConsumerPosition < 0)
{
continue;
}
var startPart = (numberOfPartsPerConsumer * myConsumerPosition) +
Math.Min(myConsumerPosition, numberOfConsumersWithExtraPart);
var numberOfParts = numberOfPartsPerConsumer + (myConsumerPosition + 1 > numberOfConsumersWithExtraPart ? 0 : 1);
if (numberOfParts <= 0)
{
Logger.WarnFormat(CultureInfo.CurrentCulture, "No broker partitions consumed by consumer thread {0} for topic {1}", consumerThreadId, item.Key);
}
else
{
for (int i = startPart; i < startPart + numberOfParts; i++)
{
var partition = curPartitions[i];
Logger.InfoFormat(CultureInfo.CurrentCulture, "{0} attempting to claim partition {1}", consumerThreadId, partition);
AddPartitionTopicInfo(currentTopicRegistry, topicDirs, partition, item.Key, consumerThreadId);
partitionOwnershipDecision.Add(new Tuple <string, string>(item.Key, partition), consumerThreadId);
}
}
}
}
if (ReflectPartitionOwnership(partitionOwnershipDecision))
{
zkClient.SlimLock.EnterWriteLock();
try
{
this.topicRegistry.Clear();
foreach (var item in currentTopicRegistry)
{
this.topicRegistry.Add(item);
}
this.UpdateFetcher(cluster);
}
finally
{
zkClient.SlimLock.ExitWriteLock();
}
this.oldPartitionsPerTopicMap = partitionsPerTopicMap;
this.oldConsumersPerTopicMap = consumersPerTopicMap;
return(true);
}
else
{
return(false);
}
}