finished test for restarting a transactional topology

Nathan Marz · Nathan Marz · commit e223f937a869 · 2012-01-26T14:50:13.000-08:00
diff --git a/src/clj/backtype/storm/testing.clj b/src/clj/backtype/storm/testing.clj
@@ -308,39 +308,46 @@
 
 (defprotocol CompletableSpout
   (exhausted? [this] "Whether all the tuples for this spout have been completed.")
-  (cleanup [this] "Cleanup any global state kept"))
+  (cleanup [this] "Cleanup any global state kept")
+  (startup [this] "Prepare the spout (globally) before starting the topology"))
 
 (extend-type FixedTupleSpout
   CompletableSpout
   (exhausted? [this]
     (= (-> this .getSourceTuples count)
        (.getCompleted this)))
   (cleanup [this]
-    (.cleanup this)))
+    (.cleanup this))
+  (startup [this]
+    ))
 
 (extend-type TransactionalSpoutCoordinator
   CompletableSpout
   (exhausted? [this]
     (exhausted? (.getSpout this)))
   (cleanup [this]
-    (cleanup (.getSpout this))
-    ))
+    (cleanup (.getSpout this)))
+  (startup [this]
+    (startup (.getSpout this))))
 
 (extend-type PartitionedTransactionalSpoutExecutor
   CompletableSpout
   (exhausted? [this]
     (exhausted? (.getPartitionedSpout this)))
   (cleanup [this]
-    (cleanup (.getPartitionedSpout this))
+    (cleanup (.getPartitionedSpout this)))
+  (startup [this]
+    (startup (.getPartitionedSpout this))
     ))
 
 (extend-type MemoryTransactionalSpout
   CompletableSpout
   (exhausted? [this]
     (.isExhaustedTuples this))
   (cleanup [this]
-    (.cleanup this)
-    ))
+    (.cleanup this))
+  (startup [this]
+    (.startup this)))
 
 (defn spout-objects [spec-map]
   (for [[_ spout-spec] spec-map]
@@ -375,7 +382,7 @@
     ))
 
 ;; TODO: mock-sources needs to be able to mock out state spouts as well
