MongoMethodDatabase.java

package org.jikesrvm.adaptive.database.methodsamples;

import java.util.Vector;
import java.util.concurrent.Semaphore;

import org.bson.Document;
import org.jikesrvm.VM;
import org.jikesrvm.adaptive.controller.Controller;
import org.jikesrvm.classloader.NormalMethod;
import org.jikesrvm.classloader.RVMClass;
import org.jikesrvm.classloader.RVMMethod;
import org.jikesrvm.compilers.common.CompiledMethod;
import org.jikesrvm.compilers.common.CompiledMethods;
import org.jikesrvm.compilers.opt.driver.CompilationPlan;
import org.jikesrvm.compilers.opt.runtimesupport.OptCompiledMethod;

import com.mongodb.MongoClient;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoDatabase;
import com.mongodb.client.model.InsertOneModel;
import com.mongodb.client.model.UpdateOneModel;
import com.mongodb.client.model.WriteModel;

import static com.mongodb.client.model.Filters.*;

public class MongoMethodDatabase extends MethodDatabase {
	private MongoClient mongo;
	private MongoDatabase aosDatabase;
	private MongoCollection<Document> aosCollection;
	private Vector<WriteModel<Document>> writeRequests;
	private BulkWriteThread bulkWriteThread;
	private Semaphore writeRequestsLock;
	private Vector<CompilationPlan> compilationPlans;
	private Semaphore compPlanLock;

	class BulkWriteThread extends Thread
	{
		private MongoCollection<Document> aosCollection;
		private Vector<WriteModel<Document>> writeRequests;
		private Semaphore writeRequestsLock;
		private Semaphore compPlanLock;
		private Vector<CompilationPlan> compilationPlans;
		
		public BulkWriteThread (MongoCollection<Document> aosCollection,
				Vector<WriteModel<Document>> writeRequests, Semaphore writeRequestsLock,
				Semaphore compPlanLock, Vector<CompilationPlan> compilationPlans)
		{
			this.writeRequestsLock = writeRequestsLock;
			this.aosCollection = aosCollection;
			this.writeRequests = writeRequests;
			this.compPlanLock = compPlanLock;
			this.compilationPlans = compilationPlans;
		}
		
		public void run ()
		{
			try {
				compPlanLock.acquire();
				for (CompilationPlan cp : compilationPlans)
				{
					_putMethodOptLevel (cp);
				}
				compilationPlans.clear();
				compPlanLock.release();
			} catch (InterruptedException e1) {
				// TODO Auto-generated catch block
				e1.printStackTrace();
			}
			
			try {
				writeRequestsLock.acquire();
				if (VM.verboseClassLoading)
					VM.sysWriteln ("Writing bulk requests");
				if (writeRequests.size () != 0)
					aosCollection.bulkWrite(writeRequests);

				if (VM.verboseClassLoading)
					VM.sysWriteln ("Wrote bulk requests");
			} catch (InterruptedException e) {
				e.printStackTrace();
			}

			writeRequests.clear();
			writeRequestsLock.release();
		}
		
		private void _putMethodOptLevel (CompilationPlan cp)
		{
			NormalMethod nm = cp.method;

			String methodFullDesc = getMethodFullDesc (nm);
			int plan = Controller.recompilationStrategy.getOptLevelForOptPlan (cp.optimizationPlan);
			Document meth = new Document();
			double counts = getCallCount (methodFullDesc);
			meth.put("methodFullDesc", methodFullDesc);
			meth.put("count", counts);
			meth.put("optLevel", plan);
			
			if (VM.verboseClassLoading)
				VM.sysWriteln ("Put Method " + methodFullDesc + " with call counts " + counts + " opt level " + plan);
			
			if (counts == 0.0)
			{
				if (VM.verboseClassLoading)
					VM.sysWriteln ("Method " + methodFullDesc + " added to MongoDB");
				try {
					writeRequestsLock.acquire();
				} catch (InterruptedException e) {
					e.printStackTrace();
				}
				writeRequests.add(new InsertOneModel<Document> (meth));
				writeRequestsLock.release();
			}
			else
			{
				if (VM.verboseClassLoading)
					VM.sysWriteln ("Method " + methodFullDesc + " updated to MongoDB");
				try {
					writeRequestsLock.acquire();
				} catch (InterruptedException e) {
					e.printStackTrace();
				}
				writeRequests.add(new UpdateOneModel <Document> (new Document ("methodFullDesc", methodFullDesc), 
						new Document ("$set", meth)));
				writeRequestsLock.release();
			}
		}
	}
	
