IoTraceAgent.java

package hu.axt.iotrace;

import static com.sun.xml.internal.ws.org.objectweb.asm.Opcodes.ACC_FINAL;
import static com.sun.xml.internal.ws.org.objectweb.asm.Opcodes.ACC_PRIVATE;
import static com.sun.xml.internal.ws.org.objectweb.asm.Opcodes.ACC_PUBLIC;
import static com.sun.xml.internal.ws.org.objectweb.asm.Opcodes.ACC_STATIC;
import static com.sun.xml.internal.ws.org.objectweb.asm.Opcodes.ACC_SUPER;
import static com.sun.xml.internal.ws.org.objectweb.asm.Opcodes.ILOAD;
import static com.sun.xml.internal.ws.org.objectweb.asm.Opcodes.INVOKESPECIAL;
import static com.sun.xml.internal.ws.org.objectweb.asm.Opcodes.INVOKESTATIC;
import static com.sun.xml.internal.ws.org.objectweb.asm.Opcodes.IRETURN;
import static com.sun.xml.internal.ws.org.objectweb.asm.Opcodes.RETURN;
import static com.sun.xml.internal.ws.org.objectweb.asm.Opcodes.V1_6;

import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.PrintStream;
import java.lang.instrument.ClassDefinition;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.net.InetAddress;
import java.util.concurrent.atomic.AtomicReference;

import sun.misc.IoTrace;

import com.sun.xml.internal.ws.org.objectweb.asm.ClassWriter;
import com.sun.xml.internal.ws.org.objectweb.asm.MethodVisitor;
import com.sun.xml.internal.ws.org.objectweb.asm.Type;

/**
 * Naive IOTrace agent implementation.
 * 
 * @author axt
 *
 */
public class IoTraceAgent {

	private static final int WRITER_THREAD_SLEEPTIME_SEC = 60;
	
	private static final String IO_LOG = "perf/io.log";
	private static final String NET_LOG = "perf/net.log";

	public static final long agentStartTime = System.nanoTime();

	public static AtomicReference<StringBuffer> io_buf = new AtomicReference<StringBuffer>(new StringBuffer());
	public static AtomicReference<StringBuffer> net_buf = new AtomicReference<StringBuffer>(new StringBuffer());

	public static PrintStream io_log;
	public static PrintStream net_log;

	public static Thread writer = new Thread() {
		@Override
		public void run() {
			io_log.println("writer started");
			io_log.flush();

			net_log.println("writer started");
			net_log.flush();
			
			for (;;) {
			
				try { Thread.sleep(1000 * WRITER_THREAD_SLEEPTIME_SEC); } catch (InterruptedException e) {}
				
				StringBuffer io_buf_old = io_buf.getAndSet(new StringBuffer());
				io_log.println(io_buf_old.toString());
				io_log.flush();
				
				StringBuffer net_buf_old = net_buf.getAndSet(new StringBuffer());
				net_log.println(net_buf_old.toString());
				net_log.flush();
			}
		}
	};

	static {
		try {
			io_log = new PrintStream(new BufferedOutputStream(new FileOutputStream(new File(IO_LOG)))); 
		} catch (IOException e) {
			e.printStackTrace();
		}
		try {
			net_log = new PrintStream(new BufferedOutputStream(new FileOutputStream(new File(NET_LOG))));
		} catch (IOException e) {
			net_log = null;
			e.printStackTrace();
		}
		
		writer.start();
	}

	
	public static class SO {
		public enum Type {
			R, W;
		}

		public SO(Type type) {
			this.startTime = System.nanoTime() - agentStartTime;
			this.type = type;
			this.threadId = Thread.currentThread().getId();
		}

		public final long startTime;
		public final long threadId;
		public final Type type;

		public long duration;
		public long bytes;
		public String addr;
		public int timeout;

		@Override
		public String toString() {
			return String.format("%d\tT%d\t%c\t%8d\t%8d\t%8d\t%s\n", startTime / 1000000, threadId, type == Type.R ? 'R' : 'W', duration / 1000, bytes, timeout, addr);
		}

	}

	public static class FO {

		public enum Type {
			R, W;
		}

