[controller] Account for max read capacity in quota change

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

@@ -155,6 +155,7 @@
 import static com.linkedin.venice.ConfigKeys.LOG_COMPACTION_SCHEDULING_ENABLED;
 import static com.linkedin.venice.ConfigKeys.LOG_COMPACTION_THREAD_COUNT;
 import static com.linkedin.venice.ConfigKeys.LOG_COMPACTION_THRESHOLD_MS;
+import static com.linkedin.venice.ConfigKeys.MAX_READ_CAPACITY;
 import static com.linkedin.venice.ConfigKeys.META_STORE_WRITER_CLOSE_CONCURRENCY;
 import static com.linkedin.venice.ConfigKeys.META_STORE_WRITER_CLOSE_TIMEOUT_MS;
 import static com.linkedin.venice.ConfigKeys.MIN_NUMBER_OF_STORE_VERSIONS_TO_PRESERVE;
@@ -578,6 +579,7 @@ public class VeniceControllerClusterConfig {
   private final boolean disableParentRequestTopicForStreamPushes;
 
   private final int defaultReadQuotaPerRouter;
+  private final long maxReadCapacityCu;
 
   private final int defaultMaxRecordSizeBytes; // default value for VeniceWriter.maxRecordSizeBytes
   private final int replicationMetadataVersion;
@@ -760,6 +762,7 @@ public VeniceControllerClusterConfig(VeniceProperties props) {
         props.getBoolean(CONTROLLER_DISABLE_PARENT_REQUEST_TOPIC_FOR_STREAM_PUSHES, false);
     this.defaultReadQuotaPerRouter =
         props.getInt(CONTROLLER_DEFAULT_READ_QUOTA_PER_ROUTER, DEFAULT_PER_ROUTER_READ_QUOTA);
+    this.maxReadCapacityCu = props.getLong(MAX_READ_CAPACITY, 100000);
     this.defaultMaxRecordSizeBytes =
         props.getInt(DEFAULT_MAX_RECORD_SIZE_BYTES, DEFAULT_MAX_RECORD_SIZE_BYTES_BACKFILL);
     if (defaultMaxRecordSizeBytes < BYTES_PER_MB) {
@@ -1248,6 +1251,10 @@ public int getDefaultReadQuotaPerRouter() {
     return defaultReadQuotaPerRouter;
   }
 
+  public long getMaxReadCapacityCu() {
+    return maxReadCapacityCu;
+  }
+
   public int getDefaultMaxRecordSizeBytes() {
     return defaultMaxRecordSizeBytes;
   }

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

@@ -5667,17 +5667,6 @@ private void internalUpdateStore(String clusterName, String storeName, UpdateSto
       }
 
       if (readQuotaInCU.isPresent()) {
-        HelixVeniceClusterResources resources = getHelixVeniceClusterResources(clusterName);
-        ZkRoutersClusterManager routersClusterManager = resources.getRoutersClusterManager();
-        int routerCount = routersClusterManager.getLiveRoutersCount();
-        VeniceControllerClusterConfig clusterConfig = getHelixVeniceClusterResources(clusterName).getConfig();
-        int defaultReadQuotaPerRouter = clusterConfig.getDefaultReadQuotaPerRouter();
-
-        if (Math.max(defaultReadQuotaPerRouter, routerCount * defaultReadQuotaPerRouter) < readQuotaInCU.get()) {
-          throw new VeniceException(
-              "Cannot update read quota for store " + storeName + " in cluster " + clusterName + ". Read quota "
-                  + readQuotaInCU.get() + " requested is more than the cluster quota.");
-        }
         setStoreReadQuota(clusterName, storeName, readQuotaInCU.get());
       }
 

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

@@ -209,6 +209,7 @@
 import com.linkedin.venice.helix.ParentHelixOfflinePushAccessor;
 import com.linkedin.venice.helix.Replica;
 import com.linkedin.venice.helix.StoragePersonaRepository;
+import com.linkedin.venice.helix.ZkRoutersClusterManager;
 import com.linkedin.venice.helix.ZkStoreConfigAccessor;
 import com.linkedin.venice.meta.BackupStrategy;
 import com.linkedin.venice.meta.BufferReplayPolicy;
@@ -2824,8 +2825,21 @@ public void updateStore(String clusterName, String storeName, UpdateStoreQueryPa
       setStore.enableWrites = writeability.map(addToUpdatedConfigList(updatedConfigsList, ENABLE_WRITES))
           .orElseGet(currStore::isEnableWrites);
 
-      setStore.readQuotaInCU = readQuotaInCU.map(addToUpdatedConfigList(updatedConfigsList, READ_QUOTA_IN_CU))
-          .orElseGet(currStore::getReadQuotaInCU);
+      if (readQuotaInCU.isPresent()) {
+        HelixVeniceClusterResources resources = getHelixVeniceClusterResources(clusterName);
+        ZkRoutersClusterManager routersClusterManager = resources.getRoutersClusterManager();
+        int routerCount = routersClusterManager.getLiveRoutersCount();
+        VeniceControllerClusterConfig clusterConfig = getHelixVeniceClusterResources(clusterName).getConfig();
+        int defaultReadQuotaPerRouter = clusterConfig.getDefaultReadQuotaPerRouter();
+        long maxReadCapacityCu = clusterConfig.getMaxReadCapacityCu();
+        if (Math.max(defaultReadQuotaPerRouter, maxReadCapacityCu) * routerCount < readQuotaInCU.get()) {
+          throw new VeniceException(
+              "Cannot update read quota for store " + storeName + " in cluster " + clusterName + ". Read quota "
+                  + readQuotaInCU.get() + " requested is more than the cluster quota.");
+        }
+        setStore.readQuotaInCU = readQuotaInCU.get();
+        updatedConfigsList.add(READ_QUOTA_IN_CU);
+      }
 
       // We need to be careful when handling currentVersion.
       // Since it is not synced between parent and local controller,

services/venice-controller/src/test/java/com/linkedin/venice/controller/TestVeniceParentHelixAdmin.java

@@ -62,6 +62,8 @@
 import com.linkedin.venice.exceptions.VeniceStoreAlreadyExistsException;
 import com.linkedin.venice.exceptions.VeniceUnsupportedOperationException;
 import com.linkedin.venice.helix.HelixReadWriteStoreRepository;
+import com.linkedin.venice.helix.StoragePersonaRepository;
+import com.linkedin.venice.helix.ZkRoutersClusterManager;
 import com.linkedin.venice.meta.BufferReplayPolicy;
 import com.linkedin.venice.meta.ConcurrentPushDetectionStrategy;
 import com.linkedin.venice.meta.HybridStoreConfigImpl;
@@ -123,6 +125,7 @@
 import org.testng.Assert;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
 
@@ -1922,6 +1925,57 @@
     assertEquals(offlineJobStatus.getExecutionStatus(), ExecutionStatus.ERROR);
   }
 
+  @DataProvider(name = "readQuotaTestCases")
+  public Object[][] readQuotaTestCases() {
+    return new Object[][] { { 100, 0, 10L, true }, // Default quota is enough
+        { 0, 100L, 10L, true }, // Default quota not enough but max router capacity is enough
+        { 0, 0, 10L, false } // Neither default quota nor max capacity is enough - should throw exception
+    };
+  }
+
+  @Test(dataProvider = "readQuotaTestCases")
+  public void testUpdateReadQuota(int defaultQuota, long maxCapacity, long requestedQuota, boolean shouldSucceed) {
+    String storeName = Utils.getUniqueString("testUpdateStore");
+    Store store = TestUtils.createTestStore(storeName, "test", System.currentTimeMillis());
+    doReturn(store).when(internalAdmin).getStore(clusterName, storeName);
+    mockResources(config, clusterName);
+    doReturn(internalAdmin.getHelixVeniceClusterResources(clusterName)).when(internalAdmin)
+        .getHelixVeniceClusterResources(clusterName);
+    when(internalAdmin.getHelixVeniceClusterResources(clusterName).getConfig()).thenReturn(config);
+    ZkRoutersClusterManager routersClusterManager = mock(ZkRoutersClusterManager.class);
+    when(internalAdmin.getHelixVeniceClusterResources(clusterName).getRoutersClusterManager())
+        .thenReturn(routersClusterManager);
+    when(routersClusterManager.getLiveRoutersCount()).thenReturn(1);
+    StoragePersonaRepository storagePersonaRepository = mock(StoragePersonaRepository.class);
+    when(internalAdmin.getHelixVeniceClusterResources(clusterName).getStoragePersonaRepository())
+        .thenReturn(storagePersonaRepository);
+
+    when(zkClient.readData(zkMetadataNodePath, null)).thenReturn(null)
+        .thenReturn(
+            AdminTopicMetadataAccessor.generateMetadataMap(
+                Optional.of(1L),
+                Optional.of(-1L),
+                Optional.of(1L),
+                Optional.of(LATEST_SCHEMA_ID_FOR_ADMIN_OPERATION)));
+
+    UpdateStoreQueryParams updateStoreQueryParams = new UpdateStoreQueryParams().setReadQuotaInCU(requestedQuota);
+    parentAdmin.initStorageCluster(clusterName);
+
+    when(config.getDefaultReadQuotaPerRouter()).thenReturn(defaultQuota);
+    when(config.getMaxReadCapacityCu()).thenReturn(maxCapacity);
+
+    if (shouldSucceed) {
+      parentAdmin.updateStore(clusterName, storeName, updateStoreQueryParams);
+      AdminOperation adminMessage = verifyAndGetSingleAdminOperation();
+      UpdateStore updateStore = (UpdateStore) adminMessage.payloadUnion;
+      Assert.assertEquals(updateStore.getReadQuotaInCU(), requestedQuota);
+    } else {
+      Assert.assertThrows(
+          VeniceException.class,
+          () -> parentAdmin.updateStore(clusterName, storeName, updateStoreQueryParams));
+    }
+  }
+
   @Test
   public void testUpdateStore() {
     String storeName = Utils.getUniqueString("testUpdateStore");
@@ -1990,7 +2044,7 @@
         .setMaxRecordSizeBytes(7777)
         .setMaxNearlineRecordSizeBytes(6666);
 
     parentAdmin.updateStore(clusterName, storeName, updateStoreQueryParams);
 
     verify(veniceWriter, times(2)).put(
         keyCaptor.capture(),