package net.vvakame.memvache;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.apache.commons.codec.digest.DigestUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import com.google.appengine.api.datastore.Key;
import com.google.appengine.api.memcache.MemcacheService;
import com.google.apphosting.api.DatastorePb.CommitResponse;
import com.google.apphosting.api.DatastorePb.DeleteRequest;
import com.google.apphosting.api.DatastorePb.GetRequest;
import com.google.apphosting.api.DatastorePb.GetResponse;
import com.google.apphosting.api.DatastorePb.GetResponse.Entity;
import com.google.apphosting.api.DatastorePb.PutRequest;
import com.google.apphosting.api.DatastorePb.PutResponse;
import com.google.apphosting.api.DatastorePb.Transaction;
import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;
import com.google.storage.onestore.v3.OnestoreEntity.Reference;
/**
* "Datastore への 単一 Entity の Get & Put の置き換え" を実装するクラス。<br>
* EntityがPutされる時は全てMemcacheに保持してDatastoreへ。<br>
* EntityがGetされる時はTx有りの時は素通し、それ以外の時はMemcacheを参照して無ければDatastoreへ。
* @author vvakame
*/
class GetPutCacheStrategy extends RpcVisitor {
private final static Log log = LogFactory.getLog(GetPutCacheStrategy.class);
// TODO Asyncを絡めて2回Getした時 pre_Get, pre_Get, post_Get, post_Get という順で動くとバグるのではないか？
/** オリジナルのリクエストが要求しているKeyの一覧 */
//List<Key> requestKeys;
Map<String, List<Key>> savedRequestKeys = new HashMap<String, List<Key>>();
/** Memcacheが持っていたEntityのキャッシュ */
//Map<Key, Entity> data;
Map<String, Map<Key, Entity>> savedCacheEntities = new HashMap<String, Map<Key, Entity>>();
Map<Long, Map<Key, Entity>> putUnderTx = new HashMap<Long, Map<Key, Entity>>();
/**
* Getを行う前の動作として、Memcacheから解決できる要素について処理を行う。<br>
* Memcacheからの不足分のみでリクエストを再構成する。<br>
* もし、Tx下であったら全てを素通しする。<br>
* @return 何も処理をしなかった場合 null を返す。キャッシュから全て済ませた場合は {@link Pair} のFirst。requestPbを再構成した時は {@link Pair} のSecond。
*/
@Override
public Pair<byte[], byte[]> pre_datastore_v3_Get(GetRequest requestPb) {
log.debug("pre_datastore_v3_Get: " + this);
if (requestPb.getTransaction().hasApp()) {
// under transaction
// 操作するEGに対してマークを付けさせるためにDatastoreに素通しする必要がある。
return null;
}
//requestKeys = PbKeyUtil.toKeys(requestPb.keys());
List<Key> theRequestKeys = PbKeyUtil.toKeys(requestPb.keys());
Map<Key, Entity> theData = new HashMap<Key, Entity>();
// Memcacheにあるものはキャッシュで済ませる
{
final MemcacheService memcache = MemvacheDelegate.getMemcache();
Map<Key, Object> all = memcache.getAll(theRequestKeys); // 存在しなかった場合Keyごと無い
theData = new HashMap<Key, Entity>();
for (Key key : all.keySet()) {
Entity entity = (Entity) all.get(key);
if (entity != null) {
theData.put(key, entity);
}
}
}
// もし全部取れた場合は Get動作を行わず結果を構成して返す。
if (theRequestKeys.size() == theData.size()) {
GetResponse responsePb = new GetResponse();
// toByteArray() を呼んだ時にNPEが発生するのを抑制するために内部的に new ArrayList() させる
responsePb.mutableEntitys();
responsePb.mutableDeferreds();
for (Key key : theRequestKeys) {
Entity entity = theData.get(key);
if (entity == null) {
theData.remove(key);
continue;
}
responsePb.addEntity(entity);
}
log.debug("all data was retrieved from memcache. finish.");
return Pair.response(responsePb.toByteArray());
}
// MemcacheにないものだけPbを再構成して投げる
for (int i = theRequestKeys.size() - 1; 0 <= i; i--) {
if (theData.containsKey(theRequestKeys.get(i))) {
requestPb.removeKey(i);
}
}
log.debug("key size: " + theRequestKeys.size() + " cache hit size: " + theData.size());
log.debug("continue to get from datastore. ");
// レスポンスのためにリクエストと紐付けてMapに持っておく
byte[] requestByte = requestPb.toByteArray();
String digest = DigestUtils.md5Hex(requestByte);
savedRequestKeys.put(digest, theRequestKeys);
savedCacheEntities.put(digest, theData);
log.debug("save data with digest: " + digest);
return Pair.request(requestByte);
}
/**
* Getを行った後の動作として、前処理で抜いた分のリクエストと実際にRPCした結果をマージし返す。<br>
* また、RPCして得られた結果についてMemcacheにキャッシュを作成する。
* @return 処理結果 or null
*/
@Override
public byte[] post_datastore_v3_Get(GetRequest requestPb, GetResponse responsePb) {
log.debug("post_datastore_v3_Get: " + this);
log.debug("requestPb: " + requestPb);
if (requestPb.getTransaction().hasApp()) {
// under transaction
return null;
}
// Memcacheに蓄える
Map<Key, Entity> newMap = new HashMap<Key, Entity>();
List<Reference> keys = requestPb.keys();
List<Entity> entitys = responsePb.entitys();
for (int i = 0; i < entitys.size(); i++) {
Key key = PbKeyUtil.toKey(keys.get(i)); // TODO SATO 並び順は保証されている？
Entity entity = entitys.get(i);
//Key key = PbKeyUtil.toKey(entity.getEntity().getKey()); TODO これだとhitしなかったときエラーになる
newMap.put(key, entity);
}
MemcacheService memcache = MemvacheDelegate.getMemcache();
memcache.putAll(newMap);
log.debug("get from datastore size: " + newMap.size());
// ここで取れてきているのはキャッシュにないヤツだけなので再構成して返す必要がある
byte[] requestByte = requestPb.toByteArray();
String digest = DigestUtils.md5Hex(requestByte);
log.debug("digest = " + digest);
log.debug("savedCacheEntities size: " + savedCacheEntities.size());
Map<Key, Entity> theData = savedCacheEntities.get(digest);
log.debug("the data is " + theData);
List<Key> theRequestKeys = savedRequestKeys.get(digest);
theData.putAll(newMap);
responsePb.clearEntity();
for (Key key : theRequestKeys) {
responsePb.addEntity(theData.get(key));
}
return responsePb.toByteArray();
}
/**
* Putを行った後の動作として、Memcacheにキャッシュを作成する。
*/
@Override
public byte[] post_datastore_v3_Put(PutRequest requestPb, PutResponse responsePb) {
Transaction tx = requestPb.getTransaction();
if (tx.hasApp()) {
// Tx下の場合はDatastoreに反映されるまで、ローカル変数に結果を保持しておく。
final long handle = tx.getHandle();
Map<Key, Entity> newMap = extractCache(requestPb, responsePb);
if (putUnderTx.containsKey(handle)) {
Map<Key, Entity> cached = putUnderTx.get(handle);
cached.putAll(newMap);
} else {
putUnderTx.put(handle, newMap);
}
} else {
MemcacheService memcache = MemvacheDelegate.getMemcache();
Map<Key, Entity> newMap = extractCache(requestPb, responsePb);
memcache.putAll(newMap);
}
return null;
}
private Map<Key, Entity> extractCache(PutRequest requestPb, PutResponse responsePb) {
Map<Key, Entity> newMap = new HashMap<Key, Entity>();
int size = requestPb.entitySize();
List<EntityProto> entitys = requestPb.entitys();
for (int i = 0; i < size; i++) {
EntityProto proto = entitys.get(i);
Reference reference = responsePb.getKey(i);
Key key = PbKeyUtil.toKey(reference);
Entity entity = new Entity();
entity.setEntity(proto);
entity.setKey(reference);
newMap.put(key, entity);
}
return newMap;
}
/**
* Deleteを行う前の動作として、とりあえずMemcacheからキャッシュを削除する。
*/
@Override
public Pair<byte[], byte[]> pre_datastore_v3_Delete(DeleteRequest requestPb) {
List<Key> keys = PbKeyUtil.toKeys(requestPb.keys());
MemcacheService memcache = MemvacheDelegate.getMemcache();
memcache.deleteAll(keys);
return null;
}
/**
* Commitを行った後の動作として、Putした時のキャッシュが存在していればMemcacheにキャッシュを作成する。
*/
@Override
public byte[] post_datastore_v3_Commit(Transaction requestPb, CommitResponse responsePb) {
final long handle = requestPb.getHandle();
if (putUnderTx.containsKey(handle)) {
Map<Key, Entity> map = putUnderTx.get(handle);
MemvacheDelegate.getMemcache().putAll(map);
return null;
} else {
return null;
}
}
/**
* Rollbackを行った後の動作として、Putした時のキャッシュが存在していればなかった事にする。
*/
@Override
public byte[] post_datastore_v3_Rollback(Transaction requestPb, CommitResponse responsePb) {
final long handle = requestPb.getHandle();
if (putUnderTx.containsKey(handle)) {
putUnderTx.remove(handle);
}
return null;
}
}