Calling real instance when intercepting method call to proxy without extending BasicLazyInitializer

Hi again,

In the autofetch project, we’re trying to call the target object when intercepting a methodcall on the proxy, but at the moment we keep on calling on the proxy object, causing a StackOverflowException. We use the following classes:

public class AutofetchLazyInitializer extends BasicLazyInitializer implements ProxyConfiguration.Interceptor {

	private static final CoreMessageLogger LOG = CoreLogging.messageLogger( AutofetchLazyInitializer.class );

	private final EntityTracker entityTracker;

	private final Class[] interfaces;

	private AutofetchLazyInitializer(
			String entityName,
			Class<?> persistentClass,
			Class<?>[] interfaces,
			Serializable id,
			Method getIdentifierMethod,
			Method setIdentifierMethod,
			CompositeType componentIdType,
			SharedSessionContractImplementor session,
			Set<Property> persistentProperties,
			boolean overridesEquals) {

		super(
				entityName,
				persistentClass,
				id,
				getIdentifierMethod,
				setIdentifierMethod,
				componentIdType,
				Objects.requireNonNull( session, "Hibernate session cannot be null" ),
				overridesEquals
		);

		this.interfaces = interfaces;

		AutofetchService autofetchService = session.getFactory()
				.getServiceRegistry()
				.getService( AutofetchService.class );
		this.entityTracker = new EntityTracker( persistentProperties, autofetchService.getExtentManager() );
	}

	@Override
	public Object intercept(Object proxy, Method thisMethod, Object[] args) throws Throwable {
		Object result = this.invoke( thisMethod, args, proxy );
		if ( result == INVOKE_IMPLEMENTATION ) {
			final String methodName = thisMethod.getName();
			if ( args.length == 0 ) {
				switch ( methodName ) {
					case "enableTracking":
						return handleEnableTracking();
					case "disableTracking":
						return handleDisableTracking();
					case "isAccessed":
						return entityTracker.isAccessed();
					default:
						break;
				}
			}
			else if ( args.length == 1 ) {
				if ( methodName.equals( "addTracker" ) && thisMethod.getParameterTypes()[0].equals(
						Statistics.class ) ) {
					return handleAddTracked( args[0] );
				}
				else if ( methodName.equals( "addTrackers" ) && thisMethod.getParameterTypes()[0].equals(
						Set.class ) ) {
					return handleAddTrackers( args[0] );
				}
				else if ( methodName
						.equals( "removeTracker" ) && thisMethod.getParameterTypes()[0].equals( Statistics.class ) ) {
					entityTracker.removeTracker( (Statistics) args[0] );
					return handleRemoveTracker( args );
				}
				else if ( methodName
						.equals( "extendProfile" ) && thisMethod.getParameterTypes()[0].equals( Statistics.class ) ) {
					return extendProfile( args );
				}
			}

			Object target = getImplementation();
			final Object returnValue;
			try {
				if ( ReflectHelper.isPublic( persistentClass, thisMethod ) ) {
					if ( !thisMethod.getDeclaringClass().isInstance( target ) ) {
						throw new ClassCastException(
								target.getClass().getName()
										+ " incompatible with "
										+ thisMethod.getDeclaringClass().getName()
						);
					}
					returnValue = thisMethod.invoke( target, args );
				}
				else {
					thisMethod.setAccessible( true );
					returnValue = thisMethod.invoke( target, args );
				}

				if ( returnValue == target ) {
					if ( returnValue.getClass().isInstance( proxy ) ) {
						return proxy;
					}
					else {
						LOG.narrowingProxy( returnValue.getClass() );
					}
				}
				return returnValue;
			}
			catch (InvocationTargetException ite) {
				throw ite.getTargetException();
			}
		}
		else {
			return result;
		}
	}

	@Override
	protected Object serializableProxy() {
		return new AutofetchSerializableProxy(
				getEntityName(),
				this.persistentClass,
				this.interfaces,
				getIdentifier(),
				( isReadOnlySettingAvailable() ?
						Boolean.valueOf( isReadOnly() ) :
						isReadOnlyBeforeAttachedToSession() ),
				this.getIdentifierMethod,
				this.setIdentifierMethod,
				this.componentIdType,
				this.entityTracker.getPersistentProperties()
		);
	}

