Utility class using JDK's tools.jar to instrument a VM at runtime (e.g. for using hibernate proxies that lazily load @lob, but without compile-time instrumentation)

HibernateClassBypassingLoadTimeWeaver.java

package gov.ne.revenue.mef.enhancement;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.instrument.classloading.LoadTimeWeaver;
import org.springframework.util.Assert;

public class HibernateClassBypassingLoadTimeWeaver implements LoadTimeWeaver {
 
	private final LoadTimeWeaver loadTimeWeaver;
 
	public HibernateClassBypassingLoadTimeWeaver(LoadTimeWeaver loadTimeWeaver) {
		Assert.notNull(loadTimeWeaver);
		this.loadTimeWeaver = loadTimeWeaver;
	}
 
	@Override
	public ClassLoader getInstrumentableClassLoader() {
		return loadTimeWeaver.getInstrumentableClassLoader();
	}
 
	@Override
	public ClassLoader getThrowawayClassLoader() {
		return loadTimeWeaver.getThrowawayClassLoader();
	}
 
	@Override
	public void addTransformer(ClassFileTransformer transformer) {
		loadTimeWeaver.addTransformer(new HibernateClassBypassingClassFileTransformer(transformer));
	}
 
	/**
	 * ClassFileTransformer decorator that suppresses processing of Hibernate
	 * classes in order to avoid potential LinkageErrors.
	 * @see org.springframework.context.annotation.LoadTimeWeavingConfiguration
	 */
	private static class HibernateClassBypassingClassFileTransformer implements ClassFileTransformer {
		private static final Logger LOGGER = LoggerFactory.getLogger(HibernateClassBypassingClassFileTransformer.class);
 
		private final ClassFileTransformer delegate;
 
		public HibernateClassBypassingClassFileTransformer(ClassFileTransformer delegate) {
			this.delegate = delegate;
		}
 
		@Override
		public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer)
				throws IllegalClassFormatException {
 
			if (className.startsWith("org.hibernate") || className.startsWith("org/hibernate") || className.startsWith("javassist")) {
				LOGGER.trace("Bypassing class: {}", className);
				return classfileBuffer;
			}
			return this.delegate.transform(loader, className, classBeingRedefined, protectionDomain, classfileBuffer);
		}
	}
}

HibernateLoadTimeWeaverAwareProcessor.java

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.config.BeanPostProcessor;
import org.springframework.context.weaving.LoadTimeWeaverAware;
import org.springframework.instrument.classloading.LoadTimeWeaver;
import org.springframework.instrument.classloading.tomcat.TomcatLoadTimeWeaver;
import org.springframework.orm.jpa.EntityManagerFactoryInfo;
import org.springframework.orm.jpa.JpaDialect;
import org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean;
import org.springframework.orm.jpa.vendor.HibernateJpaDialect;

/**
 * The default Tomcat 8 Load Time Weaving configuration will throw a
 * {@code java.lang.ClassCircularityError: org/apache/log4j/spi/ThrowableInformation}
 * in {@link JdbcConfig} when {@link LocalContainerEntityManagerFactoryBean#getObject()}
 * is called.  Therefore we jump through some hoops to avoid that.
 * @see https://jira.spring.io/browse/SPR-13886
 * @see TomcatLoadTimeWeaver
 */
public class HibernateLoadTimeWeaverAwareProcessor implements BeanPostProcessor {
	private static final Logger LOGGER = LoggerFactory.getLogger(HibernateLoadTimeWeaverAwareProcessor.class);
 
	private final LoadTimeWeaver loadTimeWeaver;
 
	public HibernateLoadTimeWeaverAwareProcessor(LoadTimeWeaver loadTimeWeaver) {
		this.loadTimeWeaver = new HibernateClassBypassingLoadTimeWeaver(loadTimeWeaver);
	}
 
