Add Router connection balancers for Avatica queries

benchmarks/src/main/java/io/druid/benchmark/ConsistentHasherBenchmark.java

@@ -0,0 +1,95 @@
+/*
+ * Licensed to Metamarkets Group Inc. (Metamarkets) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. Metamarkets licenses this file
+ * to you 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.druid.benchmark;
+
+import com.google.common.collect.Sets;
+
+import io.druid.java.util.common.StringUtils;
+import io.druid.server.router.ConsistentHasher;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.TimeUnit;
+
+@State(Scope.Benchmark)
+@Fork(value = 1)
+@Warmup(iterations = 15)
+@Measurement(iterations = 30)
+public class ConsistentHasherBenchmark
+{
+  @Param({"100000"})
+  int numIds;
+
+  ConsistentHasher hasher;
+  List<String> uuids;
+  Set<String> servers;
+
+  @Setup
+  public void setup() throws IOException
+  {
+    hasher = new ConsistentHasher(null);
+    uuids = new ArrayList<>();
+    servers = Sets.newHashSet(
+        "localhost:1",
+        "localhost:2",
+        "localhost:3",
+        "localhost:4",
+        "localhost:5",
+        "localhost:6",
+        "localhost:7",
+        "localhost:8",
+        "localhost:9",
+        "localhost:10"
+    );
+
+    for (int i = 0; i < numIds; i++) {
+      UUID uuid = UUID.randomUUID();
+      uuids.add(uuid.toString());
+    }
+
+    hasher.updateKeys(servers);
+  }
+
+  @Benchmark
+  @BenchmarkMode(Mode.AverageTime)
+  @OutputTimeUnit(TimeUnit.MICROSECONDS)
+  public void hash(Blackhole blackhole) throws Exception
+  {
+    for (String uuid : uuids) {
+      String server = hasher.findKey(StringUtils.toUtf8(uuid));
+      blackhole.consume(server);
+    }
+  }
+}

benchmarks/src/main/java/io/druid/benchmark/RendezvousHasherBenchmark.java

@@ -0,0 +1,93 @@
+/*
+ * Licensed to Metamarkets Group Inc. (Metamarkets) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. Metamarkets licenses this file
+ * to you 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.druid.benchmark;
+
+import com.google.common.collect.Sets;
+import io.druid.java.util.common.StringUtils;
+import io.druid.server.router.RendezvousHasher;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.TimeUnit;
+
+@State(Scope.Benchmark)
+@Fork(value = 1)
+@Warmup(iterations = 15)
+@Measurement(iterations = 30)
+public class RendezvousHasherBenchmark
+{
+  @Param({"100000"})
+  int numIds;
+
+  RendezvousHasher hasher;
+  List<String> uuids;
+  Set<String> servers;
+
+  @Setup
+  public void setup() throws IOException
+  {
+    hasher = new RendezvousHasher();
+    uuids = new ArrayList<>();
+    servers = Sets.newHashSet(
+        "localhost:1",
+        "localhost:2",
+        "localhost:3",
+        "localhost:4",
+        "localhost:5",
+        "localhost:6",
+        "localhost:7",
+        "localhost:8",
+        "localhost:9",
+        "localhost:10"
+    );
+
+
+    for (int i = 0; i < numIds; i++) {
+      UUID uuid = UUID.randomUUID();
+      uuids.add(uuid.toString());
+    }
+  }
+
+  @Benchmark
+  @BenchmarkMode(Mode.AverageTime)
+  @OutputTimeUnit(TimeUnit.MICROSECONDS)
+  public void hash(Blackhole blackhole) throws Exception
+  {
+    for (String uuid : uuids) {
+      String server = hasher.chooseNode(servers, StringUtils.toUtf8(uuid));
+      blackhole.consume(server);
+    }
+  }
+}

docs/content/development/router.md

@@ -75,6 +75,7 @@ The router module uses several of the default modules in [Configuration](../conf
 |`druid.router.coordinatorServiceName`|Any string.|The service discovery name of the coordinator.|druid/coordinator|
 |`druid.router.pollPeriod`|Any ISO8601 duration.|How often to poll for new rules.|PT1M|
 |`druid.router.strategies`|An ordered JSON array of objects.|All custom strategies to use for routing.|[{"type":"timeBoundary"},{"type":"priority"}]|
+|`druid.router.avatica.balancer.type`|String representing an AvaticaConnectionBalancer name|Class to use for balancing Avatica queries across brokers|rendezvousHash|
 
 Router Strategies
 -----------------
@@ -119,6 +120,40 @@ Allows defining arbitrary routing rules using a JavaScript function. The functio
 JavaScript-based functionality is disabled by default. Please refer to the Druid <a href="../development/javascript.html">JavaScript programming guide</a> for guidelines about using Druid's JavaScript functionality, including instructions on how to enable it.
 </div>
 
+
+Avatica Query Balancing
+--------------
+
+All Avatica JDBC requests with a given connection ID must be routed to the same broker, since Druid brokers do not share connection state with each other.
+
+To accomplish this, Druid provides two built-in balancers that use rendezvous hashing and consistent hashing of a request's connection ID respectively to assign requests to brokers.
+
+Note that when multiple routers are used, all routers should have identical balancer configuration to ensure that they make the same routing decisions.
+
+### Rendezvous Hash Balancer
+
+This balancer uses [Rendezvous Hashing](https://en.wikipedia.org/wiki/Rendezvous_hashing) on an Avatica request's connection ID to assign the request to a broker.
+
+To use this balancer, specify the following property:
+
+```
+druid.router.avatica.balancer.type=rendezvousHash
+```
+
+If no `druid.router.avatica.balancer` property is set, the Router will also default to using the Rendezvous Hash Balancer.
+
+### Consistent Hash Balancer
+
+This balancer uses [Consistent Hashing](https://en.wikipedia.org/wiki/Consistent_hashing) on an Avatica request's connection ID to assign the request to a broker.
+
+To use this balancer, specify the following property:
+
+```
+druid.router.avatica.balancer.type=consistentHash
+```
+
+This is a non-default implementation that is provided for experimentation purposes. The consistent hasher has longer setup times on initialization and when the set of brokers changes, but has a faster broker assignment time than the rendezous hasher when tested with 5 brokers. Benchmarks for both implementations have been provided in `ConsistentHasherBenchmark` and `RendezvousHasherBenchmark`. The consistent hasher also requires locking, while the rendezvous hasher does not.
+
 HTTP Endpoints
 --------------
 

docs/content/querying/sql.md

@@ -388,10 +388,16 @@ Table metadata is available over JDBC using `connection.getMetaData()` or by que
 ["INFORMATION_SCHEMA" tables](#retrieving-metadata). Parameterized queries (using `?` or other placeholders) don't work properly,
 so avoid those.
 
+#### Connection Stickiness
+
 Druid's JDBC server does not share connection state between brokers. This means that if you're using JDBC and have
 multiple Druid brokers, you should either connect to a specific broker, or use a load balancer with sticky sessions
 enabled.
 
+The Druid Router node provides connection stickiness when balancing JDBC requests. Please see [Router](../development/router.html) documentation for more details.
+
+Note that the non-JDBC [JSON over HTTP](#json-over-http) API is stateless and does not require stickiness.
+
 ### Connection context
 
 Druid SQL supports setting connection parameters on the client. The parameters in the table below affect SQL planning.

server/src/main/java/io/druid/server/AsyncQueryForwardingServlet.java

@@ -33,6 +33,7 @@
 import io.druid.guice.annotations.Smile;
 import io.druid.guice.http.DruidHttpClientConfig;
 import io.druid.java.util.common.DateTimes;
+import io.druid.java.util.common.IAE;
 import io.druid.query.DruidMetrics;
 import io.druid.query.GenericQueryMetricsFactory;
 import io.druid.query.Query;
@@ -62,6 +63,7 @@
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.net.URLDecoder;
+import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
@@ -78,6 +80,7 @@ public class AsyncQueryForwardingServlet extends AsyncProxyServlet implements Qu
   private static final String HOST_ATTRIBUTE = "io.druid.proxy.to.host";
   private static final String SCHEME_ATTRIBUTE = "io.druid.proxy.to.host.scheme";
   private static final String QUERY_ATTRIBUTE = "io.druid.proxy.query";
+  private static final String AVATICA_QUERY_ATTRIBUTE = "io.druid.proxy.avaticaQuery";
   private static final String OBJECTMAPPER_ATTRIBUTE = "io.druid.proxy.objectMapper";
 
   private static final int CANCELLATION_TIMEOUT_MILLIS = 500;
@@ -186,7 +189,17 @@ protected void service(HttpServletRequest request, HttpServletResponse response)
     final boolean isQueryEndpoint = request.getRequestURI().startsWith("/druid/v2")
                                     && !request.getRequestURI().startsWith("/druid/v2/sql");
 
-    if (isQueryEndpoint && HttpMethod.DELETE.is(request.getMethod())) {
+    final boolean isAvatica = request.getRequestURI().startsWith("/druid/v2/sql/avatica");
+
+    if (isAvatica) {
+      Map<String, Object> requestMap = objectMapper.readValue(request.getInputStream(), Map.class);
+      String connectionId = getAvaticaConnectionId(requestMap);
+      Server targetServer = hostFinder.findServerAvatica(connectionId);
+      byte[] requestBytes = objectMapper.writeValueAsBytes(requestMap);
+      request.setAttribute(HOST_ATTRIBUTE, targetServer.getHost());
+      request.setAttribute(SCHEME_ATTRIBUTE, targetServer.getScheme());
+      request.setAttribute(AVATICA_QUERY_ATTRIBUTE, requestBytes);
+    } else if (isQueryEndpoint && HttpMethod.DELETE.is(request.getMethod())) {
       // query cancellation request
       for (final Server server: hostFinder.getAllServers()) {
         // send query cancellation to all brokers this query may have gone to
@@ -263,6 +276,11 @@ protected void sendProxyRequest(
     proxyRequest.timeout(httpClientConfig.getReadTimeout().getMillis(), TimeUnit.MILLISECONDS);
     proxyRequest.idleTimeout(httpClientConfig.getReadTimeout().getMillis(), TimeUnit.MILLISECONDS);
 
+    byte[] avaticaQuery = (byte[]) clientRequest.getAttribute(AVATICA_QUERY_ATTRIBUTE);
+    if (avaticaQuery != null) {
+      proxyRequest.content(new BytesContentProvider(avaticaQuery));
+    }
+
     final Query query = (Query) clientRequest.getAttribute(QUERY_ATTRIBUTE);
     if (query != null) {
       final ObjectMapper objectMapper = (ObjectMapper) clientRequest.getAttribute(OBJECTMAPPER_ATTRIBUTE);
@@ -371,6 +389,18 @@ public long getInterruptedQueryCount()
     return interruptedQueryCount.get();
   }
 
+  private static String getAvaticaConnectionId(Map<String, Object> requestMap) throws IOException
+  {
+    Object connectionIdObj = requestMap.get("connectionId");
+    if (connectionIdObj == null) {
+      throw new IAE("Received an Avatica request without a connectionId.");
+    }
+    if (!(connectionIdObj instanceof String)) {
+      throw new IAE("Received an Avatica request with a non-String connectionId.");
+    }
+
+    return (String) connectionIdObj;
+  }
 
   private class MetricsEmittingProxyResponseListener extends ProxyResponseListener
   {

server/src/main/java/io/druid/server/router/AvaticaConnectionBalancer.java

@@ -0,0 +1,45 @@
+/*
+ * Licensed to Metamarkets Group Inc. (Metamarkets) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. Metamarkets licenses this file
+ * to you 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.druid.server.router;
+
+import com.fasterxml.jackson.annotation.JsonSubTypes;
+import com.fasterxml.jackson.annotation.JsonTypeInfo;
+import io.druid.client.selector.Server;
+
+import java.util.Collection;
+
+/**
+ * An AvaticaConnectionBalancer balances Avatica connections across a collection of servers.
+ */
+@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "type", defaultImpl = RendezvousHashAvaticaConnectionBalancer.class)
+@JsonSubTypes(value = {
+    @JsonSubTypes.Type(name = "rendezvousHash", value = RendezvousHashAvaticaConnectionBalancer.class),
+    @JsonSubTypes.Type(name = "consistentHash", value = ConsistentHashAvaticaConnectionBalancer.class)
+})
+public interface AvaticaConnectionBalancer
+{
+  /**
+   * @param servers Servers to balance across
+   * @param connectionId Connection ID to be balanced
+   * @return Server that connectionId should be assigned to. The process for choosing a server must be deterministic and
+   *         sticky (with a fixed set of servers, the same connectionId should always be assigned to the same server)
+   */
+  Server pickServer(Collection<Server> servers, String connectionId);
+}