Check for entity nonexistence in SqlBatchWriter (#1824)

Passing in an already-existing instance is an antipattern because it can
lead to race conditions where something else modified the object in
between when the pipeline loaded it and when you're saving it. The Write
action should only be writing new entities.

We cannot check IDs for the objects (some IDs are not autogenerated so
they might exist already). We also cannot call `insert` on the objects
because the underlying JPA `persist` call adds the input object to the
persistence context, meaning that any modifications (e.g.
updateTimestamp) are reflected in the input object. Beam doesn't allow
modification of input objects.
This commit is contained in:
gbrodman
2022-10-27 14:46:26 -04:00
committed by GitHub
parent ca60ca159f
commit 9c6c210e21
2 changed files with 55 additions and 3 deletions
@@ -14,6 +14,7 @@
package google.registry.beam.common;
import static com.google.common.base.Preconditions.checkArgument;
import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
import static org.apache.beam.sdk.values.TypeDescriptors.integers;
@@ -405,7 +406,13 @@ public final class RegistryJpaIO {
.filter(Objects::nonNull)
.collect(ImmutableList.toImmutableList());
try {
jpaTm().transact(() -> jpaTm().putAll(entities));
jpaTm()
.transact(
() -> {
// Don't modify existing objects as it could lead to race conditions
entities.forEach(this::verifyObjectNonexistence);
jpaTm().putAll(entities);
});
counter.inc(entities.size());
} catch (RuntimeException e) {
processSingly(entities);
@@ -419,7 +426,13 @@ public final class RegistryJpaIO {
private void processSingly(ImmutableList<Object> entities) {
for (Object entity : entities) {
try {
jpaTm().transact(() -> jpaTm().put(entity));
jpaTm()
.transact(
() -> {
// Don't modify existing objects as it could lead to race conditions
verifyObjectNonexistence(entity);
jpaTm().put(entity);
});
counter.inc();
} catch (RuntimeException e) {
throw new RuntimeException(toEntityKeyString(entity), e);
@@ -445,5 +458,16 @@ public final class RegistryJpaIO {
return "Non-SqlEntity: " + entity;
}
}
/** SqlBatchWriter should not re-write existing entities due to potential race conditions. */
private void verifyObjectNonexistence(Object obj) {
// We cannot rely on calling "insert" on the objects because the underlying JPA persist call
// adds the input object to the persistence context, meaning that any modifications (e.g.
// updateTimestamp) are reflected in the input object. Beam doesn't allow modification of
// input objects, so this throws an exception.
// TODO(go/non-datastore-allocateid): also check that all the objects have IDs
checkArgument(
!jpaTm().exists(obj), "Entities created in SqlBatchWriter must not already exist");
}
}
}