STORM-162: Load Aware Shuffle Grouping

Project: http://git-wip-us.apache.org/repos/asf/storm/repo
Commit: http://git-wip-us.apache.org/repos/asf/storm/commit/af64a3a2
Tree: http://git-wip-us.apache.org/repos/asf/storm/tree/af64a3a2
Diff: http://git-wip-us.apache.org/repos/asf/storm/diff/af64a3a2

Branch: refs/heads/master
Commit: af64a3a2ad951a5753c82d7acffc97518b847dcd
Parents: 7edf520
Author: Bobby Evans <[email protected]>
Authored: Tue Feb 10 09:07:21 2015 -0600
Committer: Robert (Bobby) Evans <[email protected]>
Committed: Tue Nov 3 12:17:34 2015 -0600

----------------------------------------------------------------------
 conf/defaults.yaml                              |   1 +
 .../src/clj/backtype/storm/daemon/executor.clj  |  46 ++--
 .../src/clj/backtype/storm/daemon/task.clj      |   4 +-
 .../src/clj/backtype/storm/daemon/worker.clj    |  29 +++
 .../src/clj/backtype/storm/messaging/local.clj  |  34 ++-
 storm-core/src/clj/backtype/storm/timer.clj     |  20 +-
 storm-core/src/jvm/backtype/storm/Config.java   |   7 +
 .../src/jvm/backtype/storm/grouping/Load.java   |  77 ++++++
 .../grouping/LoadAwareCustomStreamGrouping.java |  24 ++
 .../grouping/LoadAwareShuffleGrouping.java      |  76 ++++++
 .../backtype/storm/grouping/LoadMapping.java    |  64 +++++
 .../storm/grouping/ShuffleGrouping.java         |  65 ++++++
 .../backtype/storm/messaging/IConnection.java   |  16 ++
 .../backtype/storm/messaging/netty/Client.java  | 100 ++++++--
 .../storm/messaging/netty/ISaslClient.java      |  28 +++
 .../storm/messaging/netty/ISaslServer.java      |  26 +++
 .../backtype/storm/messaging/netty/IServer.java |  26 +++
 .../messaging/netty/SaslStormClientHandler.java |  37 +--
 .../backtype/storm/messaging/netty/Server.java  |  75 ++++--
 .../messaging/netty/StormClientHandler.java     |  51 +++-
 .../netty/StormClientPipelineFactory.java       |  11 +-
 .../test/clj/backtype/storm/grouping_test.clj   |  90 +++++--
 .../clj/backtype/storm/integration_test.clj     |   4 +-
 .../storm/messaging/netty_integration_test.clj  |   3 +-
 .../storm/messaging/netty_unit_test.clj         | 232 ++++++++++++++++---
 .../test/clj/backtype/storm/messaging_test.clj  |   3 +-
 26 files changed, 1005 insertions(+), 144 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/conf/defaults.yaml
----------------------------------------------------------------------
diff --git a/conf/defaults.yaml b/conf/defaults.yaml
index 160c29f..ac1b75b 100644
--- a/conf/defaults.yaml
+++ b/conf/defaults.yaml
@@ -226,6 +226,7 @@ topology.bolts.outgoing.overflow.buffer.enable: false
 topology.disruptor.wait.timeout.millis: 1000
 topology.disruptor.batch.size: 100
 topology.disruptor.batch.timeout.millis: 1
+topology.disable.loadaware: false
 
 # Configs for Resource Aware Scheduler
 topology.component.resources.onheap.memory.mb: 128.0

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/clj/backtype/storm/daemon/executor.clj
----------------------------------------------------------------------
diff --git a/storm-core/src/clj/backtype/storm/daemon/executor.clj 
b/storm-core/src/clj/backtype/storm/daemon/executor.clj
index 1bc2db4..f989654 100644
--- a/storm-core/src/clj/backtype/storm/daemon/executor.clj
+++ b/storm-core/src/clj/backtype/storm/daemon/executor.clj
@@ -34,6 +34,7 @@
   (:import [backtype.storm.daemon Shutdownable])
   (:import [backtype.storm.metric.api IMetric IMetricsConsumer$TaskInfo 
IMetricsConsumer$DataPoint StateMetric])
   (:import [backtype.storm Config Constants])
+  (:import [backtype.storm.grouping LoadAwareCustomStreamGrouping 
LoadAwareShuffleGrouping LoadMapping ShuffleGrouping])
   (:import [java.util.concurrent ConcurrentLinkedQueue])
   (:require [backtype.storm [tuple :as tuple] [thrift :as thrift]
              [cluster :as cluster] [disruptor :as disruptor] [stats :as 
stats]])
@@ -44,51 +45,53 @@
 (defn- mk-fields-grouper [^Fields out-fields ^Fields group-fields ^List 
target-tasks]
   (let [num-tasks (count target-tasks)
         task-getter (fn [i] (.get target-tasks i))]
-    (fn [task-id ^List values]
+    (fn [task-id ^List values load]
       (-> (.select out-fields group-fields values)
           tuple/list-hash-code
           (mod num-tasks)
           task-getter))))
 
-(defn- mk-shuffle-grouper [^List target-tasks]
-  (let [choices (rotating-random-range target-tasks)]
-    (fn [task-id tuple]
-      (acquire-random-range-id choices))))
-
 (defn- mk-custom-grouper [^CustomStreamGrouping grouping 
^WorkerTopologyContext context ^String component-id ^String stream-id 
target-tasks]
   (.prepare grouping context (GlobalStreamId. component-id stream-id) 
target-tasks)
-  (fn [task-id ^List values]
-    (.chooseTasks grouping task-id values)
-    ))
+  (if (instance? LoadAwareCustomStreamGrouping grouping)
+    (fn [task-id ^List values load]
+        (.chooseTasks grouping task-id values load))
+    (fn [task-id ^List values load]
+      (.chooseTasks grouping task-id values))))
+
+(defn mk-shuffle-grouper [^List target-tasks topo-conf ^WorkerTopologyContext 
context ^String component-id ^String stream-id]
+  (if (.get topo-conf TOPOLOGY-DISABLE-LOADAWARE-MESSAGING)
+    (mk-custom-grouper (ShuffleGrouping.) context component-id stream-id 
target-tasks)
+    (mk-custom-grouper (LoadAwareShuffleGrouping.) context component-id 
stream-id target-tasks)))
 
 (defn- mk-grouper
   "Returns a function that returns a vector of which task indices to send 
tuple to, or just a single task index."
-  [^WorkerTopologyContext context component-id stream-id ^Fields out-fields 
thrift-grouping ^List target-tasks]
+  [^WorkerTopologyContext context component-id stream-id ^Fields out-fields 
thrift-grouping ^List target-tasks topo-conf]
   (let [num-tasks (count target-tasks)
         random (Random.)
         target-tasks (vec (sort target-tasks))]
     (condp = (thrift/grouping-type thrift-grouping)
       :fields
         (if (thrift/global-grouping? thrift-grouping)
-          (fn [task-id tuple]
+          (fn [task-id tuple load]
             ;; It's possible for target to have multiple tasks if it reads 
multiple sources
             (first target-tasks))
           (let [group-fields (Fields. (thrift/field-grouping thrift-grouping))]
             (mk-fields-grouper out-fields group-fields target-tasks)
             ))
       :all
-        (fn [task-id tuple] target-tasks)
+        (fn [task-id tuple load] target-tasks)
       :shuffle
-        (mk-shuffle-grouper target-tasks)
+        (mk-shuffle-grouper target-tasks topo-conf context component-id 
stream-id)
       :local-or-shuffle
         (let [same-tasks (set/intersection
                            (set target-tasks)
                            (set (.getThisWorkerTasks context)))]
           (if-not (empty? same-tasks)
-            (mk-shuffle-grouper (vec same-tasks))
-            (mk-shuffle-grouper target-tasks)))
+            (mk-shuffle-grouper (vec same-tasks) topo-conf context 
component-id stream-id)
+            (mk-shuffle-grouper target-tasks topo-conf context component-id 
stream-id)))
       :none
-        (fn [task-id tuple]
+        (fn [task-id tuple load]
           (let [i (mod (.nextInt random) num-tasks)]
             (get target-tasks i)
             ))
@@ -102,7 +105,7 @@
         :direct
       )))
 
-(defn- outbound-groupings [^WorkerTopologyContext worker-context 
this-component-id stream-id out-fields component->grouping]
+(defn- outbound-groupings [^WorkerTopologyContext worker-context 
this-component-id stream-id out-fields component->grouping topo-conf]
   (->> component->grouping
        (filter-key #(-> worker-context
                         (.getComponentTasks %)
@@ -116,13 +119,13 @@
                             out-fields
                             tgrouping
                             (.getComponentTasks worker-context component)
-                            )]))
+                            topo-conf)]))
        (into {})
        (HashMap.)))
 
 (defn outbound-components
   "Returns map of stream id to component id to grouper"
-  [^WorkerTopologyContext worker-context component-id]
+  [^WorkerTopologyContext worker-context component-id topo-conf]
   (->> (.getTargets worker-context component-id)
         clojurify-structure
         (map (fn [[stream-id component->grouping]]
@@ -132,7 +135,8 @@
                   component-id
                   stream-id
                   (.getComponentOutputFields worker-context component-id 
stream-id)
-                  component->grouping)]))
+                  component->grouping
+                  topo-conf)]))
          (into {})
          (HashMap.)))
 
