HHH-7667 - Investigate expanding bytecode enhancement support

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/EnhancementException.java

@@ -0,0 +1,39 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.bytecode.enhance;
+
+import org.hibernate.HibernateException;
+
+/**
+ * @author Steve Ebersole
+ */
+public class EnhancementException extends HibernateException {
+	public EnhancementException(String message) {
+		super( message );
+	}
+
+	public EnhancementException(String message, Throwable root) {
+		super( message, root );
+	}
+}

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

@@ -0,0 +1,34 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.bytecode.enhance.spi;
+
+import java.security.ProtectionDomain;
+
+/**
+ * @author Steve Ebersole
+ */
+public interface EnhancementContext {
+	public boolean isEntityClass(String className);
+	public boolean isCompositeClass(String className);
+}

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

@@ -0,0 +1,246 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.bytecode.enhance.spi;
+
+import javax.persistence.Transient;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+
+import javassist.CannotCompileException;
+import javassist.ClassPool;
+import javassist.CtClass;
+import javassist.CtField;
+import javassist.CtNewMethod;
+import javassist.bytecode.AnnotationsAttribute;
+import javassist.bytecode.ConstPool;
+import javassist.bytecode.FieldInfo;
+import javassist.bytecode.annotation.Annotation;
+
+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.ManagedEntity;
+import org.hibernate.engine.spi.EntityEntry;
+import org.hibernate.internal.CoreMessageLogger;
+
+/**
+ * @author Steve Ebersole
+ * @author Jason Greene
+ */
+public class Enhancer  {
+	private static final CoreMessageLogger log = Logger.getMessageLogger( CoreMessageLogger.class, Enhancer.class.getName() );
+
+	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";
+
+	private final EnhancementContext enhancementContext;
+
+	private final ClassPool classPool;
+	private final CtClass managedEntityCtClass;
+	private final CtClass managedCompositeCtClass;
+	private final CtClass entityEntryCtClass;
+
+	public Enhancer(EnhancementContext enhancementContext) {
+		this.enhancementContext = enhancementContext;
+
+		this.classPool = new ClassPool( false );
+		try {
+			// add ManagedEntity contract
+			this.managedEntityCtClass = classPool.makeClass(
+					ManagedEntity.class.getClassLoader().getResourceAsStream(
+							ManagedEntity.class.getName().replace( '.', '/' ) + ".class"
+					)
+			);
+
+			// add ManagedEntity contract
+			this.managedCompositeCtClass = classPool.makeClass(
+					ManagedComposite.class.getClassLoader().getResourceAsStream(
+							ManagedComposite.class.getName().replace( '.', '/' ) + ".class"
+					)
+			);
+
+			// "add" EntityEntry
+			this.entityEntryCtClass = classPool.makeClass( EntityEntry.class.getName() );
+		}
+		catch (IOException e) {
+			throw new EnhancementException( "Could not prepare Javassist ClassPool", e );
+		}
+	}
+
+	/**
+	 * Performs the enhancement.
+	 *
+	 * @param className The name of the class whose bytecode is being enhanced.
+	 * @param originalBytes The class's original (pre-enhancement) byte code
+	 *
+	 * @return The enhanced bytecode.  Could be the same as the original bytecode if the original was
+	 * already enhanced or we could not enhance it for some reason.
+	 *
+	 * @throws EnhancementException
+	 */
+	public byte[] enhance(String className, byte[] originalBytes) throws EnhancementException {
+		final CtClass managedCtClass;
+		try {
+			managedCtClass = classPool.makeClassIfNew( new ByteArrayInputStream( originalBytes ) );
+		}
+		catch (IOException e) {
+			log.unableToBuildEnhancementMetamodel( className );
+			return originalBytes;
+		}
+
+		enhance( managedCtClass );
+
+		DataOutputStream out = null;
+		try {
+			ByteArrayOutputStream byteStream = new ByteArrayOutputStream();
+			out = new DataOutputStream( byteStream );
+			managedCtClass.toBytecode( out );
+			return byteStream.toByteArray();
+		}
+		catch (Exception e) {
+			log.unableToTransformClass( e.getMessage() );
+			throw new HibernateException( "Unable to transform class: " + e.getMessage() );
+		}
+		finally {
+			try {
+				if ( out != null ) {
+					out.close();
+				}
+			}
+			catch (IOException e) {
+				//swallow
+			}
+		}
+	}
+
+	private void enhance(CtClass managedCtClass) {
+		final String className = managedCtClass.getName();
+		log.debugf( "Enhancing %s", className );
+
+		// can't effectively enhance interfaces
+		if ( managedCtClass.isInterface() ) {
+			log.debug( "skipping enhancement : interface" );
+			return;
+		}
+
+		// skip already enhanced classes
+
+		final String[] interfaceNames = managedCtClass.getClassFile2().getInterfaces();
+		for ( String interfaceName : interfaceNames ) {
+			if ( FieldHandled.class.getName().equals( interfaceName ) ) {
+				log.debug( "skipping enhancement : already enhanced" );
+				return;
+			}
+		}
+
+		if ( enhancementContext.isEntityClass( className ) ) {
+			enhanceAsEntity( managedCtClass );
+		}
+		else if ( enhancementContext.isCompositeClass( className ) ) {
+			enhanceAsComposite( managedCtClass );
+		}
+		else {
+			log.debug( "skipping enhancement : not entity or composite" );
+		}
+	}
+
+	private void enhanceAsEntity(CtClass managedCtClass) {
+		final ConstPool constPool = managedCtClass.getClassFile().getConstPool();
+
+		// add the ManagedEntity interface
+		managedCtClass.addInterface( managedEntityCtClass );
+
+		// add field to hold EntityEntry
+		final CtField entityEntryField;
+		try {
+			entityEntryField = new CtField( entityEntryCtClass, ENTITY_ENTRY_FIELD_NAME, managedCtClass );
+			managedCtClass.addField( entityEntryField );
+		}
+		catch (CannotCompileException e) {
+			throw new EnhancementException(
+					String.format(
+							"Could not enhance entity class [%s] to add field for holding EntityEntry",
+							managedCtClass.getName()
+					),
+					e
+			);
+		}
+
+		// make that new field @Transient
+		AnnotationsAttribute annotationsAttribute = getVisibleAnnotations( entityEntryField.getFieldInfo() );
+		annotationsAttribute.addAnnotation( new Annotation( Transient.class.getName(), constPool ) );
+
+		// add the ManagedEntity#hibernate_getEntityEntry method
+		try {
+			managedCtClass.addMethod(
+					CtNewMethod.getter( ENTITY_ENTRY_GETTER_NAME, entityEntryField )
+			);
+		}
+		catch (CannotCompileException e) {
+			throw new EnhancementException(
+					String.format(
+							"Could not enhance entity class [%s] to add EntityEntry getter",
+							managedCtClass.getName()
+					),
+					e
+			);
+		}
+
+		// add the ManagedEntity#hibernate_setEntityEntry method
+		try {
+			managedCtClass.addMethod(
+					CtNewMethod.setter( ENTITY_ENTRY_SETTER_NAME, entityEntryField )
+			);
+		}
+		catch (CannotCompileException e) {
+			throw new EnhancementException(
+					String.format(
+							"Could not enhance entity class [%s] to add EntityEntry setter",
+							managedCtClass.getName()
+					),
+					e
+			);
+		}
+	}
+
+	private AnnotationsAttribute getVisibleAnnotations(FieldInfo fieldInfo) {
+		AnnotationsAttribute annotationsAttribute = (AnnotationsAttribute) fieldInfo.getAttribute( AnnotationsAttribute.visibleTag );
+		if ( annotationsAttribute == null ) {
+			annotationsAttribute = new AnnotationsAttribute( fieldInfo.getConstPool(), AnnotationsAttribute.visibleTag );
+			fieldInfo.addAttribute( annotationsAttribute );
+		}
+		return annotationsAttribute;
+	}
+
+	private void enhanceAsComposite(CtClass classFile) {
+	}
+
+
+}

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

