Implement more efficient bulk put/update in JPATM (#3237)

This is more complicated than the bulk load unfortunately.

Previously, for update/put or their plural equivalents, we'd have to run
one or two SELECT statements per entity to check if it exists already
(and for Hibernate to load the object into the context to merge it with
the input).

Both updateAll and putAll now run, per input class, one large SELECT statement
to find out which entities already exist, which also loads the entities
into the persistence context (which is basically a cache). For updateAll,
we require that all of the entities already exist. For putAll, we have a
split path
- for entities that already exist, we just call a normal merge()
- for entities that don't already exist, we run merge() while instructing
Hibernate that the object is transient (i.e. new). This means that we
get the "put" behavior (not modifying the original object) while avoid a
superfluous SELECT statement that'd check again to see if the object exists in
the DB already.

For new objects in putAll, we also use EntityManager::merge (rather than
insert/persist) so that Hibernate creates a deep copy of the entity
instead of mutating the caller's instance in place. To prevent
Hibernate's DefaultMergeEventListener from firing a redundant SELECT
statement when merging a known-new entity with a non-null ID,
TransactionInfo implements Interceptor and overrides
isTransient(Object) to return Boolean.TRUE while merging a
known-transient entity.

We add comments clarifying that "insert" will modify an entity in place
(i.e. UpdateAutoTimestamp or generated IDs) whereas "put" will not. This
has always been the case; the comments are just to clarify.

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

@@ -30,6 +30,7 @@ import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Multimaps;
+import com.google.common.collect.Sets;
 import com.google.common.collect.Streams;
 import com.google.common.flogger.FluentLogger;
 import com.google.common.flogger.StackSize;
@@ -50,6 +51,7 @@ import jakarta.persistence.FlushModeType;
 import jakarta.persistence.LockModeType;
 import jakarta.persistence.Parameter;
 import jakarta.persistence.PersistenceException;
+import jakarta.persistence.PersistenceUnitUtil;
 import jakarta.persistence.Query;
 import jakarta.persistence.TemporalType;
 import jakarta.persistence.TypedQuery;
@@ -79,7 +81,9 @@ import java.util.function.UnaryOperator;
 import java.util.regex.Pattern;
 import java.util.stream.Stream;
 import javax.annotation.Nullable;
+import org.hibernate.Interceptor;
 import org.hibernate.Session;
+import org.hibernate.SessionBuilder;
 import org.hibernate.SessionFactory;
 import org.hibernate.cfg.Environment;
 
@@ -253,20 +257,17 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     }
     TransactionInfo txnInfo = transactionInfo.get();
 
