[controller] admin tool command to clean execution IDs

clients/venice-admin-tool/src/main/java/com/linkedin/venice/AdminTool.java

@@ -262,6 +262,11 @@ public static void main(String[] args) throws Exception {
           storeResponse = queryStoreList(cmd);
           printObject(storeResponse);
           break;
+        case CLEAN_EXECUTION_IDS:
+          String cluster = getRequiredArgument(cmd, Arg.CLUSTER);
+          response = controllerClient.cleanExecutionIds(cluster);
+          printObject(response);
+          break;
         case DESCRIBE_STORE:
           storeName = getRequiredArgument(cmd, Arg.STORE, Command.DESCRIBE_STORE);
           for (String store: storeName.split(",")) {

clients/venice-admin-tool/src/main/java/com/linkedin/venice/Command.java

@@ -600,6 +600,10 @@ public enum Command {
   UPDATE_ADMIN_OPERATION_PROTOCOL_VERSION(
       "update-admin-operation-protocol-version", "Update the admin operation protocol version",
       new Arg[] { URL, CLUSTER, ADMIN_OPERATION_PROTOCOL_VERSION }
+  ),
+  CLEAN_EXECUTION_IDS(
+      "clean-execution-ids", "Clean execution ids for the deleted store from `succeededPerStore` map.",
+      new Arg[] { URL, CLUSTER }
   );
 
   private final String commandName;

internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/CleanExecutionIdsResponse.java

@@ -0,0 +1,16 @@
+package com.linkedin.venice.controllerapi;
+
+import java.util.Map;
+
+
+public class CleanExecutionIdsResponse extends ControllerResponse {
+  private Map<String, Long> cleanedExecutionIds;
+
+  public void setCleanedExecutionIds(Map<String, Long> cleanedExecutionIds) {
+    this.cleanedExecutionIds = cleanedExecutionIds;
+  }
+
+  public Map<String, Long> getCleanedExecutionIds() {
+    return this.cleanedExecutionIds;
+  }
+}

internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/ControllerClient.java

@@ -891,6 +891,11 @@ public MultiStoreResponse queryStoreList(
     return request(ControllerRoute.LIST_STORES, queryParams, MultiStoreResponse.class);
   }
 
+  public CleanExecutionIdsResponse cleanExecutionIds(String clusterName) {
+    QueryParams queryParams = newParams().add(CLUSTER, clusterName);
+    return request(ControllerRoute.CLEAN_EXECUTION_IDS, queryParams, CleanExecutionIdsResponse.class);
+  }
+
   public MultiStoreStatusResponse listStoresStatuses() {
     return request(ControllerRoute.CLUSTER_HEALTH_STORES, newParams(), MultiStoreStatusResponse.class);
   }

internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/ControllerRoute.java

@@ -153,6 +153,7 @@ public enum ControllerRoute {
   JOB("/job", HttpMethod.GET, Arrays.asList(NAME, VERSION)),
   KILL_OFFLINE_PUSH_JOB("/kill_offline_push_job", HttpMethod.POST, Collections.singletonList(TOPIC)),
   LIST_STORES("/list_stores", HttpMethod.GET, Collections.emptyList(), INCLUDE_SYSTEM_STORES),
+  CLEAN_EXECUTION_IDS("/clean_execution_ids", HttpMethod.GET, Collections.emptyList(), CLUSTER),
   LIST_CHILD_CLUSTERS("/list_child_clusters", HttpMethod.GET, Collections.emptyList()),
   LIST_NODES("/list_instances", HttpMethod.GET, Collections.emptyList()),
   CLUSTER_HEALTH_STORES("/cluster_health_stores", HttpMethod.GET, Collections.emptyList()),

internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/AdminToolE2ETest.java

@@ -9,16 +9,20 @@
 import com.linkedin.venice.AdminTool;
 import com.linkedin.venice.common.VeniceSystemStoreType;
 import com.linkedin.venice.controllerapi.ControllerClient;
+import com.linkedin.venice.controllerapi.ControllerResponse;
 import com.linkedin.venice.controllerapi.MultiStoreResponse;
 import com.linkedin.venice.controllerapi.NewStoreResponse;
 import com.linkedin.venice.controllerapi.StoreResponse;
+import com.linkedin.venice.controllerapi.UpdateStoreQueryParams;
+import com.linkedin.venice.exceptions.ErrorType;
 import com.linkedin.venice.integration.utils.ServiceFactory;
 import com.linkedin.venice.integration.utils.VeniceControllerWrapper;
 import com.linkedin.venice.integration.utils.VeniceMultiClusterWrapper;
 import com.linkedin.venice.integration.utils.VeniceMultiRegionClusterCreateOptions;
 import com.linkedin.venice.integration.utils.VeniceTwoLayerMultiRegionMultiClusterWrapper;
 import com.linkedin.venice.meta.StoreInfo;
 import com.linkedin.venice.utils.TestUtils;
+import com.linkedin.venice.utils.TestWriteUtils;
 import com.linkedin.venice.utils.Utils;
 import java.util.Arrays;
 import java.util.HashSet;
@@ -28,6 +32,7 @@
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
+import org.testng.Assert;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
@@ -299,6 +304,62 @@ public void testUpdateStoreMigrationStatus() throws Exception {
     }
   }
 
+  @Test(timeOut = TEST_TIMEOUT)
+  public void testCleanExecutionIds() throws Exception {
+    String storeName1 = Utils.getUniqueString("testCleanExecutionIds-1");
+    String storeName2 = Utils.getUniqueString("testCleanExecutionIds-2");
+    List<VeniceControllerWrapper> parentControllers = multiRegionMultiClusterWrapper.getParentControllers();
+    String clusterName = clusterNames[0];
+    String parentControllerURLs =
+        parentControllers.stream().map(VeniceControllerWrapper::getControllerUrl).collect(Collectors.joining(","));
+    ControllerClient parentControllerClient = new ControllerClient(clusterName, parentControllerURLs);
+    ControllerClient childControllerClient = ControllerClient
+        .constructClusterControllerClient(clusterName, childDatacenters.get(0).getControllerConnectString());
+
+    createStore(parentControllerClient, childControllerClient, storeName1);
+    createStore(parentControllerClient, childControllerClient, storeName2);
+
+    String[] adminToolArgs = new String[] { "--url", childControllerClient.getLeaderControllerUrl(), "--cluster",
+        clusterName, "--clean-execution-ids" };
+
+    AdminTool.main(adminToolArgs);
+
+    UpdateStoreQueryParams updateStoreParams = new UpdateStoreQueryParams();
+    updateStoreParams.setEnableReads(false).setEnableWrites(false);
+    TestWriteUtils.updateStore(storeName1, parentControllerClient, updateStoreParams);
+    parentControllerClient.deleteStore(storeName1);
+    ControllerResponse deleteStoreResponse = parentControllerClient.retryableRequest(5, c -> c.deleteStore(storeName1));
+    Assert.assertFalse(
+        deleteStoreResponse.isError(),
+        "The DeleteStoreResponse returned an error: " + deleteStoreResponse.getError());
+
+    TestUtils.waitForNonDeterministicAssertion(30, TimeUnit.SECONDS, () -> {
+      StoreResponse getStoreResponse = childControllerClient.getStore(storeName1);
+      Assert.assertEquals(getStoreResponse.getErrorType(), ErrorType.STORE_NOT_FOUND);
+    });
+
+    AdminTool.main(adminToolArgs);
+  }
+
+  private void createStore(
+      ControllerClient parentControllerClient,
+      ControllerClient childControllerClient,
+      String storeName) {
+    TestUtils.assertCommand(
+        parentControllerClient
+            .retryableRequest(5, c -> c.createNewStore(storeName, "test", "\"string\"", "\"string\"")));
+
+    TestUtils.assertCommand(
+        parentControllerClient
+            .retryableRequest(5, c -> c.emptyPush(storeName, Utils.getUniqueString("empty-push-1"), 1L)));
+
+    TestUtils.waitForNonDeterministicCompletion(
+        100,
+        TimeUnit.SECONDS,
+        () -> childControllerClient.getStore(storeName).getStore().getCurrentVersion() > 0);
+
+  }
+
   private void validateStoreMigrationStatusAcrossChildRegions(
       String storeName,
       String clusterName,

internal/venice-test-common/src/main/java/com/linkedin/venice/admin/InMemoryExecutionIdAccessor.java

@@ -3,6 +3,8 @@
 import com.linkedin.venice.controller.ExecutionIdAccessor;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
 
 
 /**
@@ -50,6 +52,22 @@ public void updateLastGeneratedExecutionId(String clusterName, Long lastGenerate
     // not used, no op.
   }
 
+  @Override
+  public Map<String, Long> cleanExecutionIdMap(String clusterName, Set<String> allStores) {
+    // initializing `executionIdsCleaned` with all the entries
+    Map<String, Long> executionIdsCleaned = new HashMap<>(executionIdMapInMem.get(clusterName));
+    Map<String, Long> executionIdsToKeep = executionIdMapInMem.get(clusterName)
+        .entrySet()
+        .parallelStream()
+        .filter(entry -> allStores.contains(entry.getKey()))
+        .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+
+    // update `executionIdsCleaned` by removing the entries that are kept
+    executionIdsCleaned.keySet().removeAll(executionIdsToKeep.keySet());
+    executionIdMapInMem.put(clusterName, executionIdsToKeep);
+    return executionIdsCleaned;
+  }
+
   @Override
   public Long incrementAndGetExecutionId(String clusterName) {
     return ++executionId;

services/venice-controller/src/main/java/com/linkedin/venice/controller/ExecutionIdAccessor.java

@@ -1,6 +1,7 @@
 package com.linkedin.venice.controller;
 
 import java.util.Map;
+import java.util.Set;
 
 
 /**
@@ -37,6 +38,8 @@ public interface ExecutionIdAccessor {
    */
   void updateLastGeneratedExecutionId(String clusterName, Long lastGeneratedExecutionId);
 
+  Map<String, Long> cleanExecutionIdMap(String clusterName, Set<String> allStores);
+
   /**
    * Read the current value from ZK and try to increment the value by 1 and write it back to ZK.
    * @return  updated execution id.

services/venice-controller/src/main/java/com/linkedin/venice/controller/ZkExecutionIdAccessor.java

@@ -11,7 +11,10 @@
 import com.linkedin.venice.utils.PathResourceRegistry;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
 import org.apache.helix.AccessOption;
 import org.apache.helix.manager.zk.ZkBaseDataAccessor;
 import org.apache.helix.zookeeper.impl.client.ZkClient;
@@ -158,6 +161,28 @@ private void updateExecutionIdMapToZk(String path, String storeName, Long lastSu
     });
   }
 
+  public Map<String, Long> cleanExecutionIdMap(String clusterName, Set<String> allStores) {
+    String path = getLastSucceededExecutionIdMapPath(clusterName);
+    AtomicReference<Map<String, Long>> executionIdsCleaned = new AtomicReference<>(new HashMap<>());
+    AtomicReference<Map<String, Long>> executionIdsToKeep = new AtomicReference<>(new HashMap<>());
+
+    HelixUtils.compareAndUpdate(zkMapAccessor, path, ZK_RETRY_COUNT, executionIdMap -> {
+      // initializing `executionIdsCleaned` with all the entries
+      executionIdsCleaned.set(executionIdMap);
+      Map<String, Long> filteredExecutionIds = executionIdMap.entrySet()
+          .parallelStream()
+          .filter(entry -> allStores.contains(entry.getKey()))
+          .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+      executionIdsToKeep.set(filteredExecutionIds);
+      return filteredExecutionIds;
+    });
+
+    // update `executionIdsCleaned` by removing the entries that are kept
+    executionIdsCleaned.get().keySet().removeAll(executionIdsToKeep.get().keySet());
+
+    return executionIdsCleaned.get();
+  }
+
   private Long getExecutionIdFromZk(String path) {
     int retry = ZK_RETRY_COUNT;
     while (retry > 0) {

services/venice-controller/src/main/java/com/linkedin/venice/controller/server/AdminSparkServer.java

@@ -11,6 +11,7 @@
 import static com.linkedin.venice.controllerapi.ControllerRoute.BACKUP_VERSION;
 import static com.linkedin.venice.controllerapi.ControllerRoute.CHECK_RESOURCE_CLEANUP_FOR_STORE_CREATION;
 import static com.linkedin.venice.controllerapi.ControllerRoute.CLEANUP_INSTANCE_CUSTOMIZED_STATES;
+import static com.linkedin.venice.controllerapi.ControllerRoute.CLEAN_EXECUTION_IDS;
 import static com.linkedin.venice.controllerapi.ControllerRoute.CLUSTER_DISCOVERY;
 import static com.linkedin.venice.controllerapi.ControllerRoute.CLUSTER_HEALTH_STORES;
 import static com.linkedin.venice.controllerapi.ControllerRoute.COMPARE_STORE;
@@ -330,6 +331,9 @@ public boolean startInner() throws Exception {
     httpService.get(
         LIST_STORES.getPath(),
         new VeniceParentControllerRegionStateHandler(admin, storesRoutes.getAllStores(admin)));
+    httpService.get(
+        CLEAN_EXECUTION_IDS.getPath(),
+        new VeniceParentControllerRegionStateHandler(admin, storesRoutes.cleanExecutionIds(admin)));
     httpService.get(
         CLUSTER_HEALTH_STORES.getPath(),
         new VeniceParentControllerRegionStateHandler(admin, storesRoutes.getAllStoresStatuses(admin)));

services/venice-controller/src/main/java/com/linkedin/venice/controller/server/StoresRoutes.java

@@ -65,8 +65,11 @@
 import com.linkedin.venice.acl.DynamicAccessController;
 import com.linkedin.venice.controller.Admin;
 import com.linkedin.venice.controller.AdminCommandExecutionTracker;
+import com.linkedin.venice.controller.VeniceHelixAdmin;
+import com.linkedin.venice.controller.VeniceParentHelixAdmin;
 import com.linkedin.venice.controller.kafka.TopicCleanupService;
 import com.linkedin.venice.controller.repush.RepushJobRequest;
+import com.linkedin.venice.controllerapi.CleanExecutionIdsResponse;
 import com.linkedin.venice.controllerapi.ClusterStaleDataAuditResponse;
 import com.linkedin.venice.controllerapi.ControllerResponse;
 import com.linkedin.venice.controllerapi.MultiStoreInfoResponse;
@@ -261,6 +264,28 @@ public void internalHandle(Request request, MultiStoreResponse veniceResponse) {
     };
   }
 
+  public Route cleanExecutionIds(Admin admin) {
+    return new VeniceRouteHandler<CleanExecutionIdsResponse>(CleanExecutionIdsResponse.class) {
+      @Override
+      public void internalHandle(Request request, CleanExecutionIdsResponse veniceResponse) {
+        String cluster = request.queryParams(CLUSTER);
+        VeniceHelixAdmin veniceHelixAdmin;
+        if (admin instanceof VeniceParentHelixAdmin) {
+          veniceHelixAdmin = ((VeniceParentHelixAdmin) admin).getVeniceHelixAdmin();
+        } else {
+          veniceHelixAdmin = (VeniceHelixAdmin) admin;
+        }
+
+        Set<String> allStores =
+            veniceHelixAdmin.getAllStores(cluster).stream().map(Store::getName).collect(Collectors.toSet());
+        Map<String, Long> executionIdsCleaned =
+            veniceHelixAdmin.getExecutionIdAccessor().cleanExecutionIdMap(cluster, allStores);
+
+        veniceResponse.setCleanedExecutionIds(executionIdsCleaned);
+      }
+    };
+  }
+
   /**
    * No ACL check; any user can try to list store statuses.
    * @see Admin#getAllStoreStatuses(String)
@@ -512,7 +537,7 @@ public void internalHandle(Request request, StoreMigrationResponse veniceRespons
   }
 
   /**
-   * @see Admin#deleteStore(String, String, int, boolean)
+   * @see Admin#deleteStore(String, String, boolean, int, boolean)
    */
   public Route deleteStore(Admin admin) {
     return new VeniceRouteHandler<TrackableControllerResponse>(TrackableControllerResponse.class) {