MemoryCounterAgent.java

package org.jblooming.tracer;

import java.util.IdentityHashMap;
import java.util.Stack;
import java.util.Map;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Field;
import java.lang.reflect.Modifier;
import java.lang.reflect.Array;

/**
 * (c) Open Lab - www.open-lab.com
 * Date: Apr 2, 2007
 * Time: 9:47:30 AM
 */
public class MemoryCounterAgent {

  private static Instrumentation instrumentation;

  /**
   * Initializes agent
   */
  public static void premain(String agentArgs,
                             Instrumentation instrumentation) {
    MemoryCounterAgent.instrumentation = instrumentation;
  }

  /**
   * Returns object size.
   */

  public static long sizeOf(Object obj) {
    if (instrumentation == null) {
      throw new IllegalStateException(
              "Instrumentation environment not initialised.");
    }
    if (isSharedFlyweight(obj)) {
      return 0;
    }
    return instrumentation.getObjectSize(obj);
  }

  /**
   * Returns deep size of object, recursively iterating over its
   * fields and superclasses.
   */

  public static long deepSizeOf(Object obj) {
    Map visited = new IdentityHashMap();
    Stack stack = new Stack();
    stack.push(obj);

    long result = 0;
    do {
      result += internalSizeOf(stack.pop(), stack, visited);
    } while (!stack.isEmpty());
    return result;
  }

  /**
   * Returns true if this is a well-known shared flyweight.
   * For example, interned Strings, Booleans and Number objects.
   */

  private static boolean isSharedFlyweight(Object obj) {
    // optimization - all of our flyweights are Comparable
    if (obj instanceof Comparable) {
      if (obj instanceof Enum) {
        return true;
      } else if (obj instanceof String) {
        return (obj == ((String) obj).intern());
      } else if (obj instanceof Boolean) {
        return (obj == Boolean.TRUE || obj == Boolean.FALSE);
      } else if (obj instanceof Integer) {
        return (obj == Integer.valueOf((Integer) obj));
      } else if (obj instanceof Short) {
        return (obj == Short.valueOf((Short) obj));
      } else if (obj instanceof Byte) {
        return (obj == Byte.valueOf((Byte) obj));
      } else if (obj instanceof Long) {
        return (obj == Long.valueOf((Long) obj));
      } else if (obj instanceof Character) {
        return (obj == Character.valueOf((Character) obj));
      }
    }
    return false;
  }

  private static boolean skipObject(Object obj, Map visited) {
    return obj == null

            || visited.containsKey(obj)
            || isSharedFlyweight(obj);
  }

  private static long internalSizeOf(
          Object obj, Stack stack, Map visited) {
    if (skipObject(obj, visited)) {
      return 0;
    }

    Class clazz = obj.getClass();
    if (clazz.isArray()) {
      addArrayElementsToStack(clazz, obj, stack);
    } else {
      // add all non-primitive fields to the stack
      while (clazz != null) {
        Field[] fields = clazz.getDeclaredFields();
        for (Field field : fields) {
          if (!Modifier.isStatic(field.getModifiers())
                  && !field.getType().isPrimitive()) {
            field.setAccessible(true);
            try {
              stack.add(field.get(obj));
            } catch (IllegalAccessException ex) {
              throw new RuntimeException(ex);
            }
          }
        }
        clazz = clazz.getSuperclass();
      }
    }
    visited.put(obj, null);
    return sizeOf(obj);
  }

  private static void addArrayElementsToStack(
          Class clazz, Object obj, Stack stack) {
    if (!clazz.getComponentType().isPrimitive()) {
      int length = Array.getLength(obj);
      for (int i = 0; i < length; i++) {
        stack.add(Array.get(obj, i));
      }
    }
  }
}