add support for upstream groups

BaragonAgentService/src/main/java/com/hubspot/baragon/agent/managers/AgentRequestManager.java

@@ -159,20 +159,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 (currentUpstream.sameAs(existingUpstream)) {
+            present = true;
+            break;
+          }
+        }
+        for (UpstreamInfo upstreamToRemove : request.getRemoveUpstreams()) {
+          if (upstreamToRemove.sameAs(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);
     }
   }

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;
     }

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;
   }

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,13 @@
 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";
+
   private final String upstream;
 
   @Size(max=250)
@@ -25,6 +27,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 +40,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 +87,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,16 +127,24 @@ public boolean equals(Object o) {
     if (!originalPath.equals(that.originalPath)) {
       return false;
     }
+    if (!group.equals(that.group)) {
+      return false;
+    }
 
     return true;
   }
 
+  public boolean sameAs(UpstreamInfo that) {
+    return upstream.equals(that.upstream) && group.equals(that.group);
+  }
+
   @Override
   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;
   }
 }

BaragonData/src/main/java/com/hubspot/baragon/data/BaragonStateDatastore.java

@@ -97,9 +97,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 +109,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 +118,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 +138,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 +152,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.sameAs(toAdd)) {
         matchingPaths.add(upstreamInfo.getOriginalPath().or(upstreamInfo.getUpstream()));
       }
     }
@@ -170,7 +161,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.sameAs(toRemove)) {
         return upstreamInfo.getOriginalPath().or(upstreamInfo.toPath());
       }
     }

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}}