/*
* Copyright (C) 2011 Google Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.google.monitoring.runtime.instrumentation;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.lang.instrument.UnmodifiableClassException;
import java.security.ProtectionDomain;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.logging.Level;
import java.util.logging.Logger;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.commons.LocalVariablesSorter;
/**
* Instruments bytecode by inserting a specified call in the constructor of a given class. This
* class is intended to be loaded by a javaagent; end-users will want to add {@link
* ConstructorCallback}s by invoking {@link #instrumentClass(Class, ConstructorCallback)}.
*/
public class ConstructorInstrumenter implements ClassFileTransformer {
// Implementation details: uses the java.lang.instrument API to
// insert an INVOKESTATIC call to a specified method directly prior to
// constructor return for the given class.
private static final Logger logger = Logger.getLogger(ConstructorInstrumenter.class.getName());
private static ConcurrentHashMap<Class<?>, List<ConstructorCallback<?>>> samplerMap =
new ConcurrentHashMap<Class<?>, List<ConstructorCallback<?>>>();
/**
* We have a read-modify-write operation when doing a put in samplerMap (above) and retransforming
* the class. This lock protects multiple threads from performing that operation concurrently.
*/
private static final Object samplerPutAtomicityLock = new Object();
/** Whether to call the samplers when subclasses of the given class are constructed. */
static boolean subclassesAlso;
// Only for package access (specifically, AllocationInstrumenter)
ConstructorInstrumenter() {}
/**
* Ensures that the given sampler will be invoked every time a constructor for class c is invoked.
*
* @param c The class to be tracked
* @param sampler the code to be invoked when an instance of c is constructed
* @throws UnmodifiableClassException if c cannot be modified.
*/
public static void instrumentClass(Class<?> c, ConstructorCallback<?> sampler)
throws UnmodifiableClassException {
// IMPORTANT: Don't forget that other threads may be accessing this
// class while this code is running. Specifically, the class may be
// executed directly after the retransformClasses is called. Thus, we need
// to be careful about what happens after the retransformClasses call.
synchronized (samplerPutAtomicityLock) {
List<ConstructorCallback<?>> list = samplerMap.get(c);
if (list == null) {
CopyOnWriteArrayList<ConstructorCallback<?>> samplerList =
new CopyOnWriteArrayList<ConstructorCallback<?>>();
samplerList.add(sampler);
samplerMap.put(c, samplerList);
Instrumentation inst = AllocationRecorder.getInstrumentation();
Class<?>[] cs = new Class<?>[1];
cs[0] = c;
inst.retransformClasses(c);
} else {
list.add(sampler);
}
}
}
/** {@inheritDoc} */
@Override
public byte[] transform(
ClassLoader loader,
String className,
Class<?> classBeingRedefined,
ProtectionDomain protectionDomain,
byte[] classfileBuffer) {
if ((classBeingRedefined == null) || (!samplerMap.containsKey(classBeingRedefined))) {
return null;
}
if (!AllocationInstrumenter.canRewriteClass(className, loader)) {
throw new RuntimeException(new UnmodifiableClassException("cannot instrument " + className));
}
return instrument(classfileBuffer, classBeingRedefined);
}
/**
* Given the bytes representing a class, add invocations of the ConstructorCallback method to the
* constructor.
*
* @param originalBytes the original <code>byte[]</code> code.
* @param classBeingRedefined the class being redefined.
* @return the instrumented <code>byte[]</code> code.
*/
public static byte[] instrument(byte[] originalBytes, Class<?> classBeingRedefined) {
try {
ClassReader cr = new ClassReader(originalBytes);
ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_FRAMES);
VerifyingClassAdapter vcw = new VerifyingClassAdapter(cw, originalBytes, cr.getClassName());
ClassVisitor adapter = new ConstructorClassAdapter(vcw, classBeingRedefined);
cr.accept(adapter, ClassReader.SKIP_FRAMES);
return vcw.toByteArray();
} catch (RuntimeException e) {
logger.log(Level.WARNING, "Failed to instrument class.", e);
throw e;
} catch (Error e) {
logger.log(Level.WARNING, "Failed to instrument class.", e);
throw e;
}
}
/** The per-method transformations to make. Really only affects the <init> methods. */
static class ConstructorMethodAdapter extends MethodVisitor {
/**
* The LocalVariablesSorter used in this adapter. Lame that it's public but the ASM architecture
* requires setting it from the outside after this AllocationMethodAdapter is fully constructed
* and the LocalVariablesSorter constructor requires a reference to this adapter. The only
* setter of this should be AllocationClassAdapter.visitMethod().
*/
public LocalVariablesSorter lvs = null;
Class<?> cl;
ConstructorMethodAdapter(MethodVisitor mv, Class<?> cl) {
super(Opcodes.ASM7, mv);
this.cl = cl;
}
/** Inserts the appropriate INVOKESTATIC call */
@Override
public void visitInsn(int opcode) {
if ((opcode == Opcodes.ARETURN)
|| (opcode == Opcodes.IRETURN)
|| (opcode == Opcodes.LRETURN)
|| (opcode == Opcodes.FRETURN)
|| (opcode == Opcodes.DRETURN)) {
throw new RuntimeException(
new UnmodifiableClassException("Constructors are supposed to return void"));
}
if (opcode == Opcodes.RETURN) {
super.visitVarInsn(Opcodes.ALOAD, 0);
super.visitMethodInsn(
Opcodes.INVOKESTATIC,
"com/google/monitoring/runtime/instrumentation/ConstructorInstrumenter",
"invokeSamplers",
"(Ljava/lang/Object;)V",
false);
}
super.visitInsn(opcode);
}
}
// Each thread remembers the last object the thread has seen.
private static final ThreadLocal<Object> lastObject = new ThreadLocal<>();
/**
* Bytecode is rewritten to invoke this method; it calls the sampler for the given class. Note
* that, unless the javaagent command line argument "subclassesAlso" is specified, it won't do
* anything if o is a subclass of the class that was supposed to be tracked.
*
* @param o the object passed to the samplers.
*/
@SuppressWarnings("unchecked")
public static void invokeSamplers(Object o) {
// Remember the last object for this thread, skip if seen already.
if (lastObject.get() == o) {
return;
}
lastObject.set(o);
Class<?> currentClass = o.getClass();
while (currentClass != null) {
List<ConstructorCallback<?>> samplers = samplerMap.get(currentClass);
if (samplers != null) {
// Leave in the @SuppressWarnings, because we define -Werror,
// and infrastructure sometimes runs with all warnings turned
// on. This would be a great place for a typesafe
// heterogeneous container, but that doesn't work with generic
// types.
for (@SuppressWarnings("rawtypes") ConstructorCallback sampler : samplers) {
sampler.sample(o);
}
// Return once the first list of registered samplers are found and
// invoked.
return;
} else {
// When subclassesAlso is not specified (default), return if no
// samplers are registered with the type of the currently-constructed
// object. Otherwise, traverse upward the class hierarchy.
if (!subclassesAlso) {
return;
}
currentClass = currentClass.getSuperclass();
}
}
}
/**
* The class that deals with per-class transformations. Basically, invokes the per-method
* transformer above if the method is an {@code <init>} method.
*/
static class ConstructorClassAdapter extends ClassVisitor {
Class<?> cl;
public ConstructorClassAdapter(ClassVisitor cv, Class<?> cl) {
super(Opcodes.ASM7, cv);
this.cl = cl;
}
/**
* For each method in the class being instrumented, <code>visitMethod</code> is called and the
* returned MethodVisitor is used to visit the method. Note that a new MethodVisitor is
* constructed for each method.
*/
@Override
public MethodVisitor visitMethod(
int access, String name, String desc, String signature, String[] exceptions) {
MethodVisitor mv = cv.visitMethod(access, name, desc, signature, exceptions);
if ((mv != null) && "<init>".equals(name)) {
ConstructorMethodAdapter aimv = new ConstructorMethodAdapter(mv, cl);
LocalVariablesSorter lvs = new LocalVariablesSorter(access, desc, aimv);
aimv.lvs = lvs;
mv = lvs;
}
return mv;
}
}
}
