[Server][FC] throws exception during store migration to enforce d2 refresh

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

@@ -7,6 +7,7 @@
 import com.linkedin.venice.helix.HelixReadOnlyLiveClusterConfigRepository;
 import com.linkedin.venice.helix.HelixReadOnlySchemaRepository;
 import com.linkedin.venice.helix.HelixReadOnlySchemaRepositoryAdapter;
+import com.linkedin.venice.helix.HelixReadOnlyStoreConfigRepository;
 import com.linkedin.venice.helix.HelixReadOnlyStoreRepository;
 import com.linkedin.venice.helix.HelixReadOnlyStoreRepositoryAdapter;
 import com.linkedin.venice.helix.HelixReadOnlyZKSharedSchemaRepository;
@@ -43,6 +44,7 @@ public class VeniceMetadataRepositoryBuilder {
   private final ICProvider icProvider;
 
   private ReadOnlyStoreRepository storeRepo;
+  private HelixReadOnlyStoreConfigRepository storeConfigRepo;
   private ReadOnlySchemaRepository schemaRepo;
   private HelixReadOnlyZKSharedSchemaRepository readOnlyZKSharedSchemaRepository;
   private ReadOnlyLiveClusterConfigRepository liveClusterConfigRepo;
@@ -79,6 +81,10 @@ public ReadOnlyStoreRepository getStoreRepo() {
     return storeRepo;
   }
 
+  public HelixReadOnlyStoreConfigRepository getStoreConfigRepo() {
+    return storeConfigRepo;
+  }
+
   public ReadOnlySchemaRepository getSchemaRepo() {
     return schemaRepo;
   }
@@ -136,6 +142,13 @@ private void initServerStoreAndSchemaRepository() {
     // Load existing store config and setup watches
     storeRepo.refresh();
 
+    storeConfigRepo = new HelixReadOnlyStoreConfigRepository(
+        zkClient,
+        adapter,
+        clusterConfig.getRefreshAttemptsForZkReconnect(),
+        clusterConfig.getRefreshIntervalForZkReconnectInMs());
+    storeConfigRepo.refresh();
+
     readOnlyZKSharedSchemaRepository = new HelixReadOnlyZKSharedSchemaRepository(
         readOnlyZKSharedSystemStoreRepository,
         zkClient,

clients/venice-client/src/test/java/com/linkedin/venice/fastclient/meta/RequestBasedMetadataTest.java

@@ -10,10 +10,13 @@
 import static com.linkedin.venice.fastclient.meta.RequestBasedMetadataTestUtils.getMockMetaData;
 import static com.linkedin.venice.utils.TestUtils.waitForNonDeterministicAssertion;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.timeout;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -26,7 +29,9 @@
 import com.linkedin.venice.client.schema.RouterBackedSchemaReader;
 import com.linkedin.venice.client.store.D2ServiceDiscovery;
 import com.linkedin.venice.client.store.transport.D2TransportClient;
+import com.linkedin.venice.client.store.transport.TransportClientResponse;
 import com.linkedin.venice.compression.CompressionStrategy;
+import com.linkedin.venice.controllerapi.D2ServiceDiscoveryResponse;
 import com.linkedin.venice.exceptions.ConfigurationException;
 import com.linkedin.venice.fastclient.ClientConfig;
 import com.linkedin.venice.serialization.avro.AvroProtocolDefinition;
@@ -370,4 +375,31 @@ public void testRequestBasedMetadataStartFailFast() throws IOException {
       Assert.assertThrows(ConfigurationException.class, requestBasedMetadata::start);
     }
   }
+
+  @Test(timeOut = TEST_TIMEOUT)
+  public void testRequestBasedMetadataOnDemandRefresh() throws IOException, InterruptedException {
+    String storeName = "testStore";
+    ClientConfig clientConfig = RequestBasedMetadataTestUtils.getMockClientConfig(storeName, false, false);
+    D2TransportClient d2TransportClient = mock(D2TransportClient.class);
+    CompletableFuture<TransportClientResponse> exceptionFuture = new CompletableFuture<>();
+    exceptionFuture.completeExceptionally(new RuntimeException("Failed to execute"));
+    doReturn(exceptionFuture).when(d2TransportClient).get(anyString());
+    D2ServiceDiscovery d2ServiceDiscovery = getMockD2ServiceDiscovery(d2TransportClient, storeName);
+    D2ServiceDiscoveryResponse d2Response = new D2ServiceDiscoveryResponse();
+    d2Response.setServerD2Service("test-service");
+    doReturn(d2Response).when(d2ServiceDiscovery).find(any(), any(), anyBoolean());
+    try (RequestBasedMetadata requestBasedMetadata = new RequestBasedMetadata(clientConfig, d2TransportClient)) {
+      RequestBasedMetadata spy = spy(requestBasedMetadata);
+      spy.setD2ServiceDiscovery(d2ServiceDiscovery);
+      // let child thread handling the start logic otherwise the main thread cannot verify the invocation times.
+      CompletableFuture.runAsync(spy::start);
+
+      // refresh would happen multiple times
+      // the first one w/o onDemond refresh and would fail due to d2 client exception
+      verify(spy, timeout(3000).atLeast(1)).updateCache(false);
+
+      // the first failed refresh triggers a onDemand refresh
+      verify(spy, timeout(3000).atLeast(1)).updateCache(true);
+    }
+  }
 }