Write one PCollection to SQL (#664)

* Write one PCollection to SQL

Defined a transform that writes a PCollection of entities to SQL using
JPA. Allows configuring parallelism level and batch size.

core/src/main/java/google/registry/beam/initsql/BeamJpaModule.java

@@ -19,7 +19,6 @@ import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.base.Strings.isNullOrEmpty;
 
 import com.google.common.base.Splitter;
-import dagger.Binds;
 import dagger.Component;
 import dagger.Lazy;
 import dagger.Module;
@@ -32,10 +31,6 @@ import google.registry.persistence.PersistenceModule;
 import google.registry.persistence.PersistenceModule.JdbcJpaTm;
 import google.registry.persistence.PersistenceModule.SocketFactoryJpaTm;
 import google.registry.persistence.transaction.JpaTransactionManager;
-import google.registry.util.Clock;
-import google.registry.util.Sleeper;
-import google.registry.util.SystemClock;
-import google.registry.util.SystemSleeper;
 import google.registry.util.UtilsModule;
 import java.io.BufferedReader;
 import java.io.IOException;
@@ -159,19 +154,10 @@ public class BeamJpaModule {
   @Provides
   @Config("beamHibernateHikariMaximumPoolSize")
   static int getBeamHibernateHikariMaximumPoolSize() {
+    // TODO(weiminyu): make this configurable. Should be equal to number of cores.
     return 4;
   }
 
-  @Module
-  interface BindModule {
-
-    @Binds
-    Sleeper sleeper(SystemSleeper sleeper);
-
-    @Binds
-    Clock clock(SystemClock clock);
-  }
-
   @Singleton
   @Component(
       modules = {

core/src/main/java/google/registry/beam/initsql/Transforms.java

@@ -19,24 +19,37 @@ import static com.google.common.base.Preconditions.checkNotNull;
 import static com.google.common.base.Preconditions.checkState;
 import static google.registry.beam.initsql.BackupPaths.getCommitLogTimestamp;
 import static google.registry.beam.initsql.BackupPaths.getExportFilePatterns;
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+import static google.registry.persistence.transaction.TransactionManagerFactory.setJpaTm;
 import static google.registry.util.DateTimeUtils.START_OF_TIME;
 import static google.registry.util.DateTimeUtils.isBeforeOrAt;
 import static java.util.Comparator.comparing;
+import static org.apache.beam.sdk.values.TypeDescriptors.integers;
 import static org.apache.beam.sdk.values.TypeDescriptors.kvs;
 import static org.apache.beam.sdk.values.TypeDescriptors.strings;
 
 import avro.shaded.com.google.common.collect.Iterators;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Streams;
+import google.registry.backup.AppEngineEnvironment;
 import google.registry.backup.CommitLogImports;
 import google.registry.backup.VersionedEntity;
+import google.registry.model.ofy.ObjectifyService;
+import google.registry.model.ofy.Ofy;
+import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.tools.LevelDbLogReader;
+import google.registry.util.SystemSleeper;
+import java.io.Serializable;
 import java.util.Collection;
 import java.util.Iterator;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.Supplier;
+import javax.persistence.OptimisticLockException;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.io.Compression;
 import org.apache.beam.sdk.io.FileIO;
@@ -47,6 +60,7 @@ import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.Flatten;
 import org.apache.beam.sdk.transforms.GroupByKey;
+import org.apache.beam.sdk.transforms.GroupIntoBatches;
 import org.apache.beam.sdk.transforms.MapElements;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
@@ -56,10 +70,12 @@ import org.apache.beam.sdk.values.PBegin;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PCollectionList;
 import org.apache.beam.sdk.values.PCollectionTuple;
+import org.apache.beam.sdk.values.PDone;
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.beam.sdk.values.TupleTagList;
 import org.apache.beam.sdk.values.TypeDescriptor;
 import org.joda.time.DateTime;
+import org.joda.time.Duration;
 
 /**
  * {@link PTransform Pipeline transforms} used in pipelines that load from both Datastore export
@@ -245,6 +261,38 @@ public final class Transforms {
                     .iterator()));
   }
 
+  /**
+   * Returns a {@link PTransform} that writes a {@link PCollection} of entities to a SQL database.
+   *
+   * @param transformId a unique ID for an instance of the returned transform
+   * @param maxWriters the max number of concurrent writes to SQL, which also determines the max
+   *     number of connection pools created
+   * @param batchSize the number of entities to write in each operation
+   * @param jpaSupplier supplier of a {@link JpaTransactionManager}
+   */
+  public static PTransform<PCollection<VersionedEntity>, PDone> writeToSql(
+      String transformId,
+      int maxWriters,
+      int batchSize,
+      SerializableSupplier<JpaTransactionManager> jpaSupplier) {
+    return new PTransform<PCollection<VersionedEntity>, PDone>() {
+      @Override
+      public PDone expand(PCollection<VersionedEntity> input) {
+        input
+            .apply(
+                "Shard data for " + transformId,
+                MapElements.into(kvs(integers(), TypeDescriptor.of(VersionedEntity.class)))
+                    .via(ve -> KV.of(ThreadLocalRandom.current().nextInt(maxWriters), ve)))
+            .apply("Batch output by shard " + transformId, GroupIntoBatches.ofSize(batchSize))
+            .apply("Write in batch for " + transformId, ParDo.of(new SqlBatchWriter(jpaSupplier)));
+        return PDone.in(input.getPipeline());
+      }
+    };
+  }
+
+  /** Interface for serializable {@link Supplier suppliers}. */
+  public interface SerializableSupplier<T> extends Supplier<T>, Serializable {}
+
   /**
    * Returns a {@link PTransform} that produces a {@link PCollection} containing all elements in the
    * given {@link Iterable}.
@@ -322,4 +370,104 @@ public final class Transforms {
       }
     }
   }
+
+  /**
+   * Writes a batch of entities to a SQL database.
+   *
+   * <p>Note that an arbitrary number of instances of this class may be created and freed in
+   * arbitrary order in a single JVM. Due to the tech debt that forced us to use a static variable
+   * to hold the {@code JpaTransactionManager} instance, we must ensure that JpaTransactionManager
+   * is not changed or torn down while being used by some instance.
+   */
+  private static class SqlBatchWriter extends DoFn<KV<Integer, Iterable<VersionedEntity>>, Void> {
+
+    private static int instanceCount = 0;
+    private static JpaTransactionManager originalJpa;
+
+    private final SerializableSupplier<JpaTransactionManager> jpaSupplier;
+
+    private transient Ofy ofy;
+    private transient SystemSleeper sleeper;
+
+    SqlBatchWriter(SerializableSupplier<JpaTransactionManager> jpaSupplier) {
+      this.jpaSupplier = jpaSupplier;
+    }
+
+    @Setup
+    public void setup() {
+      sleeper = new SystemSleeper();
+
+      ObjectifyService.initOfy();
+      ofy = ObjectifyService.ofy();
+
+      synchronized (SqlBatchWriter.class) {
+        if (instanceCount == 0) {
+          originalJpa = jpaTm();
+          setJpaTm(jpaSupplier);
+        }
+        instanceCount++;
+      }
+    }
+
+    @Teardown
+    public void teardown() {
+      synchronized (SqlBatchWriter.class) {
+        instanceCount--;
+        if (instanceCount == 0) {
+          jpaTm().teardown();
+          setJpaTm(() -> originalJpa);
+        }
+      }
+    }
+
+    @ProcessElement
+    public void processElement(@Element KV<Integer, Iterable<VersionedEntity>> kv) {
+      try (AppEngineEnvironment env = new AppEngineEnvironment()) {
+        ImmutableList<Object> ofyEntities =
+            Streams.stream(kv.getValue())
+                .map(VersionedEntity::getEntity)
+                .map(Optional::get)
+                .map(ofy::toPojo)
+                .collect(ImmutableList.toImmutableList());
+        retry(() -> jpaTm().transact(() -> jpaTm().saveNewOrUpdateAll(ofyEntities)));
+      }
+    }
+
+    // TODO(b/160632289): Enhance Retrier and use it here.
+    private void retry(Runnable runnable) {
+      int maxAttempts = 5;
+      int initialDelayMillis = 100;
+      double jitterRatio = 0.2;
+
+      for (int attempt = 0; attempt < maxAttempts; attempt++) {
+        try {
+          runnable.run();
+          return;
+        } catch (Throwable throwable) {
+          throwIfNotCausedBy(throwable, OptimisticLockException.class);
+          int sleepMillis = (1 << attempt) * initialDelayMillis;
+          int jitter =
+              ThreadLocalRandom.current().nextInt((int) (sleepMillis * jitterRatio))
+                  - (int) (sleepMillis * jitterRatio / 2);
+          sleeper.sleepUninterruptibly(Duration.millis(sleepMillis + jitter));
+        }
+      }
+    }
+
+    /**
+     * Rethrows {@code throwable} if it is not (and does not have a cause of) {@code causeType};
+     * otherwise returns with no side effects.
+     */
+    private void throwIfNotCausedBy(Throwable throwable, Class<? extends Throwable> causeType) {
+      Throwable t = throwable;
+      while (t != null) {
+        if (causeType.isInstance(t)) {
+          return;
+        }
+        t = t.getCause();
+      }
+      Throwables.throwIfUnchecked(t);
+      throw new RuntimeException(t);
+    }
+  }
 }

core/src/main/java/google/registry/persistence/transaction/JpaTransactionManager.java

@@ -25,4 +25,11 @@ public interface JpaTransactionManager extends TransactionManager {
 
   /** Deletes the entity by its id, throws exception if the entity is not deleted. */
   public abstract <T> void assertDelete(VKey<T> key);
+
+  /**
+   * Releases all resources and shuts down.
+   *
+   * <p>The errorprone check forbids injection of {@link java.io.Closeable} resources.
+   */
+  void teardown();
 }

core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java

@@ -62,6 +62,11 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     this.clock = clock;
   }
 
+  @Override
+  public void teardown() {
+    emf.close();
+  }
+
   @Override
   public EntityManager getEntityManager() {
     if (transactionInfo.get().entityManager == null) {

core/src/main/java/google/registry/persistence/transaction/TransactionManagerFactory.java

@@ -75,7 +75,12 @@ public class TransactionManagerFactory {
     return tm;
   }
 
-  /** Returns {@link JpaTransactionManager} instance. */
+  /**
+   * Returns {@link JpaTransactionManager} instance.
+   *
+   * <p>Between invocations of {@link TransactionManagerFactory#setJpaTm} every call to this method
+   * returns the same instance.
+   */
   public static JpaTransactionManager jpaTm() {
     return jpaTm.get();
   }
@@ -93,7 +98,7 @@ public class TransactionManagerFactory {
         RegistryEnvironment.get().equals(RegistryEnvironment.UNITTEST)
             || RegistryToolEnvironment.get() != null,
         "setJpamTm() should only be called by tools and tests.");
-    jpaTm = jpaTmSupplier;
+    jpaTm = Suppliers.memoize(jpaTmSupplier::get);
   }
 
   /** Sets the return of {@link #tm()} to the given instance of {@link TransactionManager}. */