		public FO(String path, Type type) {
			this.startTime = System.nanoTime() - agentStartTime;
			this.path = path.intern();
			this.type = type;
			this.threadId = Thread.currentThread().getId();
		}

		public final long startTime;
		public final long threadId;
		public final String path;
		public final Type type;
		public long duration;
		public long bytes;

		@Override
		public String toString() {
			return String.format("%d\tT%d\t%c\t%8d\t%8d\t%s\n", startTime / 1000000, threadId, type == Type.R ? 'R' : 'W', duration / 1000, bytes, path);
		}
	}


	// socket operations
	
	public static Object socketReadBegin() {
		return new SO(SO.Type.R);
	}

	public static Object socketWriteBegin() {
		return new SO(SO.Type.W);
	}

	public static void socketReadEnd(Object context, InetAddress address, int port, int timeout, long bytes) {
		handleSocketOp(context, address, port, timeout, bytes);
	}

	public static void socketWriteEnd(Object context, InetAddress address, int port, long bytes) {
		handleSocketOp(context, address, port, -1, bytes);
	}

	private static void handleSocketOp(Object context, InetAddress address, int port, int timeout, long bytes) {
		if (context != null) {
			SO so = (SO) context;
			so.duration = (System.nanoTime() - agentStartTime) - so.startTime;
			so.bytes = bytes;
			so.timeout = timeout;
			so.addr = (address.toString() + ":" + port).intern();

			if (net_buf != null) {
				net_buf.get().append(so.toString());
			}
		}
	}

	// file related operations
	public static Object fileReadBegin(String path) {
		if (path != null) { 
			return new FO(path, FO.Type.R); 
		}

		return null;
	}

	public static Object fileWriteBegin(String path) {
		if (path != null) { 
			return new FO(path, FO.Type.W); 
		}

		return null;
	}

	public static void fileReadEnd(Object context, long bytes) {
		handleFileOp(context, bytes);
	}

	public static void fileWriteEnd(Object context, long bytesWritten) {
		handleFileOp(context, bytesWritten);
	}

	private static void handleFileOp(Object context, long bytesWritten) {
		if (context != null) {
			FO fo = (FO) context;
			fo.duration = (System.nanoTime() - agentStartTime) - fo.startTime;
			fo.bytes = bytesWritten;

			if (io_buf != null && !fo.path.equals(IO_LOG) && !fo.path.equals(NET_LOG)) {
				io_buf.get().append(fo.toString());
			}
		}
	}

	public static void premain(String artuments, Instrumentation instrumentation) throws Exception {
		instrumentation.redefineClasses(new ClassDefinition(IoTrace.class, generateIoTraceClassAsm()));
	}

	// Code borrowed from OpenJDK
	// TODO add proper reference later
	private static byte[] generateIoTraceClassAsm() {

		ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_MAXS);
		cw.visit(V1_6, ACC_PUBLIC | ACC_SUPER | ACC_FINAL, "sun/misc/IoTrace", null, "java/lang/Object", null);

		for (Method om : IoTrace.class.getDeclaredMethods()) {
			
			if (!Modifier.isStatic(om.getModifiers())) continue;

			MethodVisitor mv = cw.visitMethod(ACC_PUBLIC | ACC_STATIC, om.getName(), Type.getMethodDescriptor(om), null, null);
			mv.visitCode();

			int i = 0;
			for (Type t : Type.getArgumentTypes(om)) {
				mv.visitVarInsn(t.getOpcode(ILOAD), i++);
			}

			mv.visitMethodInsn(INVOKESTATIC, "hu/axt/iotrace/IoTraceAgent", om.getName(), Type.getMethodDescriptor(om));

			mv.visitInsn(Type.getReturnType(om).getOpcode(IRETURN));
			mv.visitEnd();
		}

		MethodVisitor mv = cw.visitMethod(ACC_PRIVATE, "<init>", "()V", null, null);
		mv.visitCode();
		mv.visitMethodInsn(INVOKESPECIAL, "java/lang/Object", "<init>", "()V");
		mv.visitInsn(RETURN);
		mv.visitEnd();

		cw.visitEnd();

		return cw.toByteArray();
	}
}