package edu.brown.hstore.callbacks;
import java.util.Collection;
import org.apache.log4j.Logger;
import com.google.protobuf.RpcCallback;
import edu.brown.hstore.HStoreSite;
import edu.brown.hstore.Hstoreservice;
import edu.brown.hstore.Hstoreservice.Status;
import edu.brown.hstore.Hstoreservice.TransactionMapResponse;
import edu.brown.hstore.txns.MapReduceTransaction;
import edu.brown.hstore.util.MapReduceHelperThread;
import edu.brown.logging.LoggerUtil;
import edu.brown.logging.LoggerUtil.LoggerBoolean;
/**
* This is callback is used on the remote side of a TransactionMapRequest
* so that the network-outbound callback is not invoked until all of the partitions
* at this HStoreSite is finished with the Map phase.
* @author pavlo
*/
public class TransactionMapWrapperCallback extends BlockingRpcCallback<TransactionMapResponse, Integer> {
private static final Logger LOG = Logger.getLogger(TransactionMapWrapperCallback.class);
private static final LoggerBoolean debug = new LoggerBoolean();
private static final LoggerBoolean trace = new LoggerBoolean();
static {
LoggerUtil.attachObserver(LOG, debug, trace);
}
private TransactionMapResponse.Builder builder = null;
private MapReduceTransaction ts;
public TransactionMapResponse.Builder getBuilder() {
return builder;
}
public TransactionMapWrapperCallback(HStoreSite hstore_site) {
super(hstore_site, false);
}
public void init(MapReduceTransaction ts, RpcCallback<TransactionMapResponse> orig_callback) {
assert(this.isInitialized() == false) :
String.format("Trying to initialize %s twice! [origTs=%s, newTs=%s]",
this.getClass().getSimpleName(), this.ts, ts);
if (debug.val)
LOG.debug("Starting new " + this.getClass().getSimpleName() + " for " + ts);
this.ts = ts;
this.builder = TransactionMapResponse.newBuilder()
.setTransactionId(ts.getTransactionId())
.setStatus(Hstoreservice.Status.OK);
super.init(ts.getTransactionId(), hstore_site.getLocalPartitionIds().size(), orig_callback);
}
@Override
protected void abortCallback(Status status) {
if (debug.val)
LOG.debug(String.format("Txn #%d - Aborting %s with status %s",
this.getTransactionId(), this.getClass().getSimpleName(), status));
this.builder.setStatus(status);
Collection<Integer> localPartitions = hstore_site.getLocalPartitionIds();
for (Integer p : this.hstore_site.getLocalPartitionIds()) {
if (localPartitions.contains(p) && this.builder.getPartitionsList().contains(p) == false) {
this.builder.addPartitions(p.intValue());
}
} // FOR
this.unblockCallback();
}
@Override
protected void finishImpl() {
this.builder = null;
this.ts = null;
}
@Override
public boolean isInitialized() {
return ( this.ts != null && this.builder != null && super.isInitialized());
}
@Override
protected synchronized int runImpl(Integer partition) {
if (this.isAborted() == false)
this.builder.addPartitions(partition.intValue());
assert(this.ts != null) :
String.format("Missing MapReduceTransaction handle for txn #%d", this.ts.getTransactionId());
return 1;
}
@Override
protected void unblockCallback() {
if (debug.val) {
LOG.debug(String.format("Txn #%d - Sending %s to %s with status %s",
this.getTransactionId(),
TransactionMapResponse.class.getSimpleName(),
this.getOrigCallback().getClass().getSimpleName(),
this.builder.getStatus()));
}
assert(this.getOrigCounter() == builder.getPartitionsCount()) :
String.format("The %s for txn #%d has results from %d partitions but it was suppose to have %d.",
builder.getClass().getSimpleName(), this.getTransactionId(), builder.getPartitionsCount(), this.getOrigCounter());
assert(this.getOrigCallback() != null) :
String.format("The original callback for txn #%d is null!", this.getTransactionId());
// Get the MapReduceHelperThread object from the HStoreSite
// Pass the MapReduceTransaction handle to the helper thread to perform the shuffle operation
// Move this to be execute after the SHUFFLE phase is finished --> this.getOrigCallback().run(this.builder.build());
MapReduceHelperThread mr_helper = this.hstore_site.getMapReduceHelper();
ts.setShufflePhase();
if (debug.val)
LOG.debug(String.format("Txn #%d - I am swithing to SHUFFLE phase, go to MR_helper thread",this.getTransactionId()));
// enqueue this MapReduceTransaction to do shuffle work
mr_helper.queue(this.ts);
}
public void runOrigCallback() {
this.getOrigCallback().run(this.builder.build());
}
}