-(defnk complete-topology [cluster-map topology :mock-sources {} :storm-conf {}]
+(defnk complete-topology [cluster-map topology :mock-sources {} :storm-conf {} :cleanup-state true]
   ;; TODO: the idea of mocking for transactional topologies should be done an
   ;; abstraction level above... should have a complete-transactional-topology for this
   (let [{topology :topology capturer :capturer} (capture-topology topology)
@@ -400,7 +407,9 @@
       (when-not (extends? CompletableSpout (.getClass spout))
         (throw (RuntimeException. "Cannot complete topology unless every spout is a CompletableSpout (or mocked to be)"))
         ))
-    
+
+    (doseq [spout (spout-objects spouts)]
+      (startup spout))
     
     (submit-local-topology (:nimbus cluster-map) storm-name storm-conf topology)
     
@@ -412,10 +421,13 @@
       (.killTopology (:nimbus cluster-map) storm-name)
       (while (.assignment-info state storm-id nil)
         (simulate-wait cluster-map))
-      (doseq [spout (spout-objects spouts)]
-        (cleanup spout)))
+      (when cleanup-state
+        (doseq [spout (spout-objects spouts)]
+          (cleanup spout))))
 
-    (.getAndClearResults capturer)
+    (if cleanup-state
+      (.getAndRemoveResults capturer)
+      (.getAndClearResults capturer))
     ))
 
 (defn read-tuples
diff --git a/src/jvm/backtype/storm/testing/MemoryTransactionalSpout.java b/src/jvm/backtype/storm/testing/MemoryTransactionalSpout.java
@@ -132,6 +132,10 @@ public Map<String, Object> getComponentConfiguration() {
         return conf;
     }
     
+    public void startup() {
+        getFinishedStatuses().clear();
+    }
+    
     public void cleanup() {
         RegisteredGlobalState.clearState(_id);
         RegisteredGlobalState.clearState(_finishedPartitionsId);
diff --git a/src/jvm/backtype/storm/testing/TupleCaptureBolt.java b/src/jvm/backtype/storm/testing/TupleCaptureBolt.java
@@ -44,10 +44,16 @@ public Map<String, List<FixedTuple>> getResults() {
     public void cleanup() {
     }
     
-    public Map<String, List<FixedTuple>> getAndClearResults() {
+    public Map<String, List<FixedTuple>> getAndRemoveResults() {
         return emitted_tuples.remove(_name);
     }
 
+    public Map<String, List<FixedTuple>> getAndClearResults() {
+        Map<String, List<FixedTuple>> ret = new HashMap<String, List<FixedTuple>>(emitted_tuples.get(_name));
+        emitted_tuples.get(_name).clear();
+        return ret;
+    }
+    
     @Override
     public void declareOutputFields(OutputFieldsDeclarer declarer) {
     }
diff --git a/test/clj/backtype/storm/transactional_test.clj b/test/clj/backtype/storm/transactional_test.clj
@@ -17,8 +17,7 @@
 ;; * Test that it repeats the meta for a partitioned state (test partitioned emitter on its own)
 ;; * Test that partitioned state emits nothing for the partition if it has seen a future transaction for that partition (test partitioned emitter on its own)
 
-
-
+;; TODO: test that FailedBatchException works
 
 (defn mk-coordinator-state-changer [atom]
   (TransactionalSpoutCoordinator.
@@ -540,8 +539,52 @@
        (-> topo-info :capturer .getAndClearResults)
        ))))
 
+(deftest test-transactional-topology-restart
+  (with-simulated-time-local-cluster [cluster]
+    (letlocals
+     (bind data (mk-transactional-source))
+     (bind builder (TransactionalTopologyBuilder.
+                    "id"
+                    "spout"
+                    (MemoryTransactionalSpout. data
+                                               (Fields. ["word"])
+                                               3)
+                    2))
+
+     (-> builder
+         (.setBolt "count" (CountingCommitBolt.) 2)
+         (.globalGrouping "spout"))
+
+     (add-transactional-data data
+                             {0 [["a"]
+                                 ["b"]
+                                 ["c"]
+                                 ["d"]]
+                              1 [["d"]
+                                 ["c"]]
+                              })
+     
+     (bind results (complete-topology cluster
+                                      (.buildTopology builder)
+                                      :cleanup-state false))
+
+     (is (ms= [[5] [0] [1] [0]] (->> (read-tuples results "count")
+                                     (take 4)
+                                     (map (partial drop 1))
+                                     )))
+
+     (add-transactional-data data
+                             {0 [["a"]
+                                 ["b"]]
+                              })
+     
+     (bind results (complete-topology cluster (.buildTopology builder)))
+
+     (println results)
+     ;; need to do it this way (check for nothing transaction) because there is one transaction already saved up before that emits nothing (because of how memorytransctionalspout detects partition completion)
+     (is (ms= [[0] [0] [2] [0]] (->> (read-tuples results "count")
+                             (take 4)
+                             (map (partial drop 1))
+                             )))
+     )))
 
-
-
-;; TODO: ;; * Test that it picks up where it left off when restarting the topology
-;;      - run topology and restart it

Original file line number	Diff line number	Diff line change
`@@ -44,10 +44,16 @@ public Map<String, List<FixedTuple>> getResults() {`
`44`	`44`	`public void cleanup() {`
`45`	`45`	`}`
`46`	`46`
`47`		`- public Map<String, List<FixedTuple>> getAndClearResults() {`
	`47`	`+ public Map<String, List<FixedTuple>> getAndRemoveResults() {`
`48`	`48`	`return emitted_tuples.remove(_name);`
`49`	`49`	`}`
`50`	`50`
	`51`	`+ public Map<String, List<FixedTuple>> getAndClearResults() {`
	`52`	`+ Map<String, List<FixedTuple>> ret = new HashMap<String, List<FixedTuple>>(emitted_tuples.get(_name));`
	`53`	`+ emitted_tuples.get(_name).clear();`
	`54`	`+ return ret;`
	`55`	`+ }`
	`56`	`+`
`51`	`57`	`@Override`
`52`	`58`	`public void declareOutputFields(OutputFieldsDeclarer declarer) {`
`53`	`59`	`}`