From d000a5dff891601b805c1909949d98624ca95555 Mon Sep 17 00:00:00 2001 From: Weimin Yu Date: Thu, 13 Jun 2024 19:18:56 -0400 Subject: [PATCH] Use replica db for non-mutating epp flows (#2478) * Use replica db for non-mutating epp flows * Add a test --- .../google/registry/flows/FlowModule.java | 12 +++ .../google/registry/flows/FlowRunner.java | 42 ++++----- .../google/registry/flows/FlowModuleTest.java | 89 +++++++++++++++++++ .../google/registry/flows/FlowRunnerTest.java | 1 + 4 files changed, 123 insertions(+), 21 deletions(-) create mode 100644 core/src/test/java/google/registry/flows/FlowModuleTest.java diff --git a/core/src/main/java/google/registry/flows/FlowModule.java b/core/src/main/java/google/registry/flows/FlowModule.java index 32c0ef0c2..6e14afe9e 100644 --- a/core/src/main/java/google/registry/flows/FlowModule.java +++ b/core/src/main/java/google/registry/flows/FlowModule.java @@ -15,6 +15,7 @@ package google.registry.flows; import static com.google.common.base.Preconditions.checkState; +import static google.registry.persistence.transaction.TransactionManagerFactory.replicaTm; import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import com.google.common.base.Strings; @@ -37,6 +38,7 @@ import google.registry.model.host.HostHistory; import google.registry.model.reporting.HistoryEntry; import google.registry.persistence.IsolationLevel; import google.registry.persistence.PersistenceModule.TransactionIsolationLevel; +import google.registry.persistence.transaction.JpaTransactionManager; import java.lang.annotation.Documented; import java.util.Optional; import javax.inject.Qualifier; @@ -191,6 +193,16 @@ public class FlowModule { } } + @Provides + @FlowScope + static JpaTransactionManager provideJpaTm(Class flowClass) { + if (MutatingFlow.class.isAssignableFrom(flowClass)) { + return tm(); + } else { + return replicaTm(); + } + } + @Provides @FlowScope static ResourceCommand provideResourceCommand(EppInput eppInput) { diff --git a/core/src/main/java/google/registry/flows/FlowRunner.java b/core/src/main/java/google/registry/flows/FlowRunner.java index db3454e3d..99b903faa 100644 --- a/core/src/main/java/google/registry/flows/FlowRunner.java +++ b/core/src/main/java/google/registry/flows/FlowRunner.java @@ -14,7 +14,6 @@ package google.registry.flows; -import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import static google.registry.xml.XmlTransformer.prettyPrint; import com.google.common.flogger.FluentLogger; @@ -28,6 +27,7 @@ import google.registry.model.eppcommon.Trid; import google.registry.model.eppoutput.EppOutput; import google.registry.monitoring.whitebox.EppMetric; import google.registry.persistence.PersistenceModule.TransactionIsolationLevel; +import google.registry.persistence.transaction.JpaTransactionManager; import java.util.Optional; import javax.inject.Inject; import javax.inject.Provider; @@ -52,6 +52,8 @@ public class FlowRunner { @Inject SessionMetadata sessionMetadata; @Inject Trid trid; @Inject FlowReporter flowReporter; + @Inject JpaTransactionManager jpaTransactionManager; + @Inject FlowRunner() {} /** Runs the EPP flow, and records metrics on the given builder. */ @@ -77,26 +79,24 @@ public class FlowRunner { return EppOutput.create(flowProvider.get().run()); } try { - // TODO(mcilwain/weiminyu): Use transactReadOnly() here for TransactionalFlow and transact() - // for MutatingFlow. - return tm().transact( - isolationLevelOverride.orElse(null), - () -> { - try { - EppOutput output = EppOutput.create(flowProvider.get().run()); - if (isDryRun) { - throw new DryRunException(output); - } - if (flowClass.equals(LoginFlow.class)) { - // In LoginFlow, registrarId isn't known until after the flow executes, so save - // it then. - eppMetricBuilder.setRegistrarId(sessionMetadata.getRegistrarId()); - } - return output; - } catch (EppException e) { - throw new EppRuntimeException(e); - } - }); + return jpaTransactionManager.transact( + isolationLevelOverride.orElse(null), + () -> { + try { + EppOutput output = EppOutput.create(flowProvider.get().run()); + if (isDryRun) { + throw new DryRunException(output); + } + if (flowClass.equals(LoginFlow.class)) { + // In LoginFlow, registrarId isn't known until after the flow executes, so save + // it then. + eppMetricBuilder.setRegistrarId(sessionMetadata.getRegistrarId()); + } + return output; + } catch (EppException e) { + throw new EppRuntimeException(e); + } + }); } catch (DryRunException e) { return e.output; } catch (EppRuntimeException e) { diff --git a/core/src/test/java/google/registry/flows/FlowModuleTest.java b/core/src/test/java/google/registry/flows/FlowModuleTest.java new file mode 100644 index 000000000..331dd4aec --- /dev/null +++ b/core/src/test/java/google/registry/flows/FlowModuleTest.java @@ -0,0 +1,89 @@ +// Copyright 2024 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; + +import static com.google.common.truth.Truth.assertThat; +import static google.registry.persistence.transaction.TransactionManagerFactory.replicaTm; +import static google.registry.persistence.transaction.TransactionManagerFactory.tm; +import static google.registry.testing.DatabaseHelper.newTld; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import dagger.Component; +import google.registry.flows.FlowComponent.FlowComponentModule; +import google.registry.model.eppinput.EppInput; +import google.registry.persistence.transaction.DatabaseException; +import google.registry.persistence.transaction.JpaTestExtensions; +import google.registry.persistence.transaction.JpaTestExtensions.JpaIntegrationTestExtension; +import google.registry.persistence.transaction.JpaTransactionManager; +import google.registry.testing.EppLoader; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +/** Unit tests for {@link FlowModule}. */ +public class FlowModuleTest { + + @RegisterExtension + final JpaIntegrationTestExtension jpa = + new JpaTestExtensions.Builder().buildIntegrationTestExtension(); + + private EppInput getEppInput(String eppInputXmlFilename) throws EppException { + return new EppLoader(this, eppInputXmlFilename).getEpp(); + } + + @Test + void givenMutatingFlow_thenPrimaryTmIsUsed() throws EppException { + String eppInputXmlFilename = "domain_create.xml"; + FlowModule flowModule = + new FlowModule.Builder().setEppInput(getEppInput(eppInputXmlFilename)).build(); + JpaTransactionManager tm = + DaggerFlowModuleTest_FlowModuleTestComponent.builder() + .flowModule(flowModule) + .build() + .jpaTm(); + assertThat(tm).isEqualTo(tm()); + tm.transact(() -> tm.put(newTld("app", "ROID"))); + } + + @Test + void givenNonMutatingFlow_thenReplicaTmIsUsed() throws EppException { + String eppInputXmlFilename = "domain_info.xml"; + FlowModule flowModule = + new FlowModule.Builder().setEppInput(getEppInput(eppInputXmlFilename)).build(); + JpaTransactionManager tm = + DaggerFlowModuleTest_FlowModuleTestComponent.builder() + .flowModule(flowModule) + .build() + .jpaTm(); + assertThat(tm).isEqualTo(replicaTm()); + assertThat( + assertThrows( + DatabaseException.class, () -> tm.transact(() -> tm.put(newTld("app", "ROID"))))) + .hasMessageThat() + .contains("cannot execute INSERT in a read-only transaction"); + } + + @FlowScope + @Component(modules = {FlowModule.class, FlowComponentModule.class}) + public interface FlowModuleTestComponent { + JpaTransactionManager jpaTm(); + + @Component.Builder + interface Builder { + Builder flowModule(FlowModule flowModule); + + FlowModuleTestComponent build(); + } + } +} diff --git a/core/src/test/java/google/registry/flows/FlowRunnerTest.java b/core/src/test/java/google/registry/flows/FlowRunnerTest.java index 9b74c72a5..7461238dd 100644 --- a/core/src/test/java/google/registry/flows/FlowRunnerTest.java +++ b/core/src/test/java/google/registry/flows/FlowRunnerTest.java @@ -111,6 +111,7 @@ class FlowRunnerTest { new StatelessRequestSessionMetadata("TheRegistrar", ImmutableSet.of()); flowRunner.trid = Trid.create("client-123", "server-456"); flowRunner.flowReporter = mock(FlowReporter.class); + flowRunner.jpaTransactionManager = tm(); } @Test