@@ -240,7 +244,7 @@
      :stats (mk-executor-stats <> (sampling-rate storm-conf))
      :interval->task->metric-registry (HashMap.)
      :task->component (:task->component worker)
-     :stream->component->grouper (outbound-components worker-context 
component-id)
+     :stream->component->grouper (outbound-components worker-context 
component-id storm-conf)
      :report-error (throttled-report-error-fn <>)
      :report-error-and-die (fn [error]
                              ((:report-error <>) error)

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/clj/backtype/storm/daemon/task.clj
----------------------------------------------------------------------
diff --git a/storm-core/src/clj/backtype/storm/daemon/task.clj 
b/storm-core/src/clj/backtype/storm/daemon/task.clj
index c48fb03..7133fdf 100644
--- a/storm-core/src/clj/backtype/storm/daemon/task.clj
+++ b/storm-core/src/clj/backtype/storm/daemon/task.clj
@@ -18,6 +18,7 @@
   (:use [backtype.storm config util log])
   (:import [backtype.storm.hooks ITaskHook])
   (:import [backtype.storm.tuple Tuple TupleImpl])
+  (:import [backtype.storm.grouping LoadMapping])
   (:import [backtype.storm.generated SpoutSpec Bolt StateSpoutSpec 
StormTopology])
   (:import [backtype.storm.hooks.info SpoutAckInfo SpoutFailInfo
             EmitInfo BoltFailInfo BoltAckInfo])
@@ -119,6 +120,7 @@
 (defn mk-tasks-fn [task-data]
   (let [task-id (:task-id task-data)
         executor-data (:executor-data task-data)
+        ^LoadMapping load-mapping (:load-mapping (:worker executor-data))
         component-id (:component-id executor-data)
         ^WorkerTopologyContext worker-context (:worker-context executor-data)
         storm-conf (:storm-conf executor-data)
@@ -152,7 +154,7 @@
                (when (= :direct grouper)
                   ;;  TODO: this is wrong, need to check how the stream was 
declared
                   (throw (IllegalArgumentException. "Cannot do regular emit to 
direct stream")))
-               (let [comp-tasks (grouper task-id values)]
+               (let [comp-tasks (grouper task-id values load-mapping)]
                  (if (or (sequential? comp-tasks) (instance? Collection 
comp-tasks))
                    (.addAll out-tasks comp-tasks)
                    (.add out-tasks comp-tasks)

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/clj/backtype/storm/daemon/worker.clj
----------------------------------------------------------------------
diff --git a/storm-core/src/clj/backtype/storm/daemon/worker.clj 
b/storm-core/src/clj/backtype/storm/daemon/worker.clj
index f304bdd..453801d 100644
--- a/storm-core/src/clj/backtype/storm/daemon/worker.clj
+++ b/storm-core/src/clj/backtype/storm/daemon/worker.clj
@@ -25,6 +25,7 @@
   (:import [java.util.concurrent Executors])
   (:import [java.util ArrayList HashMap])
   (:import [backtype.storm.utils Utils TransferDrainer ThriftTopologyUtils 
WorkerBackpressureThread DisruptorQueue])
+  (:import [backtype.storm.grouping LoadMapping])
   (:import [backtype.storm.messaging TransportFactory])
   (:import [backtype.storm.messaging TaskMessage IContext IConnection 
ConnectionWithStatus ConnectionWithStatus$Status])
   (:import [backtype.storm.daemon Shutdownable])
@@ -269,6 +270,7 @@
       :topology topology
       :system-topology (system-topology! storm-conf topology)
       :heartbeat-timer (mk-halting-timer "heartbeat-timer")
+      :refresh-load-timer (mk-halting-timer "refresh-load-timer")
       :refresh-connections-timer (mk-halting-timer "refresh-connections-timer")
       :refresh-credentials-timer (mk-halting-timer "refresh-credentials-timer")
       :reset-log-levels-timer (mk-halting-timer "reset-log-levels-timer")
@@ -295,6 +297,7 @@
       :transfer-local-fn (mk-transfer-local-fn <>)
       :receiver-thread-count (get storm-conf WORKER-RECEIVER-THREAD-COUNT)
       :transfer-fn (mk-transfer-fn <>)
+      :load-mapping (LoadMapping.)
       :assignment-versions assignment-versions
       :backpressure (atom false) ;; whether this worker is going slow
       :transfer-backpressure (atom false) ;; if the transfer queue is backed-up
@@ -310,6 +313,26 @@
     [node (Integer/valueOf port-str)]
     ))
 
+(defn mk-refresh-load [worker]
+  (let [local-tasks (set (:task-ids worker))
+        remote-tasks (set/difference (worker-outbound-tasks worker) 
local-tasks)
+        short-executor-receive-queue-map (:short-executor-receive-queue-map 
worker)
+        next-update (atom 0)]
+    (fn this
+      ([]
+        (let [^LoadMapping load-mapping (:load-mapping worker)
+              local-pop (map-val (fn [queue]
+                                   (let [q-metrics (.getMetrics queue)]
+                                     (/ (double (.population q-metrics)) 
(.capacity q-metrics))))
+                                 short-executor-receive-queue-map)
+              remote-load (reduce merge (for [[np conn] 
@(:cached-node+port->socket worker)] (into {} (.getLoad conn remote-tasks))))
+              now (System/currentTimeMillis)]
+          (.setLocal load-mapping local-pop)
+          (.setRemote load-mapping remote-load)
+          (when (> now @next-update)
+            (.sendLoadMetrics (:receiver worker) local-pop)
+            (reset! next-update (+ 5000 now))))))))
+
 (defn mk-refresh-connections [worker]
   (let [outbound-tasks (worker-outbound-tasks worker)
         conf (:conf worker)
@@ -572,6 +595,7 @@
         receive-thread-shutdown (launch-receive-thread worker)
 
         refresh-connections (mk-refresh-connections worker)
+        refresh-load (mk-refresh-load worker)
 
         _ (refresh-connections nil)
 
@@ -633,6 +657,7 @@
                     (cancel-timer (:refresh-active-timer worker))
                     (cancel-timer (:executor-heartbeat-timer worker))
                     (cancel-timer (:user-timer worker))
+                    (cancel-timer (:refresh-load-timer worker))
 
                     (close-resources worker)
 
@@ -653,6 +678,7 @@
                (and
                  (timer-waiting? (:heartbeat-timer worker))
                  (timer-waiting? (:refresh-connections-timer worker))
+                 (timer-waiting? (:refresh-load-timer worker))
                  (timer-waiting? (:refresh-credentials-timer worker))
                  (timer-waiting? (:refresh-active-timer worker))
                  (timer-waiting? (:executor-heartbeat-timer worker))
@@ -689,6 +715,9 @@
                           (check-credentials-changed)
                           (if ((:storm-conf worker) 
TOPOLOGY-BACKPRESSURE-ENABLE)
                             (check-throttle-changed))))
+    ;; The jitter allows the clients to get the data at different times, and 
avoids thundering herd
+    (when-not (.get conf TOPOLOGY-DISABLE-LOADAWARE-MESSAGING)
+      (schedule-recurring-with-jitter (:refresh-load-timer worker) 0 1 500 
refresh-load))
     (schedule-recurring (:refresh-connections-timer worker) 0 (conf 
TASK-REFRESH-POLL-SECS) refresh-connections)
     (schedule-recurring (:reset-log-levels-timer worker) 0 (conf 
WORKER-LOG-LEVEL-RESET-POLL-SECS) (fn [] (reset-log-levels latest-log-config)))
     (schedule-recurring (:refresh-active-timer worker) 0 (conf 
TASK-REFRESH-POLL-SECS) (partial refresh-storm-active worker))

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/clj/backtype/storm/messaging/local.clj
----------------------------------------------------------------------
diff --git a/storm-core/src/clj/backtype/storm/messaging/local.clj 
b/storm-core/src/clj/backtype/storm/messaging/local.clj
index 4aa67ab..60a6bd2 100644
--- a/storm-core/src/clj/backtype/storm/messaging/local.clj
+++ b/storm-core/src/clj/backtype/storm/messaging/local.clj
@@ -15,13 +15,18 @@
 ;; limitations under the License.
 (ns backtype.storm.messaging.local
   (:refer-clojure :exclude [send])
-  (:use [backtype.storm log])
+  (:use [backtype.storm log util])
   (:import [backtype.storm.messaging IContext IConnection TaskMessage])
+  (:import [backtype.storm.grouping Load])
   (:import [java.util.concurrent LinkedBlockingQueue])
-  (:import [java.util Map Iterator])
+  (:import [java.util Map Iterator Collection])
   (:import [java.util Iterator ArrayList])
   (:gen-class))
 
+(defn update-load! [cached-task->load lock task->load]
+  (locking lock
+    (swap! cached-task->load merge task->load)))
+
 (defn add-queue! [queues-map lock storm-id port]
   (let [id (str storm-id "-" port)]
     (locking lock
@@ -29,7 +34,7 @@
         (swap! queues-map assoc id (LinkedBlockingQueue.))))
     (@queues-map id)))
 
-(deftype LocalConnection [storm-id port queues-map lock queue]
+(deftype LocalConnection [storm-id port queues-map lock queue task->load]
   IConnection
   (^Iterator recv [this ^int flags ^int clientId]
     (when-not queue
@@ -50,24 +55,35 @@
       (while (.hasNext iter) 
          (.put send-queue (.next iter)))
       ))
-  (^void close [this]
-    ))
+  (^void sendLoadMetrics [this ^Map taskToLoad]
+    (update-load! task->load lock taskToLoad))
+  (^Map getLoad [this ^Collection tasks]
+    (locking lock
+      (into {}
+        (for [task tasks
+              :let [load (.get @task->load task)]
+              :when (not-nil? load)]
+          ;; for now we are ignoring the connection load locally
+          [task (Load. true load 0.0)]))))
+  (^void close [this]))
 
 
 (deftype LocalContext [^{:unsynchronized-mutable true} queues-map
-                       ^{:unsynchronized-mutable true} lock]
+                       ^{:unsynchronized-mutable true} lock
+                       ^{:unsynchronized-mutable true} task->load]
   IContext
   (^void prepare [this ^Map storm-conf]
     (set! queues-map (atom {}))
+    (set! task->load (atom {}))
     (set! lock (Object.)))
   (^IConnection bind [this ^String storm-id ^int port]
-    (LocalConnection. storm-id port queues-map lock (add-queue! queues-map 
lock storm-id port)))
+    (LocalConnection. storm-id port queues-map lock (add-queue! queues-map 
lock storm-id port) task->load))
   (^IConnection connect [this ^String storm-id ^String host ^int port]
-    (LocalConnection. storm-id port queues-map lock nil))
+    (LocalConnection. storm-id port queues-map lock nil task->load))
   (^void term [this]
     ))
 
 (defn mk-context [] 
-  (let [context  (LocalContext. nil nil)]
+  (let [context  (LocalContext. nil nil nil)]
     (.prepare ^IContext context nil)
     context))

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/clj/backtype/storm/timer.clj
----------------------------------------------------------------------
diff --git a/storm-core/src/clj/backtype/storm/timer.clj 
b/storm-core/src/clj/backtype/storm/timer.clj
index 2c76ce2..b5f73f7 100644
--- a/storm-core/src/clj/backtype/storm/timer.clj
+++ b/storm-core/src/clj/backtype/storm/timer.clj
@@ -16,7 +16,7 @@
 
 (ns backtype.storm.timer
   (:import [backtype.storm.utils Time])
-  (:import [java.util PriorityQueue Comparator])
+  (:import [java.util PriorityQueue Comparator Random])
   (:import [java.util.concurrent Semaphore])
   (:use [backtype.storm util log]))
 
@@ -79,6 +79,7 @@
      :queue queue
      :active active
      :lock lock
+     :random (Random.)
      :cancel-notifier notifier}))
 
 (defn- check-active!
@@ -87,12 +88,14 @@
     (throw (IllegalStateException. "Timer is not active"))))
 
 (defnk schedule
-  [timer delay-secs afn :check-active true]
+  [timer delay-secs afn :check-active true :jitter-ms 0]
   (when check-active (check-active! timer))
   (let [id (uuid)
-        ^PriorityQueue queue (:queue timer)]
+        ^PriorityQueue queue (:queue timer)
+        end-time-ms (+ (current-time-millis) (secs-to-millis-long delay-secs))
+        end-time-ms (if (< 0 jitter-ms) (+ (.nextInt (:random timer) 
jitter-ms) end-time-ms) end-time-ms)]
     (locking (:lock timer)
-      (.add queue [(+ (current-time-millis) (secs-to-millis-long delay-secs)) 
afn id]))))
+      (.add queue [end-time-ms afn id]))))
 
 (defn schedule-recurring
   [timer delay-secs recur-secs afn]
@@ -103,6 +106,15 @@
               ; This avoids a race condition with cancel-timer.
               (schedule timer recur-secs this :check-active false))))
 