@@ -0,0 +1,34 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.engine;
+
+/**
+ * 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.
+ *
+ * @author Steve Ebersole
+ */
+public interface Managed {
+}

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

@@ -0,0 +1,32 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.engine;
+
+/**
+ * Specialized {@link Managed} contract for component/embeddable classes.
+ *
+ * @author Steve Ebersole
+ */
+public interface ManagedComposite extends Managed {
+}

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

@@ -0,0 +1,50 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.engine;
+
+import org.hibernate.engine.spi.EntityEntry;
+
+/**
+ * Specialized {@link Managed} contract for entity classes.
+ *
+ * @author Steve Ebersole
+ */
+public interface ManagedEntity extends Managed {
+	/**
+	 * Callback to get any associated EntityEntry.
+	 *
+	 * @return The EntityEntry associated with this entity instance.
+	 *
+	 * @see #hibernate_setEntityEntry
+	 */
+	public EntityEntry hibernate_getEntityEntry();
+
+	/**
+	 * Injects the EntityEntry associated with this entity instance.  The EntityEntry represents state associated
+	 * with the entity in regards to its association with a Hibernate Session.
+	 *
+	 * @param entityEntry The EntityEntry associated with this entity instance.
+	 */
+	public void hibernate_setEntityEntry(EntityEntry entityEntry);
+}

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