	public MongoMethodDatabase (int bulkMethodCount)
	{
		super (bulkMethodCount);
		
		writeRequests = new Vector<WriteModel<Document>> ();
	}
	
	@Override
	public void initialize() {
		try {
			// TODO Auto-generated method stub
			mongo = new MongoClient("localhost", 27017);
			if (VM.verboseClassLoading)
				VM.sysWriteln ("MongoClient Initialized");

			aosDatabase = mongo.getDatabase("AOSDatabase");
			if (VM.verboseClassLoading)
				VM.sysWriteln ("Mongo AOS Database Initialized");

			aosCollection = aosDatabase.getCollection("AOSCollection");
			if (VM.verboseClassLoading)
				VM.sysWriteln ("Mongo AOS Collection Initialized");
			writeRequestsLock = new Semaphore (1);
			compPlanLock = new Semaphore (1);
			
			VM.sysWriteln ("RequestsLock created");
			
			compilationPlans = new Vector<CompilationPlan> ();
		}
		catch (Exception e) {
			e.printStackTrace();
		}

	}

	public static String getMethodFullDesc (RVMMethod nm)
	{
		RVMClass declaringClass = nm.getDeclaringClass();
		String desc = "";
		if (declaringClass.getDescriptor() == null)
			VM.sysWriteln ("class desc null");
		else
			desc = declaringClass.getDescriptor().toString();
		String methodDesc = "";
		if (nm.getDescriptor() != null)
			methodDesc = nm.getDescriptor().toString();
		
		String name = "";
		if (nm.getName() == null)
			VM.sysWriteln ("name null ");
		else
			name = nm.getName().toString();
		
		return desc + name + methodDesc;
	}
	private int getOptLevel (int cmid)
	{
		CompiledMethod cm = CompiledMethods.getCompiledMethodUnchecked(cmid);

		if (cm instanceof OptCompiledMethod)
		{
			return ((OptCompiledMethod)cm).getOptLevel();
		}
		
		return -1;
	}
	
	public int getOptLevelForMethod (RVMMethod m)
	{
		if (aosCollection == null)
		{
			return -1;
		}
		
		String methodFullDesc = getMethodFullDesc (m);
		if (methodFullDesc.contains("mongodb") || methodFullDesc.contains("bson") || methodFullDesc.contains("MethodDatabase") ||
				methodFullDesc.contains("Formatter") || methodFullDesc.indexOf("Ljava/") == 0 ||
				methodFullDesc.indexOf ("Lgnu/") == 0)
			return -1;
		
		if (VM.verboseClassLoading)
			VM.sysWriteln ("Getting opt level for Method "+ methodFullDesc);
		
		Document meth = aosCollection.find(eq("methodFullDesc", methodFullDesc)).first();
		
		if (meth == null)
		{
			if (VM.verboseClassLoading)
				VM.sysWriteln ("Cannot find opt level for Method " + methodFullDesc);
			
			return -1;
		}
		
		int optLevel = ((Integer)meth.get("optLevel")).intValue();
		
		if (VM.verboseClassLoading)
			VM.sysWriteln ("Opt Level for Method " + methodFullDesc + " is " + optLevel);
		
		return optLevel;
	}

