HHH-7667 - Investigate expanding bytecode enhancement support

(cherry picked from commit 93f4fe0668)

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/spi/Enhancer.java

@@ -34,6 +34,7 @@ import javassist.ClassPool;
 import javassist.CtClass;
 import javassist.CtField;
 import javassist.CtNewMethod;
+import javassist.Modifier;
 import javassist.bytecode.AnnotationsAttribute;
 import javassist.bytecode.ConstPool;
 import javassist.bytecode.FieldInfo;
@@ -44,8 +45,8 @@ import org.jboss.logging.Logger;
 import org.hibernate.HibernateException;
 import org.hibernate.bytecode.enhance.EnhancementException;
 import org.hibernate.bytecode.internal.javassist.FieldHandled;
-import org.hibernate.engine.ManagedComposite;
+import org.hibernate.engine.spi.ManagedComposite;
-import org.hibernate.engine.ManagedEntity;
+import org.hibernate.engine.spi.ManagedEntity;
 import org.hibernate.engine.spi.EntityEntry;
 import org.hibernate.internal.CoreMessageLogger;
 
@@ -56,16 +57,27 @@ import org.hibernate.internal.CoreMessageLogger;
 public class Enhancer  {
 	private static final CoreMessageLogger log = Logger.getMessageLogger( CoreMessageLogger.class, Enhancer.class.getName() );
 
+	public static final String ENTITY_INSTANCE_GETTER_NAME = "hibernate_getEntityInstance";
+
 	public static final String ENTITY_ENTRY_FIELD_NAME = "$hibernate_entityEntryHolder";
 	public static final String ENTITY_ENTRY_GETTER_NAME = "hibernate_getEntityEntry";
 	public static final String ENTITY_ENTRY_SETTER_NAME = "hibernate_setEntityEntry";
 
+	public static final String PREVIOUS_FIELD_NAME = "$hibernate_previousManagedEntity";
+	public static final String PREVIOUS_GETTER_NAME = "hibernate_getPreviousManagedEntity";
+	public static final String PREVIOUS_SETTER_NAME = "hibernate_setPreviousManagedEntity";
+
+	public static final String NEXT_FIELD_NAME = "$hibernate_nextManagedEntity";
+	public static final String NEXT_GETTER_NAME = "hibernate_getNextManagedEntity";
+	public static final String NEXT_SETTER_NAME = "hibernate_setNextManagedEntity";
+
 	private final EnhancementContext enhancementContext;
 
 	private final ClassPool classPool;
 	private final CtClass managedEntityCtClass;
 	private final CtClass managedCompositeCtClass;
 	private final CtClass entityEntryCtClass;
+	private final CtClass objectCtClass;
 
 	public Enhancer(EnhancementContext enhancementContext) {
 		this.enhancementContext = enhancementContext;
@@ -88,6 +100,8 @@ public class Enhancer  {
 
 			// "add" EntityEntry
 			this.entityEntryCtClass = classPool.makeClass( EntityEntry.class.getName() );
+
+			this.objectCtClass = classPool.makeClass( Object.class.getName() );
 		}
 		catch (IOException e) {
 			throw new EnhancementException( "Could not prepare Javassist ClassPool", e );
@@ -177,6 +191,40 @@ public class Enhancer  {
 		// add the ManagedEntity interface
 		managedCtClass.addInterface( managedEntityCtClass );
 
+		addEntityInstanceHandling( managedCtClass, constPool );
+		addEntityEntryHandling( managedCtClass, constPool );
+		addLinkedPreviousHandling( managedCtClass, constPool );
+		addLinkedNextHandling( managedCtClass, constPool );
+	}
+
+	private void addEntityInstanceHandling(CtClass managedCtClass, ConstPool constPool) {
+		// add the ManagedEntity#hibernate_getEntityInstance method
+		try {
+			managedCtClass.addMethod(
+					CtNewMethod.make(
+							objectCtClass,
+							ENTITY_INSTANCE_GETTER_NAME,
+							new CtClass[0],
+							new CtClass[0],
+							"{ return this; }",
+							managedCtClass
+					)
+			);
+		}
+		catch (CannotCompileException e) {
+			throw new EnhancementException(
+					String.format(
+							"Could not enhance entity class [%s] to add EntityEntry getter",
+							managedCtClass.getName()
+					),
+					e
+			);
+		}
+
+		// essentially add `return this;`
+	}
+
+	private void addEntityEntryHandling(CtClass managedCtClass, ConstPool constPool) {
 		// add field to hold EntityEntry
 		final CtField entityEntryField;
 		try {
@@ -193,7 +241,8 @@ public class Enhancer  {
 			);
 		}
 
-		// make that new field @Transient
+		// make that new field transient and @Transient
+		entityEntryField.setModifiers( entityEntryField.getModifiers() | Modifier.TRANSIENT );
 		AnnotationsAttribute annotationsAttribute = getVisibleAnnotations( entityEntryField.getFieldInfo() );
 		annotationsAttribute.addAnnotation( new Annotation( Transient.class.getName(), constPool ) );
 
@@ -230,6 +279,108 @@ public class Enhancer  {
 		}
 	}
 
+	private void addLinkedPreviousHandling(CtClass managedCtClass, ConstPool constPool) {
+		// add field to hold "previous" ManagedEntity
+		final CtField previousField;
+		try {
+			previousField = new CtField( managedCtClass, PREVIOUS_FIELD_NAME, managedCtClass );
+			managedCtClass.addField( previousField );
+		}
+		catch (CannotCompileException e) {
+			throw new EnhancementException(
+					String.format(
+							"Could not enhance entity class [%s] to add field for holding previous ManagedEntity",
+							managedCtClass.getName()
+					),
+					e
+			);
+		}
+
+		// make that new field transient and @Transient
+		previousField.setModifiers( previousField.getModifiers() | Modifier.TRANSIENT );
+		AnnotationsAttribute annotationsAttribute = getVisibleAnnotations( previousField.getFieldInfo() );
+		annotationsAttribute.addAnnotation( new Annotation( Transient.class.getName(), constPool ) );
+
+		// add the "getter"
+		try {
+			managedCtClass.addMethod( CtNewMethod.getter( PREVIOUS_GETTER_NAME, previousField ) );
+		}
+		catch (CannotCompileException e) {
+			throw new EnhancementException(
+					String.format(
+							"Could not enhance entity class [%s] to add previous ManagedEntity getter",
+							managedCtClass.getName()
+					),
+					e
+			);
+		}
+
+		// add the "setter"
+		try {
+			managedCtClass.addMethod( CtNewMethod.setter( PREVIOUS_SETTER_NAME, previousField ) );
+		}
+		catch (CannotCompileException e) {
+			throw new EnhancementException(
+					String.format(
+							"Could not enhance entity class [%s] to add previous ManagedEntity setter",
+							managedCtClass.getName()
+					),
+					e
+			);
+		}
+	}
+
+	private void addLinkedNextHandling(CtClass managedCtClass, ConstPool constPool) {
+		// add field to hold "next" ManagedEntity
+		final CtField nextField;
+		try {
+			nextField = new CtField( managedCtClass, NEXT_FIELD_NAME, managedCtClass );
+			managedCtClass.addField( nextField );
+		}
+		catch (CannotCompileException e) {
+			throw new EnhancementException(
+					String.format(
+							"Could not enhance entity class [%s] to add field for holding next ManagedEntity",
+							managedCtClass.getName()
+					),
+					e
+			);
+		}
+
+		// make that new field transient and @Transient
+		nextField.setModifiers( nextField.getModifiers() | Modifier.TRANSIENT );
+		AnnotationsAttribute annotationsAttribute = getVisibleAnnotations( nextField.getFieldInfo() );
+		annotationsAttribute.addAnnotation( new Annotation( Transient.class.getName(), constPool ) );
+
+		// add the "getter"
+		try {
+			managedCtClass.addMethod( CtNewMethod.getter( NEXT_GETTER_NAME, nextField ) );
+		}
+		catch (CannotCompileException e) {
+			throw new EnhancementException(
+					String.format(
+							"Could not enhance entity class [%s] to add next ManagedEntity getter",
+							managedCtClass.getName()
+					),
+					e
+			);
+		}
+
+		// add the "setter"
+		try {
+			managedCtClass.addMethod( CtNewMethod.setter( NEXT_SETTER_NAME, nextField ) );
+		}
+		catch (CannotCompileException e) {
+			throw new EnhancementException(
+					String.format(
+							"Could not enhance entity class [%s] to add next ManagedEntity setter",
+							managedCtClass.getName()
+					),
+					e
+			);
+		}
+	}
+
 	private AnnotationsAttribute getVisibleAnnotations(FieldInfo fieldInfo) {
 		AnnotationsAttribute annotationsAttribute = (AnnotationsAttribute) fieldInfo.getAttribute( AnnotationsAttribute.visibleTag );
 		if ( annotationsAttribute == null ) {

hibernate-core/src/main/java/org/hibernate/engine/internal/EntityEntryContext.java

@@ -26,15 +26,15 @@ package org.hibernate.engine.internal;
 import java.io.IOException;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
-import java.io.Serializable;
+import java.util.IdentityHashMap;
-import java.util.LinkedHashMap;
 import java.util.Map;
 
 import org.jboss.logging.Logger;
 
+import org.hibernate.AssertionFailure;
 import org.hibernate.LockMode;
-import org.hibernate.engine.ManagedEntity;
 import org.hibernate.engine.spi.EntityEntry;
+import org.hibernate.engine.spi.ManagedEntity;
 
 /**
  * Defines a context for maintaining the relation between an entity associated with the Session ultimately owning this
@@ -43,22 +43,22 @@ import org.hibernate.engine.spi.EntityEntry;
  *         the entity->EntityEntry association is maintained in a Map within this class
  *     </li>
  *     <li>
- *         the EntityEntry is injected into the entity via it implementing the {@link org.hibernate.engine.ManagedEntity} contract,
+ *         the EntityEntry is injected into the entity via it implementing the {@link org.hibernate.engine.spi.ManagedEntity} contract,
  *         either directly or through bytecode enhancement.
  *     </li>
  * </ul>
  * <p/>
- * IMPL NOTE : This current implementation is not ideal in the {@link org.hibernate.engine.ManagedEntity} case.  The problem is that
- * the 'backingMap' is still the means to maintain ordering of the entries; but in the {@link org.hibernate.engine.ManagedEntity} case
- * the use of a Map is overkill, and double here so because of the need for wrapping the map keys.  But this is just
- * a quick prototype.
  *
  * @author Steve Ebersole
  */
 public class EntityEntryContext {
 	private static final Logger log = Logger.getLogger( EntityEntryContext.class );
 
-	private LinkedHashMap<KeyWrapper,EntityEntry> backingMap = new LinkedHashMap<KeyWrapper, EntityEntry>();
+	private transient ManagedEntity head;
+	private transient ManagedEntity tail;
+	private transient int count = 0;
+
+	private transient IdentityHashMap<Object,ManagedEntity> nonEnhancedEntityXref;
 
 	@SuppressWarnings( {"unchecked"})
 	private transient Map.Entry<Object,EntityEntry>[] reentrantSafeEntries = new Map.Entry[0];
@@ -67,193 +67,356 @@ public class EntityEntryContext {
 	public EntityEntryContext() {
 	}
 
-	/**
+	public void addEntityEntry(Object entity, EntityEntry entityEntry) {
-	 * Private constructor used during deserialization
+		// IMPORTANT!!!!!
-	 *
+		//		add is called more than once of some entities.  In such cases the first
-	 * @param backingMap The backing map re-built from the serial stream
+		//		call is simply setting up a "marker" to avoid infinite looping from reentrancy
-	 */
+		//
-	private EntityEntryContext(LinkedHashMap<KeyWrapper, EntityEntry> backingMap) {
+		// any addition (even the double one described above) should invalidate the cross-ref array
-		this.backingMap = backingMap;
-		// mark dirty so we can rebuild the 'reentrantSafeEntries'
 		dirty = true;
+
+		// determine the appropriate ManagedEntity instance to use based on whether the entity is enhanced or not.
+		// also track whether the entity was already associated with the context
+		final ManagedEntity managedEntity;
+		final boolean alreadyAssociated;
+		if ( ManagedEntity.class.isInstance( entity ) ) {
+			managedEntity = (ManagedEntity) entity;
+			alreadyAssociated = managedEntity.hibernate_getEntityEntry() != null;
+		}
+		else {
+			ManagedEntity wrapper = null;
+			if ( nonEnhancedEntityXref == null ) {
+				nonEnhancedEntityXref = new IdentityHashMap<Object, ManagedEntity>();
+			}
+			else {
+				wrapper = nonEnhancedEntityXref.get( entity );
 			}
 
-	public void clear() {
+			if ( wrapper == null ) {
-		dirty = true;
+				wrapper = new ManagedEntityImpl( entity );
-		for ( Map.Entry<KeyWrapper,EntityEntry> mapEntry : backingMap.entrySet() ) {
+				nonEnhancedEntityXref.put( entity, wrapper );
-			final Object realKey = mapEntry.getKey().getRealKey();
+				alreadyAssociated = false;
-			if ( ManagedEntity.class.isInstance( realKey ) ) {
-				( (ManagedEntity) realKey ).hibernate_setEntityEntry( null );
 			}
-		}
+			else {
-		backingMap.clear();
+				alreadyAssociated = true;
-		reentrantSafeEntries = null;
 			}
 
-	public void downgradeLocks() {
+			managedEntity = wrapper;
-		// Downgrade locks
+		}
-		for ( Map.Entry<KeyWrapper,EntityEntry> mapEntry : backingMap.entrySet() ) {
+
-			final Object realKey = mapEntry.getKey().getRealKey();
+		// associate the EntityEntry with the entity
-			final EntityEntry entityEntry = ManagedEntity.class.isInstance( realKey )
+		managedEntity.hibernate_setEntityEntry( entityEntry );
-					? ( (ManagedEntity) realKey ).hibernate_getEntityEntry()
+
-					: mapEntry.getValue();
+		if ( alreadyAssociated ) {
-			entityEntry.setLockMode( LockMode.NONE );
+			// if the entity was already associated with the context, skip the linking step.
+			return;
+		}
+
+		// finally, set up linking and count
+		if ( tail == null ) {
+			assert head == null;
+			head = managedEntity;
+			tail = head;
+			count = 1;
+		}
+		else {
+			tail.hibernate_setNextManagedEntity( managedEntity );
+			managedEntity.hibernate_setPreviousManagedEntity( tail );
+			tail = managedEntity;
+			count++;
 		}
 	}
 
 	public boolean hasEntityEntry(Object entity) {
-		return ManagedEntity.class.isInstance( entity )
+		return getEntityEntry( entity ) != null;
-				? ( (ManagedEntity) entity ).hibernate_getEntityEntry() == null
-				: backingMap.containsKey( new KeyWrapper<Object>( entity ) );
 	}
 
 	public EntityEntry getEntityEntry(Object entity) {
-		return ManagedEntity.class.isInstance( entity )
+		final ManagedEntity managedEntity;
-				? ( (ManagedEntity) entity ).hibernate_getEntityEntry()
+		if ( ManagedEntity.class.isInstance( entity ) ) {
-				: backingMap.get( new KeyWrapper<Object>( entity ) );
+			managedEntity = (ManagedEntity) entity;
+		}
+		else if ( nonEnhancedEntityXref == null ) {
+			managedEntity = null;
+		}
+		else {
+			managedEntity = nonEnhancedEntityXref.get( entity );
+		}
+
+		return managedEntity == null
+				? null
+				: managedEntity.hibernate_getEntityEntry();
 	}
 
 	public EntityEntry removeEntityEntry(Object entity) {
 		dirty = true;
+
+		final ManagedEntity managedEntity;
 		if ( ManagedEntity.class.isInstance( entity ) ) {
-			backingMap.remove( new KeyWrapper<Object>( entity ) );
+			managedEntity = (ManagedEntity) entity;
-			final EntityEntry entityEntry = ( (ManagedEntity) entity ).hibernate_getEntityEntry();
+		}
-			( (ManagedEntity) entity ).hibernate_setEntityEntry( null );
+		else if ( nonEnhancedEntityXref == null ) {
-			return entityEntry;
+			managedEntity = null;
 		}
 		else {
-			return backingMap.remove( new KeyWrapper<Object>( entity ) );
+			managedEntity = nonEnhancedEntityXref.remove( entity );
+		}
+
+		if ( managedEntity == null ) {
+			throw new AssertionFailure( "Unable to resolve entity reference to EntityEntry for removal" );
+		}
+
+		// prepare for re-linking...
+		ManagedEntity previous = managedEntity.hibernate_getPreviousManagedEntity();
+		ManagedEntity next = managedEntity.hibernate_getNextManagedEntity();
+		managedEntity.hibernate_setPreviousManagedEntity( null );
+		managedEntity.hibernate_setNextManagedEntity( null );
+
+		count--;
+
+		if ( count == 0 ) {
+			// handle as a special case...
+			head = null;
+			tail = null;
+
+			assert previous == null;
+			assert next == null;
+		}
+		else {
+			// otherwise, previous or next (or both) should be non-null
+			if ( previous == null ) {
+				// we are removing head
+				assert managedEntity == head;
+				head = next;
+			}
+			else {
+				previous.hibernate_setNextManagedEntity( next );
+			}
+
+			if ( next == null ) {
+				// we are removing tail
+				assert managedEntity == tail;
+				tail = previous;
+			}
+			else {
+				next.hibernate_setPreviousManagedEntity( previous );
 			}
 		}
 
-	public void addEntityEntry(Object entity, EntityEntry entityEntry) {
+		EntityEntry theEntityEntry = managedEntity.hibernate_getEntityEntry();
-		dirty = true;
+		managedEntity.hibernate_setEntityEntry( null );
-		if ( ManagedEntity.class.isInstance( entity ) ) {
+		return theEntityEntry;
-			backingMap.put( new KeyWrapper<Object>( entity ), null );
-			( (ManagedEntity) entity ).hibernate_setEntityEntry( entityEntry );
-		}
-		else {
-			backingMap.put( new KeyWrapper<Object>( entity ), entityEntry );
-		}
 	}
 
 	public Map.Entry<Object, EntityEntry>[] reentrantSafeEntityEntries() {
 		if ( dirty ) {
-			reentrantSafeEntries = new MapEntryImpl[ backingMap.size() ];
+			reentrantSafeEntries = new EntityEntryCrossRefImpl[count];
 			int i = 0;
-			for ( Map.Entry<KeyWrapper,EntityEntry> mapEntry : backingMap.entrySet() ) {
+			ManagedEntity managedEntity = head;
-				final Object entity = mapEntry.getKey().getRealKey();
+			while ( managedEntity != null ) {
-				final EntityEntry entityEntry = ManagedEntity.class.isInstance( entity )
+				reentrantSafeEntries[i++] = new EntityEntryCrossRefImpl(
-						? ( (ManagedEntity) entity ).hibernate_getEntityEntry()
+						managedEntity.hibernate_getEntityInstance(),
-						: mapEntry.getValue();
+						managedEntity.hibernate_getEntityEntry()
-				reentrantSafeEntries[i++] = new MapEntryImpl( entity, entityEntry );
+				);
+				managedEntity = managedEntity.hibernate_getNextManagedEntity();
 			}
 			dirty = false;
 		}
 		return reentrantSafeEntries;
 	}
 
+	public void clear() {
+		dirty = true;
+
+		ManagedEntity node = head;
+		while ( node != null ) {
+			final ManagedEntity nextNode = node.hibernate_getNextManagedEntity();
+
+			node.hibernate_setEntityEntry( null );
+			node.hibernate_setPreviousManagedEntity( null );
+			node.hibernate_setNextManagedEntity( null );
+
+			node = nextNode;
+		}
+
+		if ( nonEnhancedEntityXref != null ) {
+			nonEnhancedEntityXref.clear();
+		}
+
+		head = null;
+		tail = null;
+		count = 0;
+
+		reentrantSafeEntries = null;
+	}
+
+	public void downgradeLocks() {
+		if ( head == null ) {
+			return;
+		}
+
+		ManagedEntity node = head;
+		while ( node != null ) {
+			node.hibernate_getEntityEntry().setLockMode( LockMode.NONE );
+
+			node = node.hibernate_getNextManagedEntity();
+		}
+	}
+
 	public void serialize(ObjectOutputStream oos) throws IOException {
-		final int count = backingMap.size();
 		log.tracef( "Starting serialization of [%s] EntityEntry entries", count );
 		oos.writeInt( count );
-		for ( Map.Entry<KeyWrapper,EntityEntry> mapEntry : backingMap.entrySet() ) {
+		if ( count == 0 ) {
-			oos.writeObject( mapEntry.getKey().getRealKey() );
+			return;
-			mapEntry.getValue().serialize( oos );
+		}
+
+		ManagedEntity managedEntity = head;
+		while ( managedEntity != null ) {
+			// so we know whether or not to build a ManagedEntityImpl on deserialize
+			oos.writeBoolean( managedEntity == managedEntity.hibernate_getEntityInstance() );
+			oos.writeObject( managedEntity.hibernate_getEntityInstance() );
+			managedEntity.hibernate_getEntityEntry().serialize( oos );
+
+			managedEntity = managedEntity.hibernate_getNextManagedEntity();
 		}
 	}
 
 	public static EntityEntryContext deserialize(ObjectInputStream ois, StatefulPersistenceContext rtn) throws IOException, ClassNotFoundException {
 		final int count = ois.readInt();
 		log.tracef( "Starting deserialization of [%s] EntityEntry entries", count );
-		final LinkedHashMap<KeyWrapper,EntityEntry> backingMap = new LinkedHashMap<KeyWrapper, EntityEntry>( count );
+
+		final EntityEntryContext context = new EntityEntryContext();
+		context.count = count;
+		context.dirty = true;
+
+		if ( count == 0 ) {
+			return context;
+		}
+
+		ManagedEntity previous = null;
+
 		for ( int i = 0; i < count; i++ ) {
+			final boolean isEnhanced = ois.readBoolean();
 			final Object entity = ois.readObject();
 			final EntityEntry entry = EntityEntry.deserialize( ois, rtn );
-			backingMap.put( new KeyWrapper<Object>( entity ), entry );
+			final ManagedEntity managedEntity;
-		}
+			if ( isEnhanced ) {
-		return new EntityEntryContext( backingMap );
+				managedEntity = (ManagedEntity) entity;
-	}
-
-	/**
-	 * @deprecated Added to support (also deprecated) PersistenceContext.getEntityEntries method until it can be removed.  Safe to use for counts.
-	 *
-	 */
-	@Deprecated
-	public Map getEntityEntryMap() {
-		return backingMap;
-	}
-
-	/**
-	 * We need to base the identity on {@link System#identityHashCode(Object)} but
-	 * attempt to lazily initialize and cache this value: being a native invocation
-	 * it is an expensive value to retrieve.
-	 */
-	public static final class KeyWrapper<K> implements Serializable {
-		private final K realKey;
-		private int hash = 0;
-
-		KeyWrapper(K realKey) {
-			this.realKey = realKey;
-		}
-
-		@SuppressWarnings( {"EqualsWhichDoesntCheckParameterClass"})
-		@Override
-		public boolean equals(Object other) {
-			return realKey == ( (KeyWrapper) other ).realKey;
-		}
-
-		@Override
-		public int hashCode() {
-			if ( this.hash == 0 ) {
-				//We consider "zero" as non-initialized value
-				final int newHash = System.identityHashCode( realKey );
-				if ( newHash == 0 ) {
-					//So make sure we don't store zeros as it would trigger initialization again:
-					//any value is fine as long as we're deterministic.
-					this.hash = -1;
 			}
 			else {
-					this.hash = newHash;
+				managedEntity = new ManagedEntityImpl( entity );
+				if ( context.nonEnhancedEntityXref == null ) {
+					context.nonEnhancedEntityXref = new IdentityHashMap<Object, ManagedEntity>();
 				}
+				context.nonEnhancedEntityXref.put( entity, managedEntity );
 			}
-			return hash;
+			managedEntity.hibernate_setEntityEntry( entry );
+
+			if ( previous == null ) {
+				context.head = managedEntity;
+			}
+			else {
+				previous.hibernate_setNextManagedEntity( managedEntity );
+				managedEntity.hibernate_setPreviousManagedEntity( previous );
+			}
+
+			previous = managedEntity;
+		}
+
+		context.tail = previous;
+
+		return context;
+	}
+
+	public int getNumberOfManagedEntities() {
+		return count;
+	}
+
+	private static class ManagedEntityImpl implements ManagedEntity {
+		private final Object entityInstance;
+		private EntityEntry entityEntry;
+		private ManagedEntity previous;
+		private ManagedEntity next;
+
+		public ManagedEntityImpl(Object entityInstance) {
+			this.entityInstance = entityInstance;
 		}
 
 		@Override
-		public String toString() {
+		public Object hibernate_getEntityInstance() {
-			return realKey.toString();
+			return entityInstance;
 		}
 
-		public K getRealKey() {
+		@Override
-			return realKey;
+		public EntityEntry hibernate_getEntityEntry() {
+			return entityEntry;
+		}
+
+		@Override
+		public void hibernate_setEntityEntry(EntityEntry entityEntry) {
+			this.entityEntry = entityEntry;
+		}
+
+		@Override
+		public ManagedEntity hibernate_getNextManagedEntity() {
+			return next;
+		}
+
+		@Override
+		public void hibernate_setNextManagedEntity(ManagedEntity next) {
+			this.next = next;
+		}
+
+		@Override
+		public ManagedEntity hibernate_getPreviousManagedEntity() {
+			return previous;
+		}
+
+		@Override
+		public void hibernate_setPreviousManagedEntity(ManagedEntity previous) {
+			this.previous = previous;
 		}
 	}
 
-	private static class MapEntryImpl implements Map.Entry<Object,EntityEntry> {
+	private static class EntityEntryCrossRefImpl implements EntityEntryCrossRef {
-		private final Object key;
+		private final Object entity;
-		private EntityEntry value;
+		private EntityEntry entityEntry;
 
-		private MapEntryImpl(Object key, EntityEntry value) {
+		private EntityEntryCrossRefImpl(Object entity, EntityEntry entityEntry) {
-			this.key = key;
+			this.entity = entity;
-			this.value = value;
+			this.entityEntry = entityEntry;
+		}
+
+		@Override
+		public Object getEntity() {
+			return entity;
+		}
+
+		@Override
+		public EntityEntry getEntityEntry() {
+			return entityEntry;
 		}
 
 		@Override
 		public Object getKey() {
-			return key;
+			return getEntity();
 		}
 
 		@Override
 		public EntityEntry getValue() {
-			return value;
+			return getEntityEntry();
 		}
 
 		@Override
-		public EntityEntry setValue(EntityEntry value) {
+		public EntityEntry setValue(EntityEntry entityEntry) {
-			final EntityEntry old = this.value;
+			final EntityEntry old = this.entityEntry;
-			this.value = value;
+			this.entityEntry = entityEntry;
-			return value;
+			return old;
 		}
 	}
+
+	public static interface EntityEntryCrossRef extends Map.Entry<Object,EntityEntry> {
+		public Object getEntity();
+		public EntityEntry getEntityEntry();
+	}
 }

hibernate-core/src/main/java/org/hibernate/engine/internal/StatefulPersistenceContext.java

@@ -1152,10 +1152,14 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		return proxiesByKey;
 	}
 
+	@Override
+	public int getNumberOfManagedEntities() {
+		return entityEntryContext.getNumberOfManagedEntities();
+	}
+
 	@Override
 	public Map getEntityEntries() {
-		return entityEntryContext.getEntityEntryMap();
+		return null;
-//		return entityEntries;
 	}
 
 	@Override

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

@@ -21,12 +21,21 @@
  * 51 Franklin Street, Fifth Floor
  * Boston, MA  02110-1301  USA
  */
-package org.hibernate.engine;
+package org.hibernate.engine.spi;
 
 /**
  * Contract for classes (specifically, entities and components/embeddables) that are "managed".  Developers can
  * choose to either have their classes manually implement these interfaces or Hibernate can enhance their classes
  * to implement these interfaces via built-time or run-time enhancement.
+ * <p/>
+ * The term managed here is used to describe both:<ul>
+ *     <li>
+ *         the fact that they are known to the persistence provider (this is defined by the interface itself)
+ *     </li>
+ *     <li>
+ *         its association with Session/EntityManager (this is defined by the state exposed through the interface)
+ *     </li>
+ * </ul>
  *
  * @author Steve Ebersole
  */

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

@@ -21,7 +21,7 @@
  * 51 Franklin Street, Fifth Floor
  * Boston, MA  02110-1301  USA
  */
-package org.hibernate.engine;
+package org.hibernate.engine.spi;
 
 /**
  * Specialized {@link Managed} contract for component/embeddable classes.

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

@@ -21,18 +21,35 @@
  * 51 Franklin Street, Fifth Floor
  * Boston, MA  02110-1301  USA
  */
-package org.hibernate.engine;
+package org.hibernate.engine.spi;
-
-import org.hibernate.engine.spi.EntityEntry;
 
 /**
- * Specialized {@link Managed} contract for entity classes.
+ * Specialized {@link Managed} contract for entity classes.  Essentially provides access to information
+ * about an instance's association to a Session/EntityManager.  Specific information includes:<ul>
+ *     <li>
+ *        the association's {@link EntityEntry} (by way of {@link #hibernate_getEntityEntry()} and
+ *        {@link #hibernate_setEntityEntry}).  EntityEntry describes states, snapshots, etc.
+ *     </li>
+ *     <li>
+ *         link information.  ManagedEntity instances are part of a "linked list", thus link information
+ *         describes the next and previous entries/nodes in that ordering.  See
+ *         {@link #hibernate_getNextManagedEntity}, {@link #hibernate_setNextManagedEntity},
+ *         {@link #hibernate_getPreviousManagedEntity()}, {@link #hibernate_setPreviousManagedEntity}
+ *     </li>
+ * </ul>
  *
  * @author Steve Ebersole
  */
 public interface ManagedEntity extends Managed {
 	/**
-	 * Callback to get any associated EntityEntry.
+	 * Obtain a reference to the entity instance.
+	 *
+	 * @return The entity instance.
+	 */
+	public Object hibernate_getEntityInstance();
+
+	/**
+	 * Provides access to the associated EntityEntry.
 	 *
 	 * @return The EntityEntry associated with this entity instance.
 	 *
@@ -47,4 +64,12 @@ public interface ManagedEntity extends Managed {
 	 * @param entityEntry The EntityEntry associated with this entity instance.
 	 */
 	public void hibernate_setEntityEntry(EntityEntry entityEntry);
+
+	public ManagedEntity hibernate_getPreviousManagedEntity();
+
+	public void hibernate_setPreviousManagedEntity(ManagedEntity previous);
+
+	public ManagedEntity hibernate_getNextManagedEntity();
+
+	public void hibernate_setNextManagedEntity(ManagedEntity next);
 }

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

@@ -481,16 +481,26 @@ public interface PersistenceContext {
 	 */
 	public Map getEntitiesByKey();
 
+	/**
+	 * Provides access to the entity/EntityEntry combos associated with the persistence context in a manner that
+	 * is safe from reentrant access.  Specifically, it is safe from additions/removals while iterating.
+	 *
+	 * @return
+	 */
 	public Map.Entry<Object,EntityEntry>[] reentrantSafeEntityEntries();
 
 	/**
 	 * Get the mapping from entity instance to entity entry
 	 *
-	 * @deprecated Due to the introduction of EntityEntryContext and bytecode enhancement
+	 * @deprecated Due to the introduction of EntityEntryContext and bytecode enhancement; only valid really for
+	 * sizing, see {@link #getNumberOfManagedEntities}.  For iterating the entity/EntityEntry combos, see
+	 * {@link #reentrantSafeEntityEntries}
 	 */
 	@Deprecated
 	public Map getEntityEntries();
 
+	public int getNumberOfManagedEntities();
+
 	/**
 	 * Get the mapping from collection instance to collection entry
 	 */

hibernate-core/src/main/java/org/hibernate/event/internal/AbstractFlushingEventListener.java

@@ -120,7 +120,7 @@ public abstract class AbstractFlushingEventListener implements Serializable {
 				session.getActionQueue().numberOfInsertions(),
 				session.getActionQueue().numberOfUpdates(),
 				session.getActionQueue().numberOfDeletions(),
-				persistenceContext.getEntityEntries().size()
+				persistenceContext.getNumberOfManagedEntities()
 		);
 		LOG.debugf(
 				"Flushed: %s (re)creations, %s updates, %s removals to %s collections",

hibernate-core/src/main/java/org/hibernate/event/internal/DefaultAutoFlushEventListener.java

@@ -86,7 +86,7 @@ public class DefaultAutoFlushEventListener extends AbstractFlushingEventListener
 	private boolean flushMightBeNeeded(final EventSource source) {
 		return !source.getFlushMode().lessThan(FlushMode.AUTO) &&
 				source.getDontFlushFromFind() == 0 &&
-				( source.getPersistenceContext().getEntityEntries().size() > 0 ||
+				( source.getPersistenceContext().getNumberOfManagedEntities() > 0 ||
 						source.getPersistenceContext().getCollectionEntries().size() > 0 );
 	}
 }

hibernate-core/src/main/java/org/hibernate/event/internal/DefaultFlushEventListener.java

@@ -45,7 +45,7 @@ public class DefaultFlushEventListener extends AbstractFlushingEventListener imp
 	public void onFlush(FlushEvent event) throws HibernateException {
 		final EventSource source = event.getSession();
 		final PersistenceContext persistenceContext = source.getPersistenceContext();
-		if ( persistenceContext.getEntityEntries().size() > 0 ||
+		if ( persistenceContext.getNumberOfManagedEntities() > 0 ||
 				persistenceContext.getCollectionEntries().size() > 0 ) {
 
 			flushEverythingToExecutions(event);

hibernate-core/src/main/java/org/hibernate/stat/internal/SessionStatisticsImpl.java

@@ -41,7 +41,7 @@ public class SessionStatisticsImpl implements SessionStatistics {
 	}
 
 	public int getEntityCount() {
-		return session.getPersistenceContext().getEntityEntries().size();
+		return session.getPersistenceContext().getNumberOfManagedEntities();
 	}
 	
 	public int getCollectionCount() {

hibernate-core/src/main/java/org/hibernate/tool/enhance/EnhancementTask.java

@@ -52,7 +52,7 @@ import org.hibernate.bytecode.enhance.spi.Enhancer;
  *
  * @author Steve Ebersole
  *
- * @see org.hibernate.engine.Managed
+ * @see org.hibernate.engine.spi.Managed
  */
 public class EnhancementTask extends Task {
 	private List<FileSet> filesets = new ArrayList<FileSet>();

hibernate-core/src/test/java/org/hibernate/test/bytecode/enhancement/EnhancerTest.java

@@ -53,6 +53,7 @@ import org.hibernate.testing.junit4.BaseUnitTestCase;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 import static org.junit.Assert.fail;
 
 /**
@@ -95,6 +96,9 @@ public class EnhancerTest extends BaseUnitTestCase {
 		assertNotNull( getter.invoke( simpleEntityInstance ) );
 		setter.invoke( simpleEntityInstance, new Object[] { null } );
 		assertNull( getter.invoke( simpleEntityInstance ) );
+
+		Method entityInstanceGetter = simpleEntityClass.getMethod( Enhancer.ENTITY_INSTANCE_GETTER_NAME );
+		assertSame( simpleEntityInstance, entityInstanceGetter.invoke( simpleEntityInstance ) );
 	}
 
 	private CtClass generateCtClassForAnEntity() throws IOException, NotFoundException {

hibernate-core/src/test/java/org/hibernate/test/bytecode/enhancement/MostBasicEnhancementTest.java

@@ -31,6 +31,7 @@ import org.hibernate.testing.junit4.BaseCoreFunctionalTestCase;
 
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertSame;
 
 /**
  * @author Steve Ebersole
@@ -46,13 +47,28 @@ public class MostBasicEnhancementTest extends BaseCoreFunctionalTestCase {
 		Session s = openSession();
 		s.beginTransaction();
 		s.save( new MyEntity( 1L ) );
+		s.save( new MyEntity( 2L ) );
 		s.getTransaction().commit();
 		s.close();
 
 		s = openSession();
 		s.beginTransaction();
-		MyEntity myEntity = (MyEntity) s.load( MyEntity.class, 1L );
+		MyEntity myEntity = (MyEntity) s.get( MyEntity.class, 1L );
+		MyEntity myEntity2 = (MyEntity) s.get( MyEntity.class, 2L );
+
+		assertNotNull( myEntity.hibernate_getEntityInstance() );
+		assertSame( myEntity, myEntity.hibernate_getEntityInstance() );
 		assertNotNull( myEntity.hibernate_getEntityEntry() );
+		assertNull( myEntity.hibernate_getPreviousManagedEntity() );
+		assertNotNull( myEntity.hibernate_getNextManagedEntity() );
+
+		assertNotNull( myEntity2.hibernate_getEntityInstance() );
+		assertSame( myEntity2, myEntity2.hibernate_getEntityInstance() );
+		assertNotNull( myEntity2.hibernate_getEntityEntry() );
+		assertNotNull( myEntity2.hibernate_getPreviousManagedEntity() );
+		assertNull( myEntity2.hibernate_getNextManagedEntity() );
+
+		s.createQuery( "delete MyEntity" ).executeUpdate();
 		s.getTransaction().commit();
 		s.close();
 

hibernate-core/src/test/java/org/hibernate/test/bytecode/enhancement/MyEntity.java

@@ -27,7 +27,7 @@ import javax.persistence.Entity;
 import javax.persistence.Id;
 import javax.persistence.Transient;
 
-import org.hibernate.engine.ManagedEntity;
+import org.hibernate.engine.spi.ManagedEntity;
 import org.hibernate.engine.spi.EntityEntry;
 
 /**
@@ -37,6 +37,10 @@ import org.hibernate.engine.spi.EntityEntry;
 public class MyEntity implements ManagedEntity {
 	@Transient
 	private transient EntityEntry entityEntry;
+	@Transient
+	private transient ManagedEntity previous;
+	@Transient
+	private transient ManagedEntity next;
 
 	private Long id;
 	private String name;
@@ -65,6 +69,11 @@ public class MyEntity implements ManagedEntity {
 		this.name = name;
 	}
 
+	@Override
+	public Object hibernate_getEntityInstance() {
+		return this;
+	}
+
 	@Override
 	public EntityEntry hibernate_getEntityEntry() {
 		return entityEntry;
@@ -74,4 +83,24 @@ public class MyEntity implements ManagedEntity {
 	public void hibernate_setEntityEntry(EntityEntry entityEntry) {
 		this.entityEntry = entityEntry;
 	}
+
+	@Override
+	public ManagedEntity hibernate_getNextManagedEntity() {
+		return next;
+	}
+
+	@Override
+	public void hibernate_setNextManagedEntity(ManagedEntity next) {
+		this.next = next;
+	}
+
+	@Override
+	public ManagedEntity hibernate_getPreviousManagedEntity() {
+		return previous;
+	}
+
+	@Override
+	public void hibernate_setPreviousManagedEntity(ManagedEntity previous) {
+		this.previous = previous;
+	}
 }

hibernate-envers/src/test/java/org/hibernate/envers/test/performance/EvictAuditDataAfterCommitTest.java

@@ -3,6 +3,7 @@ package org.hibernate.envers.test.performance;
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 import org.junit.Assert;
@@ -93,8 +94,8 @@ public class EvictAuditDataAfterCommitTest extends BaseEnversFunctionalTestCase
     private void checkEmptyAuditSessionCache(Session session, String ... auditEntityNames) {
         List<String> entityNames = Arrays.asList(auditEntityNames);
         PersistenceContext persistenceContext = ((SessionImplementor) session).getPersistenceContext();
-        for (Object entry : persistenceContext.getEntityEntries().values()) {
+        for ( Map.Entry<Object,EntityEntry> entrySet : persistenceContext.reentrantSafeEntityEntries() ) {
-            EntityEntry entityEntry = (EntityEntry) entry;
+            final EntityEntry entityEntry = entrySet.getValue();
             if (entityNames.contains(entityEntry.getEntityName())) {
                 assert false : "Audit data shall not be stored in the session level cache. This causes performance issues.";
             }