Add whitelist methods to Cluster to allow only a subset of hosts to be revealed as available slots

Nathan Marz · Nathan Marz · commit a8df52dcc427 · 2012-12-16T16:13:53.000-08:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -28,6 +28,7 @@
  * Storm UI displays exception instead of blank page when there's an error rendering the page (thanks Frostman)
  * Added MultiScheme interface (thanks sritchie)
  * Added MockTridentTuple for testing (thanks emblem)
+ * Add whitelist methods to Cluster to allow only a subset of hosts to be revealed as available slots
  * Updated Trident Debug filter to take in an identifier to use when logging (thanks emblem)
  * Bug fix: Fix for bug that could cause topology to hang when ZMQ blocks sending to a worker that got reassigned
  * Bug fix: Fix deadlock bug due to variant of dining philosophers problem. Spouts now use an overflow buffer to prevent blocking and guarantee that it can consume the incoming queue of acks/fails.
diff --git a/src/jvm/backtype/storm/scheduler/Cluster.java b/src/jvm/backtype/storm/scheduler/Cluster.java
@@ -23,6 +23,8 @@ public class Cluster {
      * a map from hostname to supervisor id.
      */
     private Map<String, List<String>>        hostToId;
+    
+    private Set<String> whiteListedHosts = new HashSet<String>();
 
     public Cluster(Map<String, SupervisorDetails> supervisors, Map<String, SchedulerAssignmentImpl> assignments){
         this.supervisors = new HashMap<String, SupervisorDetails>(supervisors.size());
@@ -39,7 +41,11 @@ public Cluster(Map<String, SupervisorDetails> supervisors, Map<String, Scheduler
             this.hostToId.get(host).add(nodeId);
         }
     }
-
+    
+    public void setWhitelistedHosts(Set<String> hosts) {
+        whiteListedHosts = hosts;
+    }
+    
     /**
      * Gets all the topologies which needs scheduling.
      * 
@@ -161,6 +167,7 @@ public List<Integer> getAvailablePorts(SupervisorDetails supervisor) {
      * @return
      */
     public List<WorkerSlot> getAvailableSlots(SupervisorDetails supervisor) {
+        if(whiteListedHosts!=null && !whiteListedHosts.isEmpty() && !whiteListedHosts.contains(supervisor.host)) return new ArrayList();
         List<Integer> ports = this.getAvailablePorts(supervisor);
         List<WorkerSlot> slots = new ArrayList<WorkerSlot>(ports.size());
 