@@ -0,0 +1,259 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.engine.internal;
+
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.io.Serializable;
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+import org.jboss.logging.Logger;
+
+import org.hibernate.LockMode;
+import org.hibernate.engine.ManagedEntity;
+import org.hibernate.engine.spi.EntityEntry;
+
+/**
+ * Defines a context for maintaining the relation between an entity associated with the Session ultimately owning this
+ * EntityEntryContext instance and that entity's corresponding EntityEntry.  2 approaches are supported:<ul>
+ *     <li>
+ *         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,
+ *         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>();
+
+	@SuppressWarnings( {"unchecked"})
+	private transient Map.Entry<Object,EntityEntry>[] reentrantSafeEntries = new Map.Entry[0];
+	private transient boolean dirty = false;
+
+	public EntityEntryContext() {
+	}
+
+	/**
+	 * Private constructor used during deserialization
+	 *
+	 * @param backingMap The backing map re-built from the serial stream
+	 */
+	private EntityEntryContext(LinkedHashMap<KeyWrapper, EntityEntry> backingMap) {
+		this.backingMap = backingMap;
+		// mark dirty so we can rebuild the 'reentrantSafeEntries'
+		dirty = true;
+	}
+
+	public void clear() {
+		dirty = true;
+		for ( Map.Entry<KeyWrapper,EntityEntry> mapEntry : backingMap.entrySet() ) {
+			final Object realKey = mapEntry.getKey().getRealKey();
+			if ( ManagedEntity.class.isInstance( realKey ) ) {
+				( (ManagedEntity) realKey ).hibernate_setEntityEntry( null );
+			}
+		}
+		backingMap.clear();
+		reentrantSafeEntries = null;
+	}
+
+	public void downgradeLocks() {
+		// Downgrade locks
+		for ( Map.Entry<KeyWrapper,EntityEntry> mapEntry : backingMap.entrySet() ) {
+			final Object realKey = mapEntry.getKey().getRealKey();
+			final EntityEntry entityEntry = ManagedEntity.class.isInstance( realKey )
+					? ( (ManagedEntity) realKey ).hibernate_getEntityEntry()
+					: mapEntry.getValue();
+			entityEntry.setLockMode( LockMode.NONE );
+		}
+	}
+
+	public boolean hasEntityEntry(Object entity) {
+		return ManagedEntity.class.isInstance( entity )
+				? ( (ManagedEntity) entity ).hibernate_getEntityEntry() == null
+				: backingMap.containsKey( new KeyWrapper<Object>( entity ) );
+	}
+
+	public EntityEntry getEntityEntry(Object entity) {
+		return ManagedEntity.class.isInstance( entity )
+				? ( (ManagedEntity) entity ).hibernate_getEntityEntry()
+				: backingMap.get( new KeyWrapper<Object>( entity ) );
+	}
+
+	public EntityEntry removeEntityEntry(Object entity) {
+		dirty = true;
+		if ( ManagedEntity.class.isInstance( entity ) ) {
+			backingMap.remove( new KeyWrapper<Object>( entity ) );
+			final EntityEntry entityEntry = ( (ManagedEntity) entity ).hibernate_getEntityEntry();
+			( (ManagedEntity) entity ).hibernate_setEntityEntry( null );
+			return entityEntry;
+		}
+		else {
+			return backingMap.remove( new KeyWrapper<Object>( entity ) );
+		}
+	}
+
+	public void addEntityEntry(Object entity, EntityEntry entityEntry) {
+		dirty = true;
+		if ( ManagedEntity.class.isInstance( entity ) ) {
+			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() ];
+			int i = 0;
+			for ( Map.Entry<KeyWrapper,EntityEntry> mapEntry : backingMap.entrySet() ) {
+				final Object entity = mapEntry.getKey().getRealKey();
+				final EntityEntry entityEntry = ManagedEntity.class.isInstance( entity )
+						? ( (ManagedEntity) entity ).hibernate_getEntityEntry()
+						: mapEntry.getValue();
+				reentrantSafeEntries[i++] = new MapEntryImpl( entity, entityEntry );
+			}
+			dirty = false;
+		}
+		return reentrantSafeEntries;
+	}
+
+	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() ) {
+			oos.writeObject( mapEntry.getKey().getRealKey() );
+			mapEntry.getValue().serialize( oos );
+		}
+	}
+
+	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 );
+		for ( int i = 0; i < count; i++ ) {
+			final Object entity = ois.readObject();
+			final EntityEntry entry = EntityEntry.deserialize( ois, rtn );
+			backingMap.put( new KeyWrapper<Object>( entity ), entry );
+		}
+		return new EntityEntryContext( backingMap );
+	}
+
+	/**
+	 * @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;
+				}
+			}
+			return hash;
+		}
+
+		@Override
+		public String toString() {
+			return realKey.toString();
+		}
+
+		public K getRealKey() {
+			return realKey;
+		}
+	}
+
+	private static class MapEntryImpl implements Map.Entry<Object,EntityEntry> {
+		private final Object key;
+		private EntityEntry value;
+
+		private MapEntryImpl(Object key, EntityEntry value) {
+			this.key = key;
+			this.value = value;
+		}
+
+		@Override
+		public Object getKey() {
+			return key;
+		}
+
+		@Override
+		public EntityEntry getValue() {
+			return value;
+		}
+
+		@Override
+		public EntityEntry setValue(EntityEntry value) {
+			final EntityEntry old = this.value;
+			this.value = value;
+			return value;
+		}
+	}
+}

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

@@ -76,7 +76,6 @@ import org.hibernate.persister.entity.EntityPersister;
 import org.hibernate.pretty.MessageHelper;
 import org.hibernate.proxy.HibernateProxy;
 import org.hibernate.proxy.LazyInitializer;
-import org.hibernate.sql.Select;
 import org.hibernate.tuple.ElementWrapper;
 import org.hibernate.type.CollectionType;
 
@@ -98,7 +97,7 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	public static final Object NO_ROW = new MarkerObject( "NO_ROW" );
 
-	private static final int INIT_COLL_SIZE = 8;
+	public static final int INIT_COLL_SIZE = 8;
 
 	private SessionImplementor session;
 
@@ -108,8 +107,8 @@ public class StatefulPersistenceContext implements PersistenceContext {
 	// Loaded entity instances, by EntityUniqueKey
 	private Map<EntityUniqueKey, Object> entitiesByUniqueKey;
 
-	// Identity map of EntityEntry instances, by the entity instance
+	private EntityEntryContext entityEntryContext;
-	private Map<Object,EntityEntry> entityEntries;
+//	private Map<Object,EntityEntry> entityEntries;
 
 	// Entity proxies, by EntityKey
 	private Map<EntityKey, Object> proxiesByKey;
@@ -156,7 +155,6 @@ public class StatefulPersistenceContext implements PersistenceContext {
 	private BatchFetchQueue batchFetchQueue;
 
 
-
 	/**
 	 * Constructs a PersistentContext, bound to the given session.
 	 *
@@ -171,9 +169,10 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		proxiesByKey = new ConcurrentReferenceHashMap<EntityKey, Object>( INIT_COLL_SIZE, .75f, 1, ConcurrentReferenceHashMap.ReferenceType.STRONG, ConcurrentReferenceHashMap.ReferenceType.WEAK, null );
 		entitySnapshotsByKey = new HashMap<EntityKey, Object>( INIT_COLL_SIZE );
 
-		entityEntries = IdentityMap.instantiateSequenced( INIT_COLL_SIZE );
+		entityEntryContext = new EntityEntryContext();
+//		entityEntries = IdentityMap.instantiateSequenced( INIT_COLL_SIZE );
 		collectionEntries = IdentityMap.instantiateSequenced( INIT_COLL_SIZE );
-		parentsByChild = IdentityMap.instantiateSequenced( INIT_COLL_SIZE );
+		parentsByChild = new IdentityHashMap<Object,Object>( INIT_COLL_SIZE );
 
 		collectionsByKey = new HashMap<CollectionKey, PersistentCollection>( INIT_COLL_SIZE );
 		arrayHolders = new IdentityHashMap<Object, PersistentCollection>( INIT_COLL_SIZE );
@@ -242,7 +241,8 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		arrayHolders.clear();
 		entitiesByKey.clear();
 		entitiesByUniqueKey.clear();
-		entityEntries.clear();
+		entityEntryContext.clear();
+//		entityEntries.clear();
 		parentsByChild.clear();
 		entitySnapshotsByKey.clear();
 		collectionsByKey.clear();
@@ -293,10 +293,11 @@ public class StatefulPersistenceContext implements PersistenceContext {
 	@Override
 	public void afterTransactionCompletion() {
 		cleanUpInsertedKeysAfterTransaction();
-		// Downgrade locks
+		entityEntryContext.downgradeLocks();
-		for ( EntityEntry o : entityEntries.values() ) {
+//		// Downgrade locks
-			o.setLockMode( LockMode.NONE );
+//		for ( EntityEntry o : entityEntries.values() ) {
-		}
+//			o.setLockMode( LockMode.NONE );
+//		}
 	}
 
 	/**
@@ -455,7 +456,8 @@ public class StatefulPersistenceContext implements PersistenceContext {
 	 */
 	@Override
 	public EntityEntry getEntry(Object entity) {
-		return entityEntries.get(entity);
+		return entityEntryContext.getEntityEntry( entity );
+//		return entityEntries.get(entity);
 	}
 
 	/**
@@ -463,7 +465,8 @@ public class StatefulPersistenceContext implements PersistenceContext {
 	 */
 	@Override
 	public EntityEntry removeEntry(Object entity) {
-		return entityEntries.remove(entity);
+		return entityEntryContext.removeEntityEntry( entity );
+//		return entityEntries.remove(entity);
 	}
 
 	/**
@@ -471,7 +474,8 @@ public class StatefulPersistenceContext implements PersistenceContext {
 	 */
 	@Override
 	public boolean isEntryFor(Object entity) {
-		return entityEntries.containsKey(entity);
+		return entityEntryContext.hasEntityEntry( entity );
+//		return entityEntries.containsKey(entity);
 	}
 
 	/**
@@ -547,7 +551,9 @@ public class StatefulPersistenceContext implements PersistenceContext {
 				lazyPropertiesAreUnfetched,
 				this
 		);
-		entityEntries.put(entity, e);
+
+		entityEntryContext.addEntityEntry( entity, e );
+//		entityEntries.put(entity, e);
 
 		setHasNonReadOnlyEnties(status);
 		return e;
@@ -1148,7 +1154,8 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	@Override
 	public Map getEntityEntries() {
-		return entityEntries;
+		return entityEntryContext.getEntityEntryMap();
+//		return entityEntries;
 	}
 
 	@Override
@@ -1227,6 +1234,11 @@ public class StatefulPersistenceContext implements PersistenceContext {
 				.toString();
 	}
 
+	@Override
+	public Entry<Object,EntityEntry>[] reentrantSafeEntityEntries() {
+		return entityEntryContext.reentrantSafeEntityEntries();
+	}
+
 	/**
 	 * Search <tt>this</tt> persistence context for an associated entity instance which is considered the "owner" of
 	 * the given <tt>childEntity</tt>, and return that owner's id value.  This is performed in the scenario of a
@@ -1257,7 +1269,8 @@ public class StatefulPersistenceContext implements PersistenceContext {
 	    // try cache lookup first
 		Object parent = parentsByChild.get( childEntity );
 		if ( parent != null ) {
-			final EntityEntry entityEntry = entityEntries.get( parent );
+			final EntityEntry entityEntry = entityEntryContext.getEntityEntry( parent );
+//			final EntityEntry entityEntry = entityEntries.get( parent );
 			//there maybe more than one parent, filter by type
 			if ( 	persister.isSubclassEntityName(entityEntry.getEntityName() )
 					&& isFoundInParent( propertyName, childEntity, persister, collectionPersister, parent ) ) {
@@ -1270,7 +1283,8 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 		//not found in case, proceed
 		// iterate all the entities currently associated with the persistence context.
-		for ( Entry<Object,EntityEntry> me : IdentityMap.concurrentEntries( entityEntries ) ) {
+		for ( Entry<Object,EntityEntry> me : reentrantSafeEntityEntries() ) {
+//		for ( Entry<Object,EntityEntry> me : IdentityMap.concurrentEntries( entityEntries ) ) {
 			final EntityEntry entityEntry = me.getValue();
 			// does this entity entry pertain to the entity persister in which we are interested (owner)?
 			if ( persister.isSubclassEntityName( entityEntry.getEntityName() ) ) {
@@ -1370,7 +1384,8 @@ public class StatefulPersistenceContext implements PersistenceContext {
 	    // try cache lookup first
 	    Object parent = parentsByChild.get(childEntity);
 		if (parent != null) {
-			final EntityEntry entityEntry = entityEntries.get(parent);
+			final EntityEntry entityEntry = entityEntryContext.getEntityEntry( parent );
+//			final EntityEntry entityEntry = entityEntries.get(parent);
 			//there maybe more than one parent, filter by type
 			if ( persister.isSubclassEntityName( entityEntry.getEntityName() ) ) {
 				Object index = getIndexInParent(property, childEntity, persister, cp, parent);
@@ -1392,7 +1407,7 @@ public class StatefulPersistenceContext implements PersistenceContext {
 		}
 
 		//Not found in cache, proceed
-		for ( Entry<Object, EntityEntry> me : IdentityMap.concurrentEntries( entityEntries ) ) {
+		for ( Entry<Object, EntityEntry> me : reentrantSafeEntityEntries() ) {
 			EntityEntry ee = me.getValue();
 			if ( persister.isSubclassEntityName( ee.getEntityName() ) ) {
 				Object instance = me.getKey();
@@ -1516,8 +1531,8 @@ public class StatefulPersistenceContext implements PersistenceContext {
 
 	@Override
 	public void replaceDelayedEntityIdentityInsertKeys(EntityKey oldKey, Serializable generatedId) {
-		Object entity = entitiesByKey.remove( oldKey );
+		final Object entity = entitiesByKey.remove( oldKey );
-		EntityEntry oldEntry = entityEntries.remove( entity );
+		final EntityEntry oldEntry = entityEntryContext.removeEntityEntry( entity );
 		parentsByChild.clear();
 
 		final EntityKey newKey = session.generateEntityKey( generatedId, oldEntry.getPersister() );
@@ -1587,14 +1602,15 @@ public class StatefulPersistenceContext implements PersistenceContext {
 			oos.writeObject( entry.getValue() );
 		}
 
-		oos.writeInt( entityEntries.size() );
+		entityEntryContext.serialize( oos );
-		if ( tracing ) LOG.trace("Starting serialization of [" + entityEntries.size() + "] entityEntries entries");
+//		oos.writeInt( entityEntries.size() );
-		itr = entityEntries.entrySet().iterator();
+//		if ( tracing ) LOG.trace("Starting serialization of [" + entityEntries.size() + "] entityEntries entries");
-		while ( itr.hasNext() ) {
+//		itr = entityEntries.entrySet().iterator();
-			Map.Entry entry = ( Map.Entry ) itr.next();
+//		while ( itr.hasNext() ) {
-			oos.writeObject( entry.getKey() );
+//			Map.Entry entry = ( Map.Entry ) itr.next();
-			( ( EntityEntry ) entry.getValue() ).serialize( oos );
+//			oos.writeObject( entry.getKey() );
-		}
+//			( ( EntityEntry ) entry.getValue() ).serialize( oos );
+//		}
 
 		oos.writeInt( collectionsByKey.size() );
 		if ( tracing ) LOG.trace("Starting serialization of [" + collectionsByKey.size() + "] collectionsByKey entries");
@@ -1691,14 +1707,15 @@ public class StatefulPersistenceContext implements PersistenceContext {
 				rtn.entitySnapshotsByKey.put( EntityKey.deserialize( ois, session ), ois.readObject() );
 			}
 
-			count = ois.readInt();
+			rtn.entityEntryContext = EntityEntryContext.deserialize( ois, rtn );
-			if ( tracing ) LOG.trace("Starting deserialization of [" + count + "] entityEntries entries");
+//			count = ois.readInt();
-			rtn.entityEntries = IdentityMap.instantiateSequenced( count < INIT_COLL_SIZE ? INIT_COLL_SIZE : count );
+//			if ( tracing ) LOG.trace("Starting deserialization of [" + count + "] entityEntries entries");
-			for ( int i = 0; i < count; i++ ) {
+//			rtn.entityEntries = IdentityMap.instantiateSequenced( count < INIT_COLL_SIZE ? INIT_COLL_SIZE : count );
-				Object entity = ois.readObject();
+//			for ( int i = 0; i < count; i++ ) {
-				EntityEntry entry = EntityEntry.deserialize( ois, rtn );
+//				Object entity = ois.readObject();
-				rtn.entityEntries.put( entity, entry );
+//				EntityEntry entry = EntityEntry.deserialize( ois, rtn );
-			}
+//				rtn.entityEntries.put( entity, entry );
+//			}
 
 			count = ois.readInt();
 			if ( tracing ) LOG.trace("Starting deserialization of [" + count + "] collectionsByKey entries");

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

@@ -481,9 +481,14 @@ public interface PersistenceContext {
 	 */
 	public Map getEntitiesByKey();
 