	static HibernateProxy getProxy(
			final String entityName,
			final Class<?> persistentClass,
			final Class<?>[] interfaces,
			final Method getIdentifierMethod,
			final Method setIdentifierMethod,
			final CompositeType componentIdType,
			final Serializable id,
			final SharedSessionContractImplementor session,
			final Set<Property> persistentProperties) throws HibernateException {

		// note: interfaces is assumed to already contain HibernateProxy.class
		final AutofetchLazyInitializer lazyInitializer = new AutofetchLazyInitializer(
				entityName,
				persistentClass,
				interfaces,
				id,
				getIdentifierMethod,
				setIdentifierMethod,
				componentIdType,
				session,
				persistentProperties,
				ReflectHelper.overridesEquals( persistentClass )
		);

		final HibernateProxy proxy;
		try {
			ProxyFactory proxyFactory = Environment.getBytecodeProvider().getProxyFactoryFactory()
					.buildProxyFactory( session.getFactory() );
			proxyFactory.postInstantiate(
					entityName,
					persistentClass,
					new HashSet<>( Arrays.asList( interfaces ) ),
					getIdentifierMethod,
					setIdentifierMethod,
					componentIdType
			);
			proxy = proxyFactory.getProxy( id, session );
			( (ProxyConfiguration) proxy ).$$_hibernate_set_interceptor( lazyInitializer );
		}
		catch (Exception e) {
			String msg = LOG.bytecodeEnhancementFailed( persistentClass.getName() );
			LOG.error( msg );
			throw new HibernateException( msg, e );
		}

		( (ProxyConfiguration) proxy ).$$_hibernate_set_interceptor( lazyInitializer );

		return proxy;
	}

	private Object handleDisableTracking() {
		boolean oldValue = entityTracker.isTracking();
		this.entityTracker.setTracking( false );
		if ( !isUninitialized() ) {
			Object o = getImplementation();
			if ( o instanceof Trackable ) {
				Trackable entity = (Trackable) o;
				entity.disableTracking();
			}
		}

		return oldValue;
	}

	private Object handleEnableTracking() {
		boolean oldValue = this.entityTracker.isTracking();
		this.entityTracker.setTracking( true );

		if ( !isUninitialized() ) {
			Object o = getImplementation();
			if ( o instanceof Trackable ) {
				Trackable entity = (Trackable) o;
				entity.enableTracking();
			}
		}

		return oldValue;
	}

	private Object extendProfile(Object[] params) {
		if ( !isUninitialized() ) {
			Object o = getImplementation();
			if ( o instanceof TrackableEntity ) {
				TrackableEntity entity = (TrackableEntity) o;
				entity.extendProfile( (Statistics) params[0] );
			}
		}
		else {
			throw new IllegalStateException( "Can't call extendProfile on unloaded self." );
		}

		return null;
	}

	private Object handleRemoveTracker(Object[] params) {
		if ( !isUninitialized() ) {
			Object o = getImplementation();
			if ( o instanceof Trackable ) {
				Trackable entity = (Trackable) o;
				entity.removeTracker( (Statistics) params[0] );
			}
		}
		return null;
	}

	@SuppressWarnings("unchecked")
	private Object handleAddTrackers(Object param) {
		Set<Statistics> newTrackers = (Set<Statistics>) param;
		this.entityTracker.addTrackers( newTrackers );
		if ( !isUninitialized() ) {
			Object o = getImplementation();
			if ( o instanceof Trackable ) {
				Trackable entity = (Trackable) o;
				entity.addTrackers( newTrackers );
			}
		}

		return null;
	}

	private Object handleAddTracked(Object param) {
		this.entityTracker.addTracker( (Statistics) param );
		if ( !isUninitialized() ) {
			Object o = getImplementation();
			if ( o instanceof Trackable ) {
				Trackable entity = (Trackable) o;
				entity.addTracker( (Statistics) param );
			}
		}

		return null;
	}
}
public class AutofetchProxyFactory implements ProxyFactory {

	private static final CoreMessageLogger LOG = CoreLogging.messageLogger( AutofetchProxyFactory.class );

	private static final Class<?>[] NO_CLASSES = new Class<?>[0];

	private Class<?> proxiedClass;
	private String entityName;
	private Class<?>[] interfaces;
	private Method getIdentifierMethod;
	private Method setIdentifierMethod;
	private CompositeType componentIdType;

	private Set<org.autofetch.hibernate.Property> persistentProperties;

