extend rt versioning to system stores

Author: arjun4084346

clients/da-vinci-client/src/main/java/com/linkedin/davinci/DaVinciBackend.java

@@ -57,6 +57,7 @@
 import com.linkedin.venice.meta.Store;
 import com.linkedin.venice.meta.StoreDataChangedListener;
 import com.linkedin.venice.meta.SubscriptionBasedReadOnlyStoreRepository;
+import com.linkedin.venice.meta.SystemStoreAttributes;
 import com.linkedin.venice.meta.Version;
 import com.linkedin.venice.pubsub.api.PubSubPosition;
 import com.linkedin.venice.pushmonitor.ExecutionStatus;
@@ -308,7 +309,8 @@ public DaVinciBackend(
             ingestionService.getVeniceWriterFactory(),
             instanceName,
             valueSchemaEntry,
-            updateSchemaEntry);
+            updateSchemaEntry,
+            (this::getStore));
       }
 
       ingestionService.start();
@@ -559,6 +561,23 @@ public SubscriptionBasedReadOnlyStoreRepository getStoreRepository() {
     return storeRepository;
   }
 
+  public final Object getStore(String storeName) {
+    VeniceSystemStoreType systemStoreType = VeniceSystemStoreType.getSystemStoreType(storeName);
+    if (systemStoreType != null) {
+      String userStoreName = VeniceSystemStoreType.extractUserStoreName(storeName);
+      Store userStore = storeRepository.getStore(userStoreName);
+      Map<String, SystemStoreAttributes> systemStores = userStore.getSystemStores();
+      for (Map.Entry<String, SystemStoreAttributes> systemStoreEntries: systemStores.entrySet()) {
+        if (storeName.startsWith(systemStoreEntries.getKey())) {
+          return systemStoreEntries.getValue();
+        }
+      }
+      return null;
+    } else {
+      return storeRepository.getStore(storeName);
+    }
+  }
+
   public ObjectCacheBackend getObjectCache() {
     return cacheBackend.get();
   }

clients/da-vinci-client/src/main/java/com/linkedin/davinci/helix/HelixParticipationService.java

@@ -361,7 +361,8 @@ private void asyncStart() {
         ingestionService.getVeniceWriterFactory(),
         instance.getNodeId(),
         valueSchemaEntry,
-        updateSchemaEntry);
+        updateSchemaEntry,
+        helixReadOnlyStoreRepository::getStore);
 
     // Record replica status in Zookeeper.
     // Need to be started before connecting to ZK, otherwise some notification will not be sent by this notifier.

clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/KafkaStoreIngestionService.java