+	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
 	public Map getEntityEntries();
 
 	/**

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

@@ -144,7 +144,8 @@ public abstract class AbstractFlushingEventListener implements Serializable {
 
 		final Object anything = getAnything();
 		//safe from concurrent modification because of how concurrentEntries() is implemented on IdentityMap
-		for ( Map.Entry me : IdentityMap.concurrentEntries( persistenceContext.getEntityEntries() ) ) {
+		for ( Map.Entry<Object,EntityEntry> me : persistenceContext.reentrantSafeEntityEntries() ) {
+//		for ( Map.Entry me : IdentityMap.concurrentEntries( persistenceContext.getEntityEntries() ) ) {
 			EntityEntry entry = (EntityEntry) me.getValue();
 			Status status = entry.getStatus();
 			if ( status == Status.MANAGED || status == Status.SAVING || status == Status.READ_ONLY ) {
@@ -212,7 +213,8 @@ public abstract class AbstractFlushingEventListener implements Serializable {
 		// So this needs to be safe from concurrent modification problems.
 		// It is safe because of how IdentityMap implements entrySet()
 
-		for ( Map.Entry me : IdentityMap.concurrentEntries( persistenceContext.getEntityEntries() ) ) {
+		for ( Map.Entry<Object,EntityEntry> me : persistenceContext.reentrantSafeEntityEntries() ) {
+//		for ( Map.Entry me : IdentityMap.concurrentEntries( persistenceContext.getEntityEntries() ) ) {
 
 			// Update the status of the object and if necessary, schedule an update
 

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

@@ -0,0 +1,150 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.tool.enhance;
+
+import javax.persistence.Entity;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.FileNotFoundException;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+
+import javassist.ClassPool;
+import javassist.CtClass;
+import javassist.NotFoundException;
+import org.apache.tools.ant.BuildException;
+import org.apache.tools.ant.DirectoryScanner;
+import org.apache.tools.ant.Project;
+import org.apache.tools.ant.Task;
+import org.apache.tools.ant.types.FileSet;
+
+import org.hibernate.bytecode.enhance.spi.EnhancementContext;
+import org.hibernate.bytecode.enhance.spi.Enhancer;
+
+/**
+ * Ant task for performing build-time enhancement of entities and component/embeddable classes.
+ * <p/>
+ * IMPL NOTE : currently makes numerous assumptions, the most "horrific" being that all entities are
+ * annotated @Entity which precludes {@code hbm.xml} mappings as well as complete {@code orm.xml} mappings.  This is
+ * just a PoC though...
+ *
+ * @author Steve Ebersole
+ *
+ * @see org.hibernate.engine.Managed
+ */
+public class EnhancementTask extends Task {
+	private List<FileSet> filesets = new ArrayList<FileSet>();
+
+	public void addFileset(FileSet set) {
+		this.filesets.add( set );
+	}
+
+	@Override
+	public void execute() throws BuildException {
+		EnhancementContext enhancementContext = new EnhancementContext() {
+			@Override
+			public boolean isEntityClass(String className) {
+				// currently we only call enhance on the classes with @Entity, so here we always return true
+				return true;
+			}
+
+			@Override
+			public boolean isCompositeClass(String className) {
+				return false;
+			}
+		};
+
+		// we use the CtClass stuff here just as a simple vehicle for obtaining low level information about
+		// the class(es) contained in a file while still maintaining easy access to the underlying byte[]
+		//
+		// Enhancer also builds CtClass instances.  Might make sense to share these (ClassPool).
+		final ClassPool classPool = new ClassPool( false );
+		final List<CtClass> ctClassList = collectCtClasses( classPool );
+
+		final Enhancer enhancer = new Enhancer( enhancementContext );
+		for ( CtClass ctClass : ctClassList ) {
+			try {
+				enhancer.enhance( ctClass.getName(), ctClass.toBytecode() );
+			}
+			catch (Exception e) {
+				log(
+						"Unable to enhance class : " + ctClass.getName(),
+						e,
+						Project.MSG_WARN
+				);
+			}
+
+		}
+	}
+
+	private List<CtClass> collectCtClasses(ClassPool classPool) {
+		final List<CtClass> ctClassList = new ArrayList<CtClass>();
+
+		final Project project = getProject();
+		for ( FileSet fileSet : filesets ) {
+			final File fileSetBaseDir = fileSet.getDir( project );
+			final DirectoryScanner directoryScanner = fileSet.getDirectoryScanner( project );
+			for ( String relativeIncludedFileName : directoryScanner.getIncludedFiles() ) {
+				final File javaClassFile = new File( fileSetBaseDir, relativeIncludedFileName );
+				if ( ! javaClassFile.exists() ) {
+					continue;
+				}
+				try {
+					final CtClass ctClass = classPool.makeClass( new FileInputStream( javaClassFile ) );
+					collectCtClasses( ctClassList, ctClass );
+				}
+				catch (FileNotFoundException ignore) {
+					// should not ever happen because of explicit check above
+				}
+				catch (IOException e) {
+					throw new BuildException(
+							String.format(
+									"Error processing included file [%s : %s]",
+									fileSetBaseDir.getAbsolutePath(),
+									relativeIncludedFileName
+							),
+							e
+					);
+				}
+			}
+		}
+
+		return ctClassList;
+	}
+
+	private void collectCtClasses(List<CtClass> ctClassList, CtClass ctClass) {
+		if ( ctClass.hasAnnotation( Entity.class ) ) {
+			ctClassList.add( ctClass );
+		}
+
+		try {
+			for ( CtClass nestedCtClass : ctClass.getNestedClasses() ) {
+				collectCtClasses( ctClassList, nestedCtClass );
+			}
+		}
+		catch (NotFoundException ignore) {
+		}
+	}
+}

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