	public AutofetchProxyFactory(PersistentClass persistentClass) {
		this.proxiedClass = persistentClass.getMappedClass();
		this.persistentProperties = new HashSet<>();

		@SuppressWarnings("unchecked")
		Iterator<Property> propIter = (Iterator<Property>) persistentClass.getPropertyClosureIterator();
		while ( propIter.hasNext() ) {
			this.persistentProperties.add( new org.autofetch.hibernate.Property( propIter.next() ) );
		}
	}

	@Override
	@SuppressWarnings("unchecked")
	public void postInstantiate(
			String entityName,
			Class persistentClass,
			Set interfaces,
			Method getIdentifierMethod,
			Method setIdentifierMethod,
			CompositeType componentIdType) throws HibernateException {

		this.entityName = entityName;
		this.proxiedClass = persistentClass;
		this.interfaces = (Class<?>[]) interfaces.toArray( NO_CLASSES );
		this.getIdentifierMethod = getIdentifierMethod;
		this.setIdentifierMethod = setIdentifierMethod;
		this.componentIdType = componentIdType;
	}

	@Override
	public HibernateProxy getProxy(Serializable id, SharedSessionContractImplementor session)
			throws HibernateException {
		return AutofetchLazyInitializer.getProxy(
				entityName,
				proxiedClass,
				interfaces,
				getIdentifierMethod,
				setIdentifierMethod,
				componentIdType,
				id,
				session,
				persistentProperties
		);
	}
}

public class EntityProxyMethodHandler implements ProxyConfiguration.Interceptor, Serializable {

	private final EntityTracker entityTracker;
	private final Object proxiedObject;
	private final String proxiedClassName;

	EntityProxyMethodHandler(
			Object proxiedObject,
			String proxiedClassName,
			Set<Property> persistentProperties,
			ExtentManager extentManager) {
		this.proxiedObject = proxiedObject;
		this.proxiedClassName = proxiedClassName;
		this.entityTracker = new EntityTracker( persistentProperties, extentManager );
	}

	@Override
	@RuntimeType
	public Object intercept(@This Object instance, @Origin Method method, @AllArguments Object[] arguments)
			throws Exception {
		final String methodName = method.getName();

		if ( "toString".equals( methodName ) ) {
			return proxiedClassName + "@" + System.identityHashCode( instance );
		}
		else if ( "equals".equals( methodName ) ) {
			return proxiedObject == instance;
		}
		else if ( "hashCode".equals( methodName ) ) {
			return System.identityHashCode( instance );
		}
		else if ( arguments.length == 0 ) {
			switch ( methodName ) {
				case "disableTracking": {
					boolean oldValue = entityTracker.isTracking();
					entityTracker.setTracking( false );
					return oldValue;
				}
				case "enableTracking": {
					boolean oldValue = entityTracker.isTracking();
					entityTracker.setTracking( true );
					return oldValue;
				}
				case "isAccessed":
					return entityTracker.isAccessed();

				default:
					break;
			}
		}
		else if ( arguments.length == 1 ) {
			if ( methodName.equals( "addTracker" ) && method.getParameterTypes()[0].equals( Statistics.class ) ) {
				entityTracker.addTracker( (Statistics) arguments[0] );
				return null;
			}
			else if ( methodName.equals( "addTrackers" ) && method.getParameterTypes()[0].equals( Set.class ) ) {
				@SuppressWarnings("unchecked")
				Set<Statistics> newTrackers = (Set) arguments[0];
				entityTracker.addTrackers( newTrackers );
				return null;
			}
			else if ( methodName
					.equals( "extendProfile" ) && method.getParameterTypes()[0].equals( Statistics.class ) ) {
				entityTracker.extendProfile( (Statistics) arguments[0], instance );
				return null;
			}
			else if ( methodName
					.equals( "removeTracker" ) && method.getParameterTypes()[0].equals( Statistics.class ) ) {
				entityTracker.removeTracker( (Statistics) arguments[0] );
				return null;
			}
		}

		entityTracker.trackAccess( instance );
		
		// TODO (ammachado): Here the real instance should be used, but how?
		return method.invoke(instance, arguments);
	}
}

As you can see in the EntityProxyMethodHandler, we want to call method.invoke(instance,arguments) on the real instance, not the proxy. But we haven’t figured out a way to get access to the actual instance here. The only solution as I have thought of is to extend the basicLazyInitializer and the call getImplementation() method on the proxy, but I would rather avoid that. Does anyone know how we could access the real object in this situation?

Thanks in advance,
Erik