Add lock to ensure no failed test.

linkedin · Oct 28, 2024 · 05c8718 · 05c8718
1 parent aa18acb
commit 05c8718
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/.../src/main/java/com/linkedin/davinci/kafka/consumer/PartitionWiseKafkaConsumerService.java b/.../src/main/java/com/linkedin/davinci/kafka/consumer/PartitionWiseKafkaConsumerService.java
@@ -94,6 +94,12 @@ protected synchronized SharedKafkaConsumer pickConsumerForPartition(
 
       // Safeguard logic, avoid infinite loops for searching consumer.
       if (consumersChecked == consumerToConsumptionTask.size()) {
+        for (int consumerIdx = 0; consumerIdx < consumerToConsumptionTask.size(); consumerIdx++) {
+          LOGGER.info(
+              "Consumer: {} is consuming for: {}",
+              consumerIdx,
+              consumerToConsumptionTask.getByIndex(consumerIdx).getValue().getDestinationIdentifier(topicPartition));
+        }
         throw new VeniceException(
             "Can not find consumer for topic: " + topicPartition.getPubSubTopic().getName() + " and partition: "
                 + topicPartition.getPartitionNumber() + " from the ingestion task belonging to version topic: "
@@ -145,7 +151,7 @@ private boolean alreadySubscribedRealtimeTopicPartition(
   }
 
   @Override
-  void handleUnsubscription(SharedKafkaConsumer consumer, PubSubTopicPartition pubSubTopicPartition) {
+  synchronized void handleUnsubscription(SharedKafkaConsumer consumer, PubSubTopicPartition pubSubTopicPartition) {
     if (pubSubTopicPartition.getPubSubTopic().isRealTime()) {
       Set<PubSubConsumerAdapter> rtTopicConsumers = rtTopicPartitionToConsumerMap.get(pubSubTopicPartition);
       if (rtTopicConsumers != null) {

diff --git a/.../src/test/java/com/linkedin/davinci/kafka/consumer/KafkaConsumerServiceDelegatorTest.java b/.../src/test/java/com/linkedin/davinci/kafka/consumer/KafkaConsumerServiceDelegatorTest.java
@@ -490,10 +490,16 @@ private void verifyConsumerServiceStartConsumptionIntoDataReceiver(
   @Test(invocationCount = 5)
   public void testKafkaConsumerServiceResubscriptionConcurrency() throws Exception {
     ApacheKafkaConsumerAdapter consumer1 = mock(ApacheKafkaConsumerAdapter.class);
-    when(consumer1.hasAnySubscription()).thenReturn(true);
+    ApacheKafkaConsumerAdapter consumer2 = mock(ApacheKafkaConsumerAdapter.class);
+    ApacheKafkaConsumerAdapter consumer3 = mock(ApacheKafkaConsumerAdapter.class);
+    ApacheKafkaConsumerAdapter consumer4 = mock(ApacheKafkaConsumerAdapter.class);
+    ApacheKafkaConsumerAdapter consumer5 = mock(ApacheKafkaConsumerAdapter.class);
+    ApacheKafkaConsumerAdapter consumer6 = mock(ApacheKafkaConsumerAdapter.class);
+    // when(consumer1.hasAnySubscription()).thenReturn(true);
 
     PubSubConsumerAdapterFactory factory = mock(PubSubConsumerAdapterFactory.class);
-    when(factory.create(any(), anyBoolean(), any(), any())).thenReturn(consumer1);
+    when(factory.create(any(), anyBoolean(), any(), any()))
+        .thenReturn(consumer1, consumer2, consumer3, consumer4, consumer5, consumer6);
 
     Properties properties = new Properties();
     String testKafkaUrl = "test_kafka_url";
@@ -576,6 +582,7 @@ public void testKafkaConsumerServiceResubscriptionConcurrency() throws Exception
     Assert.assertFalse(
         raceConditionFound,
         "Found race condition in KafkaConsumerService with time passed in milliseconds: " + elapsedTime);
+    delegator.close();
   }
 
   private Runnable getResubscriptionRunnableFor(