ConsumerFetcherThread extends ShutdownableThread { public ConsumerFetcherThread(String threadName, CustomizedConsumerConfig properties, RateLimiter rateLimiter, BlockingQueue<FetchedDataChunk> chunkQueue) { super(threadName, true); this.fetchBackOffMs = properties.getFetcherThreadBackoffMs(); this.offsetMonitorMs = properties.getOffsetMonitorInterval(); this.pollTimeoutMs = properties.getPollTimeoutMs(); this.kafkaConsumer = new KafkaConsumer(properties); KafkaUReplicatorMetricsReporter.get() .registerKafkaMetrics("consumer." + threadName, kafkaConsumer.metrics()); this.rateLimiter = rateLimiter; this.chunkQueue = chunkQueue; } ConsumerFetcherThread(String threadName, CustomizedConsumerConfig properties,
      RateLimiter rateLimiter, BlockingQueue<FetchedDataChunk> chunkQueue); @Override void doWork(); void addPartitions(Map<TopicPartition, PartitionOffsetInfo> partitionAndOffsets); void removePartitions(Set<TopicPartition> topicAndPartitions); Set<TopicPartition> getTopicPartitions(); void shutdown(); @Override void awaitShutdown();  }
@Test public void testConsumerFetcherThread() throws InterruptedException, IOException { String threadName = "testCompactConsumerFetcherThread"; ConcurrentHashMap<TopicAndPartition, PartitionOffsetInfo> partitionInfoMap = new ConcurrentHashMap<>(); KafkaStarterUtils.createTopic(testTopic1, clusterZk); KafkaStarterUtils.createTopic(testTopic2, clusterZk); BlockingQueue<FetchedDataChunk> queue = new LinkedBlockingQueue<>(3); CustomizedConsumerConfig properties = new CustomizedConsumerConfig(Utils.loadProps("src/test/resources/consumer.properties")); properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, bootstrapServer); properties.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, false); properties.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest"); ConsumerFetcherThread fetcherThread = new ConsumerFetcherThread(threadName, properties, null, queue); fetcherThread.start(); TestUtils.produceMessages(bootstrapServer, testTopic1, 10); TestUtils.produceMessages(bootstrapServer, testTopic2, 10); Map<TopicPartition, PartitionOffsetInfo> addPartitions = new HashMap<>(); TopicPartition tp1 = new TopicPartition(testTopic1, 0); PartitionOffsetInfo pi1 = new PartitionOffsetInfo(tp1,0L, null); addPartitions.put(tp1, pi1); TopicPartition tp2 = new TopicPartition(testTopic2, 0); PartitionOffsetInfo pi2 = new PartitionOffsetInfo(tp2,5L, null); addPartitions.put(tp2, pi2); fetcherThread.addPartitions(addPartitions); Thread.sleep(1000); Assert.assertEquals(1, queue.remainingCapacity()); FetchedDataChunk record1 = queue.take(); FetchedDataChunk record2 = queue.take(); Assert.assertEquals(15, record1.messageSize() + record2.messageSize()); addPartitions.remove(tp1); fetcherThread.addPartitions(addPartitions); Thread.sleep(1000); Assert.assertEquals(3, queue.remainingCapacity()); com.uber.stream.ureplicator.worker.TestUtils.produceMessages(bootstrapServer, testTopic1, 10); Thread.sleep(1000); Assert.assertEquals(2, queue.remainingCapacity()); FetchedDataChunk records = queue.take(); Assert.assertEquals(10, records.messageSize()); LOGGER.info("removePartitions/addPartitions to test specify same start offset again"); Set<TopicPartition> removePartitions = new HashSet<>(); removePartitions.add(tp2); fetcherThread.removePartitions(removePartitions); Thread.sleep(1000); fetcherThread.addPartitions(addPartitions); Thread.sleep(1000); Assert.assertEquals(2, queue.remainingCapacity()); records = queue.take(); Assert.assertEquals(5, records.messageSize()); LOGGER.info("removePartitions/addPartitions to test out of range"); PartitionOffsetInfo pi3 = new PartitionOffsetInfo(tp2, 15L, null); addPartitions.remove(tp1); addPartitions.put(tp2, pi3); fetcherThread.removePartitions(removePartitions); Thread.sleep(1000); fetcherThread.addPartitions(addPartitions); Thread.sleep(1000); Assert.assertEquals(2, queue.remainingCapacity()); records = queue.take(); Assert.assertEquals(10, records.messageSize()); com.uber.stream.ureplicator.worker.TestUtils.produceMessages(bootstrapServer, testTopic1, 10); Thread.sleep(1000); Assert.assertEquals(2, queue.remainingCapacity()); records = queue.take(); Assert.assertEquals(10, records.messageSize()); fetcherThread.shutdown(); }