Merge pull request #672 from metamx/router-strat

Add more config to router strategies; fix bad test

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

@@ -19,6 +19,8 @@
 
 package io.druid.server.router;
 
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Optional;
 import com.google.common.collect.Iterables;
 import io.druid.query.Query;
@@ -27,12 +29,32 @@ import io.druid.query.Query;
  */
 public class PriorityTieredBrokerSelectorStrategy implements TieredBrokerSelectorStrategy
 {
+  private final int minPriority;
+  private final int maxPriority;
+
+  @JsonCreator
+  public PriorityTieredBrokerSelectorStrategy(
+      @JsonProperty("minPriority") Integer minPriority,
+      @JsonProperty("maxPriority") Integer maxPriority
+  )
+  {
+    this.minPriority = minPriority == null ? 0 : minPriority;
+    this.maxPriority = maxPriority == null ? 1 : maxPriority;
+  }
+
   @Override
   public Optional<String> getBrokerServiceName(TieredBrokerConfig tierConfig, Query query)
   {
     final int priority = query.getContextPriority(0);
 
-    if (priority < 0) {
+    if (priority < minPriority) {
+      return Optional.of(
+          Iterables.getLast(
+              tierConfig.getTierToBrokerMap().values(),
+              tierConfig.getDefaultBrokerServiceName()
+          )
+      );
+    } else if (priority >= maxPriority) {
       return Optional.of(
           Iterables.getFirst(
               tierConfig.getTierToBrokerMap().values(),

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

@@ -20,12 +20,15 @@
 package io.druid.server.router;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import io.druid.client.DruidServer;
 import org.joda.time.Period;
 
 import javax.validation.constraints.NotNull;
+import java.util.Arrays;
 import java.util.LinkedHashMap;
+import java.util.List;
 
 /**
  */
@@ -56,7 +59,10 @@ public class TieredBrokerConfig
 
   @JsonProperty
   @NotNull
-  private String strategies = "[{\"type\":\"timeBoundary\"},{\"type\":\"priority\"}]";
+  private List<TieredBrokerSelectorStrategy> strategies = Arrays.asList(
+      new TimeBoundaryTieredBrokerSelectorStrategy(),
+      new PriorityTieredBrokerSelectorStrategy(0, 1)
+  );
 
   // tier, <bard, numThreads>
   public LinkedHashMap<String, String> getTierToBrokerMap()
@@ -93,8 +99,8 @@ public class TieredBrokerConfig
     return pollPeriod;
   }
 
-  public String getStrategies()
+  public List<TieredBrokerSelectorStrategy> getStrategies()
   {
-    return strategies;
+    return ImmutableList.copyOf(strategies);
   }
 }

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

@@ -19,10 +19,6 @@
 
 package io.druid.server.router;
 
-import com.fasterxml.jackson.core.type.TypeReference;
-import com.fasterxml.jackson.databind.ObjectMapper;
-import com.google.api.client.util.Lists;
-import com.google.common.base.Throwables;
 import com.google.inject.Inject;
 import com.google.inject.Provider;
 
@@ -32,23 +28,12 @@ import java.util.List;
  */
 public class TieredBrokerSelectorStrategiesProvider implements Provider<List<TieredBrokerSelectorStrategy>>
 {
-  private final List<TieredBrokerSelectorStrategy> strategies = Lists.newArrayList();
+  private final List<TieredBrokerSelectorStrategy> strategies;
 
   @Inject
-  public TieredBrokerSelectorStrategiesProvider(ObjectMapper jsonMapper, TieredBrokerConfig config)
+  public TieredBrokerSelectorStrategiesProvider(TieredBrokerConfig config)
   {
-    try {
-      this.strategies.addAll(
-          (List<TieredBrokerSelectorStrategy>) jsonMapper.readValue(
-              config.getStrategies(), new TypeReference<List<TieredBrokerSelectorStrategy>>()
-          {
-          }
-          )
-      );
-    }
-    catch (Exception e) {
-      throw Throwables.propagate(e);
-    }
+    this.strategies = config.getStrategies();
   }
 
   @Override

server/src/test/java/io/druid/server/router/TieredBrokerHostSelectorTest.java

@@ -22,7 +22,6 @@ package io.druid.server.router;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.base.Supplier;
 import com.google.common.collect.ImmutableMap;
-import com.metamx.common.Pair;
 import com.metamx.http.client.HttpClient;
 import io.druid.client.DruidServer;
 import io.druid.curator.discovery.ServerDiscoveryFactory;
@@ -30,11 +29,9 @@ import io.druid.curator.discovery.ServerDiscoverySelector;
 import io.druid.guice.annotations.Global;
 import io.druid.guice.annotations.Json;
 import io.druid.query.Druids;
-import io.druid.query.TableDataSource;
 import io.druid.query.aggregation.AggregatorFactory;
 import io.druid.query.aggregation.CountAggregatorFactory;
 import io.druid.query.spec.MultipleIntervalSegmentSpec;
-import io.druid.query.timeboundary.TimeBoundaryQuery;
 import io.druid.server.coordinator.rules.IntervalLoadRule;
 import io.druid.server.coordinator.rules.Rule;
 import junit.framework.Assert;
@@ -85,7 +82,7 @@ public class TieredBrokerHostSelectorTest
           }
         },
         factory,
-        Arrays.asList(new TimeBoundaryTieredBrokerSelectorStrategy(), new PriorityTieredBrokerSelectorStrategy())
+        Arrays.asList(new TimeBoundaryTieredBrokerSelectorStrategy(), new PriorityTieredBrokerSelectorStrategy(0, 1))
     );
     EasyMock.expect(factory.createSelector(EasyMock.<String>anyObject())).andReturn(selector).atLeastOnce();
     EasyMock.replay(factory);
@@ -198,28 +195,50 @@ public class TieredBrokerHostSelectorTest
   }
 
   @Test
-    public void testPrioritySelect() throws Exception
-    {
-      String brokerName = (String) brokerSelector.select(
-          Druids.newTimeseriesQueryBuilder()
-                .dataSource("test")
-                .aggregators(Arrays.<AggregatorFactory>asList(new CountAggregatorFactory("count")))
-                .intervals(
-                    new MultipleIntervalSegmentSpec(
-                        Arrays.<Interval>asList(
-                            new Interval("2011-08-31/2011-09-01"),
-                            new Interval("2012-08-31/2012-09-01"),
-                            new Interval("2013-08-31/2013-09-01")
-                        )
-                    )
-                )
-                .context(ImmutableMap.<String, Object>of("priority", -1))
-                .build()
-      ).lhs;
+  public void testPrioritySelect() throws Exception
+  {
+    String brokerName = (String) brokerSelector.select(
+        Druids.newTimeseriesQueryBuilder()
+              .dataSource("test")
+              .aggregators(Arrays.<AggregatorFactory>asList(new CountAggregatorFactory("count")))
+              .intervals(
+                  new MultipleIntervalSegmentSpec(
+                      Arrays.<Interval>asList(
+                          new Interval("2011-08-31/2011-09-01"),
+                          new Interval("2012-08-31/2012-09-01"),
+                          new Interval("2013-08-31/2013-09-01")
+                      )
+                  )
+              )
+              .context(ImmutableMap.<String, Object>of("priority", -1))
+              .build()
+    ).lhs;
 
-      Assert.assertEquals("hotBroker", brokerName);
-    }
+    Assert.assertEquals("coldBroker", brokerName);
+  }
 
+  @Test
+  public void testPrioritySelect2() throws Exception
+  {
+    String brokerName = (String) brokerSelector.select(
+        Druids.newTimeseriesQueryBuilder()
+              .dataSource("test")
+              .aggregators(Arrays.<AggregatorFactory>asList(new CountAggregatorFactory("count")))
+              .intervals(
+                  new MultipleIntervalSegmentSpec(
+                      Arrays.<Interval>asList(
+                          new Interval("2011-08-31/2011-09-01"),
+                          new Interval("2012-08-31/2012-09-01"),
+                          new Interval("2013-08-31/2013-09-01")
+                      )
+                  )
+              )
+              .context(ImmutableMap.<String, Object>of("priority", 5))
+              .build()
+    ).lhs;
+
+    Assert.assertEquals("hotBroker", brokerName);
+  }
 
   private static class TestRuleManager extends CoordinatorRuleManager
   {