[controller][test] Add durable store update callback

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

@@ -0,0 +1,18 @@
+package com.linkedin.venice.controller;
+
+import com.linkedin.venice.meta.Store;
+
+
+/**
+ * Handles a successful store update consumed by a parent controller.
+ */
+@FunctionalInterface
+public interface StoreUpdateHandler {
+  StoreUpdateHandler NO_OP = (clusterName, store) -> {};
+
+  /**
+   * @param clusterName the cluster containing the updated store
+   * @param store a read-only snapshot of the final store state
+   */
+  void handleStoreUpdate(String clusterName, Store store);
+}

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

@@ -156,6 +156,7 @@ public static List<Class<? extends ModuleMetricEntityInterface>> getMetricEntity
   private final Optional<List<VeniceVersionLifecycleEventListener>> versionLifecycleEventListeners;
   private final Optional<List<ValueSchemaCreatedListener>> valueSchemaCreatedListeners;
   private final Optional<ExternalETLService> externalETLService;
+  private final StoreUpdateHandler storeUpdateHandler;
 
   /**
    * Allocates a new {@code VeniceController} object.
@@ -226,6 +227,7 @@ public VeniceController(VeniceControllerContext ctx) {
     this.versionLifecycleEventListeners = Optional.ofNullable(ctx.getVersionLifecycleEventListeners());
     this.valueSchemaCreatedListeners = Optional.ofNullable(ctx.getValueSchemaCreatedListeners());
     this.externalETLService = Optional.ofNullable(ctx.getExternalETLService());
+    this.storeUpdateHandler = ctx.getStoreUpdateHandler();
     this.controllerService = createControllerService();
     this.adminServer = createAdminServer(false);
     this.secureAdminServer = sslEnabled ? createAdminServer(true) : null;
@@ -262,7 +264,8 @@ private VeniceControllerService createControllerService() {
         pubSubPositionTypeRegistry,
         versionLifecycleEventListeners,
         valueSchemaCreatedListeners,
-        externalETLService);
+        externalETLService,
+        storeUpdateHandler);
     Admin admin = veniceControllerService.getVeniceHelixAdmin();
     if (multiClusterConfigs.isParent() && !(admin instanceof VeniceParentHelixAdmin)) {
       throw new VeniceException(

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

@@ -40,6 +40,7 @@ public class VeniceControllerContext {
   private List<VeniceVersionLifecycleEventListener> versionLifecycleEventListeners;
   private List<ValueSchemaCreatedListener> valueSchemaCreatedListeners;
   private ExternalETLService externalETLService;
+  private StoreUpdateHandler storeUpdateHandler;
 
   public List<VeniceProperties> getPropertiesList() {
     return propertiesList;
@@ -97,6 +98,10 @@ public ExternalETLService getExternalETLService() {
     return externalETLService;
   }
 
+  public StoreUpdateHandler getStoreUpdateHandler() {
+    return storeUpdateHandler;
+  }
+
   public VeniceControllerContext(Builder builder) {
     this.propertiesList = builder.propertiesList;
     this.metricsRepository = builder.metricsRepository;
@@ -112,6 +117,8 @@ public VeniceControllerContext(Builder builder) {
     this.versionLifecycleEventListeners = builder.versionLifecycleEventListeners;
     this.valueSchemaCreatedListeners = builder.valueSchemaCreatedListeners;
     this.externalETLService = builder.externalETLService;
+    this.storeUpdateHandler =
+        builder.storeUpdateHandler == null ? StoreUpdateHandler.NO_OP : builder.storeUpdateHandler;
   }
 
   public static class Builder {
@@ -132,6 +139,7 @@ public static class Builder {
     private List<VeniceVersionLifecycleEventListener> versionLifecycleEventListeners;
     private List<ValueSchemaCreatedListener> valueSchemaCreatedListeners;
     private ExternalETLService externalETLService;
+    private StoreUpdateHandler storeUpdateHandler;
 
     public Builder setPropertiesList(List<VeniceProperties> propertiesList) {
       this.propertiesList = propertiesList;
@@ -207,6 +215,11 @@ public Builder setExternalETLService(ExternalETLService externalETLService) {
       return this;
     }
 
+    public Builder setStoreUpdateHandler(StoreUpdateHandler storeUpdateHandler) {
+      this.storeUpdateHandler = storeUpdateHandler;
+      return this;
+    }
+
     private void addDefaultValues() {
       if (metricsRepository == null && !isMetricsRepositorySet) {
 
@@ -223,6 +236,9 @@ private void addDefaultValues() {
       if (serviceDiscoveryAnnouncers == null && !isServiceDiscoveryAnnouncerSet) {
         serviceDiscoveryAnnouncers = Collections.emptyList();
       }
+      if (storeUpdateHandler == null) {
+        storeUpdateHandler = StoreUpdateHandler.NO_OP;
+      }
     }
 
     public VeniceControllerContext build() {

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

@@ -73,6 +73,46 @@ public VeniceControllerService(
       Optional<List<VeniceVersionLifecycleEventListener>> versionLifecycleEventListeners,
       Optional<List<ValueSchemaCreatedListener>> valueSchemaCreatedListeners,
       Optional<ExternalETLService> externalETLService) {
+    this(
+        multiClusterConfigs,
+        metricsRepository,
+        sslEnabled,
+        sslConfig,
+        accessController,
+        authorizerService,
+        d2Client,
+        d2Clients,
+        routerClientConfig,
+        icProvider,
+        externalSupersetSchemaGenerator,
+        pubSubTopicRepository,
+        pubSubClientsFactory,
+        pubSubPositionTypeRegistry,
+        versionLifecycleEventListeners,
+        valueSchemaCreatedListeners,
+        externalETLService,
+        StoreUpdateHandler.NO_OP);
+  }
+
+  public VeniceControllerService(
+      VeniceControllerMultiClusterConfig multiClusterConfigs,
+      MetricsRepository metricsRepository,
+      boolean sslEnabled,
+      Optional<SSLConfig> sslConfig,
+      Optional<DynamicAccessController> accessController,
+      Optional<AuthorizerService> authorizerService,
+      D2Client d2Client,
+      Map<String, D2Client> d2Clients,
+      Optional<ClientConfig> routerClientConfig,
+      Optional<ICProvider> icProvider,
+      Optional<SupersetSchemaGenerator> externalSupersetSchemaGenerator,
+      PubSubTopicRepository pubSubTopicRepository,
+      PubSubClientsFactory pubSubClientsFactory,
+      PubSubPositionTypeRegistry pubSubPositionTypeRegistry,
+      Optional<List<VeniceVersionLifecycleEventListener>> versionLifecycleEventListeners,
+      Optional<List<ValueSchemaCreatedListener>> valueSchemaCreatedListeners,
+      Optional<ExternalETLService> externalETLService,
+      StoreUpdateHandler storeUpdateHandler) {
     this.multiClusterConfigs = multiClusterConfigs;
 
     DelegatingClusterLeaderInitializationRoutine initRoutineForPushJobDetailsSystemStore =
@@ -207,7 +247,8 @@ public VeniceControllerService(
             metricsRepository,
             pubSubClientsFactory.getConsumerAdapterFactory(),
             pubSubTopicRepository,
-            pubSubMessageDeserializer);
+            pubSubMessageDeserializer,
+            storeUpdateHandler);
         this.consumerServicesByClusters.put(cluster, adminConsumerService);
 
         this.admin.setAdminConsumerService(cluster, adminConsumerService);

services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumerService.java

@@ -3,6 +3,7 @@
 import static com.linkedin.venice.pubsub.PubSubUtil.getPubSubPositionWireFormat;
 
 import com.linkedin.venice.annotation.VisibleForTesting;
+import com.linkedin.venice.controller.StoreUpdateHandler;
 import com.linkedin.venice.controller.VeniceControllerClusterConfig;
 import com.linkedin.venice.controller.VeniceHelixAdmin;
 import com.linkedin.venice.controller.ZkAdminTopicMetadataAccessor;
@@ -50,6 +51,7 @@ public class AdminConsumerService extends AbstractVeniceService {
   private final PubSubMessageDeserializer pubSubMessageDeserializer;
   private final LogContext logContext;
   private final PubSubPositionDeserializer pubSubPositionDeserializer;
+  private final StoreUpdateHandler storeUpdateHandler;
 
   public AdminConsumerService(
       VeniceHelixAdmin admin,
@@ -58,6 +60,24 @@ public AdminConsumerService(
       PubSubConsumerAdapterFactory consumerFactory,
       PubSubTopicRepository pubSubTopicRepository,
       PubSubMessageDeserializer pubSubMessageDeserializer) {
+    this(
+        admin,
+        config,
+        metricsRepository,
+        consumerFactory,
+        pubSubTopicRepository,
+        pubSubMessageDeserializer,
+        StoreUpdateHandler.NO_OP);
+  }
+
+  public AdminConsumerService(
+      VeniceHelixAdmin admin,
+      VeniceControllerClusterConfig config,
+      MetricsRepository metricsRepository,
+      PubSubConsumerAdapterFactory consumerFactory,
+      PubSubTopicRepository pubSubTopicRepository,
+      PubSubMessageDeserializer pubSubMessageDeserializer,
+      StoreUpdateHandler storeUpdateHandler) {
     this.config = config;
     this.logContext = config.getLogContext();
     this.admin = admin;
@@ -77,6 +97,7 @@ public AdminConsumerService(
     this.consumerFactory = consumerFactory;
     this.pubSubPositionDeserializer = config.getPubSubPositionDeserializer();
     this.threadFactory = new DaemonThreadFactory("AdminConsumerService-" + config.getClusterName(), logContext);
+    this.storeUpdateHandler = storeUpdateHandler;
   }
 
   @Override
@@ -118,7 +139,8 @@ private AdminConsumptionTask getAdminConsumptionTaskForCluster(String clusterNam
         config.getAdminConsumptionCycleTimeoutMs(),
         config.getAdminConsumptionMaxWorkerThreadPoolSize(),
         pubSubTopicRepository,
-        config.getRegionName());
+        config.getRegionName(),
+        storeUpdateHandler);
   }
 
   /**

services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumptionTask.java

@@ -5,6 +5,7 @@
 import com.linkedin.venice.common.VeniceSystemStoreType;
 import com.linkedin.venice.controller.AdminTopicMetadataAccessor;
 import com.linkedin.venice.controller.ExecutionIdAccessor;
+import com.linkedin.venice.controller.StoreUpdateHandler;
 import com.linkedin.venice.controller.VeniceHelixAdmin;
 import com.linkedin.venice.controller.kafka.AdminTopicUtils;
 import com.linkedin.venice.controller.kafka.protocol.admin.AdminOperation;
@@ -262,6 +263,7 @@ public ExecutorService getExecutorService() {
    * The local region name of the controller.
    */
   private final String regionName;