@@ -0,0 +1,127 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.test.bytecode.enhancement;
+
+import javax.persistence.Entity;
+import javax.persistence.Id;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.security.ProtectionDomain;
+import java.util.Arrays;
+
+import javassist.CannotCompileException;
+import javassist.ClassPool;
+import javassist.CtClass;
+import javassist.LoaderClassPath;
+import javassist.NotFoundException;
+
+import org.hibernate.EntityMode;
+import org.hibernate.LockMode;
+import org.hibernate.bytecode.enhance.spi.EnhancementContext;
+import org.hibernate.bytecode.enhance.spi.Enhancer;
+import org.hibernate.engine.spi.EntityEntry;
+import org.hibernate.engine.spi.Status;
+
+import org.junit.Test;
+
+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.fail;
+
+/**
+ * @author Steve Ebersole
+ */
+public class EnhancerTest extends BaseUnitTestCase {
+	@Test
+	public void testEnhancement()
+			throws IOException, CannotCompileException, NotFoundException,
+			NoSuchMethodException, IllegalAccessException, InstantiationException, InvocationTargetException {
+		EnhancementContext enhancementContext = new EnhancementContext() {
+			@Override
+			public boolean isEntityClass(String className) {
+				return true;
+			}
+
+			@Override
+			public boolean isCompositeClass(String className) {
+				return false;
+			}
+		};
+		Enhancer enhancer = new Enhancer( enhancementContext );
+		CtClass anEntityCtClass = generateCtClassForAnEntity();
+		byte[] original = anEntityCtClass.toBytecode();
+		byte[] enhanced = enhancer.enhance( anEntityCtClass.getName(), original );
+		assertFalse( "entity was not enhanced", Arrays.equals( original, enhanced ) );
+
+		ClassLoader cl = new ClassLoader() { };
+		ClassPool cp2 = new ClassPool( false );
+		cp2.appendClassPath( new LoaderClassPath( cl ) );
+		CtClass enhancedCtClass = cp2.makeClass( new ByteArrayInputStream( enhanced ) );
+		Class simpleEntityClass = enhancedCtClass.toClass( cl, this.getClass().getProtectionDomain() );
+		Object simpleEntityInstance = simpleEntityClass.newInstance();
+
+		// call the new methods
+		Method setter = simpleEntityClass.getMethod( Enhancer.ENTITY_ENTRY_SETTER_NAME, EntityEntry.class );
+		Method getter = simpleEntityClass.getMethod( Enhancer.ENTITY_ENTRY_GETTER_NAME );
+		assertNull( getter.invoke( simpleEntityInstance ) );
+		setter.invoke( simpleEntityInstance, makeEntityEntry() );
+		assertNotNull( getter.invoke( simpleEntityInstance ) );
+		setter.invoke( simpleEntityInstance, new Object[] { null } );
+		assertNull( getter.invoke( simpleEntityInstance ) );
+	}
+
+	private CtClass generateCtClassForAnEntity() throws IOException, NotFoundException {
+		ClassPool cp = new ClassPool( false );
+		return cp.makeClass(
+				getClass().getClassLoader().getResourceAsStream(
+						SimpleEntity.class.getName().replace( '.', '/' ) + ".class"
+				)
+		);
+	}
+
+	private EntityEntry makeEntityEntry() {
+		return new EntityEntry(
+				Status.MANAGED,
+				null,
+				null,
+				new Long(1),
+				null,
+				LockMode.NONE,
+				false,
+				null,
+				EntityMode.POJO,
+				null,
+				false,
+				false,
+				null
+		);
+	}
+
+}

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

