Reference to dynamic proxy in a proxied class

There was an interesting question in Stackoverflow about how a Spring Bean can get a reference to the proxy created by Spring to handle transactions, Spring AOP,  Caching, Async flows etc. A reference to the proxy was required because if there is a call to itself by the proxied bean, this call would completely bypass the proxy.

Consider an InventoryService interface:

public interface InventoryService{
    public Inventory create(Inventory inventory);
    public List<Inventory> list();
    public Inventory findByVin(String vin);
    public Inventory update(Inventory inventory);
    public boolean delete(Long id);
    public Inventory compositeUpdateService(String vin, String newMake);
}

Consider also that there is a default implementation for this service, and assume that the last method compositeUpdateService, internally invokes two methods on the bean itself, like this:

public Inventory compositeUpdateService(String vin, String newMake) {
    logger.info("composite Update Service called");
    Inventory inventory = this.findByVin(vin);
    inventory.setMake(newMake);
    this.update(inventory);
    return inventory;
}

If I now create an aspect to advice any calls to InventoryService with the objective of tracking how long each method call takes, Spring AOP would create a dynamic proxy for the InventoryService bean:

However the calls to compositeUpdateService will record the time only at the level of this method, the calls that compositeUpdateService internally makes to findByVin, update is bypassing the proxy and hence will not be tracked:

A good fix for this is to use the full power of AspectJ - AspectJ would change the bytecode of all the methods of DefaultInventoryService to include the call to the advice.

A workaround that we worked out was to inject a reference to the proxy itself into the bean and instead of calling say this.findByVin and this.update, call proxy.findByVin and proxy.update!

So now how do we cleanly inject in a reference to the proxy into the bean - the solution that I offered was to create an interface to mark beans interested in their own proxies:

public interface ProxyAware<T> {
    void setProxy(T proxy);
}

The interested interface and its implementation would look like this:

public interface InventoryService extends ProxyAware<InventoryService>{
...
}

public class DefaultInventoryService implements InventoryService{ 
    ...

    private InventoryService proxy;
 @Override
 public void setProxy(InventoryService proxy) {
  this.proxy = proxy;
 }
}

and then define a BeanPostProcessor to inject in this proxy!

public class ProxyInjectingBeanPostProcessor implements BeanPostProcessor, Ordered {
 @Override
 public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException {
  return bean;
 }

 @Override
 public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
  if (AopUtils.isAopProxy((bean))){
   try {
    Object target = ((Advised)bean).getTargetSource().getTarget();

    if (target instanceof ProxyAware){
     ((ProxyAware) target).setProxy(bean);
    }
   } catch (Exception e) {
    return bean;
   }
  }
  return bean;
 }

 @Override
 public int getOrder() {
  return Integer.MAX_VALUE;
 }
}

Not the cleanest of implementations, but works!

Update:
Tomasz has pointed out a much simpler way to get hold of the proxy in the proxied class, this is by using AopContext.currentProxy() method. Spring AOP has to be instructed to set up this method, it is done through configuration this way:

<aop:aspectj-autoproxy expose-proxy="true"/>

Once this is done, in the proxied class, the reference to the proxy can be obtained by calling AopContext.currentProxy()!

Reference: http://blog.springsource.org/2012/05/23/understanding-proxy-usage-in-spring/