added cleanup method to task hook, added getThisWorkerPort to TopologyContext

Nathan Marz · Nathan Marz · commit 43a8df18f9d3 · 2012-03-15T17:41:29.000-07:00
diff --git a/src/clj/backtype/storm/daemon/task.clj b/src/clj/backtype/storm/daemon/task.clj
@@ -239,6 +239,8 @@
         (doseq [t all-threads]
           (.interrupt t)
           (.join t))
+        (doseq [hook (.getHooks user-context)]
+          (.cleanup hook))
         (.remove-task-heartbeat! storm-cluster-state storm-id task-id)
         (.disconnect storm-cluster-state)
         (.close puller)
diff --git a/src/clj/backtype/storm/daemon/worker.clj b/src/clj/backtype/storm/daemon/worker.clj
@@ -94,15 +94,17 @@
                                                (supervisor-storm-resources-path
                                                  (supervisor-stormdist-root conf storm-id))
                                                (worker-pids-root conf worker-id)
-                                               %)
+                                               %
+                                               port)
         mk-user-context #(TopologyContext. topology
                                            storm-conf
                                            task->component
                                            storm-id
                                            (supervisor-storm-resources-path
                                              (supervisor-stormdist-root conf storm-id))
                                            (worker-pids-root conf worker-id)
-                                           %)
+                                           %
+                                           port)
         mq-context (if mq-context
                      mq-context
                      (msg-loader/mk-zmq-context (storm-conf ZMQ-THREADS)
diff --git a/src/clj/backtype/storm/testing.clj b/src/clj/backtype/storm/testing.clj
@@ -542,6 +542,6 @@
         spout-spec (mk-spout-spec* (TestWordSpout.)
                                    {stream fields})
         topology (StormTopology. {component spout-spec} {} {})
-        context (TopologyContext. topology (read-storm-config) {1 component} "test-storm-id" nil nil 1)]
+        context (TopologyContext. topology (read-storm-config) {1 component} "test-storm-id" nil nil 1 nil)]
     (Tuple. context values 1 stream)
     ))
diff --git a/src/jvm/backtype/storm/hooks/BaseTaskHook.java b/src/jvm/backtype/storm/hooks/BaseTaskHook.java
@@ -13,6 +13,10 @@ public class BaseTaskHook implements ITaskHook {
     public void prepare(Map conf, TopologyContext context) {
     }
 
+    @Override
+    public void cleanup() {
+    }    
+
     @Override
     public void emit(EmitInfo info) {
     }
diff --git a/src/jvm/backtype/storm/hooks/ITaskHook.java b/src/jvm/backtype/storm/hooks/ITaskHook.java
@@ -10,6 +10,7 @@
 
 public interface ITaskHook {
     void prepare(Map conf, TopologyContext context);
+    void cleanup();
     void emit(EmitInfo info);
     void spoutAck(SpoutAckInfo info);
     void spoutFail(SpoutFailInfo info);
diff --git a/src/jvm/backtype/storm/task/TopologyContext.java b/src/jvm/backtype/storm/task/TopologyContext.java
@@ -40,10 +40,12 @@ public class TopologyContext {
     private Object _taskData = null;
     private List<ITaskHook> _hooks = new ArrayList<ITaskHook>();
     private Map _stormConf;
+    private Integer _workerPort;
     
-    public TopologyContext(StormTopology topology, Map stormConf, Map<Integer, String> taskToComponent, String stormId, String codeDir, String pidDir, Integer taskId) {
+    public TopologyContext(StormTopology topology, Map stormConf, Map<Integer, String> taskToComponent, String stormId, String codeDir, String pidDir, Integer taskId, Integer workerPort) {
         _topology = topology;
         _stormConf = stormConf;
+        _workerPort = workerPort;
         _taskToComponent = taskToComponent;
         _stormId = stormId;
         _taskId = taskId;
@@ -349,6 +351,10 @@ public int maxTopologyMessageTimeout() {
         return max;
     }
     
+    public Integer getThisWorkerPort() {
+        return _workerPort;
+    }
+    
     public void addTaskHook(ITaskHook hook) {
         hook.prepare(_stormConf, this);
         _hooks.add(hook);
diff --git a/test/clj/backtype/storm/integration_test.clj b/test/clj/backtype/storm/integration_test.clj
@@ -542,6 +542,8 @@
                   (reify backtype.storm.hooks.ITaskHook
                     (prepare [this conf context]
                       )
+                    (cleanup [this]
+                      )
                     (emit [this info]
                       (swap! emitted inc))
                     (boltAck [this info]

Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,10 @@ public class BaseTaskHook implements ITaskHook {`
`13`	`13`	`public void prepare(Map conf, TopologyContext context) {`
`14`	`14`	`}`
`15`	`15`
	`16`	`+ @Override`
	`17`	`+ public void cleanup() {`
	`18`	`+ }`
	`19`	`+`
`16`	`20`	`@Override`
`17`	`21`	`public void emit(EmitInfo info) {`
`18`	`22`	`}`
Original file line number	Diff line number	Diff line change
`@@ -542,6 +542,8 @@`
`542`	`542`	`(reify backtype.storm.hooks.ITaskHook`
`543`	`543`	`(prepare [this conf context]`
`544`	`544`	`)`
	`545`	`+ (cleanup [this]`
	`546`	`+ )`
`545`	`547`	`(emit [this info]`
`546`	`548`	`(swap! emitted inc))`
`547`	`549`	`(boltAck [this info]`