$ show google/allocation-instrumenter
This repository provides a Java agent that uses the java.lang.instrument API and ASM bytecode manipulation library to instrument every heap allocation in a Java program. It rewrites bytecodes such as NEW, NEWARRAY, ANEWARRAY, MULTIANEWARRAY, and reflection-based Array.newInstance() to insert pre- and post-allocation hooks. A user-defined 'Sampler' callback receives allocation details including count, type descriptor, object reference, and object size. Additionally, it supports constructor-level instrumentation via ConstructorInstrumenter, allowing callbacks on specific class instantiations. The agent is self-contained, handles reentrancy, caches object sizes, and carefully manages bootstrap classloader concerns through a generated Bootstrap class that dynamically appends the agent JAR to the bootclasspath. The build system uses Bazel with custom rules for assembling the agent jar, including ASM relocation via JarJar to avoid classpath conflicts, and stripping of build metadata. The codebase demonstrates expert-level bytecode transformation, robust error handling, and thorough testing covering many edge cases (e.g., multidimensional arrays, lambda allocations, metaspace exhaustion).
The LLM identified 5 files worth highlighting for reuse:
An extension of ASM's ClassWriter that provides a fallback mechanism for computing common superclasses when the JVM's class hierarchy is not fully loaded. It attempts to load class info at runtime first, and if that fails, reads the class file directly from the classpath to extract superclass and interface information. This is critical for any bytecode agent that transforms classes before they are fully loaded, preventing ClassCircularityErrors. The inner ClassInfo class reimplements the assignment logic statically. Highly reusable for other ASM-based agents.
A ClassVisitor wrapper that verifies each transformed method does not exceed the 64KB bytecode limit after instrumentation. If a method exceeds the limit, it logs a warning and returns the original unchanged class bytes, preventing the JVM from rejecting the class. This pattern is essential for any automatic code transformer that could potentially increase method size beyond the JVM spec's limit. The adapter uses ASM's CodeSizeEvaluator to measure the final size. Easily extracted for use in other agents to provide graceful degradation.
A generated class template that solves the bootstrap classloader problem for Java agents. It finds the agent JAR via ClassLoader.getSystemResource(), appends it to the bootclassloader search using Instrumentation.appendToBootstrapClassLoaderSearch(), and then invokes the real premain via reflection. This avoids requiring the exact JAR filename in the manifest's Boot-Class-Path, which is problematic when version numbers or build outputs vary. The template uses Goolge's @Generated annotation and parameterized package/class names, making it reusable across different agents. The actual premain class can be in the same JAR without compile-time dependency.
The core MethodVisitor that inserts recording calls into allocation bytecodes. Handles all allocation scenarios: NEW (deferred until <init>), NEWARRAY (primitive arrays), ANEWARRAY (object arrays), MULTIANEWARRAY (multidimensional arrays with dynamic dimension product calculation), reflection-based Array.newInstance(), clone(), and Constructor.newInstance(). It carefully manages stack states, uses local variables for scope tracking, and handles the tricky case of constructor superclass initialization tracking via outstandingAllocs counter. The pattern for computing the product of multidimensional array dimensions using bytecode loops is particularly noteworthy. This component embodies the essential complexity of the agent and is a reference implementation for bytecode rewriting patterns.
A reusable Bazel build macro for constructing Java instrumentation agents. It encapsulates the full pipeline: generating a Bootstrap class from a template, building a java_binary with proper premain manifest attributes, then using JarJar to relocate ASM and Guava classes to avoid classpath conflicts, and finally stripping build-data.properties and module-info.class files from the final deploy JAR. The rule handles all the tricky details (bootclasspath, class renaming, manifest generation) and can be reused for any Java agent project using Bazel.