Add ability to proxy to discovery server in another environment

server/pom.xml

@@ -176,5 +176,12 @@
             <artifactId>testing</artifactId>
             <scope>test</scope>
         </dependency>
+
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-core</artifactId>
+            <scope>test</scope>
+        </dependency>
+
     </dependencies>
 </project>

server/src/main/java/io/airlift/discovery/DiscoveryConfig.java

@@ -15,15 +15,24 @@
  */
 package io.airlift.discovery;
 
+import com.google.common.collect.ForwardingSet;
+import com.google.common.collect.ImmutableSet;
 import io.airlift.configuration.Config;
+import io.airlift.configuration.ConfigDescription;
 import io.airlift.units.Duration;
 
+import javax.validation.constraints.AssertTrue;
 import javax.validation.constraints.NotNull;
+import java.net.URI;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
 public class DiscoveryConfig
 {
     private Duration maxAge = new Duration(30, TimeUnit.SECONDS);
+    private StringSet proxyProxiedTypes = StringSet.of();
+    private String proxyEnvironment = null;
+    private URI proxyUri = null;
 
     @NotNull
     public Duration getMaxAge()
@@ -32,9 +41,87 @@ public Duration getMaxAge()
     }
 
     @Config("discovery.max-age")
+    @ConfigDescription("Dynamic announcement expiration")
     public DiscoveryConfig setMaxAge(Duration maxAge)
     {
         this.maxAge = maxAge;
         return this;
     }
+
+    public StringSet getProxyProxiedTypes()
+    {
+        return proxyProxiedTypes;
+    }
+
+    @Config("discovery.proxy.proxied-types")
+    @ConfigDescription("Service types to proxy (test environments only)")
+    public DiscoveryConfig setProxyProxiedTypes(StringSet proxyProxiedTypes)
+    {
+        this.proxyProxiedTypes = proxyProxiedTypes;
+        return this;
+    }
+
+    public String getProxyEnvironment()
+    {
+        return proxyEnvironment;
+    }
+
+    @Config("discovery.proxy.environment")
+    @ConfigDescription("Environment to proxy to (test environments only)")
+    public DiscoveryConfig setProxyEnvironment(String proxyEnvironment)
+    {
+        this.proxyEnvironment = proxyEnvironment;
+        return this;
+    }
+
+    public URI getProxyUri()
+    {
+        return proxyUri;
+    }
+
+    @Config("discovery.proxy.uri")
+    @ConfigDescription("Discovery server to proxy to (test environments only)")
+    public DiscoveryConfig setProxyUri(URI proxyUri)
+    {
+        this.proxyUri = proxyUri;
+        return this;
+    }
+
+    @AssertTrue(message = "discovery.proxy.environment specified if and only if any proxy types")
+    public boolean isProxyTypeAndEnvironment()
+    {
+        return proxyProxiedTypes.isEmpty() == (proxyEnvironment == null);
+    }
+
+    @AssertTrue(message = "discovery.proxy.uri specified if and only if any proxy types")
+    public boolean isProxyTypeAndUri()
+    {
+        return proxyProxiedTypes.isEmpty() == (proxyUri == null);
+    }
+
+    public static final class StringSet extends ForwardingSet<String>
+    {
+        private final Set<String> delegate;
+
+        private StringSet(Set<String> delegate)
+        {
+            this.delegate = ImmutableSet.copyOf(delegate);
+        }
+
+        public static StringSet of(String... strings)
+        {
+            return new StringSet(ImmutableSet.copyOf(strings));
+        }
+
+        public static StringSet valueOf(String string)
+        {
+            return of(string.split("\\s*,\\s*"));
+        }
+
+        @Override
+        protected Set<String> delegate()
+        {
+            return delegate;
+        }
+    }
 }

server/src/main/java/io/airlift/discovery/DiscoveryServerModule.java

@@ -34,6 +34,7 @@
 
 import static io.airlift.configuration.ConfigurationModule.bindConfig;
 import static io.airlift.discovery.client.DiscoveryBinder.discoveryBinder;
+import static io.airlift.http.client.HttpClientBinder.httpClientBinder;
 
 public class DiscoveryServerModule
         implements Module
@@ -55,6 +56,10 @@ public void configure(Binder binder)
         binder.bind(StaticStore.class).to(ReplicatedStaticStore.class).in(Scopes.SINGLETON);
         binder.install(new ReplicatedStoreModule("static", ForStaticStore.class, PersistentStore.class));
         bindConfig(binder).prefixedWith("static").to(PersistentStoreConfig.class);
