[da-vinci] Bug fix for peer discovery in DVC with multi-stores

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

@@ -1,16 +1,12 @@
 package com.linkedin.davinci.blobtransfer;
 
-import static com.linkedin.venice.client.store.ClientFactory.getTransportClient;
-
 import com.linkedin.davinci.blobtransfer.client.NettyFileTransferClient;
 import com.linkedin.davinci.blobtransfer.server.P2PBlobTransferService;
 import com.linkedin.davinci.stats.AggVersionedBlobTransferStats;
 import com.linkedin.davinci.storage.StorageEngineRepository;
 import com.linkedin.davinci.storage.StorageMetadataService;
 import com.linkedin.venice.blobtransfer.DaVinciBlobFinder;
 import com.linkedin.venice.blobtransfer.ServerBlobFinder;
-import com.linkedin.venice.client.store.AbstractAvroStoreClient;
-import com.linkedin.venice.client.store.AvroGenericStoreClientImpl;
 import com.linkedin.venice.client.store.ClientConfig;
 import com.linkedin.venice.helix.HelixCustomizedViewOfflinePushRepository;
 import com.linkedin.venice.meta.ReadOnlyStoreRepository;
@@ -24,40 +20,14 @@ public class BlobTransferUtil {
 
   /**
    * Get a P2P blob transfer manager for DaVinci Client and start it.
-   * @param p2pTransferPort, the port used by the P2P transfer server and client
+   * @param p2pTransferServerPort, the port used by the P2P transfer server
+   * @param p2pTransferClientPort, the port used by the P2P transfer client
    * @param baseDir, the base directory of the underlying storage
    * @param clientConfig, the client config to start up a transport client
    * @param storageMetadataService, the storage metadata service
    * @return the blob transfer manager
    * @throws Exception
    */
-  public static BlobTransferManager<Void> getP2PBlobTransferManagerForDVCAndStart(
-      int p2pTransferPort,
-      String baseDir,
-      ClientConfig clientConfig,
-      StorageMetadataService storageMetadataService,
-      ReadOnlyStoreRepository readOnlyStoreRepository,
-      StorageEngineRepository storageEngineRepository,
-      int maxConcurrentSnapshotUser,
-      int snapshotRetentionTimeInMin,
-      int blobTransferMaxTimeoutInMin,
-      AggVersionedBlobTransferStats aggVersionedBlobTransferStats,
-      BlobTransferUtils.BlobTransferTableFormat transferSnapshotTableFormat) {
-    return getP2PBlobTransferManagerForDVCAndStart(
-        p2pTransferPort,
-        p2pTransferPort,
-        baseDir,
-        clientConfig,
-        storageMetadataService,
-        readOnlyStoreRepository,
-        storageEngineRepository,
-        maxConcurrentSnapshotUser,
-        snapshotRetentionTimeInMin,
-        blobTransferMaxTimeoutInMin,
-        aggVersionedBlobTransferStats,
-        transferSnapshotTableFormat);
-  }
-
   public static BlobTransferManager<Void> getP2PBlobTransferManagerForDVCAndStart(
       int p2pTransferServerPort,
       int p2pTransferClientPort,
@@ -79,12 +49,10 @@ public static BlobTransferManager<Void> getP2PBlobTransferManagerForDVCAndStart(
           maxConcurrentSnapshotUser,
           snapshotRetentionTimeInMin,
           transferSnapshotTableFormat);
-      AbstractAvroStoreClient storeClient =
-          new AvroGenericStoreClientImpl<>(getTransportClient(clientConfig), false, clientConfig);
       BlobTransferManager<Void> manager = new NettyP2PBlobTransferManager(
           new P2PBlobTransferService(p2pTransferServerPort, baseDir, blobTransferMaxTimeoutInMin, blobSnapshotManager),
           new NettyFileTransferClient(p2pTransferClientPort, baseDir, storageMetadataService),
-          new DaVinciBlobFinder(storeClient),
+          new DaVinciBlobFinder(clientConfig),
           baseDir,
           aggVersionedBlobTransferStats);
       manager.start();

clients/venice-thin-client/src/main/java/com/linkedin/venice/client/store/ClientFactory.java

@@ -22,7 +22,7 @@ public class ClientFactory {
   private static Function<ClientConfig, TransportClient> configToTransportClientProviderForTests = null;
 
   // Visible for testing
-  static void setUnitTestMode() {
+  public static void setUnitTestMode() {
     unitTestMode = true;
   }
 
@@ -33,7 +33,7 @@ static void resetUnitTestMode() {
 
   // Allow for overriding with mock D2Client for unit tests. The caller must release the object to prevent side-effects
   // VisibleForTesting
-  static void setTransportClientProvider(Function<ClientConfig, TransportClient> transportClientProvider) {
+  public static void setTransportClientProvider(Function<ClientConfig, TransportClient> transportClientProvider) {
     if (!unitTestMode) {
       throw new VeniceUnsupportedOperationException("setTransportClientProvider in non-unit-test-mode");
     }

internal/venice-common/src/main/java/com/linkedin/venice/blobtransfer/DaVinciBlobFinder.java

@@ -1,12 +1,16 @@
 package com.linkedin.venice.blobtransfer;
 
+import static com.linkedin.venice.client.store.ClientFactory.getTransportClient;
 import static com.linkedin.venice.controllerapi.ControllerApiConstants.NAME;
 import static com.linkedin.venice.controllerapi.ControllerApiConstants.STORE_PARTITION;
 import static com.linkedin.venice.controllerapi.ControllerApiConstants.STORE_VERSION;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.linkedin.venice.client.store.AbstractAvroStoreClient;
+import com.linkedin.venice.client.store.AvroGenericStoreClientImpl;
+import com.linkedin.venice.client.store.ClientConfig;
 import com.linkedin.venice.utils.ObjectMapperFactory;
+import com.linkedin.venice.utils.concurrent.VeniceConcurrentHashMap;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
@@ -29,14 +33,35 @@ public class DaVinciBlobFinder implements BlobFinder {
   private static final String TYPE_BLOB_DISCOVERY = "blob_discovery";
   private static final String ERROR_DISCOVERY_MESSAGE =
       "Error finding DVC peers for blob transfer in store: %s, version: %d, partition: %d";
-  private final AbstractAvroStoreClient storeClient;
+  private final ClientConfig clientConfig;
+  private VeniceConcurrentHashMap<String, AbstractAvroStoreClient> storeToClientMap;
 
-  public DaVinciBlobFinder(AbstractAvroStoreClient storeClient) {
-    this.storeClient = storeClient;
+  public DaVinciBlobFinder(ClientConfig clientConfig) {
+    this.storeToClientMap = new VeniceConcurrentHashMap<>();
+    this.clientConfig = clientConfig;
+  }
+
+  /**
+   * Get the store client for the given store name
+   * @param storeName
+   * @return the store client
+   */
+  AbstractAvroStoreClient getStoreClient(String storeName) {
+    return storeToClientMap.computeIfAbsent(storeName, k -> {
+      // update the config with respective store name
+      ClientConfig storeClientConfig = ClientConfig.cloneConfig(clientConfig).setStoreName(storeName);
+      AbstractAvroStoreClient storeLevelClient =
+          new AvroGenericStoreClientImpl<>(getTransportClient(storeClientConfig), false, storeClientConfig);
+      storeLevelClient.start();
+      LOGGER.info("Started store client for store: {}", storeName);
+      return storeLevelClient;
+    });
   }
 
   @Override
   public BlobPeersDiscoveryResponse discoverBlobPeers(String storeName, int version, int partition) {
+    AbstractAvroStoreClient storeClient = getStoreClient(storeName);
+
     String uri = buildUriForBlobDiscovery(storeName, version, partition);
     CompletableFuture<BlobPeersDiscoveryResponse> futureResponse = CompletableFuture.supplyAsync(() -> {
       try {
@@ -85,6 +110,8 @@ private BlobPeersDiscoveryResponse handleError(
 
   @Override
   public void close() throws IOException {
-    storeClient.close();
+    for (AbstractAvroStoreClient storeClient: storeToClientMap.values()) {
+      storeClient.close();
+    }
   }
 }

internal/venice-common/src/test/java/com/linkedin/venice/blobtransfer/DaVinciBlobFinderTest.java

@@ -11,15 +11,20 @@
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.linkedin.venice.client.store.AbstractAvroStoreClient;
+import com.linkedin.venice.client.store.ClientConfig;
+import com.linkedin.venice.client.store.ClientFactory;
+import com.linkedin.venice.client.store.transport.TransportClient;
 import com.linkedin.venice.client.store.transport.TransportClientResponse;
 import com.linkedin.venice.utils.ObjectMapperFactory;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
+import java.util.function.Function;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
 import org.testng.Assert;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
@@ -35,12 +40,18 @@ public class DaVinciBlobFinderTest {
 
   @BeforeMethod
   public void setUp() {
+    ClientConfig clientConfig = ClientConfig.defaultGenericClientConfig(storeName);
+
     storeClient = mock(AbstractAvroStoreClient.class);
-    daVinciBlobFinder = new DaVinciBlobFinder(storeClient);
+    daVinciBlobFinder = spy(new DaVinciBlobFinder(clientConfig));
+
+    Mockito.doReturn(storeName).when(storeClient).getStoreName();
   }
 
   @Test
   public void testDiscoverBlobPeers_Success() {
+    Mockito.doReturn(storeClient).when(daVinciBlobFinder).getStoreClient(storeName);
+
     String responseBodyJson =
         "{\"error\":false,\"errorMessage\":\"\",\"discoveryResult\":[\"host1\",\"host2\",\"host3\"]}";
     byte[] responseBody = responseBodyJson.getBytes(StandardCharsets.UTF_8);
@@ -57,6 +68,7 @@ public void testDiscoverBlobPeers_Success() {
 
   @Test
   public void testDiscoverBlobPeers_CallsTransportClientWithCorrectURI() {
+    Mockito.doReturn(storeClient).when(daVinciBlobFinder).getStoreClient(storeName);
     String responseBodyJson =
         "{\"error\":false,\"errorMessage\":\"\",\"discoveryResult\":[\"host1\",\"host2\",\"host3\"]}";
     byte[] responseBody = responseBodyJson.getBytes(StandardCharsets.UTF_8);
@@ -78,6 +90,8 @@ public void testDiscoverBlobPeers_CallsTransportClientWithCorrectURI() {
 
   @Test
   public void testDiscoverBlobPeers_ContentDeserializationError() throws Exception {
+    Mockito.doReturn(storeClient).when(daVinciBlobFinder).getStoreClient(storeName);
+
     String responseBodyJson = "{\"error\":true,\"errorMessage\":\"some error\",\"discoveryResult\":[]}";
     byte[] responseBody = responseBodyJson.getBytes(StandardCharsets.UTF_8);
     TransportClientResponse mockResponse = new TransportClientResponse(0, null, responseBody);
@@ -98,6 +112,8 @@ public void testDiscoverBlobPeers_ContentDeserializationError() throws Exception
 
   @Test
   public void testDiscoverBlobPeers_ClientWithIncorrectUri() {
+    Mockito.doReturn(storeClient).when(daVinciBlobFinder).getStoreClient(storeName);
+
     CompletableFuture<byte[]> futureResponse = new CompletableFuture<>();
     futureResponse.completeExceptionally(new RuntimeException("Test Exception"));
     when(storeClient.getRaw(anyString())).thenReturn(futureResponse);
@@ -109,4 +125,39 @@ public void testDiscoverBlobPeers_ClientWithIncorrectUri() {
         response.getErrorMessage(),
         "Error finding DVC peers for blob transfer in store: testStore, version: 1, partition: 1");
   }
+
+  @Test
+  public void testGetStoreClient() throws IOException {
+    // set up the transport client provider used to initialize the store client
+    TransportClient transportClient1 = mock(TransportClient.class);
+    TransportClient transportClient2 = mock(TransportClient.class);
+    Function<ClientConfig, TransportClient> clientConfigTransportClientFunction = (clientConfig) -> {
+      if (clientConfig.getStoreName().equals(storeName)) {
+        return transportClient1;
+      } else if (clientConfig.getStoreName().equals("storeName2")) {
+        return transportClient2;
+      } else {
+        // Create TransportClient the regular way
+        return null;
+      }
+    };
+    ClientFactory.setUnitTestMode();
+    ClientFactory.setTransportClientProvider(clientConfigTransportClientFunction);
+
+    // ClientConfig is initialized with storeName
+    AbstractAvroStoreClient storeClient = daVinciBlobFinder.getStoreClient(storeName);
+    Assert.assertNotNull(storeClient);
+    Assert.assertEquals(storeClient.getStoreName(), storeName);
+
+    // Even if the daVinciBlobFinder is initialized at the beginning with "storeName", the getStoreClient
+    // method should be able to return a store client for "storeName2"
+    AbstractAvroStoreClient storeClient2 = daVinciBlobFinder.getStoreClient("storeName2");
+    Assert.assertNotNull(storeClient2);
+    Assert.assertEquals(storeClient2.getStoreName(), "storeName2");
+
+    daVinciBlobFinder.close();
+    // verify that the store client is closed
+    verify(transportClient1).close();
+    verify(transportClient2).close();
+  }
 }

internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DaVinciClientTest.java

@@ -1209,6 +1209,13 @@ public void testBlobP2PTransferAmongDVC() throws Exception {
         String snapshotPath = RocksDBUtils.composeSnapshotDir(dvcPath1 + "/rocksdb", storeName + "_v1", i);
         Assert.assertTrue(Files.exists(Paths.get(snapshotPath)));
       }
+
+      for (int i = 0; i < 3; i++) {
+        String partitionPath2 = RocksDBUtils.composePartitionDbDir(dvcPath2 + "/rocksdb", storeName + "_v1", i);
+        Assert.assertTrue(Files.exists(Paths.get(partitionPath2)));
+        String snapshotPath2 = RocksDBUtils.composeSnapshotDir(dvcPath2 + "/rocksdb", storeName + "_v1", i);
+        Assert.assertFalse(Files.exists(Paths.get(snapshotPath2)));
+      }
     }
   }