HHH-1775 - collection batch fetching

Conflicts:

	hibernate-core/src/main/java/org/hibernate/engine/spi/BatchFetchQueue.java

hibernate-core/src/main/java/org/hibernate/engine/spi/BatchFetchQueue.java

@@ -25,15 +25,16 @@ package org.hibernate.engine.spi;
 
 import java.io.Serializable;
 import java.util.HashMap;
-import java.util.Iterator;
 import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
 import java.util.Map;
 
+import org.jboss.logging.Logger;
+
 import org.hibernate.EntityMode;
 import org.hibernate.cache.spi.CacheKey;
 import org.hibernate.collection.spi.PersistentCollection;
-import org.hibernate.internal.util.MarkerObject;
+import org.hibernate.internal.CoreMessageLogger;
-import org.hibernate.internal.util.collections.IdentityMap;
 import org.hibernate.persister.collection.CollectionPersister;
 import org.hibernate.persister.entity.EntityPersister;
 
@@ -43,33 +44,35 @@ import org.hibernate.persister.entity.EntityPersister;
  * can be re-used as a subquery for loading owned collections.
  *
  * @author Gavin King
+ * @author Steve Ebersole
+ * @author Guenther Demetz
  */
 public class BatchFetchQueue {
+	private static final CoreMessageLogger LOG = Logger.getMessageLogger( CoreMessageLogger.class, BatchFetchQueue.class.getName() );
 
-	public static final Object MARKER = new MarkerObject( "MARKER" );
+	private final PersistenceContext context;
-
-	/**
-	 * Defines a sequence of {@link EntityKey} elements that are currently
-	 * elegible for batch-fetching.
-	 * <p/>
-	 * Even though this is a map, we only use the keys.  A map was chosen in
-	 * order to utilize a {@link LinkedHashMap} to maintain sequencing
-	 * as well as uniqueness.
-	 * <p/>
-	 * TODO : this would be better as a SequencedReferenceSet, but no such beast exists!
-	 */
-	private final Map batchLoadableEntityKeys = new LinkedHashMap(8);
 
 	/**
 	 * A map of {@link SubselectFetch subselect-fetch descriptors} keyed by the
 	 * {@link EntityKey) against which the descriptor is registered.
 	 */
-	private final Map subselectsByEntityKey = new HashMap(8);
+	private final Map<EntityKey, SubselectFetch> subselectsByEntityKey = new HashMap<EntityKey, SubselectFetch>(8);
 
 	/**
-	 * The owning persistence context.
+	 * Used to hold information about the entities that are currently eligible for batch-fetching.  Ultimately
+	 * used by {@link #getEntityBatch} to build entity load batches.
+	 * <p/>
+	 * A Map structure is used to segment the keys by entity type since loading can only be done for a particular entity
+	 * type at a time.
 	 */
-	private final PersistenceContext context;
+	private final Map <String,LinkedHashSet<EntityKey>> batchLoadableEntityKeys = new HashMap <String,LinkedHashSet<EntityKey>>(8);
+	
+	/**
+	 * Used to hold information about the collections that are currently eligible for batch-fetching.  Ultimately
+	 * used by {@link #getCollectionBatch} to build collection load batches.
+	 */
+	private final Map<String, LinkedHashMap<CollectionEntry, PersistentCollection>> batchLoadableCollections =
+			new HashMap<String, LinkedHashMap <CollectionEntry, PersistentCollection>>(8);
 
 	/**
 	 * Constructs a queue for the given context.
@@ -88,6 +91,9 @@ public class BatchFetchQueue {
 		subselectsByEntityKey.clear();
 	}
 
+
+	// sub-select support ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
 	/**
 	 * Retrieve the fetch descriptor associated with the given entity key.
 	 *
@@ -96,7 +102,7 @@ public class BatchFetchQueue {
 	 * this entity key.
 	 */
 	public SubselectFetch getSubselect(EntityKey key) {
-		return (SubselectFetch) subselectsByEntityKey.get(key);
+		return subselectsByEntityKey.get( key );
 	}
 
 	/**
@@ -106,7 +112,7 @@ public class BatchFetchQueue {
 	 * @param subquery The fetch descriptor.
 	 */
 	public void addSubselect(EntityKey key, SubselectFetch subquery) {
-		subselectsByEntityKey.put(key, subquery);
+		subselectsByEntityKey.put( key, subquery );
 	}
 
 	/**
@@ -116,7 +122,7 @@ public class BatchFetchQueue {
 	 * need to load its collections)
 	 */
 	public void removeSubselect(EntityKey key) {
-		subselectsByEntityKey.remove(key);
+		subselectsByEntityKey.remove( key );
 	}
 
 	/**
@@ -128,6 +134,9 @@ public class BatchFetchQueue {
 		subselectsByEntityKey.clear();
 	}
 
+
+	// entity batch support ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
 	/**
 	 * If an EntityKey represents a batch loadable entity, add
 	 * it to the queue.
@@ -140,7 +149,12 @@ public class BatchFetchQueue {
 	 */
 	public void addBatchLoadableEntityKey(EntityKey key) {
 		if ( key.isBatchLoadable() ) {
-			batchLoadableEntityKeys.put( key, MARKER );
+			LinkedHashSet<EntityKey> set =  batchLoadableEntityKeys.get( key.getEntityName());
+			if (set == null) {
+				set = new LinkedHashSet<EntityKey>(8);
+				batchLoadableEntityKeys.put( key.getEntityName(), set);
+			}
+			set.add(key);
 		}
 	}
 
@@ -150,69 +164,12 @@ public class BatchFetchQueue {
 	 * if necessary
 	 */
 	public void removeBatchLoadableEntityKey(EntityKey key) {
-		if ( key.isBatchLoadable() ) batchLoadableEntityKeys.remove(key);
+		if ( key.isBatchLoadable() ) {
-	}
+			LinkedHashSet<EntityKey> set =  batchLoadableEntityKeys.get( key.getEntityName());
-
+			if (set != null) {
-	/**
+				set.remove(key);
-	 * Get a batch of uninitialized collection keys for a given role
-	 *
-	 * @param collectionPersister The persister for the collection role.
-	 * @param id A key that must be included in the batch fetch
-	 * @param batchSize the maximum number of keys to return
-	 * @return an array of collection keys, of length batchSize (padded with nulls)
-	 */
-	public Serializable[] getCollectionBatch(
-			final CollectionPersister collectionPersister,
-			final Serializable id,
-			final int batchSize) {
-		Serializable[] keys = new Serializable[batchSize];
-		keys[0] = id;
-		int i = 1;
-		//int count = 0;
-		int end = -1;
-		boolean checkForEnd = false;
-		// this only works because collection entries are kept in a sequenced
-		// map by persistence context (maybe we should do like entities and
-		// keep a separate sequences set...)
-
-		for ( Map.Entry<PersistentCollection,CollectionEntry> me :
-			IdentityMap.concurrentEntries( (Map<PersistentCollection,CollectionEntry>) context.getCollectionEntries() )) {
-
-			CollectionEntry ce = me.getValue();
-			PersistentCollection collection = me.getKey();
-			if ( !collection.wasInitialized() && ce.getLoadedPersister() == collectionPersister ) {
-
-				if ( checkForEnd && i == end ) {
-					return keys; //the first key found after the given key
-				}
-
-				//if ( end == -1 && count > batchSize*10 ) return keys; //try out ten batches, max
-
-				final boolean isEqual = collectionPersister.getKeyType().isEqual(
-						id,
-						ce.getLoadedKey(),
-						collectionPersister.getFactory()
-				);
-
-				if ( isEqual ) {
-					end = i;
-					//checkForEnd = false;
-				}
-				else if ( !isCached( ce.getLoadedKey(), collectionPersister ) ) {
-					keys[i++] = ce.getLoadedKey();
-					//count++;
-				}
-
-				if ( i == batchSize ) {
-					i = 1; //end of array, start filling again from start
-					if ( end != -1 ) {
-						checkForEnd = true;
-					}
-				}
 			}
-
 		}
-		return keys; //we ran out of keys to try
 	}
 
 	/**
@@ -236,10 +193,11 @@ public class BatchFetchQueue {
 		int end = -1;
 		boolean checkForEnd = false;
 
-		Iterator iter = batchLoadableEntityKeys.keySet().iterator();
+		// TODO: this needn't exclude subclasses...
-		while ( iter.hasNext() ) {
+
-			EntityKey key = (EntityKey) iter.next();
+		LinkedHashSet<EntityKey> set =  batchLoadableEntityKeys.get( persister.getEntityName() );
-			if ( key.getEntityName().equals( persister.getEntityName() ) ) { //TODO: this needn't exclude subclasses...
+		if ( set != null ) {
+			for ( EntityKey key : set ) {
 				if ( checkForEnd && i == end ) {
 					//the first id found after the given id
 					return ids;
@@ -253,8 +211,10 @@ public class BatchFetchQueue {
 					}
 				}
 				if ( i == batchSize ) {
-					i = 1; //end of array, start filling again from start
+					i = 1; // end of array, start filling again from start
-					if (end!=-1) checkForEnd = true;
+					if ( end != -1 ) {
+						checkForEnd = true;
+					}
 				}
 			}
 		}
@@ -273,6 +233,98 @@ public class BatchFetchQueue {
 		return false;
 	}
 	
+
+	// collection batch support ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+	/**
+	 * If an CollectionEntry represents a batch loadable collection, add
+	 * it to the queue.
+	 */
+	public void addBatchLoadableCollection(PersistentCollection collection, CollectionEntry ce) {
+		final CollectionPersister persister = ce.getLoadedPersister();
+
+		LinkedHashMap<CollectionEntry, PersistentCollection> map =  batchLoadableCollections.get( persister.getRole() );
+		if ( map == null ) {
+			map = new LinkedHashMap<CollectionEntry, PersistentCollection>( 16 );
+			batchLoadableCollections.put( persister.getRole(), map );
+		}
+		map.put( ce, collection );
+	}
+	
+	/**
+	 * After a collection was initialized or evicted, we don't
+	 * need to batch fetch it anymore, remove it from the queue
+	 * if necessary
+	 */
+	public void removeBatchLoadableCollection(CollectionEntry ce) {
+		LinkedHashMap<CollectionEntry, PersistentCollection> map =  batchLoadableCollections.get( ce.getLoadedPersister().getRole() );
+		if ( map != null ) {
+			map.remove( ce );
+		}
+	}
+
+	/**
+	 * Get a batch of uninitialized collection keys for a given role
+	 *
+	 * @param collectionPersister The persister for the collection role.
+	 * @param id A key that must be included in the batch fetch
+	 * @param batchSize the maximum number of keys to return
+	 * @return an array of collection keys, of length batchSize (padded with nulls)
+	 */
+	public Serializable[] getCollectionBatch(
+			final CollectionPersister collectionPersister,
+			final Serializable id,
+			final int batchSize) {
+
+		Serializable[] keys = new Serializable[batchSize];
+		keys[0] = id;
+
+		int i = 1;
+		int end = -1;
+		boolean checkForEnd = false;
+
+		final LinkedHashMap<CollectionEntry, PersistentCollection> map =  batchLoadableCollections.get( collectionPersister.getRole() );
+		if ( map != null ) {
+			for ( Map.Entry<CollectionEntry, PersistentCollection> me : map.entrySet() ) {
+				final CollectionEntry ce = me.getKey();
+				final PersistentCollection collection = me.getValue();
+
+				if ( collection.wasInitialized() ) {
+					// should never happen
+					LOG.warn( "Encountered initialized collection in BatchFetchQueue, this should not happen." );
+					continue;
+				}
+
+				if ( checkForEnd && i == end ) {
+					return keys; //the first key found after the given key
+				}
+
+				final boolean isEqual = collectionPersister.getKeyType().isEqual(
+						id,
+						ce.getLoadedKey(),
+						collectionPersister.getFactory()
+				);
+
+				if ( isEqual ) {
+					end = i;
+					//checkForEnd = false;
+				}
+				else if ( !isCached( ce.getLoadedKey(), collectionPersister ) ) {
+					keys[i++] = ce.getLoadedKey();
+					//count++;
+				}
+
+				if ( i == batchSize ) {
+					i = 1; //end of array, start filling again from start
+					if ( end != -1 ) {
+						checkForEnd = true;
+					}
+				}
+			}
+		}
+		return keys; //we ran out of keys to try
+	}
+
 	private boolean isCached(Serializable collectionKey, CollectionPersister persister) {
 		if ( persister.hasCache() ) {
 			CacheKey cacheKey = context.getSession().generateCacheKey(