/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package backtype.storm.utils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;

import backtype.storm.messaging.IConnection;
import backtype.storm.messaging.TaskMessage;

public class TransferDrainer {

  private HashMap<Integer, ArrayList<ArrayList<TaskMessage>>> bundles = new HashMap();
  
  public void add(HashMap<Integer, ArrayList<TaskMessage>> taskTupleSetMap) {
    for (Integer key : taskTupleSetMap.keySet()) {
      
      ArrayList<ArrayList<TaskMessage>> bundle = bundles.get(key);
      if (null == bundle) {
        bundle = new ArrayList<ArrayList<TaskMessage>>();
        bundles.put(key, bundle);
      }
      
      ArrayList tupleSet = taskTupleSetMap.get(key);
      if (null != tupleSet && tupleSet.size() > 0) {
        bundle.add(tupleSet);
      }
    } 
  }
  
  public void send(HashMap<Integer, String> taskToNode, HashMap<String, IConnection> connections) {
    for (Integer task : bundles.keySet()) {
      String hostPort = taskToNode.get(task);
      if (hostPort != null) {
        IConnection connection = connections.get(hostPort);
        if (null != connection) {
          ArrayList<ArrayList<TaskMessage>> bundle = bundles.get(task);
          Iterator<TaskMessage> iter = getBundleIterator(bundle);
          if (null != iter && iter.hasNext()) {
            connection.send(iter);
          }
        }
      }
    }
  }
  
  private Iterator<TaskMessage> getBundleIterator(final ArrayList<ArrayList<TaskMessage>> bundle) {
    
    if (null == bundle) {
      return null;
    }
    
    return new Iterator<TaskMessage> () {
      
      private int offset = 0;
      private int size = 0;
      {
        for (ArrayList<TaskMessage> list : bundle) {
            size += list.size();
        }
      }
      
      private int bundleOffset = 0;
      private Iterator<TaskMessage> iter = bundle.get(bundleOffset).iterator();
      
      @Override
      public boolean hasNext() {
        if (offset < size) {
          return true;
        }
        return false;
      }

      @Override
      public TaskMessage next() {
        TaskMessage msg = null;
        if (iter.hasNext()) {
          msg = iter.next(); 
        } else {
          bundleOffset++;
          iter = bundle.get(bundleOffset).iterator();
          msg = iter.next();
        }
        if (null != msg) {
          offset++;
        }
        return msg;
      }

      @Override
      public void remove() {
        throw new RuntimeException("not supported");
      }
    };
  }
  
  public void clear() {
    bundles.clear();
  }
}
@@ -998,16 +998,55 @@
 Message;
+%0Aimport com.google.common.collect.Maps;
 %0A%0Apublic
@@ -1256,19 +1256,20 @@
 Integer 
-key
+task
  : taskT
@@ -1301,558 +1301,365 @@
     
-%0A      ArrayList%3CArrayList%3CTaskMessage%3E%3E bundle = bundles.get(key);%0A      if (null == bundle) %7B%0A        bundle = new ArrayList%3CArrayList%3CTaskMessage%3E%3E();%0A        bundles.put(key, bundle);%0A      %7D%0A      %0A      ArrayList tupleSet = taskTupleSetMap.get(key);%0A      if (null != tupleSet && tupleSet.size() %3E 0) %7B%0A        bundle.add(tupleSet);%0A      %7D%0A    %7D %0A  %7D%0A  %0A  public void send(HashMap%3CInteger, String%3E taskToNode, HashMap%3CString, IConnection%3E connections) %7B%0A    for (Integer task : bundles.keySet()) %7B%0A      String hostPort = taskToNode.get(task);
+addListRefToMap(this.bundles, task, taskTupleSetMap.get(task));%0A    %7D%0A  %7D%0A  %0A  public void send(HashMap%3CInteger, String%3E taskToNode, HashMap%3CString, IConnection%3E connections) %7B%0A    HashMap%3CString, ArrayList%3CArrayList%3CTaskMessage%3E%3E%3E bundleMapByDestination = groupBundleByDestination(taskToNode);%0A%0A    for (String hostPort : bundleMapByDestination.keySet()) %7B
 %0A   
@@ -1842,18 +1842,37 @@
 ndle
-s.get(task
+MapByDestination.get(hostPort
 );%0A 
@@ -2055,26 +2055,972 @@
 %7D%0A    %7D%0A  %7D%0A
+%0A
   
+private HashMap%3CString, ArrayList%3CArrayList%3CTaskMessage%3E%3E%3E groupBundleByDestination(HashMap%3CInteger, String%3E taskToNode) %7B%0A    HashMap%3CString, ArrayList%3CArrayList%3CTaskMessage%3E%3E%3E bundleMap = Maps.newHashMap();%0A    for (Integer task : this.bundles.keySet()) %7B%0A      String hostPort = taskToNode.get(task);%0A      if (hostPort != null) %7B%0A        for (ArrayList%3CTaskMessage%3E chunk : this.bundles.get(task)) %7B%0A          addListRefToMap(bundleMap, hostPort, chunk);%0A        %7D%0A      %7D%0A    %7D%0A    return bundleMap;%0A  %7D%0A%0A  private %3CT%3E void addListRefToMap(HashMap%3CT, ArrayList%3CArrayList%3CTaskMessage%3E%3E%3E bundles,%0A                                   T key, ArrayList%3CTaskMessage%3E tuples) %7B%0A    ArrayList%3CArrayList%3CTaskMessage%3E%3E bundle = bundles.get(key);%0A%0A    if (null == bundle) %7B%0A      bundle = new ArrayList%3CArrayList%3CTaskMessage%3E%3E();%0A      bundles.put(key, bundle);%0A    %7D%0A%0A    if (null != tuples && tuples.size() %3E 0) %7B%0A      bundle.add(tuples);%0A    %7D%0A  %7D%0A
 %0A  private I

Change TransferDrainer to re-group msg by destination when sending