+  private final StoreUpdateHandler storeUpdateHandler;
 
   public AdminConsumptionTask(
       String clusterName,
@@ -279,6 +281,42 @@ public AdminConsumptionTask(
       int maxWorkerThreadPoolSize,
       PubSubTopicRepository pubSubTopicRepository,
       String regionName) {
+    this(
+        clusterName,
+        consumer,
+        remoteConsumptionEnabled,
+        remoteKafkaServerUrl,
+        admin,
+        adminTopicMetadataAccessor,
+        executionIdAccessor,
+        isParentController,
+        stats,
+        adminTopicReplicationFactor,
+        minInSyncReplicas,
+        processingCycleTimeoutInMs,
+        maxWorkerThreadPoolSize,
+        pubSubTopicRepository,
+        regionName,
+        StoreUpdateHandler.NO_OP);
+  }
+
+  public AdminConsumptionTask(
+      String clusterName,
+      PubSubConsumerAdapter consumer,
+      boolean remoteConsumptionEnabled,
+      Optional<String> remoteKafkaServerUrl,
+      VeniceHelixAdmin admin,
+      AdminTopicMetadataAccessor adminTopicMetadataAccessor,
+      ExecutionIdAccessor executionIdAccessor,
+      boolean isParentController,
+      AdminConsumptionStats stats,
+      int adminTopicReplicationFactor,
+      Optional<Integer> minInSyncReplicas,
+      long processingCycleTimeoutInMs,
+      int maxWorkerThreadPoolSize,
+      PubSubTopicRepository pubSubTopicRepository,
+      String regionName,
+      StoreUpdateHandler storeUpdateHandler) {
     this.clusterName = clusterName;
     this.pubSubAdminTopic = pubSubTopicRepository.getTopic(AdminTopicUtils.getTopicNameFromClusterName(clusterName));
     this.adminTopicPartition = new PubSubTopicPartitionImpl(pubSubAdminTopic, AdminTopicUtils.ADMIN_TOPIC_PARTITION_ID);
@@ -312,6 +350,7 @@ public AdminConsumptionTask(
         new DaemonThreadFactory(String.format("Venice-Admin-Execution-Task-%s", clusterName), admin.getLogContext()));
     this.undelegatedRecords = new LinkedList<>();
     this.regionName = regionName;
+    this.storeUpdateHandler = storeUpdateHandler;
     this.storeRetryCountMap = new ConcurrentHashMap<>();
 
     if (remoteConsumptionEnabled) {
@@ -569,7 +608,8 @@ private void executeMessagesAndCollectResults() throws InterruptedException {
             isParentController,
             stats,
             regionName,
-            inflightThreadsByStore);
+            inflightThreadsByStore,
+            storeUpdateHandler);
         // Check if there is previously created scheduled task still occupying one thread from the pool.
         if (storesWithScheduledTask.add(storeName)) {
           // Log the store name and the position of the task being added into the task list

services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminExecutionTask.java

@@ -7,6 +7,7 @@
 import com.linkedin.venice.common.VeniceSystemStoreUtils;
 import com.linkedin.venice.compression.CompressionStrategy;
 import com.linkedin.venice.controller.ExecutionIdAccessor;
+import com.linkedin.venice.controller.StoreUpdateHandler;
 import com.linkedin.venice.controller.VeniceHelixAdmin;
 import com.linkedin.venice.controller.kafka.protocol.admin.AbortMigration;
 import com.linkedin.venice.controller.kafka.protocol.admin.AddVersion;
@@ -53,6 +54,7 @@
 import com.linkedin.venice.meta.IngestionPauseMode;
 import com.linkedin.venice.meta.LifecycleHooksRecord;
 import com.linkedin.venice.meta.LifecycleHooksRecordImpl;
+import com.linkedin.venice.meta.ReadOnlyStore;
 import com.linkedin.venice.meta.StorageMode;
 import com.linkedin.venice.meta.Store;
 import com.linkedin.venice.meta.VeniceETLStrategy;
@@ -95,6 +97,7 @@ public class AdminExecutionTask implements Callable<Void> {
   private final long lastPersistedExecutionId;
 
   private final ConcurrentHashMap<String, AtomicInteger> inflightThreadsByStore;
+  private final StoreUpdateHandler storeUpdateHandler;
 
   AdminExecutionTask(
       Logger LOGGER,
@@ -109,6 +112,36 @@ public class AdminExecutionTask implements Callable<Void> {
       AdminConsumptionStats stats,
       String regionName,
       ConcurrentHashMap<String, AtomicInteger> inflightThreadsByStore) {
+    this(
+        LOGGER,
+        clusterName,
+        storeName,
+        lastSucceededExecutionIdMap,
+        lastPersistedExecutionId,
+        internalTopic,
+        admin,
+        executionIdAccessor,
+        isParentController,
+        stats,
+        regionName,
+        inflightThreadsByStore,
+        StoreUpdateHandler.NO_OP);
+  }
+
+  AdminExecutionTask(
+      Logger LOGGER,
+      String clusterName,
+      String storeName,
+      ConcurrentHashMap<String, Long> lastSucceededExecutionIdMap,
+      long lastPersistedExecutionId,
+      Queue<AdminOperationWrapper> internalTopic,
+      VeniceHelixAdmin admin,
+      ExecutionIdAccessor executionIdAccessor,
+      boolean isParentController,
+      AdminConsumptionStats stats,
+      String regionName,
+      ConcurrentHashMap<String, AtomicInteger> inflightThreadsByStore,
+      StoreUpdateHandler storeUpdateHandler) {
     this.LOGGER = LOGGER;
     this.clusterName = clusterName;
     this.storeName = storeName;
@@ -121,6 +154,7 @@ public class AdminExecutionTask implements Callable<Void> {
     this.stats = stats;
     this.regionName = regionName;
     this.inflightThreadsByStore = inflightThreadsByStore;
+    this.storeUpdateHandler = storeUpdateHandler;
   }
 
   @Override
@@ -250,6 +284,7 @@ private void processMessage(AdminOperation adminOperation) {
           lastSucceededExecutionId);
       return;
     }
+    boolean storeUpdated = false;
     try {
       switch (AdminMessageType.valueOf(adminOperation)) {
         case STORE_CREATION:
@@ -287,6 +322,7 @@ private void processMessage(AdminOperation adminOperation) {
           break;
         case UPDATE_STORE:
           handleSetStore((UpdateStore) adminOperation.payloadUnion);
+          storeUpdated = true;
           break;
         case DELETE_STORE:
           handleDeleteStore((DeleteStore) adminOperation.payloadUnion);
@@ -354,6 +390,11 @@ private void processMessage(AdminOperation adminOperation) {
           AdminMessageType.valueOf(adminOperation),
           e.getMessage());
     }
+    if (storeUpdated && isParentController) {
+      Store finalStore = admin.getStore(clusterName, storeName).cloneStore();
+      // Invoke before advancing checkpoints so callback failures leave the admin operation eligible for retry.
+      storeUpdateHandler.handleStoreUpdate(clusterName, new ReadOnlyStore(finalStore));
+    }
     executionIdAccessor.updateLastSucceededExecutionIdMap(clusterName, storeName, adminOperation.executionId);
     lastSucceededExecutionIdMap.put(storeName, adminOperation.executionId);
   }