@@ -393,7 +393,8 @@ public KafkaStoreIngestionService(
           zkSharedSchemaRepository.get(),
           pubSubTopicRepository,
           serverConfig.getMetaStoreWriterCloseTimeoutInMS(),
-          serverConfig.getMetaStoreWriterCloseConcurrency());
+          serverConfig.getMetaStoreWriterCloseConcurrency(),
+          storeName -> metadataRepo.getStore(storeName));
       metadataRepo.registerStoreDataChangedListener(new StoreDataChangedListener() {
         @Override
         public void handleStoreDeleted(Store store) {
@@ -1486,7 +1487,8 @@ public InternalDaVinciRecordTransformerConfig getInternalRecordTransformerConfig
     return storeNameToInternalRecordTransformerConfig.get(storeName);
   }
 
-  public void attemptToPrintIngestionInfoFor(String storeName, Integer version, Integer partition, String regionName) {
+  public void attemptToPrintIngestionInfoFor(Store store, Integer version, Integer partition, String regionName) {
+    String storeName = store.getName();
     try {
       PubSubTopic versionTopic = pubSubTopicRepository.getTopic(Version.composeKafkaTopic(storeName, version));
       StoreIngestionTask storeIngestionTask = getStoreIngestionTask(versionTopic.getName());
@@ -1508,7 +1510,7 @@ public void attemptToPrintIngestionInfoFor(String storeName, Integer version, In
       String infoPrefix = "isCurrentVersion: " + (storeIngestionTask.isCurrentVersion()) + "\n";
       if (storeIngestionTask.isHybridMode() && partitionConsumptionState.isEndOfPushReceived()
           && partitionConsumptionState.getLeaderFollowerState() == LeaderFollowerStateType.LEADER) {
-        ingestingTopic = pubSubTopicRepository.getTopic(Utils.composeRealTimeTopic(storeName));
+        ingestingTopic = pubSubTopicRepository.getTopic(Utils.getRealTimeTopicName(store));
       }
       PubSubTopicPartition ingestingTopicPartition = new PubSubTopicPartitionImpl(ingestingTopic, partition);
 

clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTask.java

@@ -4342,9 +4342,12 @@ private void waitUntilValueSchemaAvailable(int schemaId) throws InterruptedExcep
         // cluster these metastore writes could be spiky
         if (metaStoreWriter != null && !VeniceSystemStoreType.META_STORE.isSystemStore(storeName)) {
           String metaStoreName = VeniceSystemStoreType.META_STORE.getSystemStoreName(storeName);
-          PubSubTopic metaStoreRT = pubSubTopicRepository.getTopic(Utils.composeRealTimeTopic(metaStoreName));
-          if (getTopicManager(localKafkaServer).containsTopicWithRetries(metaStoreRT, 5)) {
-            metaStoreWriter.writeInUseValueSchema(storeName, versionNumber, schemaId);
+          Store metaStore = metaStoreWriter.storeResolver.apply(metaStoreName);
+          if (metaStore != null) {
+            PubSubTopic metaStoreRT = pubSubTopicRepository.getTopic(Utils.getRealTimeTopicName(metaStore));
+            if (getTopicManager(localKafkaServer).containsTopicWithRetries(metaStoreRT, 5)) {
+              metaStoreWriter.writeInUseValueSchema(storeName, versionNumber, schemaId);
+            }
           }
         }
         return;

clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/ingestion/heartbeat/HeartbeatMonitoringService.java

@@ -647,7 +647,7 @@ void checkAndMaybeLogHeartbeatDelayMap(
                   heartbeatTs,
                   currentTimestamp);
               getKafkaStoreIngestionService().attemptToPrintIngestionInfoFor(
-                  storeName.getKey(),
+                  metadataRepository.getStore(storeName.getKey()),
                   version.getKey(),
                   partition.getKey(),
                   region.getKey());

clients/da-vinci-client/src/test/java/com/linkedin/davinci/DaVinciBackendTest.java

@@ -11,6 +11,7 @@
 import static com.linkedin.venice.utils.DataProviderUtils.BOOLEAN;
 import static com.linkedin.venice.utils.DataProviderUtils.allPermutationGenerator;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.mock;
@@ -34,6 +35,7 @@
 import com.linkedin.venice.client.store.ClientFactory;
 import com.linkedin.venice.exceptions.DiskLimitExhaustedException;
 import com.linkedin.venice.exceptions.VeniceException;
+import com.linkedin.venice.helix.ZkClientFactory;
 import com.linkedin.venice.meta.ClusterInfoProvider;
 import com.linkedin.venice.meta.ReadOnlySchemaRepository;
 import com.linkedin.venice.meta.SubscriptionBasedReadOnlyStoreRepository;
@@ -43,16 +45,19 @@
 import com.linkedin.venice.schema.writecompute.DerivedSchemaEntry;
 import com.linkedin.venice.serialization.avro.SchemaPresenceChecker;
 import com.linkedin.venice.service.ICProvider;
+import com.linkedin.venice.stats.ZkClientStatusStats;
 import com.linkedin.venice.utils.VeniceProperties;
 import io.tehuti.metrics.MetricsRepository;
 import java.util.Optional;
 import java.util.Properties;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import org.apache.helix.zookeeper.impl.client.ZkClient;
 import org.mockito.MockedConstruction;
 import org.mockito.MockedStatic;
 import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
@@ -67,6 +72,14 @@ public class DaVinciBackendTest {
   private MockedConstruction<VeniceMetadataRepositoryBuilder> mockMetadataBuilder;
   private MockedConstruction<SchemaPresenceChecker> mockSchemaPresenceChecker;
 
+  @BeforeClass
+  public void init() {
+    MockedStatic<ZkClientFactory> mockZkFactory = mockStatic(ZkClientFactory.class);
+    ZkClient mockZkClient = mock(ZkClient.class);
+    mockZkFactory.when(() -> ZkClientFactory.newZkClient(anyString())).thenReturn(mockZkClient);
+    doNothing().when(mockZkClient).subscribeStateChanges(any(ZkClientStatusStats.class));
+  }
+
   @BeforeMethod
   public void setUp() throws Exception {
     ClientConfig clientConfig = new ClientConfig(STORE_NAME).setVeniceURL("http://localhost:7777")