pendingDeposits)
throws IOException {
@@ -715,6 +715,12 @@ public static void main(String[] args) throws IOException, ClassNotFoundExceptio
PipelineOptionsFactory.register(RdePipelineOptions.class);
RdePipelineOptions options =
PipelineOptionsFactory.fromArgs(args).withValidation().as(RdePipelineOptions.class);
+ // We need to self allocate the IDs because the pipeline creates EPP resources from history
+ // entries and projects them to watermark. These buildable entities would otherwise request an
+ // ID from datastore, which Beam does not have access to. The IDs are not included in the
+ // deposits or are these entities persisted back to the database, so it is OK to use a self
+ // allocated ID to get around the limitations of beam.
+ options.setUseSelfAllocatedId(true);
RegistryPipelineOptions.validateRegistryPipelineOptions(options);
options.setIsolationOverride(TransactionIsolationLevel.TRANSACTION_READ_COMMITTED);
DaggerRdePipeline_RdePipelineComponent.builder().options(options).build().rdePipeline().run();
diff --git a/core/src/main/java/google/registry/model/IdService.java b/core/src/main/java/google/registry/model/IdService.java
index 8074c0583e6..0a6040c0b80 100644
--- a/core/src/main/java/google/registry/model/IdService.java
+++ b/core/src/main/java/google/registry/model/IdService.java
@@ -17,58 +17,114 @@
import static com.google.common.base.Preconditions.checkState;
import com.google.appengine.api.datastore.DatastoreServiceFactory;
-import com.google.common.annotations.VisibleForTesting;
+import com.google.common.flogger.FluentLogger;
import google.registry.beam.common.RegistryPipelineWorkerInitializer;
import google.registry.config.RegistryEnvironment;
import google.registry.model.annotations.DeleteAfterMigration;
import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Supplier;
/**
- * Allocates a globally unique {@link Long} number to use as an Ofy {@code @Id}.
+ * Allocates a {@link long} to use as a {@code @Id}, (part) of the primary SQL key for an entity.
*
- * In non-test, non-beam environments the Id is generated by Datastore, otherwise it's from an
- * atomic long number that's incremented every time this method is called.
+ *
Normally, the ID is globally unique and allocated by Datastore. It is possible to override
+ * this behavior by providing an ID supplier, such as in unit tests, where a self-allocated ID based
+ * on a monotonically increasing atomic {@link long} is used. Such an ID supplier can also be used
+ * in other scenarios, such as in a Beam pipeline to get around the limitation of Beam's inability
+ * to use GAE SDK to access Datastore. The override should be used with great care lest it results
+ * in irreversible data corruption.
+ *
+ * @see #setIdSupplier(Supplier)
*/
@DeleteAfterMigration
public final class IdService {
- /**
- * A placeholder String passed into DatastoreService.allocateIds that ensures that all ids are
- * initialized from the same id pool.
- */
- private static final String APP_WIDE_ALLOCATION_KIND = "common";
+ private static final FluentLogger logger = FluentLogger.forEnclosingClass();
+
+ private IdService() {}
+
+ private static Supplier idSupplier =
+ RegistryEnvironment.UNITTEST.equals(RegistryEnvironment.get())
+ ? SelfAllocatedIdSupplier.getInstance()
+ : DatastoreIdSupplier.getInstance();
/**
- * Counts of used ids for use in unit tests or Beam.
+ * Provides a {@link Supplier} of ID that overrides the default.
+ *
+ * Currently, the only use case for an override is in the Beam pipeline, where access to
+ * Datastore is not possible through the App Engine API. As such, the setter explicitly checks if
+ * the runtime is Beam.
*
- *
Note that one should only use self-allocate Ids in Beam for entities whose Ids are not
- * important and are not persisted back to the database, i. e. nowhere the uniqueness of the ID is
- * required.
+ *
Because the provided supplier is not guaranteed to be globally unique and compatible with
+ * existing IDs in the database, one should proceed with great care. It is safe to use an
+ * arbitrary supplier when the resulting IDs are not significant and not persisted back to the
+ * database, i.e. the IDs are only required by the {@link Buildable} contract but are not used in
+ * any meaningful way. One example is the RDE pipeline where we project EPP resource entities from
+ * history entries to watermark time, which are then marshalled into XML elements in the RDE
+ * deposits, where the IDs are omitted.
*/
- private static final AtomicLong nextSelfAllocatedId = new AtomicLong(1); // ids cannot be zero
-
- private static final boolean isSelfAllocated() {
- return RegistryEnvironment.UNITTEST.equals(RegistryEnvironment.get())
- || "true".equals(System.getProperty(RegistryPipelineWorkerInitializer.PROPERTY, "false"));
+ public static void setIdSupplier(Supplier idSupplier) {
+ checkState(
+ "true".equals(System.getProperty(RegistryPipelineWorkerInitializer.PROPERTY, "false")),
+ "Can only set ID supplier in a Beam pipeline");
+ logger.atWarning().log("Using ID supplier override!");
+ IdService.idSupplier = idSupplier;
}
/** Allocates an id. */
- // TODO(b/201547855): Find a way to allocate a unique ID without datastore.
public static long allocateId() {
- return isSelfAllocated()
- ? nextSelfAllocatedId.getAndIncrement()
- : DatastoreServiceFactory.getDatastoreService()
- .allocateIds(APP_WIDE_ALLOCATION_KIND, 1)
- .iterator()
- .next()
- .getId();
+ return idSupplier.get();
}
- /** Resets the global self-allocated id counter (i.e. sets the next id to 1). */
- @VisibleForTesting
- public static void resetSelfAllocatedId() {
- checkState(
- isSelfAllocated(), "Can only call resetSelfAllocatedId() in unit tests or Beam pipelines");
- nextSelfAllocatedId.set(1); // ids cannot be zero
+ // TODO(b/201547855): Find a way to allocate a unique ID without datastore.
+ private static class DatastoreIdSupplier implements Supplier {
+
+ private static final DatastoreIdSupplier INSTANCE = new DatastoreIdSupplier();
+
+ /**
+ * A placeholder String passed into {@code DatastoreService.allocateIds} that ensures that all
+ * IDs are initialized from the same ID pool.
+ */
+ private static final String APP_WIDE_ALLOCATION_KIND = "common";
+
+ public static DatastoreIdSupplier getInstance() {
+ return INSTANCE;
+ }
+
+ @Override
+ public Long get() {
+ return DatastoreServiceFactory.getDatastoreService()
+ .allocateIds(APP_WIDE_ALLOCATION_KIND, 1)
+ .iterator()
+ .next()
+ .getId();
+ }
+ }
+
+ /**
+ * An ID supplier that allocates an ID from a monotonically increasing atomic {@link long}.
+ *
+ * The generated IDs are only unique within the same JVM. It is not suitable for production use
+ * unless in cases the IDs are not significant.
+ */
+ public static class SelfAllocatedIdSupplier implements Supplier {
+
+ private static final SelfAllocatedIdSupplier INSTANCE = new SelfAllocatedIdSupplier();
+
+ /** Counts of used ids for self allocating IDs. */
+ private static final AtomicLong nextSelfAllocatedId = new AtomicLong(1); // ids cannot be zero
+
+ public static SelfAllocatedIdSupplier getInstance() {
+ return INSTANCE;
+ }
+
+ @Override
+ public Long get() {
+ return nextSelfAllocatedId.getAndIncrement();
+ }
+
+ public void reset() {
+ nextSelfAllocatedId.set(1);
+ }
}
}
diff --git a/core/src/test/java/google/registry/testing/AppEngineExtension.java b/core/src/test/java/google/registry/testing/AppEngineExtension.java
index 0cf37bef697..090d98eb33d 100644
--- a/core/src/test/java/google/registry/testing/AppEngineExtension.java
+++ b/core/src/test/java/google/registry/testing/AppEngineExtension.java
@@ -40,7 +40,7 @@
import com.google.common.io.Files;
import com.googlecode.objectify.Key;
import com.googlecode.objectify.ObjectifyFilter;
-import google.registry.model.IdService;
+import google.registry.model.IdService.SelfAllocatedIdSupplier;
import google.registry.model.ofy.ObjectifyService;
import google.registry.model.registrar.Registrar;
import google.registry.model.registrar.Registrar.State;
@@ -441,7 +441,7 @@ public void setUp() throws Exception {
ObjectifyService.initOfy();
// Reset id allocation in ObjectifyService so that ids are deterministic in tests.
- IdService.resetSelfAllocatedId();
+ SelfAllocatedIdSupplier.getInstance().reset();
this.ofyTestEntities.forEach(AppEngineExtension::register);
}