Use replica db for non-mutating epp flows (#2478)

* Use replica db for non-mutating epp flows

* Add a test

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.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 @@ static Class<? extends Flow> provideFlowClass(EppInput eppInput) {
  }
  }
 
+ @Provides
+ @FlowScope
+ static JpaTransactionManager provideJpaTm(Class<? extends Flow> flowClass) {
+ if (MutatingFlow.class.isAssignableFrom(flowClass)) {
+ return tm();
+ } else {
+ return replicaTm();
+ }
+ }
+
  @Provides
  @FlowScope
  static ResourceCommand provideResourceCommand(EppInput eppInput) {

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.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 EppOutput run(final EppMetric.Builder eppMetricBuilder) throws EppExcepti
  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) {

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();
+ }
+ }
+}

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

@@ -111,6 +111,7 @@ void beforeEach() {
  new StatelessRequestSessionMetadata("TheRegistrar", ImmutableSet.of());
  flowRunner.trid = Trid.create("client-123", "server-456");
  flowRunner.flowReporter = mock(FlowReporter.class);
+ flowRunner.jpaTransactionManager = tm();
  }
 
  @Test