	@Override
	public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException {
		if (bean instanceof LoadTimeWeaverAware && bean instanceof EntityManagerFactoryInfo) {
			JpaDialect jpaDialect = ((EntityManagerFactoryInfo) bean).getJpaDialect();
			boolean isHibernate = jpaDialect instanceof HibernateJpaDialect;
			LOGGER.info("EMF {} has Hibernate dialect: {}", bean, isHibernate);
			if (isHibernate) {
				LOGGER.info("Injecting custom LTW: {}", loadTimeWeaver);
				((LoadTimeWeaverAware) bean).setLoadTimeWeaver(loadTimeWeaver);
			}
		}
		return bean;
	}
 
	@Override
	public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
		return bean;
	}
}

InstrumentationUtil.java

import java.lang.management.ManagementFactory;
import java.lang.management.RuntimeMXBean;
import java.net.URISyntaxException;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.EnableLoadTimeWeaving;
import org.springframework.instrument.InstrumentationSavingAgent;

import com.sun.tools.attach.VirtualMachine;

/**
 * Instrument the current VM using the classes from the runtime's tools.jar
 
 */
public class InstrumentationUtil {
	private static Logger logger = LoggerFactory.getLogger(InstrumentationUtil.class);
	
	private static Boolean isInstrumented = null;

	@Deprecated
	public static synchronized void instrument() throws InstrumentationException {
		String systemProp = System.getProperty("mef.instrument");
		if(systemProp == null){
			logger.info("instrumentation disabled because -Dmef.instrument is unset");
			return;
		} else if(systemProp.equals("false")){
			logger.info("instrumentation disabled because -Dmef.instrument=false");
			return;
		}
		
		if(isInstrumented()){
			logger.info("VM already instrumented!");
			return;
		}	
		
		String jar;
		try {
			
			jar = InstrumentationSavingAgent.class.getProtectionDomain().getCodeSource().getLocation().toURI().getPath();
			
			if(System.getProperty("os.name").toLowerCase().contains("win")){
				if(jar.startsWith("/")){
					// sometimes the jar comes back as /c:/path/to/jarfile.jar
					// but we need the real path
					jar = jar.substring(1);
				}	
			}
			
			System.err.println("--------------------------------------------------------------------------------");
			System.err.println(jar);
			System.err.println("--------------------------------------------------------------------------------");
			
			logger.info("Dynamically instrumented VM using {}", jar);
		} catch (URISyntaxException e) {
			e.printStackTrace();
			throw new InstrumentationException(e);
		}
	
		String nameOfRunningVM = ManagementFactory.getRuntimeMXBean().getName();
		logger.info("Instrumenting {}", nameOfRunningVM);
		int p = nameOfRunningVM.indexOf('@');
		String pid = nameOfRunningVM.substring(0, p);
		try {
			VirtualMachine vm = VirtualMachine.attach(pid);
			vm.loadAgent(jar, "");
			vm.detach();
			System.out.println("--------------------------------------------------------------------------------");
			System.out.println(pid + " instrumented with " + jar);
			System.out.println("--------------------------------------------------------------------------------");
			isInstrumented = true;
		} catch (Exception e) {
			e.printStackTrace();
			throw new InstrumentationException(e);
		} 
		
	}

	/**
	 * If you instrument by some other mechanism (e.g. in a Tomcat 8 container
	 * with {@link EnableLoadTimeWeaving}), then kindly set this property
	 * so that {@link JdbcConfig} knows to set 
	 * {@link org.hibernate.jpa.AvailableSettings.ENHANCER_ENABLE_LAZY_INITIALIZATION}
	 * in JPA Properties.
	 * @param b
	 */
	public static void setInstrumented(boolean b){
		isInstrumented = b ;
	}