	public static String getMethodFullDesc (int cmid)
	{
		CompiledMethod cm = CompiledMethods.getCompiledMethodUnchecked(cmid);
		
		return getMethodFullDesc (cm.method);
	}
	@Override
	public void incrementCallCount(int cmid) {
		// TODO Auto-generated method stub
		
		if (aosCollection == null)
		{
			return;
		}
		
		String methodFullDesc = getMethodFullDesc(cmid);

		double count = getCallCount (methodFullDesc);
		int optLevel = getOptLevel (cmid);
		
		if (VM.verboseClassLoading)
			VM.sysWriteln ("Method " + cmid + " with call counts " + count + " will be incremented");
		try {
			writeRequestsLock.acquire();
		} catch (InterruptedException e) {
			e.printStackTrace();
		}
		
		if (count != 0.0)
		{
			writeRequests.add(new UpdateOneModel <Document> (new Document ("methodFullDesc", methodFullDesc), 
				new Document ("$set", new Document ("methodFullDesc", methodFullDesc).append("count", new Double (count + 1.0)).append ("optLevel", optLevel))));
		}
		else
		{
			writeRequests.add(new InsertOneModel<Document> (new Document ("methodFullDesc", methodFullDesc).append ("count", new Double (1.0)).append ("optLevel", optLevel)));
		}
		
		writeRequestsLock.release();
		if (VM.verboseClassLoading)
			VM.sysWriteln ("Method " + cmid + " incremented to MongoDB");
		
		bulkWrite ();
	}

	@Override
	public void putCallCount(int cmid, double counts) {
		// TODO Auto-generated method stub
		//MongoDB implementation
		if (aosCollection == null)
		{
			return;
		}
		
		String methodFullDesc = getMethodFullDesc (cmid);
		int optLevel = getOptLevel (cmid);
		
		Document meth = new Document();
		meth.put("methodFullDesc", methodFullDesc);
		meth.put("count", counts);
		meth.put("optLevel", optLevel);
		
		if (VM.verboseClassLoading)
			VM.sysWriteln ("Put Method " + cmid + " with call counts " + counts);
		
		if (getCallCount (methodFullDesc) == 0.0)
		{
			if (VM.verboseClassLoading)
				VM.sysWriteln ("Method " + cmid + " added to MongoDB");
			try {
				writeRequestsLock.acquire();
			} catch (InterruptedException e) {
				e.printStackTrace();
			}
			writeRequests.add(new InsertOneModel<Document> (meth));
			writeRequestsLock.release();
		}
		else
		{
			if (VM.verboseClassLoading)
				VM.sysWriteln ("Method " + cmid + " updated to MongoDB");
			try {
				writeRequestsLock.acquire();
			} catch (InterruptedException e) {
				e.printStackTrace();
			}
			writeRequests.add(new UpdateOneModel <Document> (new Document ("methodFullDesc", methodFullDesc), 
					new Document ("$set", new Document ("methodFullDesc", methodFullDesc).append("count", counts))));
			writeRequestsLock.release();
		}
		bulkWrite ();
	}
	
	private double getCallCount (String methodFullDesc)
	{
		Document meth = aosCollection.find(eq("methodFullDesc", methodFullDesc)).first();
		if (meth == null)
		{
			return 0.0;
		}
		
		return ((Double)meth.get("count")).doubleValue();
	}
	@Override
	public double getCallCount(int cmid) {
		// TODO Auto-generated method stub
		if (aosCollection == null)
		{
			return 0.0;
		}
		
		String methodFullDesc = getMethodFullDesc (cmid);

		if (VM.verboseClassLoading)
			VM.sysWriteln ("Getting call count for Method "+ cmid);
		
		double counts = getCallCount (methodFullDesc);
		
		if (counts == 0.0 && VM.verboseClassLoading)
			VM.sysWriteln ("Cannot find call count for Method " + cmid);
		
		if (VM.verboseClassLoading)
			VM.sysWriteln ("Call count for Method " + cmid + " is " + counts);
		
		return counts;
	}
	
	@Override
	public void putMethodOptLevel (CompilationPlan cp)
	{
		compilationPlans.add(cp);
	}
	
	private void bulkWrite ()
	{
		if (writeRequests.size() >= bulkUpdateCount)
		{
			if (bulkWriteThread != null && bulkWriteThread.isAlive())
			{
				try {
					bulkWriteThread.join();
				} catch (InterruptedException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				}
			}
			
			bulkWriteThread = new BulkWriteThread(aosCollection, writeRequests, 
					writeRequestsLock, compPlanLock, compilationPlans);
			bulkWriteThread.start();
		}
	}
}