@@ -0,0 +1,63 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.test.bytecode.enhancement;
+
+import org.hibernate.Session;
+
+import org.junit.Test;
+
+import org.hibernate.testing.junit4.BaseCoreFunctionalTestCase;
+
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+
+/**
+ * @author Steve Ebersole
+ */
+public class MostBasicEnhancementTest extends BaseCoreFunctionalTestCase {
+	@Override
+	protected Class<?>[] getAnnotatedClasses() {
+		return new Class[] { MyEntity.class };
+	}
+
+	@Test
+	public void testIt() {
+		Session s = openSession();
+		s.beginTransaction();
+		s.save( new MyEntity( 1L ) );
+		s.getTransaction().commit();
+		s.close();
+
+		s = openSession();
+		s.beginTransaction();
+		MyEntity myEntity = (MyEntity) s.load( MyEntity.class, 1L );
+		assertNotNull( myEntity.hibernate_getEntityEntry() );
+		s.getTransaction().commit();
+		s.close();
+
+		assertNull( myEntity.hibernate_getEntityEntry() );
+	}
+
+
+}

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

@@ -0,0 +1,77 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.test.bytecode.enhancement;
+
+import javax.persistence.Entity;
+import javax.persistence.Id;
+import javax.persistence.Transient;
+
+import org.hibernate.engine.ManagedEntity;
+import org.hibernate.engine.spi.EntityEntry;
+
+/**
+ * @author Steve Ebersole
+ */
+@Entity
+public class MyEntity implements ManagedEntity {
+	@Transient
+	private transient EntityEntry entityEntry;
+
+	private Long id;
+	private String name;
+
+	public MyEntity() {
+	}
+
+	public MyEntity(Long id) {
+		this.id = id;
+	}
+
+	@Id
+	public Long getId() {
+		return id;
+	}
+
+	public void setId(Long id) {
+		this.id = id;
+	}
+
+	public String getName() {
+		return name;
+	}
+
+	public void setName(String name) {
+		this.name = name;
+	}
+
+	@Override
+	public EntityEntry hibernate_getEntityEntry() {
+		return entityEntry;
+	}
+
+	@Override
+	public void hibernate_setEntityEntry(EntityEntry entityEntry) {
+		this.entityEntry = entityEntry;
+	}
+}

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

