Made use of StoreName rather than String in the VenicePathParser's Re…

…tryManager maps.
linkedin · Dec 21, 2024 · bbc9f53 · bbc9f53
1 parent 0e44d33
commit bbc9f53
Showing 1 changed file with 11 additions and 8 deletions.
diff --git a/services/venice-router/src/main/java/com/linkedin/venice/router/api/VenicePathParser.java b/services/venice-router/src/main/java/com/linkedin/venice/router/api/VenicePathParser.java
@@ -25,6 +25,7 @@
 import com.linkedin.venice.meta.RetryManager;
 import com.linkedin.venice.meta.Store;
 import com.linkedin.venice.meta.StoreDataChangedListener;
+import com.linkedin.venice.meta.StoreName;
 import com.linkedin.venice.meta.StoreVersionName;
 import com.linkedin.venice.read.RequestType;
 import com.linkedin.venice.router.VeniceRouterConfig;
@@ -121,9 +122,10 @@ public class VenicePathParser implements ExtendedResourcePathParser<VenicePath,
   private final CompressorFactory compressorFactory;
   private final MetricsRepository metricsRepository;
   private final ScheduledExecutorService retryManagerScheduler;
-  private final Map<String, RetryManager> routerSingleKeyRetryManagers;
-  private final Map<String, RetryManager> routerMultiKeyRetryManagers;
+  private final Map<StoreName, RetryManager> routerSingleKeyRetryManagers;
+  private final Map<StoreName, RetryManager> routerMultiKeyRetryManagers;
   private final NameRepository nameRepository = new NameRepository();
+  /** Outer map of {client's desired compression strategy -> inner map of {store-version -> decompressor}} */
   private final EnumMap<CompressionStrategy, Map<StoreVersionName, VeniceResponseDecompressor>> decompressorMaps;
 
   public VenicePathParser(
@@ -145,29 +147,30 @@ public VenicePathParser(
     }
     this.storeRepository.registerStoreDataChangedListener(new StoreDataChangedListener() {
       @Override
-      public void handleStoreDeleted(String storeName) {
+      public void handleStoreDeleted(String storeNameString) {
+        StoreName storeName = nameRepository.getStoreName(storeNameString);
         routerSingleKeyRetryManagers.remove(storeName);
         routerMultiKeyRetryManagers.remove(storeName);
         cleanDecompressorMaps(storeName, storeVersionName -> true);
       }
 
       @Override
       public void handleStoreChanged(Store store) {
-        String storeName = store.getName();
+        StoreName storeName = nameRepository.getStoreName(store.getName());
         IntSet upToDateVersionsSet = store.getVersionNumbers();
         cleanDecompressorMaps(
             storeName,
             storeVersionName -> !upToDateVersionsSet.contains(storeVersionName.getVersionNumber()));
       }
 
-      private void cleanDecompressorMaps(String storeName, Function<StoreVersionName, Boolean> criteriaForRemoval) {
+      private void cleanDecompressorMaps(StoreName storeName, Function<StoreVersionName, Boolean> criteriaForRemoval) {
         for (Map<StoreVersionName, VeniceResponseDecompressor> decompressorMap: decompressorMaps.values()) {
           // remove out dated versions (if any) from the map
           Iterator<StoreVersionName> storeVersionNameIterator = decompressorMap.keySet().iterator();
           StoreVersionName storeVersionName;
           while (storeVersionNameIterator.hasNext()) {
             storeVersionName = storeVersionNameIterator.next();
-            if (storeVersionName.getStoreName().equals(storeName)) {
+            if (storeVersionName.getStore().equals(storeName)) {
               if (criteriaForRemoval.apply(storeVersionName)) {
                 storeVersionNameIterator.remove();
               }
@@ -212,7 +215,7 @@ public VenicePath parseResourceUri(String uri, BasicFullHttpRequest fullHttpRequ
 
       if (VeniceRouterUtils.isHttpGet(method)) {
         RetryManager singleKeyRetryManager = routerSingleKeyRetryManagers.computeIfAbsent(
-            storeName,
+            storeVersionName.getStore(),
             ignored -> new RetryManager(
                 metricsRepository,
                 SINGLE_KEY_RETRY_MANAGER_STATS_PREFIX + storeName,
@@ -231,7 +234,7 @@ public VenicePath parseResourceUri(String uri, BasicFullHttpRequest fullHttpRequ
             responseDecompressor);
       } else if (VeniceRouterUtils.isHttpPost(method)) {
         RetryManager multiKeyRetryManager = routerMultiKeyRetryManagers.computeIfAbsent(
-            storeName,
+            storeVersionName.getStore(),
             ignored -> new RetryManager(
                 metricsRepository,
                 MULTI_KEY_RETRY_MANAGER_STATS_PREFIX + storeName,