Merge pull request #179 from HubSpot/upstream-groups

add support for upstream groups
HubSpot · May 6, 2016 · 72ae63e · 72ae63e
2 parents 8797a63 + 42ac7b4
commit 72ae63e
Show file tree

Hide file tree

Showing 6 changed files with 103 additions and 65 deletions.
diff --git a/...gonAgentService/src/main/java/com/hubspot/baragon/agent/managers/AgentRequestManager.java b/...gonAgentService/src/main/java/com/hubspot/baragon/agent/managers/AgentRequestManager.java
@@ -1,14 +1,5 @@
 package com.hubspot.baragon.agent.managers;
 
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.List;
-import java.util.Random;
-import java.util.concurrent.atomic.AtomicReference;
-
-import javax.ws.rs.core.Response;
-
 import com.google.common.base.Optional;
 import com.google.inject.Inject;
 import com.google.inject.Singleton;
@@ -30,6 +21,13 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicReference;
+import javax.ws.rs.core.Response;
+
 @Singleton
 public class AgentRequestManager {
   private static final Logger LOG = LoggerFactory.getLogger(AgentRequestManager.class);
@@ -159,20 +157,27 @@ private ServiceContext getApplyContext(BaragonRequest request) throws Exception
       return new ServiceContext(request.getLoadBalancerService(), request.getReplaceUpstreams(), System.currentTimeMillis(), true);
     } else {
       List<UpstreamInfo> upstreams = new ArrayList<>();
-      upstreams.addAll(stateDatastore.getUpstreams(request.getLoadBalancerService().getServiceId()));
-
-      List<UpstreamInfo> toRemove = new ArrayList<>();
-      for (UpstreamInfo removeUpstreamInfo : request.getRemoveUpstreams()) {
-        for (UpstreamInfo upstreamInfo : upstreams) {
-          if (upstreamInfo.getUpstream().equals(removeUpstreamInfo.getUpstream())) {
-            toRemove.add(upstreamInfo);
+      upstreams.addAll(request.getAddUpstreams());
+      for (UpstreamInfo existingUpstream : stateDatastore.getUpstreams(request.getLoadBalancerService().getServiceId())) {
+        boolean present = false;
+        boolean toRemove = false;
+        for (UpstreamInfo currentUpstream : upstreams) {
+          if (UpstreamInfo.upstreamAndGroupMatches(currentUpstream, existingUpstream)) {
+            present = true;
+            break;
           }
         }
+        for (UpstreamInfo upstreamToRemove : request.getRemoveUpstreams()) {
+          if (UpstreamInfo.upstreamAndGroupMatches(upstreamToRemove, existingUpstream)) {
+            toRemove = true;
+            break;
+          }
+        }
+        if (!present && !toRemove) {
+          upstreams.add(existingUpstream);
+        }
       }
 
-      upstreams.removeAll(toRemove);
-      upstreams.addAll(request.getAddUpstreams());
-
       return new ServiceContext(request.getLoadBalancerService(), upstreams, System.currentTimeMillis(), true);
     }
   }

diff --git a/BaragonCore/src/main/java/com/hubspot/baragon/models/BaragonRequest.java b/BaragonCore/src/main/java/com/hubspot/baragon/models/BaragonRequest.java
@@ -122,7 +122,7 @@ private List<UpstreamInfo> addRequestId(List<UpstreamInfo> upstreams, String req
 
   private UpstreamInfo addRequestId(UpstreamInfo upstream, String requestId) {
     if (!upstream.getRequestId().isPresent()) {
-      return new UpstreamInfo(upstream.getUpstream(), Optional.of(requestId), upstream.getRackId());
+      return new UpstreamInfo(upstream.getUpstream(), Optional.of(requestId), upstream.getRackId(), Optional.of(upstream.getGroup()));
     } else {
       return upstream;
     }

diff --git a/BaragonCore/src/main/java/com/hubspot/baragon/models/ServiceContext.java b/BaragonCore/src/main/java/com/hubspot/baragon/models/ServiceContext.java
@@ -1,17 +1,21 @@
 package com.hubspot.baragon.models;
 
-import java.util.Collection;
-import java.util.Collections;
-
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Objects;
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Map;
 
 @JsonIgnoreProperties(ignoreUnknown = true)
 public class ServiceContext {
   private final BaragonService service;
   private final Collection<UpstreamInfo> upstreams;
+  private final Map<String, Collection<UpstreamInfo>> upstreamGroups;
   private final Long timestamp;
   private final boolean present;
   private final boolean rootPath;
@@ -26,6 +30,16 @@ public ServiceContext(@JsonProperty("service") BaragonService service,
     this.upstreams = Objects.firstNonNull(upstreams, Collections.<UpstreamInfo>emptyList());
     this.present = present;
     this.rootPath = service.getServiceBasePath().equals("/");
+
+    if (!this.upstreams.isEmpty()) {
+      final Multimap<String, UpstreamInfo> upstreamGroupsMultimap = ArrayListMultimap.create();
+      for (UpstreamInfo upstream : this.upstreams) {
+        upstreamGroupsMultimap.put(upstream.getGroup(), upstream);
+      }
+      this.upstreamGroups = upstreamGroupsMultimap.asMap();
+    } else {
+      this.upstreamGroups = Collections.emptyMap();
+    }
   }
 
   public BaragonService getService() {
@@ -36,6 +50,10 @@ public Collection<UpstreamInfo> getUpstreams() {
     return upstreams;
   }
 
+  public Map<String, Collection<UpstreamInfo>> getUpstreamGroups() {
+    return upstreamGroups;
+  }
+
   public Long getTimestamp() {
     return timestamp;
   }

diff --git a/BaragonCore/src/main/java/com/hubspot/baragon/models/UpstreamInfo.java b/BaragonCore/src/main/java/com/hubspot/baragon/models/UpstreamInfo.java
@@ -1,8 +1,5 @@
 package com.hubspot.baragon.models;
 
-import javax.validation.constraints.Pattern;
-import javax.validation.constraints.Size;
-
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
@@ -12,8 +9,17 @@
 import com.google.common.base.Strings;
 import com.google.common.io.BaseEncoding;
 
+import javax.validation.constraints.Pattern;
+import javax.validation.constraints.Size;
+
 @JsonIgnoreProperties( ignoreUnknown = true )
 public class UpstreamInfo {
+  public static final String DEFAULT_GROUP = "default";
+
+  public static boolean upstreamAndGroupMatches(UpstreamInfo a, UpstreamInfo b) {
+    return a.getUpstream().equals(b.getUpstream()) && a.getGroup().equals(b.getGroup());
+  }
+
   private final String upstream;
 
   @Size(max=250)
@@ -25,6 +31,10 @@ public class UpstreamInfo {
   private final String rackId;
   private final Optional<String> originalPath;
 
+  @Size(max=100)
+  @Pattern(regexp = "^$|[^\\s|]+", message = "cannot contain whitespace, '/', or '|'")
+  private final String group;
+
   @JsonCreator
   public static UpstreamInfo fromString(String value) {
     String[] split = value.split("\\|", -1);
@@ -34,27 +44,34 @@ public static UpstreamInfo fromString(String value) {
       String upstream = new String(BaseEncoding.base64Url().decode(split[0]), Charsets.UTF_8);
       Optional<String> requestId = split.length > 1 && !split[1].equals("") ? Optional.of(split[1]) : Optional.<String>absent();
       Optional<String> rackId = split.length > 2 && !split[2].equals("") ? Optional.of(split[2]) : Optional.<String>absent();
-      return new UpstreamInfo(upstream, requestId, rackId, Optional.of(value));
+      Optional<String> group = split.length > 3 && !split[3].equals("") ? Optional.of(split[3]) : Optional.<String>absent();
+      return new UpstreamInfo(upstream, requestId, rackId, Optional.of(value), group);
     }
   }
 
   public static UpstreamInfo fromUnEncodedString(String upstream) {
-    return new UpstreamInfo(upstream, Optional.<String>absent(), Optional.<String>absent(), Optional.of(upstream));
+    return new UpstreamInfo(upstream, Optional.<String>absent(), Optional.<String>absent(), Optional.of(upstream), Optional.<String>absent());
+  }
+
+  public UpstreamInfo(String upstream, Optional<String> requestId, Optional<String> rackId) {
+    this(upstream, requestId, rackId, Optional.<String>absent(), Optional.<String>absent());
   }
 
-  public UpstreamInfo (String upstream, Optional<String> requestId, Optional<String> rackId) {
-    this(upstream, requestId, rackId, Optional.<String>absent());
+  public UpstreamInfo(String upstream, Optional<String> requestId, Optional<String> rackId, Optional<String> group) {
+    this(upstream, requestId, rackId, Optional.<String>absent(), group);
   }
 
   @JsonCreator
   public UpstreamInfo(@JsonProperty("upstream") String upstream,
                       @JsonProperty("requestId") Optional<String> requestId,
                       @JsonProperty("rackId") Optional<String> rackId,
-                      @JsonProperty("originalPath") Optional<String> originalPath) {
+                      @JsonProperty("originalPath") Optional<String> originalPath,
+                      @JsonProperty("group") Optional<String> group) {
     this.upstream = upstream;
     this.requestId = requestId.or("");
     this.rackId = rackId.or("");
     this.originalPath = originalPath;
+    this.group = group.or(DEFAULT_GROUP);
   }
 
   public String getUpstream() {
@@ -74,13 +91,17 @@ public Optional<String> getOriginalPath() {
     return originalPath;
   }
 
+  public String getGroup() {
+    return group;
+  }
+
   @Override
   public String toString() {
     return upstream;
   }
 
   public String toPath() {
-    return String.format("%s|%s|%s", sanitizeUpstream(upstream), requestId, rackId);
+    return String.format("%s|%s|%s|%s", sanitizeUpstream(upstream), requestId, rackId, group);
   }
 
   protected String sanitizeUpstream(String name) {
@@ -110,6 +131,9 @@ public boolean equals(Object o) {
     if (!originalPath.equals(that.originalPath)) {
       return false;
     }
+    if (!group.equals(that.group)) {
+      return false;
+    }
 
     return true;
   }
@@ -119,7 +143,8 @@ public int hashCode() {
     int result = upstream.hashCode();
     result = 31 * result + requestId.hashCode();
     result = 31 * result + rackId.hashCode();
-    result = 31 *result + originalPath.hashCode();
+    result = 31 * result + originalPath.hashCode();
+    result = 31 * result + group.hashCode();
     return result;
   }
 }
diff --git a/BaragonData/src/main/java/com/hubspot/baragon/data/BaragonStateDatastore.java b/BaragonData/src/main/java/com/hubspot/baragon/data/BaragonStateDatastore.java
@@ -1,21 +1,5 @@
 package com.hubspot.baragon.data;
 
-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Map.Entry;
-
-import org.apache.curator.framework.CuratorFramework;
-import org.apache.curator.framework.api.transaction.CuratorTransactionFinal;
-import org.apache.curator.utils.ZKPaths;
-import org.apache.zookeeper.data.Stat;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
 import com.codahale.metrics.annotation.Timed;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.base.Function;
@@ -31,6 +15,21 @@
 import com.hubspot.baragon.models.BaragonServiceState;
 import com.hubspot.baragon.models.UpstreamInfo;
 import com.hubspot.baragon.utils.ZkParallelFetcher;
+import org.apache.curator.framework.CuratorFramework;
+import org.apache.curator.framework.api.transaction.CuratorTransactionFinal;
+import org.apache.curator.utils.ZKPaths;
+import org.apache.zookeeper.data.Stat;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
 
 @Singleton
 public class BaragonStateDatastore extends AbstractDataStore {
@@ -97,9 +96,8 @@ public void updateService(BaragonRequest request) throws Exception {
       createNode(SERVICES_FORMAT);
     }
 
-    Collection<UpstreamInfo> currentUpstreams = getUpstreams(request.getLoadBalancerService().getServiceId());
-
     String serviceId = request.getLoadBalancerService().getServiceId();
+    Collection<UpstreamInfo> currentUpstreams = getUpstreams(serviceId);
     String servicePath = String.format(SERVICE_FORMAT, serviceId);
     CuratorTransactionFinal transaction;
     if (nodeExists(servicePath)) {
@@ -110,7 +108,7 @@ public void updateService(BaragonRequest request) throws Exception {
 
     List<String> pathsToDelete = new ArrayList<>();
     if (!request.getReplaceUpstreams().isEmpty()) {
-      for (UpstreamInfo upstreamInfo : getUpstreams(serviceId)) {
+      for (UpstreamInfo upstreamInfo : currentUpstreams) {
         String removePath = String.format(UPSTREAM_FORMAT, serviceId, getRemovePath(currentUpstreams, upstreamInfo));
         if (nodeExists(removePath)) {
           pathsToDelete.add(removePath);
@@ -119,15 +117,7 @@ public void updateService(BaragonRequest request) throws Exception {
       }
       for (UpstreamInfo upstreamInfo : request.getReplaceUpstreams()) {
         String addPath = String.format(UPSTREAM_FORMAT, serviceId, upstreamInfo.toPath());
-        List<String> matchingUpstreamPaths = matchingUpstreamPaths(currentUpstreams, upstreamInfo);
-        for (String matchingPath : matchingUpstreamPaths) {
-          String fullPath = String.format(UPSTREAM_FORMAT, serviceId, matchingPath);
-          if (nodeExists(fullPath) && !pathsToDelete.contains(fullPath)) {
-            pathsToDelete.add(fullPath);
-            transaction.delete().forPath(fullPath);
-          }
-        }
-        if (!nodeExists(addPath)) {
+        if (!nodeExists(addPath) || pathsToDelete.contains(addPath)) {
           transaction.create().forPath(addPath).and();
         }
       }
@@ -147,7 +137,7 @@ public void updateService(BaragonRequest request) throws Exception {
           if (nodeExists(fullPath) && !pathsToDelete.contains(fullPath) && !fullPath.equals(addPath)) {
             LOG.info(String.format("Deleting %s", fullPath));
             pathsToDelete.add(fullPath);
-            transaction.delete().forPath(fullPath);
+            transaction.delete().forPath(fullPath).and();
           }
         }
         if (!nodeExists(addPath) || pathsToDelete.contains(addPath)) {
@@ -161,7 +151,7 @@ public void updateService(BaragonRequest request) throws Exception {
   private List<String> matchingUpstreamPaths(Collection<UpstreamInfo> currentUpstreams, UpstreamInfo toAdd) {
     List<String> matchingPaths = new ArrayList<>();
     for (UpstreamInfo upstreamInfo : currentUpstreams) {
-      if (upstreamInfo.getUpstream().equals(toAdd.getUpstream())) {
+      if (UpstreamInfo.upstreamAndGroupMatches(upstreamInfo, toAdd)) {
         matchingPaths.add(upstreamInfo.getOriginalPath().or(upstreamInfo.getUpstream()));
       }
     }
@@ -170,7 +160,7 @@ private List<String> matchingUpstreamPaths(Collection<UpstreamInfo> currentUpstr
 
   private String getRemovePath(Collection<UpstreamInfo> currentUpstreams, UpstreamInfo toRemove) {
     for (UpstreamInfo upstreamInfo : currentUpstreams) {
-      if (upstreamInfo.getUpstream().equals(toRemove.getUpstream())) {
+      if (UpstreamInfo.upstreamAndGroupMatches(upstreamInfo, toRemove)) {
         return upstreamInfo.getOriginalPath().or(upstreamInfo.toPath());
       }
     }

diff --git a/BaragonUI/app/templates/serviceDetail.hbs b/BaragonUI/app/templates/serviceDetail.hbs
@@ -124,7 +124,7 @@
                     {{#each service.splitUpstreams}}
                         <div class="col-md-6">
                             {{#each this}}
-                                <h4>{{upstream}} {{#if rackId}}({{ rackId }}){{/if}}{{#if ../../config.allowEdit}}<a class="pull-left" data-action="removeUpstream"><span class="glyphicon glyphicon-remove inactive" data-upstream="{{ upstream }}"></span></a>{{/if}}</h4>
+                                <h4>{{upstream}} {{#if rackId}}({{ rackId }}){{/if}}{{#if group}}({{{ group }}}){{/if}}{{#if ../../config.allowEdit}}<a class="pull-left" data-action="removeUpstream"><span class="glyphicon glyphicon-remove inactive" data-upstream="{{ upstream }}"></span></a>{{/if}}</h4>
                             {{/each}}
                         </div>
                     {{/each}}