+(defn schedule-recurring-with-jitter
+  [timer delay-secs recur-secs jitter-ms afn]
+  (schedule timer
+            delay-secs
+            (fn this []
+              (afn)
+              ; This avoids a race condition with cancel-timer.
+              (schedule timer recur-secs this :check-active false :jitter-ms 
jitter-ms))))
+
 (defn cancel-timer
   [timer]
   (check-active! timer)

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/jvm/backtype/storm/Config.java
----------------------------------------------------------------------
diff --git a/storm-core/src/jvm/backtype/storm/Config.java 
b/storm-core/src/jvm/backtype/storm/Config.java
index 73e41cb..eab1575 100644
--- a/storm-core/src/jvm/backtype/storm/Config.java
+++ b/storm-core/src/jvm/backtype/storm/Config.java
@@ -239,6 +239,13 @@ public class Config extends HashMap<String, Object> {
     public static final String TOPOLOGY_TUPLE_SERIALIZER = 
"topology.tuple.serializer";
 
     /**
+     * Disable load aware grouping support.
+     */
+    @isBoolean
+    @NotNull
+    public static final String TOPOLOGY_DISABLE_LOADAWARE_MESSAGING = 
"topology.disable.loadaware.messaging";
+
+    /**
      * Try to serialize all tuples, even for local transfers.  This should 
only be used
      * for testing, as a sanity check that all of your tuples are setup 
properly.
      */

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/jvm/backtype/storm/grouping/Load.java
----------------------------------------------------------------------
diff --git a/storm-core/src/jvm/backtype/storm/grouping/Load.java 
b/storm-core/src/jvm/backtype/storm/grouping/Load.java
new file mode 100644
index 0000000..a10c791
--- /dev/null
+++ b/storm-core/src/jvm/backtype/storm/grouping/Load.java
@@ -0,0 +1,77 @@
+/**
+ * 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.grouping;
+
+import java.io.Serializable;
+
+/**
+ * Represents the load that a Bolt is currently under to help in
+ * deciding where to route a tuple, to help balance the load.
+ */
+public class Load {
+    private boolean hasMetrics = false;
+    private double boltLoad = 0.0; //0 no load to 1 fully loaded
+    private double connectionLoad = 0.0; //0 no load to 1 fully loaded
+
+    /**
+     * Create a new load
+     * @param hasMetrics have metrics been reported yet?
+     * @param boltLoad the load as reported by the bolt 0.0 no load 1.0 fully 
loaded
+     * @param connectionLoad the load as reported by the connection to the 
bolt 0.0 no load 1.0 fully loaded.
+     */
+    public Load(boolean hasMetrics, double boltLoad, double connectionLoad) {
+        this.hasMetrics = hasMetrics;
+        this.boltLoad = boltLoad;
+        this.connectionLoad = connectionLoad;
+    }
+
+    /**
+     * @return true if metrics have been reported so far.
+     */
+    public boolean hasMetrics() {
+        return hasMetrics;
+    }
+
+    /**
+     * @return the load as reported by the bolt.
+     */
+    public double getBoltLoad() {
+        return boltLoad;
+    }
+
+    /**
+     * @return the load as reported by the connection
+     */
+    public double getConnectionLoad() {
+        return connectionLoad;
+    }
+
+    /**
+     * @return the load that is a combination of sub loads.
+     */
+    public double getLoad() {
+        if (!hasMetrics) {
+            return 1.0;
+        }
+        return connectionLoad > boltLoad ? connectionLoad : boltLoad;
+    }
+
+    public String toString() {
+        return "[:load "+boltLoad+" "+connectionLoad+"]";
+    }
+}

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/jvm/backtype/storm/grouping/LoadAwareCustomStreamGrouping.java
----------------------------------------------------------------------
diff --git 
a/storm-core/src/jvm/backtype/storm/grouping/LoadAwareCustomStreamGrouping.java 
b/storm-core/src/jvm/backtype/storm/grouping/LoadAwareCustomStreamGrouping.java
new file mode 100644
index 0000000..f39273f
--- /dev/null
+++ 
b/storm-core/src/jvm/backtype/storm/grouping/LoadAwareCustomStreamGrouping.java
@@ -0,0 +1,24 @@
+/**
+ * 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.grouping;
+
+import java.util.List;
+
+public interface LoadAwareCustomStreamGrouping extends CustomStreamGrouping {
+   List<Integer> chooseTasks(int taskId, List<Object> values, LoadMapping 
load);
+}

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/jvm/backtype/storm/grouping/LoadAwareShuffleGrouping.java
----------------------------------------------------------------------
diff --git 
a/storm-core/src/jvm/backtype/storm/grouping/LoadAwareShuffleGrouping.java 
b/storm-core/src/jvm/backtype/storm/grouping/LoadAwareShuffleGrouping.java
new file mode 100644
index 0000000..a19b585
--- /dev/null
+++ b/storm-core/src/jvm/backtype/storm/grouping/LoadAwareShuffleGrouping.java
@@ -0,0 +1,76 @@
+/**
+ * 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.grouping;
+
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Random;
+
+import backtype.storm.generated.GlobalStreamId;
+import backtype.storm.task.WorkerTopologyContext;
+
+public class LoadAwareShuffleGrouping implements 
LoadAwareCustomStreamGrouping, Serializable {
+    private Random random;
+    private List<Integer>[] rets;
+    private int[] targets;
+    private int[] loads;
+    private int total;
+    private long lastUpdate = 0;
+
+    @Override
+    public void prepare(WorkerTopologyContext context, GlobalStreamId stream, 
List<Integer> targetTasks) {
+        random = new Random();
+        rets = (List<Integer>[])new List<?>[targetTasks.size()];
+        targets = new int[targetTasks.size()];
+        for (int i = 0; i < targets.length; i++) {
+            rets[i] = Arrays.asList(targetTasks.get(i));
+            targets[i] = targetTasks.get(i);
+        }
+        loads = new int[targets.length];
+    }
+
+    @Override
+    public List<Integer> chooseTasks(int taskId, List<Object> values) {
+        throw new RuntimeException("NOT IMPLEMENTED");
+    }
+
+    @Override
+    public List<Integer> chooseTasks(int taskId, List<Object> values, 
LoadMapping load) {
+        if ((lastUpdate + 1000) < System.currentTimeMillis()) {
+            int local_total = 0;
+            for (int i = 0; i < targets.length; i++) {
+                int val = (int)(101 - (load.get(targets[i]) * 100));
+                loads[i] = val;
+                local_total += val;
+            }
+            total = local_total;
+            lastUpdate = System.currentTimeMillis();
+        }
+        int selected = random.nextInt(total);
+        int sum = 0;
+        for (int i = 0; i < targets.length; i++) {
+            sum += loads[i];
+            if (selected < sum) {
+                return rets[i];
+            }
+        }
+        return rets[rets.length-1];
+    }
+}

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/jvm/backtype/storm/grouping/LoadMapping.java
----------------------------------------------------------------------
diff --git a/storm-core/src/jvm/backtype/storm/grouping/LoadMapping.java 
b/storm-core/src/jvm/backtype/storm/grouping/LoadMapping.java
new file mode 100644
index 0000000..abd8983
--- /dev/null
+++ b/storm-core/src/jvm/backtype/storm/grouping/LoadMapping.java
@@ -0,0 +1,64 @@
+/**
+ * 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.grouping;
+
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.Map;
+import java.util.HashMap;
+
+/**
+ * Holds a list of the current loads
+ */
+public class LoadMapping {
+    private static final Load NOT_CONNECTED = new Load(false, 1.0, 1.0);
+    private final AtomicReference<Map<Integer,Load>> _local = new 
AtomicReference<Map<Integer,Load>>(new HashMap<Integer,Load>());
+    private final AtomicReference<Map<Integer,Load>> _remote = new 
AtomicReference<Map<Integer,Load>>(new HashMap<Integer,Load>());
+
+    public void setLocal(Map<Integer, Double> local) {
+        Map<Integer, Load> newLocal = new HashMap<Integer, Load>();
+        if (local != null) {
+          for (Map.Entry<Integer, Double> entry: local.entrySet()) {
+            newLocal.put(entry.getKey(), new Load(true, entry.getValue(), 
0.0));
+          }
+        }
+        _local.set(newLocal);
+    }
+
+    public void setRemote(Map<Integer, Load> remote) {
+        if (remote != null) {
+          _remote.set(new HashMap<Integer, Load>(remote));
+        } else {
+          _remote.set(new HashMap<Integer, Load>());
+        }
+    }
+
+    public Load getLoad(int task) {
+        Load ret = _local.get().get(task);
+        if (ret == null) {
+          ret = _remote.get().get(task);
+        }
+        if (ret == null) {
+          ret = NOT_CONNECTED;
+        }
+        return ret;
+    }
+
+    public double get(int task) {
+        return getLoad(task).getLoad();
+    }
+}

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/jvm/backtype/storm/grouping/ShuffleGrouping.java
----------------------------------------------------------------------
diff --git a/storm-core/src/jvm/backtype/storm/grouping/ShuffleGrouping.java 
b/storm-core/src/jvm/backtype/storm/grouping/ShuffleGrouping.java
new file mode 100644
index 0000000..01755bb
--- /dev/null
+++ b/storm-core/src/jvm/backtype/storm/grouping/ShuffleGrouping.java
@@ -0,0 +1,65 @@
+/**
+ * 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.grouping;
+
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import backtype.storm.generated.GlobalStreamId;
+import backtype.storm.task.WorkerTopologyContext;
+
+public class ShuffleGrouping implements CustomStreamGrouping, Serializable {
+    private Random random;
+    private ArrayList<List<Integer>> choices;
+    private AtomicInteger current;
+
+    @Override
+    public void prepare(WorkerTopologyContext context, GlobalStreamId stream, 
List<Integer> targetTasks) {
+        random = new Random();
+        choices = new ArrayList<List<Integer>>(targetTasks.size());
+        for (Integer i: targetTasks) {
+            choices.add(Arrays.asList(i));
+        }
+        Collections.shuffle(choices, random);
+        current = new AtomicInteger(0);
+    }
+
+    @Override
+    public List<Integer> chooseTasks(int taskId, List<Object> values) {
+        int rightNow;
+        int size = choices.size();
+        while (true) {
+            rightNow = current.incrementAndGet();
+            if (rightNow < size) {
+                return choices.get(rightNow);
+            } else if (rightNow == size) {
+                current.set(0);
+                //This should be thread safe so long as ArrayList does not 
have any internal state that can be messed up by multi-treaded access.
+                Collections.shuffle(choices, random);
+                return choices.get(0);
+            }
+            //race condition with another thread, and we lost
+            // try again
+        }
+    }
+}

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/jvm/backtype/storm/messaging/IConnection.java
----------------------------------------------------------------------
diff --git a/storm-core/src/jvm/backtype/storm/messaging/IConnection.java 
b/storm-core/src/jvm/backtype/storm/messaging/IConnection.java
index ead4935..88f0c0a 100644
--- a/storm-core/src/jvm/backtype/storm/messaging/IConnection.java
+++ b/storm-core/src/jvm/backtype/storm/messaging/IConnection.java
@@ -17,7 +17,10 @@
  */
 package backtype.storm.messaging;
 
+import backtype.storm.grouping.Load;
+import java.util.Collection;
 import java.util.Iterator;
+import java.util.Map;
 
 public interface IConnection {   
     
@@ -27,6 +30,12 @@ public interface IConnection {
      * @return
      */
     public Iterator<TaskMessage> recv(int flags, int clientId);
+
+    /**
+     * Send load metrics to all downstream connections.
+     * @param taskToLoad a map from the task id to the load for that task.
+     */
+    public void sendLoadMetrics(Map<Integer, Double> taskToLoad);
     
     /**
      * send a message with taskId and payload
@@ -43,6 +52,13 @@ public interface IConnection {
     public void send(Iterator<TaskMessage> msgs);
     
     /**
+     * Get the current load for the given tasks
+     * @param tasks the tasks to look for.
+     * @return a Load for each of the tasks it knows about.
+     */
+    public Map<Integer, Load> getLoad(Collection<Integer> tasks);
+
+    /**
      * close this connection
      */
     public void close();

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/jvm/backtype/storm/messaging/netty/Client.java
----------------------------------------------------------------------
diff --git a/storm-core/src/jvm/backtype/storm/messaging/netty/Client.java 
b/storm-core/src/jvm/backtype/storm/messaging/netty/Client.java
index 2149c0d..b9151f3 100644
--- a/storm-core/src/jvm/backtype/storm/messaging/netty/Client.java
+++ b/storm-core/src/jvm/backtype/storm/messaging/netty/Client.java
@@ -17,7 +17,19 @@
  */
 package backtype.storm.messaging.netty;
 
+import java.net.InetSocketAddress;
+import java.net.SocketAddress;
+import java.util.Iterator;
+import java.util.Collection;
+import java.util.Map;
+import java.util.HashMap;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.lang.InterruptedException;
+
 import backtype.storm.Config;
+import backtype.storm.grouping.Load;
 import backtype.storm.messaging.ConnectionWithStatus;
 import backtype.storm.messaging.TaskMessage;
 import backtype.storm.metric.api.IStatefulObject;
@@ -34,17 +46,11 @@ import org.jboss.netty.util.TimerTask;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.net.InetSocketAddress;
-import java.net.SocketAddress;
+
 import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.Iterator;
 import java.util.List;
-import java.util.Map;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
-import java.util.concurrent.atomic.AtomicReference;
 
 import static com.google.common.base.Preconditions.checkState;
 
@@ -60,7 +66,7 @@ import static com.google.common.base.Preconditions.checkState;
  *     - Note: The current implementation drops any messages that are being 
enqueued for sending if the connection to
  *       the remote destination is currently unavailable.
  */
-public class Client extends ConnectionWithStatus implements IStatefulObject {
+public class Client extends ConnectionWithStatus implements IStatefulObject, 
ISaslClient {
     private static final long PENDING_MESSAGES_FLUSH_TIMEOUT_MS = 600000L;
     private static final long PENDING_MESSAGES_FLUSH_INTERVAL_MS = 1000L;
 
@@ -73,6 +79,7 @@ public class Client extends ConnectionWithStatus implements 
IStatefulObject {
     private final ClientBootstrap bootstrap;
     private final InetSocketAddress dstAddress;
     protected final String dstAddressPrefixedName;
+    private volatile Map<Integer, Double> serverLoad = null;
 
     /**
      * The channel used for all write operations from this client to the 
remote destination.
@@ -104,6 +111,10 @@ public class Client extends ConnectionWithStatus 
implements IStatefulObject {
      */
     private final AtomicLong pendingMessages = new AtomicLong(0);
 
+    /**
+     * Whether the SASL channel is ready.
+     */
+    private final AtomicBoolean saslChannelReady = new AtomicBoolean(false);
 
     /**
      * This flag is set to true if and only if a client instance is being 
closed.
@@ -125,6 +136,8 @@ public class Client extends ConnectionWithStatus implements 
IStatefulObject {
         this.scheduler = scheduler;
         this.context = context;
         int bufferSize = 
Utils.getInt(stormConf.get(Config.STORM_MESSAGING_NETTY_BUFFER_SIZE));
+        // if SASL authentication is disabled, saslChannelReady is initialized 
as true; otherwise false
+        
saslChannelReady.set(!Utils.getBoolean(stormConf.get(Config.STORM_MESSAGING_NETTY_AUTHENTICATION),
 false));
         LOG.info("creating Netty Client, connecting to {}:{}, bufferSize: {}", 
host, port, bufferSize);
         int messageBatchSize = 
Utils.getInt(stormConf.get(Config.STORM_NETTY_MESSAGE_BATCH_SIZE), 262144);
 
@@ -134,19 +147,19 @@ public class Client extends ConnectionWithStatus 
implements IStatefulObject {
         retryPolicy = new StormBoundedExponentialBackoffRetry(minWaitMs, 
maxWaitMs, maxReconnectionAttempts);
 
         // Initiate connection to remote destination
-        bootstrap = createClientBootstrap(factory, bufferSize);
+        bootstrap = createClientBootstrap(factory, bufferSize, stormConf);
         dstAddress = new InetSocketAddress(host, port);
         dstAddressPrefixedName = prefixedName(dstAddress);
         scheduleConnect(NO_DELAY_MS);
         batcher = new MessageBuffer(messageBatchSize);
     }
 
-    private ClientBootstrap createClientBootstrap(ChannelFactory factory, int 
bufferSize) {
+    private ClientBootstrap createClientBootstrap(ChannelFactory factory, int 
bufferSize, Map stormConf) {
         ClientBootstrap bootstrap = new ClientBootstrap(factory);
         bootstrap.setOption("tcpNoDelay", true);
         bootstrap.setOption("sendBufferSize", bufferSize);
         bootstrap.setOption("keepAlive", true);
-        bootstrap.setPipelineFactory(new StormClientPipelineFactory(this));
+        bootstrap.setPipelineFactory(new StormClientPipelineFactory(this, 
stormConf));
         return bootstrap;
     }
 
@@ -158,7 +171,7 @@ public class Client extends ConnectionWithStatus implements 
IStatefulObject {
     }
 
     /**
-     * We will retry connection with exponential back-off policy
+     * Enqueue a task message to be sent to server
      */
     private void scheduleConnect(long delayMs) {
         scheduler.newTimeout(new Connect(dstAddress), delayMs, 
TimeUnit.MILLISECONDS);
@@ -190,7 +203,11 @@ public class Client extends ConnectionWithStatus 
implements IStatefulObject {
         } else if (!connectionEstablished(channelRef.get())) {
             return Status.Connecting;
         } else {
-            return Status.Ready;
+            if (saslChannelReady.get()) {
+                return Status.Ready;
+            } else {
+                return Status.Connecting; // need to wait until sasl channel 
is also ready
+            }
         }
     }
 
@@ -205,6 +222,11 @@ public class Client extends ConnectionWithStatus 
implements IStatefulObject {
     }
 
     @Override
+    public void sendLoadMetrics(Map<Integer, Double> taskToLoad) {
+        throw new RuntimeException("Client connection should not send load 
metrics");
+    }
+
+    @Override
     public void send(int taskId, byte[] payload) {
         TaskMessage msg = new TaskMessage(taskId, payload);
         List<TaskMessage> wrapper = new ArrayList<TaskMessage>(1);
@@ -283,11 +305,14 @@ public class Client extends ConnectionWithStatus 
implements IStatefulObject {
         }
     }
 
+    public InetSocketAddress getDstAddress() {
+        return dstAddress;
+    }
+
     private boolean hasMessages(Iterator<TaskMessage> msgs) {
         return msgs != null && msgs.hasNext();
     }
 
-
     private void dropMessages(Iterator<TaskMessage> msgs) {
         // We consume the iterator by traversing and thus "emptying" it.
         int msgCount = iteratorSize(msgs);
@@ -383,15 +408,13 @@ public class Client extends ConnectionWithStatus 
implements IStatefulObject {
                     break;
                 }
                 Thread.sleep(PENDING_MESSAGES_FLUSH_INTERVAL_MS);
-            }
-            catch (InterruptedException e) {
+            } catch (InterruptedException e) {
                 break;
             }
         }
 
     }
 
-
     private void closeChannel() {
         Channel channel = channelRef.get();
         if (channel != null) {
@@ -400,9 +423,29 @@ public class Client extends ConnectionWithStatus 
implements IStatefulObject {
         }
     }
 
+    void setLoadMetrics(Map<Integer, Double> taskToLoad) {
+        this.serverLoad = taskToLoad;
+    }
+
+    @Override
+    public Map<Integer, Load> getLoad(Collection<Integer> tasks) {
+        Map<Integer, Double> loadCache = serverLoad;
+        Map<Integer, Load> ret = new HashMap<Integer, Load>();
+        if (loadCache != null) {
+            double clientLoad = Math.min(pendingMessages.get(), 1024)/1024.0;
+            for (Integer task : tasks) {
+                Double found = loadCache.get(task);
+                if (found != null) {
+                    ret.put(task, new Load(true, found, clientLoad));
+                }
+            }
+        }
+        return ret;
+    }
+
     @Override
     public Object getState() {
-        LOG.info("Getting metrics for client connection to {}", 
dstAddressPrefixedName);
+        LOG.debug("Getting metrics for client connection to {}", 
dstAddressPrefixedName);
         HashMap<String, Object> ret = new HashMap<String, Object>();
         ret.put("reconnects", totalConnectionAttempts.getAndSet(0));
         ret.put("sent", messagesSent.getAndSet(0));
@@ -416,10 +459,28 @@ public class Client extends ConnectionWithStatus 
implements IStatefulObject {
         return ret;
     }
 
-    public Map getStormConf() {
+    public Map getConfig() {
         return stormConf;
     }
 
+    /** ISaslClient interface **/
+    public void channelConnected(Channel channel) {
+//        setChannel(channel);
+    }
+
+    public void channelReady() {
+        saslChannelReady.set(true);
+    }
+
+    public String name() {
+        return (String)stormConf.get(Config.TOPOLOGY_NAME);
+    }
+
+    public String secretKey() {
+        return SaslUtils.getSecretKey(stormConf);
+    }
+    /** end **/
+
     private String srcAddressName() {
         String name = null;
         Channel channel = channelRef.get();
@@ -513,5 +574,4 @@ public class Client extends ConnectionWithStatus implements 
IStatefulObject {
             }
         }
     }
-
 }

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/jvm/backtype/storm/messaging/netty/ISaslClient.java
----------------------------------------------------------------------
diff --git a/storm-core/src/jvm/backtype/storm/messaging/netty/ISaslClient.java 
b/storm-core/src/jvm/backtype/storm/messaging/netty/ISaslClient.java
new file mode 100644
index 0000000..57dcfe8
--- /dev/null
+++ b/storm-core/src/jvm/backtype/storm/messaging/netty/ISaslClient.java
@@ -0,0 +1,28 @@
+/**
+ * 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.messaging.netty;
+
+import org.jboss.netty.channel.Channel;
+import backtype.storm.Config;
+
+public interface ISaslClient {
+    void channelConnected(Channel channel);
+    void channelReady();
+    String name();
+    String secretKey();
+}

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/jvm/backtype/storm/messaging/netty/ISaslServer.java
----------------------------------------------------------------------
diff --git a/storm-core/src/jvm/backtype/storm/messaging/netty/ISaslServer.java 
b/storm-core/src/jvm/backtype/storm/messaging/netty/ISaslServer.java
new file mode 100644
index 0000000..4203dcc
--- /dev/null
+++ b/storm-core/src/jvm/backtype/storm/messaging/netty/ISaslServer.java
@@ -0,0 +1,26 @@
+/**
+ * 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.messaging.netty;
+
+import org.jboss.netty.channel.Channel;
+
+public interface ISaslServer extends IServer {
+    String name();
+    String secretKey();
+    void authenticated(Channel c);
+}

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/jvm/backtype/storm/messaging/netty/IServer.java
----------------------------------------------------------------------
diff --git a/storm-core/src/jvm/backtype/storm/messaging/netty/IServer.java 
b/storm-core/src/jvm/backtype/storm/messaging/netty/IServer.java
new file mode 100644
index 0000000..d046492
--- /dev/null
+++ b/storm-core/src/jvm/backtype/storm/messaging/netty/IServer.java
@@ -0,0 +1,26 @@
+/**
+ * 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.messaging.netty;
+
+import org.jboss.netty.channel.Channel;
+
+public interface IServer {
+    void channelConnected(Channel c);
+    void received(Object message, String remote, Channel channel) throws 
InterruptedException;
+    void closeChannel(Channel c);
+}

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/jvm/backtype/storm/messaging/netty/SaslStormClientHandler.java
----------------------------------------------------------------------
diff --git 
a/storm-core/src/jvm/backtype/storm/messaging/netty/SaslStormClientHandler.java 
b/storm-core/src/jvm/backtype/storm/messaging/netty/SaslStormClientHandler.java
index 12b466c..2a5ae99 100644
--- 
a/storm-core/src/jvm/backtype/storm/messaging/netty/SaslStormClientHandler.java
+++ 
b/storm-core/src/jvm/backtype/storm/messaging/netty/SaslStormClientHandler.java
@@ -28,19 +28,18 @@ import org.jboss.netty.channel.SimpleChannelUpstreamHandler;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import backtype.storm.Config;
-
 public class SaslStormClientHandler extends SimpleChannelUpstreamHandler {
 
     private static final Logger LOG = LoggerFactory
             .getLogger(SaslStormClientHandler.class);
-    private Client client;
+
+    private ISaslClient client;
     long start_time;
     /** Used for client or server's token to send or receive from each other. 
*/
     private byte[] token;
-    private String topologyName;
+    private String name;
 
-    public SaslStormClientHandler(Client client) throws IOException {
+    public SaslStormClientHandler(ISaslClient client) throws IOException {
         this.client = client;
         start_time = System.currentTimeMillis();
         getSASLCredentials();
@@ -51,9 +50,7 @@ public class SaslStormClientHandler extends 
SimpleChannelUpstreamHandler {
             ChannelStateEvent event) {
         // register the newly established channel
         Channel channel = ctx.getChannel();
-
-        LOG.info("Connection established from " + channel.getLocalAddress()
-                + " to " + channel.getRemoteAddress());
+        client.channelConnected(channel);
 
         try {
             SaslNettyClient saslNettyClient = 
SaslNettyClientState.getSaslNettyClient
@@ -62,10 +59,11 @@ public class SaslStormClientHandler extends 
SimpleChannelUpstreamHandler {
             if (saslNettyClient == null) {
                 LOG.debug("Creating saslNettyClient now " + "for channel: "
                         + channel);
-                saslNettyClient = new SaslNettyClient(topologyName, token);
+                saslNettyClient = new SaslNettyClient(name, token);
                 SaslNettyClientState.getSaslNettyClient.set(channel,
                         saslNettyClient);
             }
+            LOG.debug("Sending SASL_TOKEN_MESSAGE_REQUEST");
             channel.write(ControlMessage.SASL_TOKEN_MESSAGE_REQUEST);
         } catch (Exception e) {
             LOG.error("Failed to authenticate with server " + "due to error: ",
@@ -96,7 +94,7 @@ public class SaslStormClientHandler extends 
SimpleChannelUpstreamHandler {
             ControlMessage msg = (ControlMessage) event.getMessage();
             if (msg == ControlMessage.SASL_COMPLETE_REQUEST) {
                 LOG.debug("Server has sent us the SaslComplete "
-                        + "message. Allowing normal work to proceed.");
+                          + "message. Allowing normal work to proceed.");
 
                 if (!saslNettyClient.isComplete()) {
                     LOG.error("Server returned a Sasl-complete message, "
@@ -106,6 +104,8 @@ public class SaslStormClientHandler extends 
SimpleChannelUpstreamHandler {
                             + "we can tell, we are not authenticated yet.");
                 }
                 ctx.getPipeline().remove(this);
+                this.client.channelReady();
+
                 // We call fireMessageReceived since the client is allowed to
                 // perform this request. The client's request will now proceed
                 // to the next pipeline component namely StormClientHandler.
@@ -116,7 +116,7 @@ public class SaslStormClientHandler extends 
SimpleChannelUpstreamHandler {
         SaslMessageToken saslTokenMessage = (SaslMessageToken) event
                 .getMessage();
         LOG.debug("Responding to server's token of length: "
-                + saslTokenMessage.getSaslToken().length);
+                  + saslTokenMessage.getSaslToken().length);
 
         // Generate SASL response (but we only actually send the response if
         // it's non-null.
@@ -127,17 +127,18 @@ public class SaslStormClientHandler extends 
SimpleChannelUpstreamHandler {
             // (if not, warn), and return without sending a response back to 
the
             // server.
             LOG.debug("Response to server is null: "
-                    + "authentication should now be complete.");
+                      + "authentication should now be complete.");
             if (!saslNettyClient.isComplete()) {
                 LOG.warn("Generated a null response, "
                         + "but authentication is not complete.");
                 throw new Exception("Server reponse is null, but as far as "
                         + "we can tell, we are not authenticated yet.");
             }
+            this.client.channelReady();
             return;
         } else {
             LOG.debug("Response to server token has length:"
-                    + responseToServer.length);
+                      + responseToServer.length);
         }
         // Construct a message containing the SASL response and send it to the
         // server.
@@ -146,12 +147,14 @@ public class SaslStormClientHandler extends 
SimpleChannelUpstreamHandler {
     }
 
     private void getSASLCredentials() throws IOException {
-        topologyName = (String) 
this.client.getStormConf().get(Config.TOPOLOGY_NAME);
-        String secretKey = SaslUtils.getSecretKey(this.client.getStormConf());
+        String secretKey;
+        name = client.name();
+        secretKey = client.secretKey();
+
         if (secretKey != null) {
             token = secretKey.getBytes();
         }
-        LOG.debug("SASL credentials for storm topology " + topologyName
+        LOG.debug("SASL credentials for storm topology " + name
                 + " is " + secretKey);
     }
-}
\ No newline at end of file
+}

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/jvm/backtype/storm/messaging/netty/Server.java
----------------------------------------------------------------------
diff --git a/storm-core/src/jvm/backtype/storm/messaging/netty/Server.java 
b/storm-core/src/jvm/backtype/storm/messaging/netty/Server.java
index e984144..7011b42 100644
--- a/storm-core/src/jvm/backtype/storm/messaging/netty/Server.java
+++ b/storm-core/src/jvm/backtype/storm/messaging/netty/Server.java
@@ -24,11 +24,13 @@ import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Collection;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executors;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ThreadFactory;
+import java.io.IOException;
 
 import org.jboss.netty.bootstrap.ServerBootstrap;
 import org.jboss.netty.channel.Channel;
@@ -36,17 +38,19 @@ import org.jboss.netty.channel.ChannelFactory;
 import org.jboss.netty.channel.group.ChannelGroup;
 import org.jboss.netty.channel.group.DefaultChannelGroup;
 import org.jboss.netty.channel.socket.nio.NioServerSocketChannelFactory;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import backtype.storm.Config;
+import backtype.storm.grouping.Load;
 import backtype.storm.messaging.ConnectionWithStatus;
-import backtype.storm.messaging.IConnection;
 import backtype.storm.messaging.TaskMessage;
 import backtype.storm.metric.api.IStatefulObject;
+import backtype.storm.serialization.KryoValuesSerializer;
 import backtype.storm.utils.Utils;
 
-class Server extends ConnectionWithStatus implements IStatefulObject {
+class Server extends ConnectionWithStatus implements IStatefulObject, 
ISaslServer {
 
     private static final Logger LOG = LoggerFactory.getLogger(Server.class);
     @SuppressWarnings("rawtypes")
@@ -56,7 +60,6 @@ class Server extends ConnectionWithStatus implements 
IStatefulObject {
     private final AtomicInteger messagesDequeued = new AtomicInteger(0);
     private final AtomicInteger[] pendingMessages;
     
-    
     // Create multiple queues for incoming messages. The size equals the 
number of receiver threads.
     // For message which is sent to same task, it will be stored in the same 
queue to preserve the message order.
     private LinkedBlockingQueue<ArrayList<TaskMessage>>[] message_queue;
@@ -71,12 +74,13 @@ class Server extends ConnectionWithStatus implements 
IStatefulObject {
        
     private volatile boolean closing = false;
     List<TaskMessage> closeMessage = Arrays.asList(new TaskMessage(-1, null));
-    
+    private KryoValuesSerializer _ser;
     
     @SuppressWarnings("rawtypes")
     Server(Map storm_conf, int port) {
         this.storm_conf = storm_conf;
         this.port = port;
+        _ser = new KryoValuesSerializer(storm_conf);
         
         queueCount = 
Utils.getInt(storm_conf.get(Config.WORKER_RECEIVER_THREAD_COUNT), 1);
         roundRobinQueueId = 0;
@@ -94,8 +98,8 @@ class Server extends ConnectionWithStatus implements 
IStatefulObject {
         int backlog = 
Utils.getInt(storm_conf.get(Config.STORM_MESSAGING_NETTY_SOCKET_BACKLOG), 500);
         int maxWorkers = 
Utils.getInt(storm_conf.get(Config.STORM_MESSAGING_NETTY_SERVER_WORKER_THREADS));
 
-        ThreadFactory bossFactory = new NettyRenameThreadFactory(name() + 
"-boss");
-        ThreadFactory workerFactory = new NettyRenameThreadFactory(name() + 
"-worker");
+        ThreadFactory bossFactory = new NettyRenameThreadFactory(netty_name() 
+ "-boss");
+        ThreadFactory workerFactory = new 
NettyRenameThreadFactory(netty_name() + "-worker");
         
         if (maxWorkers > 0) {
             factory = new 
NioServerSocketChannelFactory(Executors.newCachedThreadPool(bossFactory), 
@@ -105,7 +109,7 @@ class Server extends ConnectionWithStatus implements 
IStatefulObject {
                 Executors.newCachedThreadPool(workerFactory));
         }
         
-        LOG.info("Create Netty Server " + name() + ", buffer_size: " + 
buffer_size + ", maxWorkers: " + maxWorkers);
+        LOG.info("Create Netty Server " + netty_name() + ", buffer_size: " + 
buffer_size + ", maxWorkers: " + maxWorkers);
         
         bootstrap = new ServerBootstrap(factory);
         bootstrap.setOption("child.tcpNoDelay", true);
@@ -181,7 +185,6 @@ class Server extends ConnectionWithStatus implements 
IStatefulObject {
         }
     }
 
-
     /**
      * enqueue a received message 
      * @throws InterruptedException
@@ -216,14 +219,12 @@ class Server extends ConnectionWithStatus implements 
IStatefulObject {
         if ((flags & 0x01) == 0x01) {
             //non-blocking
             ret = message_queue[queueId].poll();
-        }
-        else {
+        } else {
             try {
                 ArrayList<TaskMessage> request = message_queue[queueId].take();
                 LOG.debug("request to be processed: {}", request);
                 ret = request;
-            }
-            catch (InterruptedException e) {
+            } catch (InterruptedException e) {
                 LOG.info("exception within msg receiving", e);
                 ret = null;
             }
@@ -249,7 +250,7 @@ class Server extends ConnectionWithStatus implements 
IStatefulObject {
      * close a channel
      * @param channel
      */
-    protected void closeChannel(Channel channel) {
+    public void closeChannel(Channel channel) {
         channel.close().awaitUninterruptibly();
         allChannels.remove(channel);
     }
@@ -265,15 +266,33 @@ class Server extends ConnectionWithStatus implements 
IStatefulObject {
         }
     }
 
+    @Override
+    public void sendLoadMetrics(Map<Integer, Double> taskToLoad) {
+        try {
+            MessageBatch mb = new MessageBatch(1);
+            mb.add(new TaskMessage(-1, 
_ser.serialize(Arrays.asList((Object)taskToLoad))));
+            allChannels.write(mb);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public Map<Integer, Load> getLoad(Collection<Integer> tasks) {
+        throw new RuntimeException("Server connection cannot get load");
+    }
+
+    @Override
     public void send(int task, byte[] message) {
         throw new UnsupportedOperationException("Server connection should not 
send any messages");
     }
     
+    @Override
     public void send(Iterator<TaskMessage> msgs) {
       throw new UnsupportedOperationException("Server connection should not 
send any messages");
     }
        
-    public String name() {
+    public String netty_name() {
       return "Netty-server-localhost-" + port;
     }
 
@@ -306,7 +325,7 @@ class Server extends ConnectionWithStatus implements 
IStatefulObject {
     }
 
     public Object getState() {
-        LOG.info("Getting metrics for server on port {}", port);
+        LOG.debug("Getting metrics for server on port {}", port);
         HashMap<String, Object> ret = new HashMap<String, Object>();
         ret.put("dequeuedMessages", messagesDequeued.getAndSet(0));
         ArrayList<Integer> pending = new 
ArrayList<Integer>(pendingMessages.length);
@@ -330,8 +349,30 @@ class Server extends ConnectionWithStatus implements 
IStatefulObject {
         return ret;
     }
 
-    @Override public String toString() {
-       return String.format("Netty server listening on port %s", port);
+    /** Implementing IServer. **/
+    public void channelConnected(Channel c) {
+        addChannel(c);
+    }
+
+    public void received(Object message, String remote, Channel channel)  
throws InterruptedException {
+        List<TaskMessage>msgs = (List<TaskMessage>)message;
+        enqueue(msgs, remote);
+    }
+
+    public String name() {
+        return (String)storm_conf.get(Config.TOPOLOGY_NAME);
+    }
+
+    public String secretKey() {
+        return SaslUtils.getSecretKey(storm_conf);
     }
 
+    public void authenticated(Channel c) {
+        return;
+    }
+
+    @Override
+    public String toString() {
+        return String.format("Netty server listening on port %s", port);
+    }
 }

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/jvm/backtype/storm/messaging/netty/StormClientHandler.java
----------------------------------------------------------------------
diff --git 
a/storm-core/src/jvm/backtype/storm/messaging/netty/StormClientHandler.java 
b/storm-core/src/jvm/backtype/storm/messaging/netty/StormClientHandler.java
index 2d25001..877b6d8 100644
--- a/storm-core/src/jvm/backtype/storm/messaging/netty/StormClientHandler.java
+++ b/storm-core/src/jvm/backtype/storm/messaging/netty/StormClientHandler.java
@@ -17,7 +17,20 @@
  */
 package backtype.storm.messaging.netty;
 
+import backtype.storm.messaging.TaskMessage;
+import backtype.storm.serialization.KryoValuesDeserializer;
+
 import java.net.ConnectException;
+import java.util.Map;
+import java.util.List;
+import java.io.IOException;
+
+import org.jboss.netty.channel.Channel;
+import org.jboss.netty.channel.ChannelHandlerContext;
+import org.jboss.netty.channel.ChannelStateEvent;
+import org.jboss.netty.channel.ExceptionEvent;
+import org.jboss.netty.channel.MessageEvent;
+import org.jboss.netty.channel.SimpleChannelUpstreamHandler;
 
 import org.jboss.netty.channel.*;
 import org.slf4j.Logger;
@@ -26,9 +39,41 @@ import org.slf4j.LoggerFactory;
 public class StormClientHandler extends SimpleChannelUpstreamHandler  {
     private static final Logger LOG = 
LoggerFactory.getLogger(StormClientHandler.class);
     private Client client;
-    
-    StormClientHandler(Client client) {
+    private KryoValuesDeserializer _des;
+
+    StormClientHandler(Client client, Map conf) {
         this.client = client;
+        _des = new KryoValuesDeserializer(conf);
+    }
+
+    @Override
+    public void messageReceived(ChannelHandlerContext ctx, MessageEvent event) 
{
+        //examine the response message from server
+        Object message = event.getMessage();
+        if (message instanceof ControlMessage) {
+            ControlMessage msg = (ControlMessage)message;
+            if (msg==ControlMessage.FAILURE_RESPONSE) {
+                LOG.info("failure response:{}", msg);
+            }
+        } else if (message instanceof List) {
+            try {
+                //This should be the metrics, and there should only be one of 
them
+                List<TaskMessage> list = (List<TaskMessage>)message;
+                if (list.size() < 1) throw new RuntimeException("Didn't see 
enough load metrics ("+client.getDstAddress()+") "+list);
+                if (list.size() != 1) LOG.warn("Messages are not being 
delivered fast enough, got "+list.size()+" metrics messages at 
once("+client.getDstAddress()+")");
+                TaskMessage tm = ((List<TaskMessage>)message).get(list.size() 
- 1);
+                if (tm.task() != -1) throw new RuntimeException("Metrics 
messages are sent to the system task ("+client.getDstAddress()+") "+tm);
+                List metrics = _des.deserialize(tm.message());
+                if (metrics.size() < 1) throw new RuntimeException("No metrics 
data in the metrics message ("+client.getDstAddress()+") "+metrics);
+                if (!(metrics.get(0) instanceof Map)) throw new 
RuntimeException("The metrics did not have a map in the first slot 
("+client.getDstAddress()+") "+metrics);
+                client.setLoadMetrics((Map<Integer, Double>)metrics.get(0));
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        } else {
+            throw new RuntimeException("Don't know how to handle a message of 
type "
+                                       + message + " (" + 
client.getDstAddress() + ")");
+        }
     }
 
     @Override
@@ -40,7 +85,7 @@ public class StormClientHandler extends 
SimpleChannelUpstreamHandler  {
     public void exceptionCaught(ChannelHandlerContext ctx, ExceptionEvent 
event) {
         Throwable cause = event.getCause();
         if (!(cause instanceof ConnectException)) {
-            LOG.info("Connection failed " + client.dstAddressPrefixedName, 
cause);
+            LOG.info("Connection to "+client.getDstAddress()+" failed:", 
cause);
         }
     }
 }

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/src/jvm/backtype/storm/messaging/netty/StormClientPipelineFactory.java
----------------------------------------------------------------------
diff --git 
a/storm-core/src/jvm/backtype/storm/messaging/netty/StormClientPipelineFactory.java
 
b/storm-core/src/jvm/backtype/storm/messaging/netty/StormClientPipelineFactory.java
index 4be06cd..6158eef 100644
--- 
a/storm-core/src/jvm/backtype/storm/messaging/netty/StormClientPipelineFactory.java
+++ 
b/storm-core/src/jvm/backtype/storm/messaging/netty/StormClientPipelineFactory.java
@@ -22,12 +22,15 @@ import org.jboss.netty.channel.ChannelPipelineFactory;
 import org.jboss.netty.channel.Channels;
 
 import backtype.storm.Config;
+import java.util.Map;
 
 class StormClientPipelineFactory implements ChannelPipelineFactory {
     private Client client;
+    private Map conf;
 
-    StormClientPipelineFactory(Client client) {
+    StormClientPipelineFactory(Client client, Map conf) {
         this.client = client;
+        this.conf = conf;
     }
 
     public ChannelPipeline getPipeline() throws Exception {
@@ -39,15 +42,15 @@ class StormClientPipelineFactory implements 
ChannelPipelineFactory {
         // Encoder
         pipeline.addLast("encoder", new MessageEncoder());
 
-        boolean isNettyAuth = (Boolean) 
this.client.getStormConf().get(Config.STORM_MESSAGING_NETTY_AUTHENTICATION);
+        boolean isNettyAuth = (Boolean) conf
+                .get(Config.STORM_MESSAGING_NETTY_AUTHENTICATION);
         if (isNettyAuth) {
             // Authenticate: Removed after authentication completes
             pipeline.addLast("saslClientHandler", new SaslStormClientHandler(
                     client));
         }
         // business logic.
-        pipeline.addLast("handler", new StormClientHandler(client));
-
+        pipeline.addLast("handler", new StormClientHandler(client, conf));
         return pipeline;
     }
 }

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/test/clj/backtype/storm/grouping_test.clj
----------------------------------------------------------------------
diff --git a/storm-core/test/clj/backtype/storm/grouping_test.clj 
b/storm-core/test/clj/backtype/storm/grouping_test.clj
index 2bfa066..f914591 100644
--- a/storm-core/test/clj/backtype/storm/grouping_test.clj
+++ b/storm-core/test/clj/backtype/storm/grouping_test.clj
@@ -17,27 +17,85 @@
   (:use [clojure test])
   (:import [backtype.storm.testing TestWordCounter TestWordSpout 
TestGlobalCount TestAggregatesCounter TestWordBytesCounter NGrouping]
            [backtype.storm.generated JavaObject JavaObjectArg])
-  (:use [backtype.storm testing clojure])
-  (:use [backtype.storm.daemon common])
+  (:import [backtype.storm.grouping LoadMapping])
+  (:use [backtype.storm testing clojure log config])
+  (:use [backtype.storm.daemon common executor])
   (:require [backtype.storm [thrift :as thrift]]))
 
 (deftest test-shuffle
+ (let [shuffle-fn (mk-shuffle-grouper [(int 1) (int 2)] 
{TOPOLOGY-DISABLE-LOADAWARE-MESSAGING true} nil "comp" "stream")
+       num-messages 100000
+       min-prcnt (int (* num-messages 0.49))
+       max-prcnt (int (* num-messages 0.51))
+       data [1 2]
+       freq (frequencies (for [x (range 0 num-messages)] (shuffle-fn (int 1) 
data nil)))
+       load1 (.get freq [(int 1)])
+       load2 (.get freq [(int 2)])]
+    (log-message "FREQ:" freq)
+    (is (>= load1 min-prcnt))
+    (is (<= load1 max-prcnt))
+    (is (>= load2 min-prcnt))
+    (is (<= load2 max-prcnt))))
+
+(deftest test-shuffle-load-even
+ (let [shuffle-fn (mk-shuffle-grouper [(int 1) (int 2)] {} nil "comp" "stream")
+       num-messages 100000
+       min-prcnt (int (* num-messages 0.49))
+       max-prcnt (int (* num-messages 0.51))
+       load (LoadMapping.)
+       _ (.setLocal load {(int 1) 0.0 (int 2) 0.0})
+       data [1 2]
+       freq (frequencies (for [x (range 0 num-messages)] (shuffle-fn (int 1) 
data load)))
+       load1 (.get freq [(int 1)])
+       load2 (.get freq [(int 2)])]
+    (log-message "FREQ:" freq)
+    (is (>= load1 min-prcnt))
+    (is (<= load1 max-prcnt))
+    (is (>= load2 min-prcnt))
+    (is (<= load2 max-prcnt))))
+
+(deftest test-shuffle-load-uneven
+ (let [shuffle-fn (mk-shuffle-grouper [(int 1) (int 2)] {} nil "comp" "stream")
+       num-messages 100000
+       min1-prcnt (int (* num-messages 0.32))
+       max1-prcnt (int (* num-messages 0.34))
+       min2-prcnt (int (* num-messages 0.65))
+       max2-prcnt (int (* num-messages 0.67))
+       load (LoadMapping.)
+       _ (.setLocal load {(int 1) 0.5 (int 2) 0.0})
+       data [1 2]
+       freq (frequencies (for [x (range 0 num-messages)] (shuffle-fn (int 1) 
data load)))
+       load1 (.get freq [(int 1)])
+       load2 (.get freq [(int 2)])]
+    (log-message "FREQ:" freq)
+    (is (>= load1 min1-prcnt))
+    (is (<= load1 max1-prcnt))
+    (is (>= load2 min2-prcnt))
+    (is (<= load2 max2-prcnt))))
+
+(deftest test-field
   (with-simulated-time-local-cluster [cluster :supervisors 4]
-    (let [topology (thrift/mk-topology
-                    {"1" (thrift/mk-spout-spec (TestWordSpout. true) 
:parallelism-hint 4)}
-                    {"2" (thrift/mk-bolt-spec {"1" :shuffle} (TestGlobalCount.)
-                                            :parallelism-hint 6)
+    (let [spout-phint 4
+          bolt-phint 6
+          topology (thrift/mk-topology
+                    {"1" (thrift/mk-spout-spec (TestWordSpout. true)
+                                               :parallelism-hint spout-phint)}
+                    {"2" (thrift/mk-bolt-spec {"1" ["word"]}
+                                              (TestWordBytesCounter.)
+                                              :parallelism-hint bolt-phint)
                      })
-          results (complete-topology cluster
-                                     topology
-                                     ;; important for test that
-                                     ;; #tuples = multiple of 4 and 6
-                                     :mock-sources {"1" (->> [["a"] ["b"]]
-                                                             (repeat 12)
-                                                             (apply concat))})]
-      (is (ms= (apply concat (repeat 6 [[1] [2] [3] [4]]))
-               (read-tuples results "2")))
-      )))
+          results (complete-topology
+                    cluster
+                    topology
+                    :mock-sources {"1" (->> [[(.getBytes "a")]
+                                             [(.getBytes "b")]]
+                                            (repeat (* spout-phint bolt-phint))
+                                            (apply concat))})]
+      (is (ms= (apply concat
+                      (for [value '("a" "b")
+                            sum (range 1 (inc (* spout-phint bolt-phint)))]
+                        [[value sum]]))
+               (read-tuples results "2"))))))
 
 (deftest test-field
   (with-simulated-time-local-cluster [cluster :supervisors 4]

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/test/clj/backtype/storm/integration_test.clj
----------------------------------------------------------------------
diff --git a/storm-core/test/clj/backtype/storm/integration_test.clj 
b/storm-core/test/clj/backtype/storm/integration_test.clj
index ca17d6e..cc0208d 100644
--- a/storm-core/test/clj/backtype/storm/integration_test.clj
+++ b/storm-core/test/clj/backtype/storm/integration_test.clj
@@ -62,13 +62,13 @@
   (with-simulated-time-local-cluster [cluster :supervisors 4]
     (let [topology (thrift/mk-topology
                     {"1" (thrift/mk-spout-spec (TestWordSpout. true))}
-                    {"2" (thrift/mk-bolt-spec {"1" :shuffle} emit-task-id
+                    {"2" (thrift/mk-bolt-spec {"1" :all} emit-task-id
                       :parallelism-hint 3
                       :conf {TOPOLOGY-TASKS 6})
                      })
           results (complete-topology cluster
                                      topology
-                                     :mock-sources {"1" [["a"] ["a"] ["a"] 
["a"] ["a"] ["a"]]})]
+                                     :mock-sources {"1" [["a"]]})]
       (is (ms= [[0] [1] [2] [3] [4] [5]]
                (read-tuples results "2")))
       )))

http://git-wip-us.apache.org/repos/asf/storm/blob/af64a3a2/storm-core/test/clj/backtype/storm/messaging/netty_integration_test.clj
----------------------------------------------------------------------
diff --git 
a/storm-core/test/clj/backtype/storm/messaging/netty_integration_test.clj 
b/storm-core/test/clj/backtype/storm/messaging/netty_integration_test.clj
index a918a4b..8081ca5 100644
--- a/storm-core/test/clj/backtype/storm/messaging/netty_integration_test.clj
+++ b/storm-core/test/clj/backtype/storm/messaging/netty_integration_test.clj
@@ -54,5 +54,4 @@
                                                          ["a"] ["b"]
                                                          ]}
                                      )]
-      (is (ms= (apply concat (repeat 6 [[1] [2] [3] [4]]))
-               (read-tuples results "2"))))))
+        (is (= (* 6 4) (.size (read-tuples results "2")))))))

Reply via email to