	@Deprecated
	public static synchronized boolean isInstrumented(){
		if(isInstrumented == null){
			RuntimeMXBean runtimeMxBean = ManagementFactory.getRuntimeMXBean();
			runtimeMxBean.getInputArguments().forEach(arg -> {
				if (arg.startsWith("-javaagent") && arg.contains("spring-instrument")) {
					isInstrumented = true;
				}
			});
		}
		
		// still?  
		if(isInstrumented == null){
			return false;
		}
		
		return isInstrumented;
	}
	
}

pom.xml

<project>	
  <!-- typically put this in the parent project, if it exists -->
<dependencyManagement>
		<dependencies>
      <!-- to allow runtime loading of the bytecode instrumenter -->
			<!-- this is required if you want to dynamically load -->
			<!-- the spring-instrumentation jar at runtime -->
			<dependency>
				<groupId>com.sun</groupId>
				<artifactId>tools</artifactId>
				<version>1.8.0</version>
				<scope>system</scope>
				<systemPath>${java.home}/../lib/tools.jar</systemPath>
			</dependency>
    </dependencies>
</dependencyManagement>

  <!-- then in each module, the actual jar -->
<dependencies>
		<!-- to allow runtime loading of the bytecode instrumenter -->
		<dependency>
			<groupId>com.sun</groupId>
			<artifactId>tools</artifactId>
			<scope>system</scope>
			<systemPath>${java.home}/../lib/tools.jar</systemPath>
		</dependency>  
</dependencies>
</project>

SomeConfig.java

@Configuration
class SomeConfig {
  @Inject
  Optional<LoadTimeWeaver> loadTimeWeaver;
@Bean
	public EntityManagerFactory entityManagerFactory(DataSource dataSource, JpaVendorAdapter vendorAdapter, Optional<LoadTimeWeaver> loadTimeWeaver) {
		LocalContainerEntityManagerFactoryBean factory = new LocalContainerEntityManagerFactoryBean();
		Map<String, Object> props = new HashMap<>();
//		props.put(AvailableSettings.USE_QUERY_CACHE, true);
//		props.put(AvailableSettings.JMX_ENABLED, true);
		props.put(AvailableSettings.STATEMENT_BATCH_SIZE, 50);
		props.put(AvailableSettings.ORDER_INSERTS, true);
		props.put(AvailableSettings.USE_MINIMAL_PUTS, true);
	
		// set the logger org.hibernate.stat=DEBUG to see hibernate statistics
		props.put(AvailableSettings.GENERATE_STATISTICS, false);
		
		// RUNTIME ENHANCEMENT FOR LAZY PROPERTY LOADING
		// to use, specify -Djavaagent:c:\path\to\spring-instrument-4.2.3.RELEASE.jar
		// as an argument to the VM
		// google "Spring Bytecode Instrumentation" for more info
		
		// we do this to allow lazy property initialization for the
		// xml fields on SubmissionArchive (which should eventually be removed)
	
		if(loadTimeWeaver.isPresent()){
			props.put("hibernate.enhancer.enableLazyInitialization", true);
			
			if(isTomcat()){
//				HibernateClassBypassingLoadTimeWeaver hackWeaver = new HibernateClassBypassingLoadTimeWeaver(loadTimeWeaver.get());
//				factory.setLoadTimeWeaver(hackWeaver);
//				logger.info("Hibernate class enhancement using Tomcat enhancement hack.");
				logger.warn("Load time weaver excluded because this is Tomcat.  Hopefully this was compiled into mef-data instead.");
			} else {
				factory.setLoadTimeWeaver(loadTimeWeaver.get());
				logger.info("Hibernate class enhancement enabled.");
			}
		} else {
			logger.info("Runtime Hibernate class enhancement disabled.  No LoadTimeWeaver present.");
		}	
		
		factory.getJpaPropertyMap().putAll(props);
		factory.setJpaVendorAdapter(vendorAdapter);
		factory.setPackagesToScan(PACKAGES_TO_SCAN);
		factory.setDataSource(dataSource);
		factory.afterPropertiesSet();
		EntityManagerFactory f = factory.getObject();
		return f;
	}
}