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.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<? 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.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) {

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 @@ class FlowRunnerTest {
         new StatelessRequestSessionMetadata("TheRegistrar", ImmutableSet.of());
     flowRunner.trid = Trid.create("client-123", "server-456");
     flowRunner.flowReporter = mock(FlowReporter.class);
+    flowRunner.jpaTransactionManager = tm();
   }
 
   @Test