cache drpc clients to eliminate overhead

Nathan Marz · Nathan Marz · commit f39c0bd185fe · 2011-10-22T01:05:39.000-07:00
diff --git a/src/jvm/backtype/storm/drpc/ReturnResults.java b/src/jvm/backtype/storm/drpc/ReturnResults.java
@@ -1,7 +1,6 @@
 package backtype.storm.drpc;
 
 import backtype.storm.Config;
-import backtype.storm.ILocalDRPC;
 import backtype.storm.generated.DistributedRPC;
 import backtype.storm.task.OutputCollector;
 import backtype.storm.task.TopologyContext;
@@ -10,6 +9,9 @@
 import backtype.storm.tuple.Tuple;
 import backtype.storm.utils.DRPCClient;
 import backtype.storm.utils.ServiceRegistry;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import org.apache.thrift7.TException;
 import org.json.simple.JSONValue;
@@ -20,6 +22,8 @@ public class ReturnResults implements IRichBolt {
     OutputCollector _collector;
     boolean local;
 
+    Map<List, DRPCClient> _clients = new HashMap<List, DRPCClient>();
+
     public void prepare(Map stormConf, TopologyContext context, OutputCollector collector) {
         _collector = collector;
         local = stormConf.get(Config.STORM_CLUSTER_MODE).equals("local");
@@ -30,21 +34,26 @@ public void execute(Tuple input) {
         String returnInfo = (String) input.getValue(1);
         if(returnInfo!=null) {
             Map retMap = (Map) JSONValue.parse(returnInfo);
-            String host = (String) retMap.get("host");
-            int port = (int) ((Long) retMap.get("port")).longValue();
+            final String host = (String) retMap.get("host");
+            final int port = (int) ((Long) retMap.get("port")).longValue();                                   
             String id = (String) retMap.get("id");
             DistributedRPC.Iface client;
             if(local) {
                 client = (DistributedRPC.Iface) ServiceRegistry.getService(host);
             } else {
-                client = new DRPCClient(host, port);
+                List server = new ArrayList() {{
+                    add(host);
+                    add(port);
+                }};
+            
+                if(!_clients.containsKey(server)) {
+                    _clients.put(server, new DRPCClient(host, port));
+                }
+                client = _clients.get(server);
             }
                 
             try {
                 client.result(id, result);
-                if(!local) {
-                    ((DRPCClient)client).close();                   
-                }
                 _collector.ack(input);
             } catch(TException e) {
                 _collector.fail(input);
@@ -53,7 +62,9 @@ public void execute(Tuple input) {
     }    
 
     public void cleanup() {
-
+        for(DRPCClient c: _clients.values()) {
+            c.close();
+        }
     }
 
     public void declareOutputFields(OutputFieldsDeclarer declarer) {