@@ -0,0 +1,53 @@
+/*
+ * Hibernate, Relational Persistence for Idiomatic Java
+ *
+ * Copyright (c) 2012, Red Hat Inc. or third-party contributors as
+ * indicated by the @author tags or express copyright attribution
+ * statements applied by the authors.  All third-party contributions are
+ * distributed under license by Red Hat Inc.
+ *
+ * This copyrighted material is made available to anyone wishing to use, modify,
+ * copy, or redistribute it subject to the terms and conditions of the GNU
+ * Lesser General Public License, as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
+ * or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
+ * for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this distribution; if not, write to:
+ * Free Software Foundation, Inc.
+ * 51 Franklin Street, Fifth Floor
+ * Boston, MA  02110-1301  USA
+ */
+package org.hibernate.test.bytecode.enhancement;
+
+import javax.persistence.Entity;
+import javax.persistence.Id;
+
+/**
+ * @author Steve Ebersole
+ */
+@Entity
+public class SimpleEntity {
+	private Long id;
+	private String name;
+
+	@Id
+	public Long getId() {
+		return id;
+	}
+
+	public void setId(Long id) {
+		this.id = id;
+	}
+
+	public String getName() {
+		return name;
+	}
+
+	public void setName(String name) {
+		this.name = name;
+	}
+}

