/// <summary>
/// Produce the given number of messages, create a consumer with the given offset policy,
/// then reset the offset to the given value and consume until we get no new messages.
/// </summary>
/// <param name="numMessages"></param>
/// <param name="resetTo"></param>
/// <param name="offset"></param>
/// <returns>The count of messages received.</returns>
public int ResetAndConsume(int numMessages, string resetTo, long offset)
{
TestUtils.WaitUntilLeaderIsElectedOrChanged(this.ZkClient, Topic, 0, 1000);
var producer = TestUtils.CreateProducer(
TestUtils.GetBrokerListFromConfigs(Configs), new DefaultEncoder(), new StringEncoder());
for (var i = 0; i < numMessages; i++)
{
producer.Send(new KeyedMessage<string, byte[]>(Topic, Topic, Encoding.UTF8.GetBytes("test")));
}
TestUtils.WaitUntilMetadataIsPropagated(this.Servers, Topic, 0, 1000);
// update offset in zookeeper for consumer to jump "forward" in time
var dirs = new ZKGroupTopicDirs(Group, Topic);
var consumerConfig = TestUtils.CreateConsumerProperties(ZkConnect, Group, TestConsumer);
consumerConfig.AutoOffsetReset = resetTo;
consumerConfig.ConsumerTimeoutMs = 2000;
consumerConfig.FetchWaitMaxMs = 0;
TestUtils.UpdateConsumerOffset(consumerConfig, dirs.ConsumerOffsetDir + "/" + "0", offset);
Logger.InfoFormat("Update consumer offset to {0}", offset);
var consumerConnector = Consumer.Create(consumerConfig);
var messagesStream = consumerConnector.CreateMessageStreams(new Dictionary<string, int> { { Topic, 1 } })[Topic].First();
var received = 0;
var iter = messagesStream.GetEnumerator();
try
{
for (var i = 0; i < numMessages; i++)
{
iter.MoveNext(); // will throw a timeout exception if the message isn't there
received++;
}
}
catch (ConsumerTimeoutException)
{
Logger.InfoFormat("consumer timeout out after receiving {0} messages", received);
}
finally
{
producer.Dispose();
consumerConnector.Shutdown();
}
return received;
}
public ZookeeperConsumerConnectorTest()
{
this.dirs = new ZKGroupTopicDirs(Group, Topic);
}
/// <summary>
/// Commits the offsets of all messages consumed so far.
/// </summary>
public void CommitOffsets()
{
this.EnsuresNotDisposed();
if (this.zkClient == null)
{
return;
}
foreach (KeyValuePair<string, IDictionary<Partition, PartitionTopicInfo>> topic in topicRegistry)
{
var topicDirs = new ZKGroupTopicDirs(this.config.GroupId, topic.Key);
foreach (KeyValuePair<Partition, PartitionTopicInfo> partition in topic.Value)
{
var newOffset = partition.Value.GetConsumeOffset();
try
{
ZkUtils.UpdatePersistentPath(zkClient, topicDirs.ConsumerOffsetDir + "/" + partition.Value.Partition.Name, newOffset.ToString());
}
catch (Exception ex)
{
Logger.WarnFormat(CultureInfo.CurrentCulture, "exception during CommitOffsets: {0}", ex);
}
if (Logger.IsDebugEnabled)
{
Logger.DebugFormat(CultureInfo.CurrentCulture, "Commited offset {0} for topic {1}", newOffset, partition);
}
}
}
}
private long ResetConsumerOffsets(string topic, Partition partition)
{
long offset;
switch (this.config.AutoOffsetReset)
{
case OffsetRequest.SmallestTime:
offset = OffsetRequest.EarliestTime;
break;
case OffsetRequest.LargestTime:
offset = OffsetRequest.LatestTime;
break;
default:
return -1;
}
var request = new OffsetRequest(topic, partition.PartId, offset, 1);
var offsets = this.simpleConsumer.GetOffsetsBefore(request);
var topicDirs = new ZKGroupTopicDirs(this.config.GroupId, topic);
Logger.InfoFormat(CultureInfo.CurrentCulture, "updating partition {0} with {1} offset {2}", partition.Name, offset == OffsetRequest.EarliestTime ? "earliest" : "latest", offsets[0]);
ZkUtils.UpdatePersistentPath(this.zkClient, topicDirs.ConsumerOffsetDir + "/" + partition.Name, offsets[0].ToString());
return offsets[0];
}
public static string GetConsumerPartitionOwnerPath(string group, string topic, int partition)
{
var topicDirs = new ZKGroupTopicDirs(group, topic);
return topicDirs.ConsumerOwnerDir + "/" + partition;
}
internal static string GetConsumerPartitionOffsetPath(string group, string topic, string partition)
{
var topicDirs = new ZKGroupTopicDirs(group, topic);
return(topicDirs.ConsumerOffsetDir + "/" + partition);
}
private void ReleasePartitionOwnership()
{
foreach (KeyValuePair<string, IDictionary<Partition, PartitionTopicInfo>> item in topicRegistry)
{
var topicDirs = new ZKGroupTopicDirs(this.config.GroupId, item.Key);
foreach (var partition in item.Value.Keys)
{
string znode = topicDirs.ConsumerOwnerDir + "/" + partition.Name;
ZkUtils.DeletePath(zkClient, znode);
if (Logger.IsDebugEnabled)
{
Logger.DebugFormat(CultureInfo.CurrentCulture, "Consumer {0} releasing {1}", this.consumerIdString, znode);
}
}
}
}
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 ProcessPartition(ZKGroupTopicDirs topicDirs, string partition, string topic, string consumerThreadId)
{
var partitionOwnerPath = topicDirs.ConsumerOwnerDir + "/" + partition;
try
{
ZkUtils.CreateEphemeralPathExpectConflict(zkClient, partitionOwnerPath, consumerThreadId);
}
catch (KeeperException.NodeExistsException)
{
//// The node hasn't been deleted by the original owner. So wait a bit and retry.
Logger.InfoFormat(CultureInfo.CurrentCulture, "waiting for the partition ownership to be deleted: {0}", partition);
return false;
}
AddPartitionTopicInfo(topicDirs, partition, topic, consumerThreadId);
return true;
}
private void AddPartitionTopicInfo(ZKGroupTopicDirs topicDirs, string partitionString, string topic, string consumerThreadId)
{
var partition = Partition.ParseFrom(partitionString);
var partTopicInfoMap = this.topicRegistry[topic];
var znode = topicDirs.ConsumerOffsetDir + "/" + partition.Name;
var offsetString = this.zkClient.ReadData<string>(znode, true);
long offset = string.IsNullOrEmpty(offsetString) ? 0 : long.Parse(offsetString, CultureInfo.InvariantCulture);
var queue = this.queues[new Tuple<string, string>(topic, consumerThreadId)];
var partTopicInfo = new PartitionTopicInfo(
topic,
partition.BrokerId,
partition,
queue,
offset,
offset,
this.config.FetchSize);
partTopicInfoMap.Add(partition, partTopicInfo);
if (Logger.IsDebugEnabled)
{
Logger.DebugFormat(CultureInfo.CurrentCulture, "{0} selected new offset {1}", partTopicInfo, offset);
}
}
internal static string GetConsumerPartitionOffsetPath(string group, string topic, string partition)
{
var topicDirs = new ZKGroupTopicDirs(group, topic);
return topicDirs.ConsumerOffsetDir + "/" + partition;
}
public static string GetConsumerPartitionOwnerPath(string group, string topic, int partition)
{
var topicDirs = new ZKGroupTopicDirs(group, topic);
return(topicDirs.ConsumerOwnerDir + "/" + partition);
}
