make MemoryTransactionalSpout work on a cluster

Nathan Marz · Nathan Marz · commit c0dae31f8ace · 2012-05-24T15:28:18.000-07:00
diff --git a/src/jvm/backtype/storm/testing/MemoryTransactionalSpout.java b/src/jvm/backtype/storm/testing/MemoryTransactionalSpout.java
@@ -26,13 +26,15 @@ public class MemoryTransactionalSpout implements IPartitionedTransactionalSpout<
     private String _finishedPartitionsId;
     private int _takeAmt;
     private Fields _outFields;
+    private Map<Integer, List<List<Object>>> _initialPartitions;
     
     public MemoryTransactionalSpout(Map<Integer, List<List<Object>>> partitions, Fields outFields, int takeAmt) {
         _id = RegisteredGlobalState.registerState(partitions);
         Map<Integer, Boolean> finished = Collections.synchronizedMap(new HashMap<Integer, Boolean>());
         _finishedPartitionsId = RegisteredGlobalState.registerState(finished);
         _takeAmt = takeAmt;
         _outFields = outFields;
+        _initialPartitions = partitions;
     }
     
     public boolean isExhaustedTuples() {
@@ -149,8 +151,10 @@ public void cleanup() {
         RegisteredGlobalState.clearState(_finishedPartitionsId);
     }
     
-    private Map<Integer, List<List<Object>>> getQueues() {
-        return (Map<Integer, List<List<Object>>>) RegisteredGlobalState.getState(_id);
+    private Map<Integer, List<List<Object>>> getQueues() {   
+        Map<Integer, List<List<Object>>> ret = (Map<Integer, List<List<Object>>>) RegisteredGlobalState.getState(_id);
+        if(ret!=null) return ret;
+        else return _initialPartitions;
     }
     
     private Map<Integer, Boolean> getFinishedStatuses() {
