PooledStorage.java

package edu.berkeley.lipstick.storage;

import java.io.InputStream;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

public class PooledStorage implements IStorage {
    AtomicInteger csCount;
    AtomicLong numReads;
    AtomicLong numWrites;
    Class storageClass;
    ConcurrentLinkedQueue<IStorage> pool;
    ConcurrentLinkedQueue<IStorage> allStorage;

    boolean open;

    public PooledStorage (Class storageClass) {
        this.storageClass = storageClass;
    }

    public void open() throws Exception {
        open = true;
        csCount = new AtomicInteger(0);
        numReads = new AtomicLong(0);
        numWrites = new AtomicLong(0);
        pool = new ConcurrentLinkedQueue<IStorage>();
        allStorage = new ConcurrentLinkedQueue<IStorage>();
    }

    IStorage getStorage() throws Exception {
        IStorage ret = pool.poll();

        if(ret == null) {
            csCount.incrementAndGet();
            ret = (IStorage) (storageClass.newInstance());
            ret.open();
            allStorage.add(ret);
        }

        return ret;
    }

    void returnStorage(IStorage cs) {
        pool.add(cs);
    }

    public void close() throws Exception {
        open = false;
        int closed = 0;
        // todo: this really isn't safe...
        while(closed < csCount.get()) {
            IStorage cs = pool.poll();
            if(cs != null) {
                cs.close();
                closed++;
            }
        }
    }

    public byte[] get(String key) throws Exception {
        return get(key, false);
    }

    public void put(String key, byte[] value, long timestamp) throws Exception {
        put(key, value, timestamp, false);
    }

    public byte[] get(String key, boolean recordLat) throws Exception {
        if(!open)
            throw new Exception("pools closed");

        numReads.incrementAndGet();

        IStorage cs = getStorage();
        byte[] ret = cs.get(key, recordLat);
        returnStorage(cs);
        return ret;
    }

    public void put(String key, byte[] value, long timestamp, boolean recordLat) throws Exception {
        if(!open)
            throw new Exception("pools closed");

        numWrites.incrementAndGet();

        IStorage cs = getStorage();
        cs.put(key, value, timestamp, recordLat);
        returnStorage(cs);
    }

    public long getBytesWritten() {
        long ret = 0;
        for(IStorage cs : allStorage) {
            ret += cs.getBytesWritten();
        }

        return ret;
    }

    public long getBytesRead() {
        long ret = 0;
        for(IStorage cs : allStorage) {
            ret += cs.getBytesRead();
        }

        return ret;
    }

    public long getReadLatency() {
        long ret = 0;
        for(IStorage cs : allStorage) {
            ret += cs.getReadLatency();
        }

        return ret;
    }

    public long getWriteLatency() {
        long ret = 0;
        for(IStorage cs : allStorage) {
            ret += cs.getWriteLatency();
        }

        return ret;
    }

        public long getNumReads() {
            return numReads.get();
    }

    public long getNumWrites() {
        return numWrites.get();
    }
}