From d0138b59c0abaefa530560030b7a4918371c31e4 Mon Sep 17 00:00:00 2001 From: Gus Brodman Date: Tue, 25 Aug 2026 13:10:04 -0400 Subject: [PATCH] Add a flow quota manager and throttling This uses the already-existing quota manager to acquire quota, if necessary, before running a flow. If we can't acquire quota, we throw a command use EPP exception. Note: this is unused so far, as you can see from the FlowQuotaModule (we provide an empty list of flow classes to throttle). This is intentionally done to limit the scope of the commit. We inject the flow-quota-parameters classes because we wish to throttle requests before any possible transaction is started -- for transactional flows, we start the transaction before actually instantiating the flow class. --- .../google/registry/flows/FlowRunner.java | 9 + .../flows/quota/FlowQuotaManager.java | 81 ++++++++ .../registry/flows/quota/FlowQuotaModule.java | 37 ++++ .../flows/quota/FlowQuotaParameters.java | 48 +++++ .../registry/module/RegistryComponent.java | 2 + .../registry/flows/EppTestComponent.java | 15 ++ .../google/registry/flows/FlowRunnerTest.java | 63 ++++++ .../flows/quota/FlowQuotaManagerTest.java | 180 ++++++++++++++++++ .../module/TestRegistryComponent.java | 2 + 9 files changed, 437 insertions(+) create mode 100644 core/src/main/java/google/registry/flows/quota/FlowQuotaManager.java create mode 100644 core/src/main/java/google/registry/flows/quota/FlowQuotaModule.java create mode 100644 core/src/main/java/google/registry/flows/quota/FlowQuotaParameters.java create mode 100644 core/src/test/java/google/registry/flows/quota/FlowQuotaManagerTest.java diff --git a/core/src/main/java/google/registry/flows/FlowRunner.java b/core/src/main/java/google/registry/flows/FlowRunner.java index b93191a8ff2..b1018f23b3b 100644 --- a/core/src/main/java/google/registry/flows/FlowRunner.java +++ b/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."); diff --git a/core/src/main/java/google/registry/flows/quota/FlowQuotaManager.java b/core/src/main/java/google/registry/flows/quota/FlowQuotaManager.java new file mode 100644 index 00000000000..d0f99f459f5 --- /dev/null +++ b/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. + * + *

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, FlowQuotaParameters> parametersMap; + + public static FlowQuotaManager create( + QuotaManager quotaManager, ImmutableList allParameters) { + ImmutableMap, FlowQuotaParameters> parametersMap = + allParameters.stream() + .collect(ImmutableMap.toImmutableMap(FlowQuotaParameters::getFlowClass, p -> p)); + return new FlowQuotaManager(quotaManager, parametersMap); + } + + private FlowQuotaManager( + QuotaManager quotaManager, + ImmutableMap, 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 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"); + } + } +} diff --git a/core/src/main/java/google/registry/flows/quota/FlowQuotaModule.java b/core/src/main/java/google/registry/flows/quota/FlowQuotaModule.java new file mode 100644 index 00000000000..ae21ceab5ce --- /dev/null +++ b/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 jedis) { + QuotaManager quotaManager = + jedis.isPresent() ? new ValkeyQuotaManager(jedis.get(), "flow") : new NoopQuotaManager(); + return FlowQuotaManager.create(quotaManager, ImmutableList.of()); + } +} diff --git a/core/src/main/java/google/registry/flows/quota/FlowQuotaParameters.java b/core/src/main/java/google/registry/flows/quota/FlowQuotaParameters.java new file mode 100644 index 00000000000..d3eed47f187 --- /dev/null +++ b/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. + * + *

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 getFlowClass(); + + /** + * The ID of the quota, e.g. a domain-name-registrar-ID combo. + * + *

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(); +} diff --git a/core/src/main/java/google/registry/module/RegistryComponent.java b/core/src/main/java/google/registry/module/RegistryComponent.java index cc9806115b7..0e3f041484c 100644 --- a/core/src/main/java/google/registry/module/RegistryComponent.java +++ b/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, diff --git a/core/src/test/java/google/registry/flows/EppTestComponent.java b/core/src/test/java/google/registry/flows/EppTestComponent.java index 8fa0779dab9..f90a2e1314b 100644 --- a/core/src/test/java/google/registry/flows/EppTestComponent.java +++ b/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 { diff --git a/core/src/test/java/google/registry/flows/FlowRunnerTest.java b/core/src/test/java/google/registry/flows/FlowRunnerTest.java index 4fdc99147ad..00ee9906352 100644 --- a/core/src/test/java/google/registry/flows/FlowRunnerTest.java +++ b/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 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); + } + } } diff --git a/core/src/test/java/google/registry/flows/quota/FlowQuotaManagerTest.java b/core/src/test/java/google/registry/flows/quota/FlowQuotaManagerTest.java new file mode 100644 index 00000000000..31b7f5cbe2d --- /dev/null +++ b/core/src/test/java/google/registry/flows/quota/FlowQuotaManagerTest.java @@ -0,0 +1,180 @@ +// 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 static com.google.common.truth.Truth.assertThat; +import static google.registry.testing.EppExceptionSubject.assertAboutEppExceptions; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.Mockito.mock; + +import com.google.common.collect.ImmutableList; +import google.registry.flows.Flow; +import google.registry.flows.quota.FlowQuotaManager.TooManyRequestsException; +import google.registry.model.eppinput.EppInput; +import google.registry.model.eppoutput.EppOutput.ResponseOrGreeting; +import google.registry.quota.ValkeyQuotaManager; +import io.github.ss_bhatt.testcontainers.valkey.ValkeyContainer; +import java.time.Duration; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; +import redis.clients.jedis.HostAndPort; +import redis.clients.jedis.RedisClient; + +/** Tests for {@link FlowQuotaManager} backed by Valkey. */ +@Testcontainers +class FlowQuotaManagerTest { + + @Container private static final ValkeyContainer valkey = new ValkeyContainer(); + + private RedisClient jedis; + private ValkeyQuotaManager quotaManager; + private final EppInput eppInput = mock(EppInput.class); + + @BeforeEach + void setUp() { + jedis = + RedisClient.builder() + .hostAndPort(new HostAndPort(valkey.getHost(), valkey.getFirstMappedPort())) + .build(); + jedis.flushAll(); + quotaManager = new ValkeyQuotaManager(jedis, "flow"); + } + + @Test + void testAcquireQuota_noConfiguredParameters_noOp() { + FlowQuotaManager manager = FlowQuotaManager.create(quotaManager, ImmutableList.of()); + assertDoesNotThrow(() -> manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar")); + assertThat(jedis.keys("*")).isEmpty(); + } + + @Test + void testAcquireQuota_unconfiguredFlow_noOp() { + FlowQuotaParameters paramA = + createParameters(FlowA.class, "quota-a", 10, Duration.ofMinutes(1)); + FlowQuotaManager manager = FlowQuotaManager.create(quotaManager, ImmutableList.of(paramA)); + assertDoesNotThrow(() -> manager.acquireQuota(FlowB.class, eppInput, "TheRegistrar")); + assertThat(jedis.keys("*")).isEmpty(); + } + + @Test + void testAcquireQuota_success() throws Exception { + FlowQuotaParameters paramA = createParameters(FlowA.class, "quota-a", 2, Duration.ofMinutes(1)); + FlowQuotaManager manager = FlowQuotaManager.create(quotaManager, ImmutableList.of(paramA)); + + manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar"); + assertThat(jedis.get("flow:TheRegistrar:quota-a")).isEqualTo("1"); + + manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar"); + assertThat(jedis.get("flow:TheRegistrar:quota-a")).isEqualTo("0"); + } + + @Test + void testAcquireQuota_exceeded_throwsTooManyRequestsException() throws Exception { + FlowQuotaParameters paramA = createParameters(FlowA.class, "quota-a", 2, Duration.ofMinutes(1)); + FlowQuotaManager manager = FlowQuotaManager.create(quotaManager, ImmutableList.of(paramA)); + + manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar"); + manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar"); + + TooManyRequestsException thrown = + assertThrows( + TooManyRequestsException.class, + () -> manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar")); + assertThat(thrown).hasMessageThat().contains("Too many requests"); + assertAboutEppExceptions().that(thrown).marshalsToXml(); + } + + @Test + void testAcquireQuota_releaseQuotaAfterDuration() throws Exception { + FlowQuotaParameters paramA = createParameters(FlowA.class, "quota-a", 2, Duration.ofMillis(10)); + FlowQuotaManager manager = FlowQuotaManager.create(quotaManager, ImmutableList.of(paramA)); + + manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar"); + Thread.sleep(30); + assertDoesNotThrow(() -> manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar")); + } + + @Test + void testAcquireQuota_isolatedByRegistrarId() throws Exception { + FlowQuotaParameters paramA = createParameters(FlowA.class, "quota-a", 1, Duration.ofMinutes(1)); + FlowQuotaManager manager = FlowQuotaManager.create(quotaManager, ImmutableList.of(paramA)); + + manager.acquireQuota(FlowA.class, eppInput, "RegistrarA"); + assertThrows( + TooManyRequestsException.class, + () -> manager.acquireQuota(FlowA.class, eppInput, "RegistrarA")); + + // RegistrarB has independent quota + assertDoesNotThrow(() -> manager.acquireQuota(FlowA.class, eppInput, "RegistrarB")); + assertThat(jedis.get("flow:RegistrarA:quota-a")).isEqualTo("0"); + assertThat(jedis.get("flow:RegistrarB:quota-a")).isEqualTo("0"); + } + + @Test + void testAcquireQuota_multipleFlowParameters() throws Exception { + FlowQuotaParameters paramA = createParameters(FlowA.class, "quota-a", 2, Duration.ofMinutes(1)); + FlowQuotaParameters paramB = createParameters(FlowB.class, "quota-b", 5, Duration.ofMinutes(5)); + FlowQuotaManager manager = + FlowQuotaManager.create(quotaManager, ImmutableList.of(paramA, paramB)); + + manager.acquireQuota(FlowA.class, eppInput, "TheRegistrar"); + manager.acquireQuota(FlowB.class, eppInput, "TheRegistrar"); + + assertThat(jedis.get("flow:TheRegistrar:quota-a")).isEqualTo("1"); + assertThat(jedis.get("flow:TheRegistrar:quota-b")).isEqualTo("4"); + } + + static class FlowA implements Flow { + @Override + public ResponseOrGreeting run() { + return null; + } + } + + static class FlowB implements Flow { + @Override + public ResponseOrGreeting run() { + return null; + } + } + + private static FlowQuotaParameters createParameters( + Class flowClass, String quotaPrefix, int maxQuota, Duration window) { + return new FlowQuotaParameters() { + @Override + public Class getFlowClass() { + return flowClass; + } + + @Override + public String getQuotaId(EppInput input, String registrarId) { + return registrarId + ":" + quotaPrefix; + } + + @Override + public int getMaxQuotaAllowed() { + return maxQuota; + } + + @Override + public Duration getWindowDuration() { + return window; + } + }; + } +} diff --git a/core/src/test/java/google/registry/module/TestRegistryComponent.java b/core/src/test/java/google/registry/module/TestRegistryComponent.java index 2147ce4d3f8..ad7e1a0e2c0 100644 --- a/core/src/test/java/google/registry/module/TestRegistryComponent.java +++ b/core/src/test/java/google/registry/module/TestRegistryComponent.java @@ -28,6 +28,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.GmailModule; import google.registry.groups.GroupsModule; import google.registry.groups.GroupssettingsModule; @@ -60,6 +61,7 @@ CustomLogicFactoryModule.class, DomainDeletionTimeCacheModule.class, DriveModule.class, + FlowQuotaModule.class, GmailModule.class, GroupsModule.class, GroupssettingsModule.class,