internal void Unassign(MockConsumer mockConsumer)
{
if (consumers.ContainsKey(mockConsumer.Name))
{
lock (rebalanceLock)
{
var c = consumers[mockConsumer.Name];
var pList = c.TopicPartitionsOffset.Select(f => new TopicPartitionOffset(f.Topic, f.Partition, f.OffsetComitted)).ToList();
c.RebalanceListener?.PartitionsRevoked(c.Consumer, pList);
// Rebalance on other consumer in the same group
var otherConsumers = consumerGroups[mockConsumer.MemberId].Where(i => consumers.ContainsKey(i)).Select(i => consumers[i]).Where(i => !i.Name.Equals(mockConsumer.Name)).ToList();
if (otherConsumers.Count > 0)
{
int partEach = c.Partitions.Count / otherConsumers.Count;
int modulo = c.Partitions.Count % otherConsumers.Count;
int j = 0;
for (int i = 0; i < otherConsumers.Count; ++i)
{
List <TopicPartition> parts = null;
if (i == otherConsumers.Count - 1)
{
parts = c.Partitions.GetRange(j, j + partEach + modulo);
}
else
{
parts = c.Partitions.GetRange(j, j + partEach);
}
otherConsumers[i].Partitions.AddRange(parts);
foreach (var k in parts)
{
if (!otherConsumers[i].TopicPartitionsOffset.Any(m => m.Topic.Equals(k.Topic) && m.Partition.Equals(k.Partition)))
{
otherConsumers[i].TopicPartitionsOffset.Add(new MockTopicPartitionOffset
{
OffsetComitted = 0,
OffsetConsumed = 0,
Partition = k.Partition,
Topic = k.Topic
});
}
}
otherConsumers[i].RebalanceListener?.PartitionsAssigned(otherConsumers[i].Consumer, otherConsumers[i].Partitions);
otherConsumers[i].Assigned = true;
j += partEach;
}
}
c.Partitions.Clear();
c.Assigned = false;
}
}
else
{
throw new StreamsException($"Consumer {mockConsumer.Name} unknown !");
}
}
public IConsumer <byte[], byte[]> GetConsumer(ConsumerConfig config, IConsumerRebalanceListener rebalanceListener)
{
var consumer = new MockConsumer(config.GroupId, config.ClientId);
consumer.SetRebalanceListener(rebalanceListener);
return(consumer);
}
internal ConsumeResult <byte[], byte[]> Consume(MockConsumer mockConsumer, TimeSpan timeout)
{
foreach (var t in mockConsumer.Subscription)
{
CreateTopic(t);
}
DateTime dt = DateTime.Now;
ConsumeResult <byte[], byte[]> result = null;
if (consumers.ContainsKey(mockConsumer.Name))
{
var c = consumers[mockConsumer.Name];
if (!c.Assigned)
{
lock (rebalanceLock)
{
NeedRebalance();
}
}
lock (rebalanceLock)
{
foreach (var p in c.Partitions)
{
if (timeout != TimeSpan.Zero && (dt + timeout) < DateTime.Now)
{
break;
}
var topic = topics[p.Topic];
var offset = c.TopicPartitionsOffset.FirstOrDefault(t => t.Topic.Equals(p.Topic) && t.Partition.Equals(p.Partition));
var record = topic.GetMessage(p.Partition, offset.OffsetConsumed);
if (record != null)
{
result = new ConsumeResult <byte[], byte[]>
{
Offset = offset.OffsetConsumed,
Topic = p.Topic,
Partition = p.Partition,
Message = new Message <byte[], byte[]> {
Key = record.Key, Value = record.Value
}
};
++offset.OffsetConsumed;
break;
}
}
}
return(result);
}
else
{
throw new StreamsException($"Consumer {mockConsumer.Name} unknown !");
}
}
internal void Unsubscribe(MockConsumer mockConsumer)
{
if (consumers.ContainsKey(mockConsumer.Name))
{
var c = consumers[mockConsumer.Name];
Unassign(mockConsumer);
c.Topics.Clear();
}
}
internal List <TopicPartitionOffset> Comitted(MockConsumer mockConsumer)
{
var c = consumers[mockConsumer.Name];
List <TopicPartitionOffset> list = new List <TopicPartitionOffset>();
foreach (var p in c.Partitions)
{
var offset = c.TopicPartitionsOffset.FirstOrDefault(t => t.Topic.Equals(p.Topic) && t.Partition.Equals(p.Partition));
if (offset != null)
{
list.Add(new TopicPartitionOffset(new TopicPartition(p.Topic, p.Partition), new Offset(offset.OffsetComitted)));
}
}
return(list);
}
internal List <TopicPartitionOffset> Commit(MockConsumer mockConsumer)
{
if (consumers.ContainsKey(mockConsumer.Name))
{
var c = consumers[mockConsumer.Name];
foreach (var p in c.TopicPartitionsOffset)
{
p.OffsetComitted = p.OffsetConsumed;
}
return(c.TopicPartitionsOffset.Select(t => new TopicPartitionOffset(new TopicPartition(t.Topic, t.Partition), t.OffsetComitted)).ToList());
}
else
{
throw new StreamsException($"Consumer {mockConsumer.Name} unknown !");
}
}
internal void Commit(MockConsumer mockConsumer, IEnumerable <TopicPartitionOffset> offsets)
{
if (consumers.ContainsKey(mockConsumer.Name))
{
var c = consumers[mockConsumer.Name];
foreach (var o in offsets)
{
var p = c.TopicPartitionsOffset.FirstOrDefault(t => t.Topic.Equals(o.Topic) && t.Partition.Equals(o.Partition));
if (p != null)
{
p.OffsetConsumed = o.Offset.Value;
p.OffsetComitted = o.Offset.Value;
}
}
}
else
{
throw new StreamsException($"Consumer {mockConsumer.Name} unknown !");
}
}
internal void SubscribeTopic(MockConsumer consumer, IEnumerable <string> topics)
{
foreach (var t in topics)
{
CreateTopic(t);
}
if (!consumers.ContainsKey(consumer.Name))
{
var cons = new MockConsumerInformation
{
GroupId = consumer.MemberId,
Name = consumer.Name,
Consumer = consumer,
Topics = new List <string>(topics),
RebalanceListener = consumer.Listener,
Partitions = new List <TopicPartition>(),
TopicPartitionsOffset = new List <MockConsumerInformation.MockTopicPartitionOffset>()
};
consumers.Add(consumer.Name, cons);
if (consumerGroups.ContainsKey(consumer.MemberId))
{
consumerGroups[consumer.MemberId].Add(consumer.Name);
}
else
{
consumerGroups.Add(consumer.MemberId, new List <string> {
consumer.Name
});
}
}
else
{
throw new StreamsException($"Client {consumer.Name} already subscribe topic. Please call unsucribe before");
}
}
internal ConsumeResult <byte[], byte[]> Consume(MockConsumer mockConsumer, CancellationToken cancellationToken) => Consume(mockConsumer, TimeSpan.FromSeconds(10));
internal void Assign(MockConsumer mockConsumer, IEnumerable <TopicPartition> topicPartitions)
{
foreach (var t in topicPartitions)
{
CreateTopic(t.Topic);
}
if (consumers.ContainsKey(mockConsumer.Name))
{
var c = consumers[mockConsumer.Name];
if (c.Partitions.Count == 0)
{
lock (rebalanceLock)
{
List <MockConsumerInformation> customerToRebalance = new List <MockConsumerInformation>();
foreach (var p in topicPartitions)
{
var info = consumers.Select(kp => kp.Value)
.Where(i => i.GroupId.Equals(mockConsumer.MemberId))
.FirstOrDefault(i => i.Partitions.Contains(p));
if (info != null && !customerToRebalance.Contains(info))
{
customerToRebalance.Add(info);
}
}
foreach (var cus in customerToRebalance)
{
var parts = cus.Partitions.Join(topicPartitions, (t) => t, (t) => t, (t1, t2) => t1);
var pList = cus.TopicPartitionsOffset.Select(f => new TopicPartitionOffset(f.Topic, f.Partition, f.OffsetComitted)).ToList();
cus.RebalanceListener?.PartitionsRevoked(cus.Consumer, pList);
foreach (var j in parts)
{
cus.Partitions.Remove(j);
}
cus.RebalanceListener?.PartitionsAssigned(cus.Consumer, cus.Partitions);
cus.Assigned = true;
}
c.Partitions = new List <TopicPartition>(topicPartitions);
foreach (var k in topicPartitions)
{
if (!c.TopicPartitionsOffset.Any(m => m.Topic.Equals(k.Topic) && m.Partition.Equals(k.Partition)))
{
c.TopicPartitionsOffset.Add(new MockTopicPartitionOffset
{
OffsetComitted = 0,
OffsetConsumed = 0,
Partition = k.Partition,
Topic = k.Topic
});
}
}
c.RebalanceListener?.PartitionsAssigned(c.Consumer, c.Partitions);
c.Assigned = true;
}
}
else
{
throw new StreamsException($"Consumer {mockConsumer.Name} was already assigned partitions. Please call unassigne before !");
}
}
else
{
throw new StreamsException($"Consumer {mockConsumer.Name} unknown !");
}
}
public MockWrappedConsumerRebalanceListener(IConsumerRebalanceListener wrapped, MockConsumer consumer)
{
this.wrapped = wrapped;
this.consumer = consumer;
}
