Add a flow quota manager and throttling

core/src/main/java/google/registry/flows/FlowRunner.java

@@ -23,8 +23,10 @@
 import google.registry.flows.FlowModule.RegistrarId;
 import google.registry.flows.FlowModule.Superuser;
 import google.registry.flows.FlowModule.Transactional;
+import google.registry.flows.quota.FlowQuotaManager;
 import google.registry.flows.session.LoginFlow;
 import google.registry.model.eppcommon.Trid;
+import google.registry.model.eppinput.EppInput;
 import google.registry.model.eppoutput.EppOutput;
 import google.registry.monitoring.whitebox.EppMetric;
 import google.registry.persistence.PersistenceModule.TransactionIsolationLevel;
@@ -56,6 +58,8 @@ public class FlowRunner {
   @Inject Trid trid;
   @Inject FlowReporter flowReporter;
   @Inject JpaTransactionManager jpaTransactionManager;
+  @Inject EppInput eppInput;
+  @Inject FlowQuotaManager flowQuotaManager;
 
   @Inject FlowRunner() {}
 
@@ -80,6 +84,11 @@ public EppOutput run(final EppMetric.Builder eppMetricBuilder) throws EppExcepti
     eppMetricBuilder.setCommandNameFromFlow(flowClass.getSimpleName());
     final StopwatchLogger stopwatch = new StopwatchLogger();
 
+    // First, acquire quota if necessary
+    if (!isSuperuser) {
+      flowQuotaManager.acquireQuota(flowClass, eppInput, registrarId);
+    }
+
     // We may already be in a transaction, e.g., when invoked by DeleteExpiredDomainsAction.
     if (!isTransactional || jpaTransactionManager.inTransaction()) {
       stopwatch.tick("We're in transaction, running the flow now.");

core/src/main/java/google/registry/flows/quota/FlowQuotaManager.java

@@ -0,0 +1,81 @@
+// Copyright 2026 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.flows.quota;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.flogger.FluentLogger;
+import google.registry.flows.EppException;
+import google.registry.flows.Flow;
+import google.registry.model.eppinput.EppInput;
+import google.registry.quota.QuotaManager;
+import java.time.Duration;
+import javax.annotation.concurrent.ThreadSafe;
+
+/**
+ * Quota management for EPP flows using Redis/Valkey.
+ *
+ * <p>This is provided the base {@link QuotaManager} as well as a list of {@link
+ * FlowQuotaParameters}. The latter is used to define per-flow throttling characteristics so that we
+ * can check quota before the flow class itself is instantiated.
+ */
+@ThreadSafe
+public class FlowQuotaManager {
+
+  private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
+  private final QuotaManager quotaManager;
+  private final ImmutableMap<Class<? extends Flow>, FlowQuotaParameters> parametersMap;
+
+  public static FlowQuotaManager create(
+      QuotaManager quotaManager, ImmutableList<FlowQuotaParameters> allParameters) {
+    ImmutableMap<Class<? extends Flow>, FlowQuotaParameters> parametersMap =
+        allParameters.stream()
+            .collect(ImmutableMap.toImmutableMap(FlowQuotaParameters::getFlowClass, p -> p));
+    return new FlowQuotaManager(quotaManager, parametersMap);
+  }
+
+  private FlowQuotaManager(
+      QuotaManager quotaManager,
+      ImmutableMap<Class<? extends Flow>, FlowQuotaParameters> parametersMap) {
+    this.quotaManager = quotaManager;
+    this.parametersMap = parametersMap;
+  }
+
+  /** Acquires one unit of quota from the quota manager. Throws an exception on failure. */
+  public void acquireQuota(Class<? extends Flow> flowClass, EppInput eppInput, String registrarId)
+      throws TooManyRequestsException {
+    FlowQuotaParameters parameters = parametersMap.get(flowClass);
+    if (parameters == null) {
+      return;
+    }
+    String quotaId = parameters.getQuotaId(eppInput, registrarId);
+    int maxQuotaAllowed = parameters.getMaxQuotaAllowed();
+    Duration windowDuration = parameters.getWindowDuration();
+    if (!quotaManager.acquireQuota(quotaId, maxQuotaAllowed, windowDuration)) {
+      logger.atWarning().log(
+          "Failed to acquire quota for flow %s, registrar %s",
+          flowClass.getSimpleName(), registrarId);
+      throw new TooManyRequestsException();
+    }
+  }
+
+  /** Too many requests too quickly. */
+  public static class TooManyRequestsException extends EppException.CommandUseErrorException {
+    public TooManyRequestsException() {
+      super("Too many requests");
+    }
+  }
+}

core/src/main/java/google/registry/flows/quota/FlowQuotaModule.java

@@ -0,0 +1,37 @@
+// Copyright 2026 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.flows.quota;
+
+import com.google.common.collect.ImmutableList;
+import dagger.Module;
+import dagger.Provides;
+import google.registry.quota.NoopQuotaManager;
+import google.registry.quota.QuotaManager;
+import google.registry.quota.ValkeyQuotaManager;
+import jakarta.inject.Singleton;
+import java.util.Optional;
+import redis.clients.jedis.UnifiedJedis;
+
+@Module
+public class FlowQuotaModule {
+
+  @Provides
+  @Singleton
+  static FlowQuotaManager provideFlowQuotaManager(Optional<UnifiedJedis> jedis) {
+    QuotaManager quotaManager =
+        jedis.isPresent() ? new ValkeyQuotaManager(jedis.get(), "flow") : new NoopQuotaManager();
+    return FlowQuotaManager.create(quotaManager, ImmutableList.of());
+  }
+}

core/src/main/java/google/registry/flows/quota/FlowQuotaParameters.java

@@ -0,0 +1,48 @@
+// Copyright 2026 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.flows.quota;
+
+import google.registry.flows.Flow;
+import google.registry.model.eppinput.EppInput;
+import java.time.Duration;
+
+/**
+ * Provider of information about how we should throttle / have quota for a given flow type.
+ *
+ * <p>Normally, we'd wish to define these in the flow class itself by using some interface that
+ * might be named something like "ThrottlingFlow". However, the flow class itself is not
+ * instantiated until a transaction has been opened (if the flow is a transactional flow). This is
+ * suboptimal, given that one of the primary purposes of throttling is to reduce database load. As a
+ * result, we define the throttling parameters outside the flow class itself so the throttling /
+ * quota can be checked before the class is actually instantiated.
+ */
+public interface FlowQuotaParameters {
+
+  /** The flow that is being throttled. */
+  Class<? extends Flow> getFlowClass();
+
+  /**
+   * The ID of the quota, e.g. a domain-name-registrar-ID combo.
+   *
+   * <p>Note: the registrar ID may be empty if we're not authenticated yet, like a LoginFlow.
+   */
+  String getQuotaId(EppInput eppInput, String registrarId);
+
+  /** The maximum number of tokens / quota for each ID for this flow. */
+  int getMaxQuotaAllowed();
+
+  /** The fixed window over which to throttle requests with the same ID. */
+  Duration getWindowDuration();
+}

core/src/main/java/google/registry/module/RegistryComponent.java

@@ -32,6 +32,7 @@
 import google.registry.flows.ServerTridProviderModule;
 import google.registry.flows.custom.CustomLogicFactoryModule;
 import google.registry.flows.domain.DomainDeletionTimeCacheModule;
+import google.registry.flows.quota.FlowQuotaModule;
 import google.registry.groups.DirectoryModule;
 import google.registry.groups.GmailModule;
 import google.registry.groups.GroupsModule;
@@ -68,6 +69,7 @@
       DirectoryModule.class,
       DomainDeletionTimeCacheModule.class,
       DriveModule.class,
+      FlowQuotaModule.class,
       GmailModule.class,
       GroupsModule.class,
       GroupssettingsModule.class,

core/src/test/java/google/registry/flows/EppTestComponent.java

@@ -14,6 +14,7 @@
 
 package google.registry.flows;
 
+import com.google.common.collect.ImmutableList;
 import dagger.Component;
 import dagger.Module;
 import dagger.Provides;
@@ -27,7 +28,9 @@
 import google.registry.flows.custom.TestCustomLogicFactory;
 import google.registry.flows.domain.DomainDeletionTimeCache;
 import google.registry.flows.domain.DomainFlowTmchUtils;
+import google.registry.flows.quota.FlowQuotaManager;
 import google.registry.monitoring.whitebox.EppMetric;
+import google.registry.quota.NoopQuotaManager;
 import google.registry.request.Modules.GsonModule;
 import google.registry.request.RequestScope;
 import google.registry.request.lock.LockHandler;
@@ -60,6 +63,7 @@ class FakesAndMocksModule {
     private FakeLockHandler lockHandler;
     private Sleeper sleeper;
     private CloudTasksHelper cloudTasksHelper;
+    private FlowQuotaManager flowQuotaManager;
 
     public CloudTasksHelper getCloudTasksHelper() {
       return cloudTasksHelper;
@@ -69,6 +73,10 @@ public EppMetric.Builder getMetricBuilder() {
       return metricBuilder;
     }
 
+    public FlowQuotaManager getFlowQuotaManager() {
+      return flowQuotaManager;
+    }
+
     public static FakesAndMocksModule create(FakeClock clock) {
       FakesAndMocksModule instance = new FakesAndMocksModule();
       CloudTasksHelper cloudTasksHelper = new CloudTasksHelper(clock);
@@ -82,6 +90,8 @@ public static FakesAndMocksModule create(FakeClock clock) {
       instance.metricBuilder = EppMetric.builderForRequest(clock);
       instance.lockHandler = new FakeLockHandler(true);
       instance.cloudTasksHelper = cloudTasksHelper;
+      instance.flowQuotaManager =
+          FlowQuotaManager.create(new NoopQuotaManager(), ImmutableList.of());
       return instance;
     }
 
@@ -134,6 +144,11 @@ ServerTridProvider provideServerTridProvider() {
     DomainDeletionTimeCache provideDomainDeletionTimeCache() {
       return DomainDeletionTimeCache.create();
     }
+
+    @Provides
+    FlowQuotaManager provideFlowQuotaManager() {
+      return flowQuotaManager;
+    }
   }
 
   class FakeServerTridProvider implements ServerTridProvider {

core/src/test/java/google/registry/flows/FlowRunnerTest.java

@@ -21,27 +21,40 @@
 import static google.registry.testing.TestLogHandlerUtils.findFirstLogMessageByPrefix;
 import static google.registry.util.DateTimeUtils.START_INSTANT;
 import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import com.google.common.base.Joiner;
 import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedMap;
 import com.google.common.net.InetAddresses;
 import com.google.common.testing.TestLogHandler;
 import google.registry.flows.certs.CertificateChecker;
+import google.registry.flows.quota.FlowQuotaManager;
+import google.registry.flows.quota.FlowQuotaManager.TooManyRequestsException;
+import google.registry.flows.quota.FlowQuotaParameters;
 import google.registry.model.eppcommon.Trid;
+import google.registry.model.eppinput.EppInput;
 import google.registry.model.eppoutput.EppOutput.ResponseOrGreeting;
 import google.registry.model.eppoutput.EppResponse;
 import google.registry.monitoring.whitebox.EppMetric;
 import google.registry.persistence.PersistenceModule.TransactionIsolationLevel;
 import google.registry.persistence.transaction.JpaTestExtensions;
 import google.registry.persistence.transaction.JpaTestExtensions.JpaIntegrationTestExtension;
+import google.registry.quota.NoopQuotaManager;
+import google.registry.quota.QuotaManager;
 import google.registry.testing.FakeClock;
 import google.registry.testing.FakeHttpSession;
 import google.registry.util.JdkLoggerConfig;
+import java.time.Duration;
 import java.time.Instant;
 import java.util.List;
 import java.util.Optional;
@@ -112,6 +125,9 @@ void beforeEach() {
     flowRunner.trid = Trid.create("client-123", "server-456");
     flowRunner.flowReporter = mock(FlowReporter.class);
     flowRunner.jpaTransactionManager = tm();
+    flowRunner.eppInput = mock(EppInput.class);
+    flowRunner.flowQuotaManager =
+        FlowQuotaManager.create(new NoopQuotaManager(), ImmutableList.of());
   }
 
   @Test
@@ -224,4 +240,51 @@ void testRun_loggingStatement_complexEppInput() throws Exception {
     String xml = Joiner.on('\n').join(lines.subList(3, lines.size() - 4));
     assertThat(xml).isEqualTo(sanitizedDomainCreateXml);
   }
+
+  @Test
+  void testRun_quotaExceeded_throwsException() {
+    QuotaManager mockQuotaManager = mock(QuotaManager.class);
+    when(mockQuotaManager.acquireQuota(
+            eq("TheRegistrar:test-quota-id"), anyInt(), any(Duration.class)))
+        .thenReturn(false);
+    flowRunner.flowQuotaManager =
+        FlowQuotaManager.create(mockQuotaManager, ImmutableList.of(new TestFlowQuotaParameters()));
+
+    assertThrows(TooManyRequestsException.class, () -> flowRunner.run(eppMetricBuilder));
+  }
+
+  @Test
+  void testRun_quotaAvailable_succeeds() throws Exception {
+    QuotaManager mockQuotaManager = mock(QuotaManager.class);
+    when(mockQuotaManager.acquireQuota(
+            eq("TheRegistrar:test-quota-id"), anyInt(), any(Duration.class)))
+        .thenReturn(true);
+    flowRunner.flowQuotaManager =
+        FlowQuotaManager.create(mockQuotaManager, ImmutableList.of(new TestFlowQuotaParameters()));
+
+    flowRunner.run(eppMetricBuilder);
+    verify(mockQuotaManager).acquireQuota("TheRegistrar:test-quota-id", 10, Duration.ofMinutes(1));
+  }
+
+  private static class TestFlowQuotaParameters implements FlowQuotaParameters {
+    @Override
+    public Class<? extends Flow> getFlowClass() {
+      return TestCommandFlow.class;
+    }
+
+    @Override
+    public String getQuotaId(EppInput eppInput, String registrarId) {
+      return registrarId + ":test-quota-id";
+    }
+
+    @Override
+    public int getMaxQuotaAllowed() {
+      return 10;
+    }
+
+    @Override
+    public Duration getWindowDuration() {
+      return Duration.ofMinutes(1);
+    }
+  }
 }