package com.onpositive.gae.profiling.rpc;
import java.beans.PropertyDescriptor;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import com.google.appengine.api.datastore.Blob;
import com.google.appengine.api.datastore.Text;
import com.google.apphosting.api.DatastorePb;
import com.google.apphosting.api.DatastorePb.Cost;
import com.google.apphosting.api.DatastorePb.GetResponse;
import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;
import com.google.storage.onestore.v3.OnestoreEntity.Path;
import com.google.storage.onestore.v3.OnestoreEntity.Property;
import com.google.storage.onestore.v3.OnestoreEntity.PropertyValue;
import com.google.storage.onestore.v3.OnestoreEntity.Reference;
import com.google.storage.onestore.v3.OnestoreEntity.Path.Element;
import com.google.storage.onestore.v3.OnestoreEntity.Property.Meaning;
import com.onpositive.gae.profiling.rpc.DataStoreCallStorageAttachment.GetRequestInfo;
import com.onpositive.gae.profiling.rpc.DataStoreCallStorageAttachment.PutRequestInfo;
import com.onpositive.gae.profiling.rpc.RPCAttachment.RequestInfo;
import com.onpositive.traps.ICallStorageAttachment;
import com.onpositive.traps.Profiler;
public class DataStoreCallProcessor implements IRPCCallProcessor {
private HashMap<String, ICallInterceptor> interceptors = new HashMap<String, ICallInterceptor>();
public HashMap<String, ICallInterceptor> getInterceptors() {
return interceptors;
}
static DataStoreCallStorageAttachment attachment = null;
static HashMap<RequestInfo, RequestInfo> infoMap = null;
static {
init();
}
private static void init() {
ICallStorageAttachment attachment2 = Profiler
.getAttachment(DataStoreCallStorageAttachment.class);
if (attachment2 == null) {
Profiler.putAttachment(new DataStoreCallStorageAttachment());
}
attachment = (DataStoreCallStorageAttachment) Profiler
.getAttachment(DataStoreCallStorageAttachment.class);
infoMap = attachment.infoMap;
}
public RequestInfo storedInfo(RequestInfo inf) {
RequestInfo requestInfo = infoMap.get(inf);
if (requestInfo == null) {
infoMap.put(inf, inf);
return inf;
}
return requestInfo;
}
public void recordInfo(RequestInfo inf) {
RequestInfo requestInfo = infoMap.get(inf);
if (requestInfo != null) {
inf.merge(requestInfo);
}
infoMap.put(inf, inf);
if (infoMap.size() > 10) {
long min = Long.MAX_VALUE;
RequestInfo mi = null;
for (RequestInfo i : infoMap.keySet()) {
if (i.cost < min) {
mi = i;
min = i.cost;
}
}
if (mi != null) {
infoMap.remove(mi);
}
}
}
public class RemoveInterCeptor implements ICallInterceptor {
private RequestInfo requestInfo;
private int cost;
public int getCost() {
return cost;
}
public void postCall(byte[] result) {
DatastorePb.DeleteResponse deleteReq = new DatastorePb.DeleteResponse();
deleteReq.mergeFrom(result);
}
public void preCall(byte[] call) {
DatastorePb.DeleteRequest deleteReq = new DatastorePb.DeleteRequest();
deleteReq.mergeFrom(call);
Exception exception = new Exception();
exception.fillInStackTrace();
StackTraceElement[] stackTrace = exception.getStackTrace();
requestInfo = new DataStoreCallStorageAttachment.DefaultRequestInfo(
"Delete Request", stackTrace);
requestInfo.count = 1;
cost = calculateCosts(deleteReq);
//System.out.println("Cost REMOVE:" + cost);
recordInfo(requestInfo);
}
private int calculateCosts(DatastorePb.DeleteRequest req) {
Iterator<Reference> ref = req.keyIterator();
int cost = 0;
while (ref.hasNext()) {
cost += 25;
Reference reff = ref.next();
PropertyDescriptor[] prd = reff.getPropertyDescriptors();
/*
* for (PropertyDescriptor prde : prd) { if
* (!Blob.class.isAssignableFrom(prde.getPropertyType()) &&
* !Text.class.isAssignableFrom(prde .getPropertyType())) { cost
* +=20; } }
*/}
return cost;
}
public void stat(int delta) {
storedInfo(requestInfo).cost += delta;
}
}
public class GetInterCeptor implements ICallInterceptor {
private GetRequestInfo requestInfo;
private int cost;
public int getCost() {
return cost;
}
public synchronized void postCall(byte[] result) {
GetResponse pr = new GetResponse();
pr.mergeFrom(result);
Exception exception = new Exception();
exception.fillInStackTrace();
StackTraceElement[] stackTrace = exception.getStackTrace();
requestInfo = new GetRequestInfo("Get Request", stackTrace);
int entitySize = pr.entitySize();
int size = pr.entitySize();
cost = size * 10;
if (requestInfo != null) {
requestInfo.entityCount += entitySize;
requestInfo.count++;
recordInfo(requestInfo);
}
// cost = -1;
//System.out.println("Cost GET:" + cost);
}
public synchronized void preCall(byte[] call) {
DatastorePb.GetRequest putRequest = new DatastorePb.GetRequest();
putRequest.mergeFrom(call);
}
public void stat(int delta) {
storedInfo(requestInfo).cost += delta;
}
}
public class PutInterceptor implements ICallInterceptor {
private int cost;
public int getCost() {
return cost;
}
PutRequestInfo requestInfo;
public synchronized void postCall(byte[] result) {
DatastorePb.PutResponse putRequest = new DatastorePb.PutResponse();
putRequest.mergeFrom(result);
Cost cost = putRequest.getCost();
if (requestInfo != null) {
updateInfo(requestInfo, cost);
recordInfo(requestInfo);
}
}
public synchronized void preCall(byte[] call) {
DatastorePb.PutRequest putRequest = new DatastorePb.PutRequest();
putRequest.mergeFrom(call);
putRequest.entitySize();
List<EntityProto> entitys = putRequest.entitys();
Exception exception = new Exception();
exception.fillInStackTrace();
StackTraceElement[] stackTrace = exception.getStackTrace();
HashSet<String> types = new HashSet<String>();
for (EntityProto p : entitys) {
Path entityGroup = p.getEntityGroup();
List<Element> elements = entityGroup.elements();
for (Element e : elements) {
types.add(e.getType());
}
}
this.cost = calculateCosts(putRequest);
requestInfo = new PutRequestInfo("Put:" + types, stackTrace);
}
public void stat(int delta) {
storedInfo(requestInfo).cost += delta;
}
private int calculateCosts(DatastorePb.PutRequest res) {
Iterator<EntityProto> i = res.entityIterator();
int cost = 0;
while (i.hasNext()) {
cost += 25;
EntityProto ep = i.next();
Iterator<Property> pi = ep.propertyIterator();
while (pi.hasNext()) {
Property p = pi.next();
Meaning m = p.getMeaningEnum();
if (m != null) {
if (m.compareTo(Meaning.BLOB) != 0
&& m.compareTo(Meaning.TEXT) != 0) {
cost += 20;
}
}
}
}
return cost;
}
}
public class QueryInterceptor implements ICallInterceptor {
private RequestInfo requestInfo;
private int cost;
public void postCall(byte[] result) {
DatastorePb.QueryResult resultQ = new DatastorePb.QueryResult();
resultQ.mergeFrom(result);
Exception exception = new Exception();
exception.fillInStackTrace();
StackTraceElement[] stackTrace = exception.getStackTrace();
requestInfo = new DataStoreCallStorageAttachment.DefaultRequestInfo(
"Query Request", stackTrace);
requestInfo.count = 1;
cost = calculateCosts(resultQ);
//System.out.println("Cost Query:" + cost);
recordInfo(requestInfo);
}
private int calculateCosts(DatastorePb.QueryResult res) {
Iterator<EntityProto> i = res.resultIterator();
int cost = 0;
while (i.hasNext()) {
cost += 25;
EntityProto ep = i.next();
Iterator<Property> pi = ep.propertyIterator();
while (pi.hasNext()) {
Property p = pi.next();
Meaning m = p.getMeaningEnum();
if (m != null) {
if (m.compareTo(Meaning.BLOB) != 0
&& m.compareTo(Meaning.TEXT) != 0) {
cost += 20;
}
}
}
}
return cost;
}
public void preCall(byte[] call) {
DatastorePb.Query query = new DatastorePb.Query();
query.mergeFrom(call);
}
public void stat(int delta) {
if (requestInfo != null) {
storedInfo(requestInfo).cost += delta;
}
requestInfo = null;
}
public int getCost() {
return cost;
}
}
public class CommitInterceptor implements ICallInterceptor {
private RequestInfo requestInfo;
public void postCall(byte[] result) {
}
public void preCall(byte[] call) {
Exception exception = new Exception();
exception.fillInStackTrace();
StackTraceElement[] stackTrace = exception.getStackTrace();
requestInfo = new DataStoreCallStorageAttachment.DefaultRequestInfo(
"Commit Request", stackTrace);
requestInfo.count = 1;
recordInfo(requestInfo);
}
public void stat(int delta) {
if (requestInfo != null) {
storedInfo(requestInfo).cost += delta;
}
requestInfo = null;
}
}
public DataStoreCallProcessor() {
interceptors.put("Get", new GetInterCeptor());
interceptors.put("Put", new PutInterceptor());
interceptors.put("RunQuery", new QueryInterceptor());
interceptors.put("Commit", new CommitInterceptor());
interceptors.put("Commit", new CommitInterceptor());
interceptors.put("Delete", new RemoveInterCeptor());
}
public String getPackage() {
return "datastore_v3";
}
public void postCall(String name, byte[] result) {
ICallInterceptor iCallInterceptor = interceptors.get(name);
if (iCallInterceptor != null) {
iCallInterceptor.postCall(result);
}
}
ICallInterceptor interceptor;
public void preCall(String name, byte[] data) {
ICallInterceptor iCallInterceptor = interceptors.get(name);
if (iCallInterceptor != null) {
iCallInterceptor.preCall(data);
interceptor = iCallInterceptor;
}
}
private void updateInfo(PutRequestInfo requestInfo, Cost cost) {
if (cost != null) {
int entityWrites = cost.getEntityWrites();
int entityWriteBytes = cost.getEntityWriteBytes();
int indexWriteBytes = cost.getIndexWriteBytes();
int indexWrites = cost.getIndexWrites();
requestInfo.entityWrites += entityWrites;
requestInfo.entityWritesBytes += entityWriteBytes;
requestInfo.indexWriteBytes += indexWriteBytes;
requestInfo.indexWrites += indexWrites;
}
requestInfo.count++;
}
public static void clear() {
init();
}
public void stat(int delta) {
if (interceptor != null) {
interceptor.stat(delta);
}
}
}