ProcedureNameFilter.java

package edu.brown.workload.filters;

import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.log4j.Logger;
import org.voltdb.VoltProcedure;
import org.voltdb.catalog.CatalogType;

import edu.brown.workload.AbstractTraceElement;
import edu.brown.workload.TransactionTrace;

/**
 * Filter TransactionTrace records based on their procedure name
 * @author pavlo
 */
public class ProcedureNameFilter extends Filter {
    private static final Logger LOG = Logger.getLogger(ProcedureNameFilter.class);
    private static final boolean d = LOG.isDebugEnabled();
    
    public static final String INCLUDE_ALL = "*";
    public static final Integer INCLUDE_UNLIMITED = Integer.MAX_VALUE;

    /** If set to true, then we will increase includedTxn by the txn's weights */
    private final boolean weighted;
    
    /**
     * Whitelist counters
     * When includeFinished set is the same size as the includeTotals, then we will HALT
     */
    private final Map<String, AtomicInteger> includeCounters = new HashMap<String, AtomicInteger>();
    private final Set<String> includeFinished = new HashSet<String>();
    private final Map<String, Integer> includeTotals = new HashMap<String, Integer>();
    
    /**
     * Blacklist
     */
    private final Set<String> exclude = new HashSet<String>();
    
    public ProcedureNameFilter(boolean weighted) {
        super();
        this.weighted = weighted;
    }
    
    /**
     * Add the given catalog object names to this filter's whitelist
     * @param names
     */
    public ProcedureNameFilter include(String...names) {
        for (String name : names) {
            this.include(name, INCLUDE_UNLIMITED);
        }
        return (this);
    }
    
    @SuppressWarnings("unchecked")
    public ProcedureNameFilter include(Class<? extends VoltProcedure>...procClasses) {
        for (Class<? extends VoltProcedure> procClass : procClasses) {
            this.include(procClass.getSimpleName(), INCLUDE_UNLIMITED);
        }
        return (this);
    }
    
    public ProcedureNameFilter include(String name, int limit) {
        this.includeCounters.put(name, new AtomicInteger(limit));
        this.includeTotals.put(name, limit);
        return (this);
    }
    
    public Map<String, Integer> getProcIncludes() {
        return (Collections.unmodifiableMap(this.includeTotals));
    }
    
    /**
     * Return the set of procedure names that we want to include
     * @return
     */
    public Set<String> getProcedureNames() {
        return (this.includeCounters.keySet());
    }
    
    /**
     * Add the given catalog object names to this filter's blacklist
     * @param names
     */
    public ProcedureNameFilter exclude(String...names) {
        for (String proc_name : names) {
            this.exclude.add(proc_name);
            this.includeCounters.remove(proc_name);
            this.includeTotals.remove(proc_name);
        } // FOR
        return (this);
    }
    
    @Override
    public String debugImpl() {
        return String.format("%s[include=%s / exclude=%s / finished=%s]",
                             this.getClass().getSimpleName(),
                             this.includeCounters, this.exclude, this.includeFinished);
    }
    
    @Override
    protected void resetImpl() {
        for (Entry<String, AtomicInteger> e : this.includeCounters.entrySet()) {
            assert(this.includeTotals.containsKey(e.getKey())) : "Missing '" + e.getKey() + "'";
            e.getValue().set(this.includeTotals.get(e.getKey()));
        } // FOR
        this.includeFinished.clear();
    }
    
    @Override
    public FilterResult filter(AbstractTraceElement<? extends CatalogType> element) {
        FilterResult result = FilterResult.ALLOW;
        if (element instanceof TransactionTrace) {
            final String name = element.getCatalogItemName();
            
            // BLACKLIST
            if (this.exclude.contains(name)) {
                result = FilterResult.SKIP;
                
            // WHITELIST
            } else if (this.includeCounters.isEmpty() == false) {
                if (d) LOG.debug("Checking whether " + name + " is in whitelist [total=" + this.includeCounters.size() + ", finished=" + this.includeFinished.size() + "]");
                
                // If the HALT countdown is zero, then we know that we have all of the procs that
                // we want for this workload
                if (this.includeFinished.size() == this.includeCounters.size()) {
                    result = FilterResult.HALT;
                    
                // Check whether this guy is allowed and keep count of how many we've added
                } else if (this.includeCounters.containsKey(name) && this.includeFinished.contains(name) == false) {
                    int weight = (this.weighted ? element.getWeight() : 1);
                    int count = this.includeCounters.get(name).getAndAdd(-1*weight);
                    
                    // If the counter goes to zero, then we have seen enough of this procedure
                    // Reset its counter back to 1 so that the next time we see it it will always get skipped
                    if (count <= 0) {
                        result = FilterResult.SKIP;
                        this.includeFinished.add(name);
                        if (d) LOG.debug("Transaction '" + name + "' has exhausted count. Skipping...");
                    } else if (d) {
                        LOG.debug("Transaction '" + name + "' is allowed [remaining=" + count + "]");
                    }
                // Not allowed. Just skip...
                } else {
                    result = FilterResult.SKIP;
                }
            }
        }
        return (result);
    }
}