VoldemortClient.java

package com.yahoo.ycsb.db;

import java.util.HashMap;
import java.util.Set;
import java.util.Vector;
import java.util.Map.Entry;

import org.apache.log4j.Logger;

import voldemort.client.ClientConfig;
import voldemort.client.SocketStoreClientFactory;
import voldemort.client.StoreClient;
import voldemort.versioning.VectorClock;
import voldemort.versioning.Versioned;

import com.yahoo.ycsb.DB;
import com.yahoo.ycsb.DBException;
import com.yahoo.ycsb.ByteIterator;
import com.yahoo.ycsb.StringByteIterator;


public class VoldemortClient extends DB {

	private StoreClient<String, HashMap<String, String>> storeClient;
	private SocketStoreClientFactory socketFactory;
	private String storeName;
    private final Logger logger = Logger.getLogger(VoldemortClient.class);
    
    public static final int OK = 0;
    public static final int ERROR = -1;
    public static final int NOT_FOUND = -2;
    
    /**
     * Initialize the DB layer. This accepts all properties allowed by the Voldemort client.
     * A store maps to a table.
     * Required : bootstrap_urls
     * Additional property : store_name -> to preload once, should be same as -t <table>
     * 
     * {@linktourl http://project-voldemort.com/javadoc/client/voldemort/client/ClientConfig.html}
     */
	public void init() throws DBException {
		ClientConfig clientConfig = new ClientConfig(getProperties());
		socketFactory = new SocketStoreClientFactory(clientConfig);
		
		// Retrieve store name
		storeName = getProperties().getProperty("store_name", "usertable");
		
		// Use store name to retrieve client
		storeClient = socketFactory.getStoreClient(storeName);
		if ( storeClient == null )
			throw new DBException("Unable to instantiate store client");
		
	}
	
	public void cleanup() throws DBException {
		socketFactory.close();
	}
	
	@Override
	public int delete(String table, String key) {
		if ( checkStore(table) == ERROR ) {
			return ERROR;
		}
		
		if ( storeClient.delete(key) )
			return OK;
		else
			return ERROR;
	}

	@Override
	public int insert(String table, String key, HashMap<String, ByteIterator> values) {
		if ( checkStore(table) == ERROR ) {
			return ERROR;
		}
		storeClient.put(key, (HashMap<String,String>)StringByteIterator.getStringMap(values));
		return OK;
	}

	@Override
	public int read(String table, String key, Set<String> fields,
			HashMap<String, ByteIterator> result) {
		if ( checkStore(table) == ERROR ) {
			return ERROR;
		}
		
		Versioned<HashMap<String, String>> versionedValue = storeClient.get(key);
		
		if ( versionedValue == null )
			return NOT_FOUND;
		
		if ( fields != null ) {
			for (String field : fields) {
				String val = versionedValue.getValue().get(field);
				if ( val != null )
				    result.put(field, new StringByteIterator(val));
			}
		} else {
			StringByteIterator.putAllAsByteIterators(result, versionedValue.getValue());
		}
		return OK;
	}

	@Override
	public int scan(String table, String startkey, int recordcount,
			Set<String> fields, Vector<HashMap<String, ByteIterator>> result) {
		logger.warn("Voldemort does not support Scan semantics");
		return OK;
	}

	@Override
	public int update(String table, String key, HashMap<String, ByteIterator> values) {
		if ( checkStore(table) == ERROR ) {
			return ERROR;
		}
		
		Versioned<HashMap<String, String>> versionedValue = storeClient.get(key);
		HashMap<String, String> value = new HashMap<String, String>();
		VectorClock version;
		if ( versionedValue != null ) {
			version = ((VectorClock) versionedValue.getVersion()).incremented(0, 1);
			value = versionedValue.getValue();
			for (Entry<String, ByteIterator> entry : values.entrySet()) {
				value.put(entry.getKey(), entry.getValue().toString());
			}
		} else {
			version = new VectorClock();
			StringByteIterator.putAllAsStrings(value, values);
		}
		
		storeClient.put(key, Versioned.value(value, version));
		return OK;
	}
	
	private int checkStore(String table) {
		if ( table.compareTo(storeName) != 0 ) {
			try {
				storeClient = socketFactory.getStoreClient(table);
				if ( storeClient == null ) {
					logger.error("Could not instantiate storeclient for " + table);
					return ERROR;
				}
				storeName = table;
			} catch ( Exception e ) {
				return ERROR;
			}
		}
		return OK;
	}

}