-    txnInfo.entityManager =
-        logSqlStatements
-            ? emf.unwrap(SessionFactory.class)
-                .withOptions()
-                .statementInspector(
-                    new UnaryOperator<String>() {
-                      @Override
-                      public String apply(String s) {
-                        logger.atInfo().log(SQL_STATEMENT_LOG_SENTINEL_FORMAT, s);
-                        return s;
-                      }
-                    })
-                .openSession()
-            : emf.createEntityManager();
+    SessionBuilder sessionBuilder =
+        emf.unwrap(SessionFactory.class).withOptions().interceptor(txnInfo);
+    if (logSqlStatements) {
+      sessionBuilder.statementInspector(
+          (UnaryOperator<String>)
+              s -> {
+                logger.atInfo().log(SQL_STATEMENT_LOG_SENTINEL_FORMAT, s);
+                return s;
+              });
+    }
+    txnInfo.entityManager = sessionBuilder.openSession();
     if (readOnly) {
       // Disable Hibernate's dirty object check on flushing, it has become more aggressive in v6.
       txnInfo.entityManager.unwrap(Session.class).setDefaultReadOnly(true);
@@ -359,16 +360,20 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   }
 
   /**
-   * Inserts an object into the database.
+   * Inserts a new object into the database, throwing an exception if it already exists.
+   *
+   * <p>This method delegates to {@link EntityManager#persist} and therefore <b>modifies {@code
+   * entity} in place</b> (e.g., assigning {@link jakarta.persistence.GeneratedValue} IDs, setting
+   * auto-timestamps, and wrapping collections).
    *
    * <p>If {@code entity} has an auto-generated identity field (i.e., a field annotated with {@link
    * jakarta.persistence.GeneratedValue}), the caller must not assign a value to this field,
    * otherwise Hibernate would mistake the entity as detached and raise an error.
    *
    * <p>The practical implication of the above is that when inserting such an entity using a
-   * retriable transaction , the entity should be instantiated inside the transaction body. A failed
-   * attempt may still assign and ID to the entity, therefore reusing the same entity would cause
-   * retries to fail.
+   * retriable transaction, the entity should be instantiated inside the transaction body. A failed
+   * attempt may still assign an ID to the entity, therefore reusing the same entity would cause
+   * retries to fail; otherwise, prefer {@link #put}.
    */
   @Override
   public void insert(Object entity) {
@@ -386,9 +391,15 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   @Override
   public void insertAll(ImmutableObject... entities) {
-    insertAll(ImmutableSet.copyOf(entities));
+    insertAll(ImmutableList.copyOf(entities));
   }
 
+  /**
+   * Persists a new object or updates an existing object in the database.
+   *
+   * <p>Unlike {@link #insert}, this method delegates to {@link EntityManager#merge} to make a deep
+   * copy and <b>never modifies {@code entity} in place</b>.
+   */
   @Override
   public void put(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
@@ -399,32 +410,73 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   @Override
   public void putAll(ImmutableObject... entities) {
     checkArgumentNotNull(entities, "entities must be specified");
-    assertInTransaction();
-    for (Object entity : entities) {
-      put(entity);
-    }
+    putAll(ImmutableList.copyOf(entities));
   }
 
   @Override
   public void putAll(ImmutableCollection<?> entities) {
     checkArgumentNotNull(entities, "entities must be specified");
     assertInTransaction();
-    entities.forEach(this::put);
+    // Group entities by class so we can batch-load existing entities per concrete type via
+    // Session::findMultiple.
+    ImmutableListMultimap<Class<?>, ?> entitiesByClass =
+        Multimaps.index(entities, Object::getClass);
+    // Pre-warm the Hibernate first-level cache (persistence context) in one batched SELECT per
+    // entity class and record which entities already exist in the database. Without this, calling
+    // EntityManager::merge on detached entities would execute an individual SELECT per entity.
+    Set<Object> existingEntities = Collections.newSetFromMap(new IdentityHashMap<>());
+    for (Class<?> clazz : entitiesByClass.keySet()) {
+      existingEntities.addAll(
+          findExistingEntitiesAndLoadContext(clazz, entitiesByClass.get(clazz)));
+    }
+    for (Object entity : entities) {
+      if (existingEntities.contains(entity)) {
+        // The entity already exists in the persistence context (loaded from the DB above);
+        // updateObject merges into the pre-warmed persistence context without issuing a SELECT.
+        transactionInfo.get().updateObject(entity);
+      } else {
+        // The entity does not exist in the DB yet (either id == null or findMultiple returned
+        // null). By calling "merge" while explicitly telling Hibernate that the object is
+        // transient, we get a deep copy (so, not modifying the immutable original) while still
+        // avoiding issuing an additional SELECT.
+        transactionInfo.get().mergeTransientObject(entity);
+      }
+    }
   }
 
   @Override
   public void update(Object entity) {
     checkArgumentNotNull(entity, "entity must be specified");
-    assertInTransaction();
-    checkArgument(exists(entity), "Given entity does not exist");
-    transactionInfo.get().updateObject(entity);
+    updateAll(ImmutableList.of(entity));
   }
 
   @Override
   public void updateAll(ImmutableCollection<?> entities) {
     checkArgumentNotNull(entities, "entities must be specified");
     assertInTransaction();
-    entities.forEach(this::update);
+    ImmutableListMultimap<Class<?>, ?> entitiesByClass =
+        Multimaps.index(entities, Object::getClass);
+    // Instead of merging the entities one by one, we load the entities in one go then merge them.
+    // If we go one-by-one, we execute one or two SELECT statements per entity (one to check
+    // existence, one to load into the persistence context). Instead, we run one large SELECT
+    // statement to check all the entities' existences. This has the extra benefit of loading all
+    // the entities into the persistence context (cache) at once instead of having to load the
+    // entities into the cache one by one.
+    for (Class<?> clazz : entitiesByClass.keySet()) {
+      ImmutableList<?> entitiesForThisClass = entitiesByClass.get(clazz);
+      ImmutableList<Object> existingEntities =
+          findExistingEntitiesAndLoadContext(clazz, entitiesForThisClass);
+      // All entities should be present and accounted for after loading the persistence context
+      if (existingEntities.size() != entitiesForThisClass.size()) {
+        throw new IllegalArgumentException(
+            String.format(
+                "Entity/entities passed to updateAll do not already exist: %s",
+                Sets.difference(
+                    ImmutableSet.copyOf(entitiesForThisClass),
+                    ImmutableSet.copyOf(existingEntities))));
+      }
+    }
+    entities.forEach(transactionInfo.get()::updateObject);
   }
 
   @Override
@@ -806,11 +858,54 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     return entity;
   }
 
-  private static class TransactionInfo {
+  /**
+   * Finds existing entities and pre-warms the persistence context (cache).
+   *
+   * <p>This first filters out entities that don't have IDs (i.e. new objects with generated IDs).
+   * Then we call Session::findMultiple to find existing entities while also populating the
+   * persistence context (cache). This means that later calls to merge() won't need to SELECT, as
+   * the object is already loaded.
+   */
+  private ImmutableList<Object> findExistingEntitiesAndLoadContext(
+      Class<?> clazz, ImmutableList<?> entitiesWithThisClass) {
+    Session session = getEntityManager().unwrap(Session.class);
+    PersistenceUnitUtil persistenceUnitUtil = emf.getPersistenceUnitUtil();
+    ImmutableList.Builder<Object> entitiesWithIdsBuilder = new ImmutableList.Builder<>();
+    ImmutableList.Builder<Object> idsBuilder = new ImmutableList.Builder<>();
+    for (Object entity : entitiesWithThisClass) {
+      Object id = persistenceUnitUtil.getIdentifier(entity);
+      // Skip entities with null IDs (e.g., unpersisted @GeneratedValue entities) since they
+      // cannot exist in the database yet.
+      if (id != null) {
+        entitiesWithIdsBuilder.add(entity);
+        idsBuilder.add(id);
+      }
+    }
+    ImmutableList<Object> entitiesWithIds = entitiesWithIdsBuilder.build();
+    ImmutableList<Object> ids = idsBuilder.build();
+    checkArgument(
+        ids.size() == ImmutableSet.copyOf(ids).size(),
+        "Multiple entities of type %s with the same ID",
+        clazz.getSimpleName());
+    ImmutableList.Builder<Object> existingEntitiesBuilder = new ImmutableList.Builder<>();
+    if (!ids.isEmpty()) {
+      // Session::findMultiple's return value includes nulls for missing entities
+      List<?> loaded = session.findMultiple(clazz, ids);
+      for (int i = 0; i < ids.size(); i++) {
+        if (loaded.get(i) != null) {
+          existingEntitiesBuilder.add(entitiesWithIds.get(i));
+        }
+      }
+    }
+    return existingEntitiesBuilder.build();
+  }
+
+  private static class TransactionInfo implements Interceptor {
     EntityManager entityManager;
     boolean inTransaction = false;
     Instant transactionTime;
     Supplier<Long> idProvider;
+    @Nullable Object transientEntityToMerge;
 
     // The set of entity objects that have been either persisted (via insert()) or merged (via
     // put()/update()). If the entity manager returns these as a result of a find() or query
@@ -818,6 +913,12 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     // them to not be saved to the database -- so we throw an exception instead.
     Set<Object> objectsToSave = Collections.newSetFromMap(new IdentityHashMap<>());
 
+    @Override
+    public Boolean isTransient(Object entity) {
+      // "null" indicates "we don't know if this is transient"
+      return (entity != null && entity == transientEntityToMerge) ? Boolean.TRUE : null;
+    }
+
     /** Start a new transaction. */
     private void start(Clock clock, Supplier<Long> idProvider) {
       checkArgumentNotNull(clock);
@@ -830,6 +931,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
       idProvider = null;
       inTransaction = false;
       transactionTime = null;
+      transientEntityToMerge = null;
       objectsToSave = Collections.newSetFromMap(new IdentityHashMap<>());
       if (entityManager != null) {
         // Close this EntityManager just let the connection pool be able to reuse it, it doesn't
@@ -839,13 +941,39 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
       }
     }
 
-    /** Does the full "update" on an object including all internal housekeeping. */
+    /**
+     * Merges {@code object} into the session (making a deep copy without modifying {@code object}
+     * in place) and records the managed copy for internal housekeeping.
+     */
     private void updateObject(Object object) {
       Object merged = entityManager.merge(object);
       objectsToSave.add(merged);
     }
 
-    /** Does the full "insert" on a new object including all internal housekeeping. */
+    /**
+     * Merges a known-transient (new) object into the session without firing a database SELECT.
+     *
+     * <p>Hibernate performs deep copies on objects when calling "merge" which we want in order to
+     * avoid modifying the original object. However, "merge" issues an additional SELECT statement
+     * unless we explicitly tell Hibernate the object is transient. By setting {@link
+     * #transientEntityToMerge}, {@link #isTransient(Object)} returns {@code Boolean.TRUE} when
+     * Hibernate's {@code DefaultMergeEventListener} inspects the entity, causing Hibernate to take
+     * the {@code entityIsTransient} path (deep-copying the entity and scheduling an INSERT) without
+     * executing a {@code SELECT} query even when the entity has a pre-assigned non-null ID.
+     */
+    private void mergeTransientObject(Object object) {
+      transientEntityToMerge = object;
+      try {
+        updateObject(object);
+      } finally {
+        transientEntityToMerge = null;
+      }
+    }
+
+    /**
+     * Persists a new {@code object} directly in the session (modifying {@code object} in place) and
+     * records it for internal housekeeping.
+     */
     private void insertObject(Object object) {
       entityManager.persist(object);
       objectsToSave.add(object);

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

@@ -131,31 +131,77 @@ public interface TransactionManager {
   /** Returns the Instant associated with the start of this particular transaction attempt. */
   Instant getTxTime();
 
-  /** Persists a new entity in the database, throws exception if the entity already exists. */
+  /**
+   * Persists a new entity in the database, throwing an exception if the entity already exists.
+   *
+   * <p>Note: This method uses {@link jakarta.persistence.EntityManager#persist} and <b>modifies the
+   * input {@code entity} in place</b> (e.g., assigning {@link jakarta.persistence.GeneratedValue}
+   * IDs, setting auto-timestamps, and wrapping collections). When used inside a retriable
+   * transaction, {@code entity} should be instantiated inside the transaction body so retries do
+   * not reuse a mutated instance; otherwise, prefer {@link #put}.
+   */
   void insert(Object entity);
 
-  /** Persists all new entities in the database, throws exception if any entity already exists. */
+  /**
+   * Persists all new entities in the database, throwing an exception if any entity already exists.
+   *
+   * <p>Like {@link #insert}, this method <b>modifies the input entities in place</b>.
+   */
   void insertAll(ImmutableCollection<?> entities);
 
-  /** Persists all new entities in the database, throws exception if any entity already exists. */
+  /**
+   * Persists all new entities in the database, throwing an exception if any entity already exists.
+   *
+   * <p>Like {@link #insert}, this method <b>modifies the input entities in place</b>.
+   */
   void insertAll(ImmutableObject... entities);
 
-  /** Persists a new entity or update the existing entity in the database. */
+  /**
+   * Persists a new entity or updates an existing entity in the database.
+   *
+   * <p>Unlike {@link #insert}, this method uses {@link jakarta.persistence.EntityManager#merge} to
+   * make a deep copy and <b>never modifies the input {@code entity} in place</b>.
+   */
   void put(Object entity);
 
-  /** Persists all new entities or updates the existing entities in the database. */
+  /**
+   * Persists all new entities or updates existing entities in the database.
+   *
+   * <p>Unlike {@link #insertAll}, this method uses {@link jakarta.persistence.EntityManager#merge}
+   * to make a deep copy and <b>never modifies the input entities in place</b>.
+   */
   void putAll(ImmutableObject... entities);
 
-  /** Persists all new entities or updates the existing entities in the database. */
+  /**
+   * Persists all new entities or updates existing entities in the database.
+   *
+   * <p>Unlike {@link #insertAll}, this method uses {@link jakarta.persistence.EntityManager#merge}
+   * to make a deep copy and <b>never modifies the input entities in place</b>.
+   */
   void putAll(ImmutableCollection<?> entities);
 
-  /** Updates an entity in the database, throws exception if the entity does not exist. */
+  /**
+   * Updates an existing entity in the database, throwing an exception if it does not exist.
+   *
+   * <p>This method uses {@link jakarta.persistence.EntityManager#merge} to make a deep copy and
+   * <b>never modifies the input {@code entity} in place</b>.
+   */
   void update(Object entity);
 
-  /** Updates all entities in the database, throws exception if any entity does not exist. */
+  /**
+   * Updates all existing entities in the database, throwing an exception if any does not exist.
+   *
+   * <p>This method uses {@link jakarta.persistence.EntityManager#merge} to make a deep copy and
+   * <b>never modifies the input entities in place</b>.
+   */
   void updateAll(ImmutableCollection<?> entities);
 
-  /** Updates all entities in the database, throws exception if any entity does not exist. */
+  /**
+   * Updates all existing entities in the database, throwing an exception if any does not exist.
+   *
+   * <p>This method uses {@link jakarta.persistence.EntityManager#merge} to make a deep copy and
+   * <b>never modifies the input entities in place</b>.
+   */
   void updateAll(ImmutableObject... entities);
 
   /** Returns whether the given entity with same ID exists. */

core/src/test/java/google/registry/beam/resave/ResaveAllEppResourcesPipelineTest.java

@@ -15,6 +15,7 @@
 package google.registry.beam.resave;
 
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static com.google.common.collect.Iterables.getOnlyElement;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.testing.DatabaseHelper.createTld;
@@ -27,10 +28,11 @@ import static google.registry.testing.DatabaseHelper.persistDomainWithPendingTra
 import static google.registry.testing.DatabaseHelper.persistNewRegistrars;
 import static google.registry.util.DateTimeUtils.minusDays;
 import static google.registry.util.DateTimeUtils.plusYears;
+import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import com.google.common.collect.ImmutableCollection;
 import google.registry.beam.TestPipelineExtension;
 import google.registry.model.EppResource;
 import google.registry.model.domain.Domain;
@@ -138,6 +140,7 @@ public class ResaveAllEppResourcesPipelineTest {
   }
 
   @Test
+  @SuppressWarnings("unchecked")
   void testPipeline_fastOnlySavesChanged() {
     Instant now = fakeClock.now();
     persistDomainWithDependentResources("renewed", "tld", now, now, plusYears(now, 1));
@@ -145,14 +148,16 @@ public class ResaveAllEppResourcesPipelineTest {
     // Spy the transaction manager so we can be sure we're only saving the renewed domain
     JpaTransactionManager spy = spy(tm());
     TransactionManagerFactory.setJpaTm(() -> spy);
-    ArgumentCaptor<Domain> domainPutCaptor = ArgumentCaptor.forClass(Domain.class);
+    ArgumentCaptor<ImmutableCollection<Domain>> domainPutCaptor =
+        ArgumentCaptor.forClass(ImmutableCollection.class);
     runPipeline();
     // We should only be attempting to put the one changed domain into the DB
-    verify(spy).put(domainPutCaptor.capture());
-    assertThat(domainPutCaptor.getValue().getDomainName()).isEqualTo("renewed.tld");
+    verify(spy).putAll(domainPutCaptor.capture());
+    assertThat(getOnlyElement(domainPutCaptor.getValue()).getDomainName()).isEqualTo("renewed.tld");
   }
 
   @Test
+  @SuppressWarnings("unchecked")
   void testPipeline_notFastResavesAll() {
     options.setFast(false);
     Instant now = fakeClock.now();
@@ -163,12 +168,14 @@ public class ResaveAllEppResourcesPipelineTest {
     // Spy the transaction manager so we can be sure we're attempting to save everything
     JpaTransactionManager spy = spy(tm());
     TransactionManagerFactory.setJpaTm(() -> spy);
-    ArgumentCaptor<EppResource> eppResourcePutCaptor = ArgumentCaptor.forClass(EppResource.class);
+    ArgumentCaptor<ImmutableCollection<EppResource>> eppResourcePutCaptor =
+        ArgumentCaptor.forClass(ImmutableCollection.class);
     runPipeline();
     // We should be attempting to put both domains in, even the unchanged one
-    verify(spy, times(2)).put(eppResourcePutCaptor.capture());
+    verify(spy, atLeastOnce()).putAll(eppResourcePutCaptor.capture());
     assertThat(
             eppResourcePutCaptor.getAllValues().stream()
+                .flatMap(ImmutableCollection::stream)
                 .map(EppResource::getRepoId)
                 .collect(toImmutableSet()))
         .containsExactly(renewed.getRepoId(), nonRenewed.getRepoId());

core/src/test/java/google/registry/model/common/DnsRefreshRequestTest.java

@@ -49,6 +49,17 @@ public class DnsRefreshRequestTest extends EntityTestCase {
     assertThat(requests.get(0).id).isNotNull();
   }
 
+  @Test
+  void testPutAll_leavesTransientIdNull() {
+    assertThat(request.id).isNull();
+    tm().transact(() -> tm().putAll(request));
+    assertThat(request.id).isNull();
+    ImmutableList<DnsRefreshRequest> requests = loadAllOf(DnsRefreshRequest.class);
+    assertThat(requests).hasSize(1);
+    assertThat(requests.get(0).id).isNotNull();
+    assertAboutImmutableObjects().that(requests.get(0)).isEqualExceptFields(request, "id");
+  }
+
   @Test
   void testNullValues() {
     // type

core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerImplTest.java

@@ -54,6 +54,8 @@ import jakarta.persistence.PersistenceException;
 import jakarta.persistence.RollbackException;
 import java.io.Serializable;
 import java.util.NoSuchElementException;
+import org.hibernate.SessionFactory;
+import org.hibernate.stat.Statistics;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import org.junit.jupiter.api.function.Executable;
@@ -480,6 +482,76 @@ class JpaTransactionManagerImplTest {
         .containsExactlyElementsIn(moreEntities);
   }
 
+  @Test
+  void putAll_mixedEntityTypes_succeeds() {
+    persistResource(theEntity);
+    TestEntity updatedTheEntity = new TestEntity("theEntity", "foo_updated");
+    TestEntity newEntity = new TestEntity("newEntity", "new_data");
+    tm().transact(
+            () -> tm().putAll(ImmutableList.of(updatedTheEntity, compoundIdEntity, newEntity)));
+    assertThat(tm().transact(() -> tm().loadByKey(theEntityKey))).isEqualTo(updatedTheEntity);
+    assertThat(tm().transact(() -> tm().loadByKey(compoundIdEntityKey)))
+        .isEqualTo(compoundIdEntity);
+    assertThat(tm().transact(() -> tm().loadByKey(VKey.create(TestEntity.class, "newEntity"))))
+        .isEqualTo(newEntity);
+  }
+
+  @Test
+  void putAll_varargs_succeeds() {
+    persistResource(theEntity);
+    TestEntity updatedTheEntity = new TestEntity("theEntity", "foo_updated");
+    tm().transact(() -> tm().putAll(updatedTheEntity, compoundIdEntity));
+    assertThat(tm().transact(() -> tm().loadByKey(theEntityKey))).isEqualTo(updatedTheEntity);
+    assertThat(tm().transact(() -> tm().loadByKey(compoundIdEntityKey)))
+        .isEqualTo(compoundIdEntity);
+  }
+
+  @Test
+  void putAll_duplicateIds_throws() {
+    assertThat(
+            assertThrows(
+                IllegalArgumentException.class,
+                () ->
+                    tm().transact(
+                            () ->
+                                tm().putAll(
+                                        new TestEntity("entity1", "foo"),
+                                        new TestEntity("entity1", "bar")))))
+        .hasMessageThat()
+        .contains("Multiple entities of type TestEntity with the same ID");
+  }
+
+  @Test
+  void putAll_deepCopiesNewEntitiesWithoutExtraSelects() {
+    persistResource(theEntity);
+    TestEntity updatedTheEntity = new TestEntity("theEntity", "foo_updated");
+    TestEntity newEntity1 = new TestEntity("newEntity1", "data1");
+    TestEntity newEntity2 = new TestEntity("newEntity2", "data2");
+    tm().transact(
+            () -> {
+              Statistics stats =
+                  tm().getEntityManager()
+                      .getEntityManagerFactory()
+                      .unwrap(SessionFactory.class)
+                      .getStatistics();
+              stats.setStatisticsEnabled(true);
+              stats.clear();
+              tm().putAll(updatedTheEntity, newEntity1, newEntity2);
+              // Only 1 SQL statement (the single batched findMultiple SELECT) should have been
+              // prepared so far; merging newEntity1 and newEntity2 must not trigger extra SELECTs.
+              assertThat(stats.getPrepareStatementCount()).isEqualTo(1);
+              assertThat(tm().getEntityManager().contains(updatedTheEntity)).isFalse();
+              assertThat(tm().getEntityManager().contains(newEntity1)).isFalse();
+              assertThat(tm().getEntityManager().contains(newEntity2)).isFalse();
+              newEntity1.data = "mutated_in_memory";
+            });
+    assertThat(tm().transact(() -> tm().loadByKey(theEntityKey))).isEqualTo(updatedTheEntity);
+    assertThat(tm().transact(() -> tm().loadByKey(VKey.create(TestEntity.class, "newEntity1"))))
+        .isEqualTo(new TestEntity("newEntity1", "data1"));
+    assertThat(tm().transact(() -> tm().loadByKey(VKey.create(TestEntity.class, "newEntity2"))))
+        .isEqualTo(new TestEntity("newEntity2", "data2"));
+  }
+
   @Test
   void update_succeeds() {
     persistResource(theEntity);
@@ -518,11 +590,54 @@ class JpaTransactionManagerImplTest {
             new TestEntity("entity1", "foo_updated"),
             new TestEntity("entity2", "bar_updated"),
             new TestEntity("entity3", "qux_updated"));
-    tm().transact(() -> tm().updateAll(updated));
+    tm().transact(
+            () -> {
+              Statistics stats =
+                  tm().getEntityManager()
+                      .getEntityManagerFactory()
+                      .unwrap(SessionFactory.class)
+                      .getStatistics();
+              stats.setStatisticsEnabled(true);
+              stats.clear();
+              tm().updateAll(updated);
+              // Only 1 SQL statement (the single batched findMultiple SELECT) should have been
+              // prepared before flush; merging the entities into the pre-warmed persistence context
+              // must not trigger extra per-entity SELECTs.
+              assertThat(stats.getPrepareStatementCount()).isEqualTo(1);
+            });
     assertThat(tm().transact(() -> tm().loadAllOf(TestEntity.class)))
         .containsExactlyElementsIn(updated);
   }
 
+  @Test
+  void updateAll_mixedEntityTypes_succeeds() {
+    persistResource(theEntity);
+    persistResource(compoundIdEntity);
+    TestEntity updatedTheEntity = new TestEntity("theEntity", "foo_updated");
+    TestCompoundIdEntity updatedCompoundEntity =
+        new TestCompoundIdEntity("compoundIdEntity", 10, "bar_updated");
+    tm().transact(() -> tm().updateAll(updatedTheEntity, updatedCompoundEntity));
+    assertThat(tm().transact(() -> tm().loadByKey(theEntityKey))).isEqualTo(updatedTheEntity);
+    assertThat(tm().transact(() -> tm().loadByKey(compoundIdEntityKey)))
+        .isEqualTo(updatedCompoundEntity);
+  }
+
+  @Test
+  void updateAll_duplicateIds_throws() {
+    persistResource(theEntity);
+    assertThat(
+            assertThrows(
+                IllegalArgumentException.class,
+                () ->
+                    tm().transact(
+                            () ->
+                                tm().updateAll(
+                                        new TestEntity("theEntity", "foo_updated"),
+                                        new TestEntity("theEntity", "bar_updated")))))
+        .hasMessageThat()
+        .contains("Multiple entities of type TestEntity with the same ID");
+  }
+
   @Test
   void updateAll_rollsBackWhenFailure() {
     persistResources(moreEntities);