package com.yahoo.ycsb.db;
import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.Properties;
import java.util.Set;
import java.util.Vector;
import org.apache.thrift.TException;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TFramedTransport;
import org.apache.thrift.transport.TSocket;
import org.apache.thrift.transport.TTransport;
import com.yahoo.mapkeeper.BinaryResponse;
import com.yahoo.mapkeeper.MapKeeper;
import com.yahoo.mapkeeper.Record;
import com.yahoo.mapkeeper.RecordListResponse;
import com.yahoo.mapkeeper.ResponseCode;
import com.yahoo.mapkeeper.ScanOrder;
import com.yahoo.ycsb.ByteIterator;
import com.yahoo.ycsb.DB;
import com.yahoo.ycsb.StringByteIterator;
import com.yahoo.ycsb.workloads.CoreWorkload;
public class MapKeeperClient extends DB {
private static final String HOST = "mapkeeper.host";
private static final String HOST_DEFAULT = "localhost";
private static final String PORT = "mapkeeper.port";
private static final String PORT_DEFAULT = "9090";
MapKeeper.Client c;
boolean writeallfields;
static boolean initteddb = false;
private synchronized static void initDB(Properties p, MapKeeper.Client c) throws TException {
if(!initteddb) {
initteddb = true;
c.addMap(p.getProperty(CoreWorkload.TABLENAME_PROPERTY, CoreWorkload.TABLENAME_PROPERTY_DEFAULT));
}
}
public void init() {
String host = getProperties().getProperty(HOST, HOST_DEFAULT);
int port = Integer.parseInt(getProperties().getProperty(PORT, PORT_DEFAULT));
TTransport tr = new TFramedTransport(new TSocket(host, port));
TProtocol proto = new TBinaryProtocol(tr);
c = new MapKeeper.Client(proto);
try {
tr.open();
initDB(getProperties(), c);
} catch(TException e) {
throw new RuntimeException(e);
}
writeallfields = Boolean.parseBoolean(getProperties().getProperty(CoreWorkload.WRITE_ALL_FIELDS_PROPERTY,
CoreWorkload.WRITE_ALL_FIELDS_PROPERTY_DEFAULT));
}
ByteBuffer encode(HashMap<String, ByteIterator> values) {
int len = 0;
for(String k : values.keySet()) {
len += (k.length() + 1 + values.get(k).bytesLeft() + 1);
}
byte[] array = new byte[len];
int i = 0;
for(String k : values.keySet()) {
for(int j = 0; j < k.length(); j++) {
array[i] = (byte)k.charAt(j);
i++;
}
array[i] = '\t'; // XXX would like to use sane delimiter (null, 254, 255, ...) but java makes this nearly impossible
i++;
ByteIterator v = values.get(k);
i = v.nextBuf(array, i);
array[i] = '\t';
i++;
}
array[array.length-1] = 0;
ByteBuffer buf = ByteBuffer.wrap(array);
buf.rewind();
return buf;
}
void decode(Set<String> fields, String tups, HashMap<String, ByteIterator> tup) {
String[] tok = tups.split("\\t");
if(tok.length == 0) { throw new IllegalStateException("split returned empty array!"); }
for(int i = 0; i < tok.length; i+=2) {
if(fields == null || fields.contains(tok[i])) {
if(tok.length < i+2) { throw new IllegalStateException("Couldn't parse tuple <" + tups + "> at index " + i); }
if(tok[i] == null || tok[i+1] == null) throw new NullPointerException("Key is " + tok[i] + " val is + " + tok[i+1]);
tup.put(tok[i], new StringByteIterator(tok[i+1]));
}
}
if(tok.length == 0) {
System.err.println("Empty tuple: " + tups);
}
}
int ycsbThriftRet(BinaryResponse succ, ResponseCode zero, ResponseCode one) {
return ycsbThriftRet(succ.responseCode, zero, one);
}
int ycsbThriftRet(ResponseCode rc, ResponseCode zero, ResponseCode one) {
return
rc == zero ? 0 :
rc == one ? 1 : 2;
}
ByteBuffer bufStr(String str) {
ByteBuffer buf = ByteBuffer.wrap(str.getBytes());
return buf;
}
String strResponse(BinaryResponse buf) {
return new String(buf.value.array());
}
@Override
public int read(String table, String key, Set<String> fields,
HashMap<String, ByteIterator> result) {
try {
ByteBuffer buf = bufStr(key);
BinaryResponse succ = c.get(table, buf);
int ret = ycsbThriftRet(
succ,
ResponseCode.RecordExists,
ResponseCode.RecordNotFound);
if(ret == 0) {
decode(fields, strResponse(succ), result);
}
return ret;
} catch(TException e) {
e.printStackTrace();
return 2;
}
}
@Override
public int scan(String table, String startkey, int recordcount,
Set<String> fields, Vector<HashMap<String, ByteIterator>> result) {
try {
//XXX what to pass in for nulls / zeros?
RecordListResponse res = c.scan(table, ScanOrder.Ascending, bufStr(startkey), true, null, false, recordcount, 0);
int ret = ycsbThriftRet(res.responseCode, ResponseCode.Success, ResponseCode.ScanEnded);
if(ret == 0) {
for(Record r : res.records) {
HashMap<String, ByteIterator> tuple = new HashMap<String, ByteIterator>();
// Note: r.getKey() and r.getValue() call special helper methods that trim the buffer
// to an appropriate length, and memcpy it to a byte[]. Trying to manipulate the ByteBuffer
// directly leads to trouble.
tuple.put("key", new StringByteIterator(new String(r.getKey())));
decode(fields, new String(r.getValue())/*strBuf(r.bufferForValue())*/, tuple);
result.add(tuple);
}
}
return ret;
} catch(TException e) {
e.printStackTrace();
return 2;
}
}
@Override
public int update(String table, String key,
HashMap<String, ByteIterator> values) {
try {
if(!writeallfields) {
HashMap<String, ByteIterator> oldval = new HashMap<String, ByteIterator>();
read(table, key, null, oldval);
for(String k: values.keySet()) {
oldval.put(k, values.get(k));
}
values = oldval;
}
ResponseCode succ = c.update(table, bufStr(key), encode(values));
return ycsbThriftRet(succ, ResponseCode.RecordExists, ResponseCode.RecordNotFound);
} catch(TException e) {
e.printStackTrace();
return 2;
}
}
@Override
public int insert(String table, String key,
HashMap<String, ByteIterator> values) {
try {
int ret = ycsbThriftRet(c.insert(table, bufStr(key), encode(values)), ResponseCode.Success, ResponseCode.RecordExists);
return ret;
} catch(TException e) {
e.printStackTrace();
return 2;
}
}
@Override
public int delete(String table, String key) {
try {
return ycsbThriftRet(c.remove(table, bufStr(key)), ResponseCode.Success, ResponseCode.RecordExists);
} catch(TException e) {
e.printStackTrace();
return 2;
}
}
}