[controller][test] Add durable store update callback

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

@@ -0,0 +1,197 @@
+package com.linkedin.venice.controller;
+
+import static com.linkedin.venice.ConfigKeys.CONTROLLER_AUTO_MATERIALIZE_DAVINCI_PUSH_STATUS_SYSTEM_STORE;
+import static com.linkedin.venice.ConfigKeys.CONTROLLER_AUTO_MATERIALIZE_META_SYSTEM_STORE;
+import static com.linkedin.venice.controllerapi.ControllerApiConstants.READ_QUOTA_IN_CU;
+
+import com.linkedin.venice.controllerapi.ControllerClient;
+import com.linkedin.venice.controllerapi.ControllerResponse;
+import com.linkedin.venice.controllerapi.NewStoreResponse;
+import com.linkedin.venice.controllerapi.UpdateStoreQueryParams;
+import com.linkedin.venice.exceptions.VeniceRetriableException;
+import com.linkedin.venice.integration.utils.ServiceFactory;
+import com.linkedin.venice.integration.utils.VeniceControllerWrapper;
+import com.linkedin.venice.integration.utils.VeniceMultiRegionClusterCreateOptions;
+import com.linkedin.venice.integration.utils.VeniceTwoLayerMultiRegionMultiClusterWrapper;
+import com.linkedin.venice.meta.Store;
+import com.linkedin.venice.meta.StoreInfo;
+import com.linkedin.venice.utils.TestUtils;
+import com.linkedin.venice.utils.Time;
+import com.linkedin.venice.utils.Utils;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Properties;
+import java.util.Set;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+
+public class StoreUpdateHandlerIntegrationTest {
+  private static final long TEST_TIMEOUT_MS = 2 * Time.MS_PER_MINUTE;
+  private static final long UPDATED_READ_QUOTA = 1234;
+
+  @Test(timeOut = TEST_TIMEOUT_MS)
+  public void testStoreUpdateHandlerRetriesWithFinalReadOnlySnapshot() throws InterruptedException {
+    String storeName = Utils.getUniqueString("store-update-handler");
+    String originalOwner = "test-owner";
+    RetryingStoreUpdateHandler storeUpdateHandler = new RetryingStoreUpdateHandler(storeName, UPDATED_READ_QUOTA);
+    AtomicInteger childHandlerInvocationCount = new AtomicInteger();
+
+    Properties parentControllerProperties = new Properties();
+    parentControllerProperties.setProperty(CONTROLLER_AUTO_MATERIALIZE_META_SYSTEM_STORE, Boolean.FALSE.toString());
+    parentControllerProperties
+        .setProperty(CONTROLLER_AUTO_MATERIALIZE_DAVINCI_PUSH_STATUS_SYSTEM_STORE, Boolean.FALSE.toString());
+    parentControllerProperties.put(VeniceControllerWrapper.STORE_UPDATE_HANDLER, storeUpdateHandler);
+    Properties childControllerProperties = new Properties();
+    childControllerProperties.setProperty(CONTROLLER_AUTO_MATERIALIZE_META_SYSTEM_STORE, Boolean.FALSE.toString());
+    childControllerProperties
+        .setProperty(CONTROLLER_AUTO_MATERIALIZE_DAVINCI_PUSH_STATUS_SYSTEM_STORE, Boolean.FALSE.toString());
+    childControllerProperties.put(
+        VeniceControllerWrapper.STORE_UPDATE_HANDLER,
+        (StoreUpdateHandler) (clusterName, store, updatedConfigs) -> childHandlerInvocationCount.incrementAndGet());
+
+    VeniceMultiRegionClusterCreateOptions options =
+        new VeniceMultiRegionClusterCreateOptions.Builder().numberOfRegions(1)
+            .numberOfClusters(1)
+            .numberOfParentControllers(1)
+            .numberOfChildControllers(1)
+            .numberOfServers(0)
+            .numberOfRouters(0)
+            .replicationFactor(1)
+            .parentControllerProperties(parentControllerProperties)
+            .childControllerProperties(childControllerProperties)
+            .build();
+
+    try (VeniceTwoLayerMultiRegionMultiClusterWrapper venice =
+        ServiceFactory.getVeniceTwoLayerMultiRegionMultiClusterWrapper(options)) {
+      String clusterName = venice.getClusterNames()[0];
+      String childControllerUrl = venice.getChildRegions().get(0).getControllerConnectString();
+      try (
+          ControllerClient parentControllerClient =
+              new ControllerClient(clusterName, venice.getControllerConnectString());
+          ControllerClient childControllerClient = new ControllerClient(clusterName, childControllerUrl)) {
+        NewStoreResponse newStoreResponse =
+            parentControllerClient.createNewStore(storeName, originalOwner, "\"string\"", "\"string\"");
+        Assert.assertFalse(newStoreResponse.isError(), newStoreResponse.getError());
+        TestUtils.waitForNonDeterministicAssertion(
+            30,
+            TimeUnit.SECONDS,
+            () -> Assert.assertFalse(childControllerClient.getStore(storeName).isError()));
+
+        ControllerResponse updateStoreResponse = parentControllerClient
+            .updateStore(storeName, new UpdateStoreQueryParams().setReadQuotaInCU(UPDATED_READ_QUOTA));
+        Assert.assertFalse(updateStoreResponse.isError(), updateStoreResponse.getError());
+        Assert.assertTrue(
+            storeUpdateHandler.awaitSuccessfulInvocation(30, TimeUnit.SECONDS),
+            "The store update handler did not succeed after its first-attempt failure");
+
+        Store callbackStore = storeUpdateHandler.getLatestStore();
+        Assert.assertTrue(storeUpdateHandler.getInvocationCount() >= 2);
+        Assert.assertEquals(storeUpdateHandler.getLatestClusterName(), clusterName);
+        Assert.assertEquals(callbackStore.getName(), storeName);
+        Assert.assertEquals(callbackStore.getOwner(), originalOwner);
+        Assert.assertEquals(callbackStore.getReadQuotaInCU(), UPDATED_READ_QUOTA);
+        Assert.assertTrue(storeUpdateHandler.receivedOnlyReadOnlyStores());
+        Assert.assertTrue(storeUpdateHandler.receivedOnlyImmutableUpdatedConfigs());
+        Assert.assertTrue(
+            storeUpdateHandler.getReceivedUpdatedConfigs()
+                .stream()
+                .allMatch(updatedConfigs -> updatedConfigs.equals(Collections.singleton(READ_QUOTA_IN_CU))));
+
+        String barrierOwner = "owner-after-update";
+        ControllerResponse setOwnerResponse = parentControllerClient.setStoreOwner(storeName, barrierOwner);
+        Assert.assertFalse(setOwnerResponse.isError(), setOwnerResponse.getError());
+        TestUtils.waitForNonDeterministicAssertion(30, TimeUnit.SECONDS, () -> {
+          StoreInfo childStore = childControllerClient.getStore(storeName).getStore();
+          Assert.assertEquals(childStore.getOwner(), barrierOwner);
+          Assert.assertEquals(childStore.getReadQuotaInCU(), UPDATED_READ_QUOTA);
+        });
+
+        StoreInfo parentStore = parentControllerClient.getStore(storeName).getStore();
+        Assert.assertEquals(parentStore.getReadQuotaInCU(), UPDATED_READ_QUOTA);
+        Assert.assertTrue(storeUpdateHandler.getInvocationCount() >= 2);
+        Assert.assertEquals(childHandlerInvocationCount.get(), 0);
+      }
+    }
+  }
+
+  private static final class RetryingStoreUpdateHandler implements StoreUpdateHandler {
+    private final String targetStoreName;
+    private final long targetReadQuota;
+    private final AtomicInteger invocationCount = new AtomicInteger();
+    private final AtomicReference<String> latestClusterName = new AtomicReference<>();
+    private final AtomicReference<Store> latestStore = new AtomicReference<>();
+    private final AtomicBoolean receivedOnlyReadOnlyStores = new AtomicBoolean(true);
+    private final AtomicBoolean receivedOnlyImmutableUpdatedConfigs = new AtomicBoolean(true);
+    private final CopyOnWriteArrayList<Set<String>> receivedUpdatedConfigs = new CopyOnWriteArrayList<>();
+    private final CountDownLatch successfulInvocation = new CountDownLatch(1);
+
+    private RetryingStoreUpdateHandler(String targetStoreName, long targetReadQuota) {
+      this.targetStoreName = targetStoreName;
+      this.targetReadQuota = targetReadQuota;
+    }
+
+    @Override
+    public void handleStoreUpdate(String clusterName, Store store, Set<String> updatedConfigs) {
+      if (!targetStoreName.equals(store.getName()) || store.getReadQuotaInCU() != targetReadQuota) {
+        return;
+      }
+
+      latestClusterName.set(clusterName);
+      latestStore.set(store);
+      receivedUpdatedConfigs.add(updatedConfigs);
+      try {
+        store.setOwner("unexpected-mutation");
+        receivedOnlyReadOnlyStores.set(false);
+      } catch (UnsupportedOperationException expected) {
+        // Expected for callback snapshots.
+      }
+      try {
+        updatedConfigs.add("unexpected-config");
+        receivedOnlyImmutableUpdatedConfigs.set(false);
+      } catch (UnsupportedOperationException expected) {
+        // Expected for callback config sets.
+      }
+
+      if (invocationCount.incrementAndGet() == 1) {
+        throw new VeniceRetriableException("Expected first-attempt store update handler failure");
+      }
+      successfulInvocation.countDown();
+    }
+
+    private boolean awaitSuccessfulInvocation(long timeout, TimeUnit unit) throws InterruptedException {
+      return successfulInvocation.await(timeout, unit);
+    }
+
+    private int getInvocationCount() {
+      return invocationCount.get();
+    }
+
+    private String getLatestClusterName() {
+      return latestClusterName.get();
+    }
+
+    private Store getLatestStore() {
+      return latestStore.get();
+    }
+
+    private boolean receivedOnlyReadOnlyStores() {
+      return receivedOnlyReadOnlyStores.get();
+    }
+
+    private boolean receivedOnlyImmutableUpdatedConfigs() {
+      return receivedOnlyImmutableUpdatedConfigs.get();
+    }
+
+    private List<Set<String>> getReceivedUpdatedConfigs() {
+      return new ArrayList<>(receivedUpdatedConfigs);
+    }
+  }
+}

internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/VeniceControllerWrapper.java

@@ -68,6 +68,7 @@
 import com.linkedin.venice.acl.VeniceComponent;
 import com.linkedin.venice.client.store.ClientConfig;
 import com.linkedin.venice.controller.Admin;
+import com.linkedin.venice.controller.StoreUpdateHandler;
 import com.linkedin.venice.controller.VeniceController;
 import com.linkedin.venice.controller.VeniceControllerContext;
 import com.linkedin.venice.controller.VeniceHelixAdmin;
@@ -117,6 +118,7 @@ public class VeniceControllerWrapper extends ProcessWrapper {
   public static final String PARENT_D2_SERVICE_NAME = "ParentController";
 
   public static final String SUPERSET_SCHEMA_GENERATOR = "SupersetSchemaGenerator";
+  public static final String STORE_UPDATE_HANDLER = "StoreUpdateHandler";
 
   public static final double DEFAULT_STORAGE_ENGINE_OVERHEAD_RATIO = 0.85d;
 
@@ -421,6 +423,11 @@ static StatefulServiceProvider<VeniceControllerWrapper> generateService(VeniceCo
       if (passedSupersetSchemaGenerator instanceof SupersetSchemaGenerator) {
         supersetSchemaGenerator = Optional.of((SupersetSchemaGenerator) passedSupersetSchemaGenerator);
       }
+      Optional<StoreUpdateHandler> storeUpdateHandler = Optional.empty();
+      Object passedStoreUpdateHandler = options.getExtraProperties().get(STORE_UPDATE_HANDLER);
+      if (passedStoreUpdateHandler instanceof StoreUpdateHandler) {
+        storeUpdateHandler = Optional.of((StoreUpdateHandler) passedStoreUpdateHandler);
+      }
       Map<String, D2Client> d2Clients = options.getD2Clients();
       VeniceControllerContext ctx = new VeniceControllerContext.Builder().setPropertiesList(propertiesList)
           .setMetricsRepository(metricsRepository)
@@ -431,6 +438,7 @@ static StatefulServiceProvider<VeniceControllerWrapper> generateService(VeniceCo
           .setRouterClientConfig(consumerClientConfig.orElse(null))
           .setExternalSupersetSchemaGenerator(supersetSchemaGenerator.orElse(null))
           .setAccessController(options.getDynamicAccessController())
+          .setStoreUpdateHandler(storeUpdateHandler.orElse(null))
           .build();
       VeniceController veniceController = new VeniceController(ctx);
       return new VeniceControllerWrapper(

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

@@ -0,0 +1,37 @@
+package com.linkedin.venice.controller;
+
+import com.linkedin.venice.meta.Store;
+import java.util.Set;
+
+
+/**
+ * Handles a successful store update consumed by a parent controller.
+ */
+@FunctionalInterface
+public interface StoreUpdateHandler {
+  StoreUpdateHandler NO_OP = new StoreUpdateHandler() {
+    @Override
+    public void handleStoreUpdate(String clusterName, Store store, Set<String> updatedConfigs) {
+    }
+
+    @Override
+    public boolean isNoOp() {
+      return true;
+    }
+  };
+
+  /**
+   * @param clusterName the cluster containing the updated store
+   * @param store a read-only snapshot of the final store state
+   * @param updatedConfigs the immutable set of config keys copied from the durable UPDATE_STORE message; the set is
+   *        stable across retries of the same admin operation
+   */
+  void handleStoreUpdate(String clusterName, Store store, Set<String> updatedConfigs);
+
+  /**
+   * @return whether callback-specific work, including fetching the final store snapshot, should be skipped
+   */
+  default boolean isNoOp() {
+    return false;
+  }
+}

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() {