From 836f127682c53b3e47b101e37e8c19a9e3796f29 Mon Sep 17 00:00:00 2001
From: Stephen Salinas <ssalinas@hubspot.com>
Date: Thu, 10 Dec 2015 14:22:21 -0500
Subject: [PATCH] take a list of string values for each reserved attribute

---
 .../singularity/config/SingularityConfiguration.java     | 7 ++++---
 .../singularity/mesos/SingularitySlaveAndRackHelper.java | 9 ++++++---
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/SingularityService/src/main/java/com/hubspot/singularity/config/SingularityConfiguration.java b/SingularityService/src/main/java/com/hubspot/singularity/config/SingularityConfiguration.java
index a9057dad7f..0dad6c45fb 100644
--- a/SingularityService/src/main/java/com/hubspot/singularity/config/SingularityConfiguration.java
+++ b/SingularityService/src/main/java/com/hubspot/singularity/config/SingularityConfiguration.java
@@ -1,6 +1,7 @@
 package com.hubspot.singularity.config;
 
 import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 
@@ -204,7 +205,7 @@ public class SingularityConfiguration extends Configuration {
   private AuthConfiguration authConfiguration = new AuthConfiguration();
 
   @NotNull
-  private Map<String, String> reserveSlavesWithAttributes = Collections.emptyMap();
+  private Map<String, List<String>> reserveSlavesWithAttributes = Collections.emptyMap();
 
   @JsonProperty("graphite")
   @NotNull
@@ -834,11 +835,11 @@ public void setHistoryPurgingConfiguration(HistoryPurgingConfiguration historyPu
     this.historyPurgingConfiguration = historyPurgingConfiguration;
   }
 
-  public Map<String, String> getReserveSlavesWithAttributes() {
+  public Map<String, List<String>> getReserveSlavesWithAttributes() {
     return reserveSlavesWithAttributes;
   }
 
-  public void setReserveSlavesWithAttrbiutes(Map<String, String> reserveSlavesWithAttributes) {
+  public void setReserveSlavesWithAttrbiutes(Map<String, List<String>> reserveSlavesWithAttributes) {
     this.reserveSlavesWithAttributes = reserveSlavesWithAttributes;
   }
 
diff --git a/SingularityService/src/main/java/com/hubspot/singularity/mesos/SingularitySlaveAndRackHelper.java b/SingularityService/src/main/java/com/hubspot/singularity/mesos/SingularitySlaveAndRackHelper.java
index 711ea7070d..6721fbc195 100644
--- a/SingularityService/src/main/java/com/hubspot/singularity/mesos/SingularitySlaveAndRackHelper.java
+++ b/SingularityService/src/main/java/com/hubspot/singularity/mesos/SingularitySlaveAndRackHelper.java
@@ -1,6 +1,7 @@
 package com.hubspot.singularity.mesos;
 
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 import javax.inject.Singleton;
@@ -95,9 +96,11 @@ public Map<String, String> getTextAttributes(Offer offer) {
   public Map<String, String> reservedSlaveAttributes(Offer offer) {
     Map<String, String> reservedAttributes = new HashMap<>();
     Map<String, String> offerTextAttributes = getTextAttributes(offer);
-    for (Map.Entry<String, String> entry : configuration.getReserveSlavesWithAttributes().entrySet()) {
-      if (offerTextAttributes.containsKey(entry.getKey()) && offerTextAttributes.get(entry.getKey()).equals(entry.getValue())) {
-        reservedAttributes.put(entry.getKey(), entry.getValue());
+    for (Map.Entry<String, List<String>> entry : configuration.getReserveSlavesWithAttributes().entrySet()) {
+      for (String attr : entry.getValue()) {
+        if (offerTextAttributes.containsKey(entry.getKey()) && offerTextAttributes.get(entry.getKey()).equals(attr)) {
+          reservedAttributes.put(entry.getKey(), attr);
+        }
       }
     }
     return reservedAttributes;