+
+        // proxy announcements
+        httpClientBinder(binder).bindAsyncHttpClient("discovery.proxy", ForProxyStore.class);
+        binder.bind(ProxyStore.class).in(Scopes.SINGLETON);
     }
 
     @Singleton

server/src/main/java/io/airlift/discovery/DynamicAnnouncementResource.java

@@ -15,6 +15,7 @@
  */
 package io.airlift.discovery;
 
+import com.google.common.base.Joiner;
 import com.google.common.base.Objects;
 import io.airlift.node.NodeInfo;
 
@@ -24,39 +25,57 @@
 import javax.ws.rs.PUT;
 import javax.ws.rs.Path;
 import javax.ws.rs.PathParam;
-import javax.ws.rs.core.Context;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
-import javax.ws.rs.core.UriInfo;
+import java.util.HashSet;
+import java.util.Set;
 
 import static java.lang.String.format;
 import static javax.ws.rs.core.Response.Status.ACCEPTED;
 import static javax.ws.rs.core.Response.Status.BAD_REQUEST;
+import static javax.ws.rs.core.Response.Status.FORBIDDEN;
 import static javax.ws.rs.core.Response.Status.NOT_FOUND;
 
 @Path("/v1/announcement/{node_id}")
 public class DynamicAnnouncementResource
 {
     private final NodeInfo nodeInfo;
     private final DynamicStore dynamicStore;
+    private Set<String> proxyTypes;
 
     @Inject
-    public DynamicAnnouncementResource(DynamicStore dynamicStore, NodeInfo nodeInfo)
+    public DynamicAnnouncementResource(DynamicStore dynamicStore, NodeInfo nodeInfo, DiscoveryConfig discoveryConfig)
     {
         this.dynamicStore = dynamicStore;
         this.nodeInfo = nodeInfo;
+        proxyTypes = discoveryConfig.getProxyProxiedTypes();
     }
 
     @PUT
     @Consumes(MediaType.APPLICATION_JSON)
-    public Response put(@PathParam("node_id") Id<Node> nodeId, @Context UriInfo uriInfo, DynamicAnnouncement announcement)
+    public Response put(@PathParam("node_id") Id<Node> nodeId, DynamicAnnouncement announcement)
     {
         if (!nodeInfo.getEnvironment().equals(announcement.getEnvironment())) {
             return Response.status(BAD_REQUEST)
                     .entity(format("Environment mismatch. Expected: %s, Provided: %s", nodeInfo.getEnvironment(), announcement.getEnvironment()))
                     .build();
         }
 
+        if (!proxyTypes.isEmpty()) {
+            Set<String> forbiddenTypes = new HashSet<>();
+            for (DynamicServiceAnnouncement serviceAnnouncement : announcement.getServiceAnnouncements()) {
+                String type = serviceAnnouncement.getType();
+                if (proxyTypes.contains(type)) {
+                    forbiddenTypes.add(type);
+                }
+            }
+            if (!forbiddenTypes.isEmpty()) {
+                return Response.status(FORBIDDEN)
+                        .entity(format("Cannot announce proxied type%s %s", forbiddenTypes.size() == 1 ? "" : "s", Joiner.on(',').join(forbiddenTypes)))
+                        .build();
+            }
+        }
+
         String location = Objects.firstNonNull(announcement.getLocation(), "/somewhere/" + nodeId.toString());
 
         DynamicAnnouncement announcementWithLocation = DynamicAnnouncement.copyOf(announcement)

server/src/main/java/io/airlift/discovery/ForProxyStore.java

@@ -0,0 +1,31 @@
+/*
+ * Copyright 2013 Proofpoint, Inc.
+ *
+ * Licensed 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 io.airlift.discovery;
+
+import com.google.inject.BindingAnnotation;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+
+import static java.lang.annotation.ElementType.PARAMETER;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+
+@Retention(RUNTIME)
+@Target(PARAMETER)
+@BindingAnnotation
+public @interface ForProxyStore
+{
+}

server/src/main/java/io/airlift/discovery/Main.java

@@ -58,7 +58,6 @@ public static void main(String[] args)
         }
         catch (Exception e) {
             log.error(e);
-            // Cassandra prevents the vm from shutting down on its own
             System.exit(1);
         }
     }