hibernate-core/src/test/java/org/hibernate/test/legacy/FumTest.java

@@ -780,8 +780,8 @@ public class FumTest extends LegacyTestCase {
 		s.beginTransaction();
 
 		simple = (Simple) s.get( Simple.class, Long.valueOf(10) );
-		assertTrue("Not same parent instances", check.getName().equals( simple.getName() ) );
+		assertEquals( "Not same parent instances", check.getName(), simple.getName() );
-		assertTrue("Not same child instances", check.getOther().getName().equals( other.getName() ) );
+		assertEquals( "Not same child instances", check.getOther().getName(), other.getName() );
 
 		simple.setName("My updated name");
 
@@ -803,8 +803,8 @@ public class FumTest extends LegacyTestCase {
 		s.beginTransaction();
 
 		simple = (Simple) s.get( Simple.class, Long.valueOf( 10 ) );
-		assertTrue("Not same parent instances", check.getName().equals( simple.getName() ) );
+		assertEquals( "Not same parent instances", check.getName(), simple.getName() );
-		assertTrue("Not same child instances", check.getOther().getName().equals( other.getName() ) );
+		assertEquals( "Not same child instances", check.getOther().getName(), other.getName() );
 
 		// Now, lets delete across serialization...
 		s.delete(simple);
@@ -856,7 +856,7 @@ public class FumTest extends LegacyTestCase {
 		s.beginTransaction();
 		fum = (Fum) s.load( Fum.class, fum.getId() );
 
-		assertTrue("the Fum.friends did not get saved", fum.getFriends().size() == 2);
+		assertEquals( "the Fum.friends did not get saved", 2, fum.getFriends().size() );
 
 		fum.setFriends(null);
 		s.getTransaction().commit();

hibernate-entitymanager/src/main/java/org/hibernate/ejb/event/EJB3PostUpdateEventListener.java

@@ -68,8 +68,7 @@ public class EJB3PostUpdateEventListener
 	}
 
 	private void handlePostUpdate(Object entity, EventSource source) {
-		EntityEntry entry = (EntityEntry) source.getPersistenceContext()
+		EntityEntry entry = (EntityEntry) source.getPersistenceContext().getEntry( entity );
-				.getEntityEntries().get(entity);
 		// mimic the preUpdate filter
 		if ( Status.DELETED != entry.getStatus()) {
 			callbackHandler.postUpdate(entity);