Add CachingCostBalancerStrategy

benchmarks/src/main/java/io/druid/server/coordinator/CachingCostBalancerStrategyBenchmark.java

@@ -0,0 +1,114 @@
+/*
+ * 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.coordinator;
+
+import io.druid.server.coordinator.cost.SegmentsCostCache;
+import io.druid.timeline.DataSegment;
+import org.joda.time.DateTime;
+import org.joda.time.Interval;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Random;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.TimeUnit;
+
+@State(Scope.Benchmark)
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MILLISECONDS)
+@Fork(1)
+public class CachingCostBalancerStrategyBenchmark
+{
+  private static final int NUMBER_OF_SEGMENTS = 100000;
+  private static final int NUMBER_OF_QUERIES = 500;
+
+  private static final long DAYS_IN_MONTH = 30;
+
+  private final DateTime referenceTime = new DateTime("2014-01-01T00:00:00");
+  private final List<DataSegment> segments = new ArrayList<>();
+  private final List<DataSegment> segmentQueries = new ArrayList<>();
+  private final int seed = ThreadLocalRandom.current().nextInt();
+
+  private SegmentsCostCache segmentsCostCache;
+
+  @Setup
+  public void createSegments()
+  {
+    Random random = new Random(seed);
+    SegmentsCostCache.Builder prototype = SegmentsCostCache.builder();
+    for (int i = 0; i < NUMBER_OF_SEGMENTS; ++i) {
+      DataSegment segment = createSegment(random.nextInt((int)TimeUnit.DAYS.toHours(DAYS_IN_MONTH)));
+      segments.add(segment);
+      prototype.addSegment(segment);
+    }
+    segmentsCostCache = prototype.build();
+    for (int i = 0; i < NUMBER_OF_QUERIES; ++i) {
+      DataSegment segment = createSegment(random.nextInt((int)TimeUnit.DAYS.toHours(DAYS_IN_MONTH)));
+      segmentQueries.add(segment);
+    }
+
+    System.out.println("GENERATING SEGMENTS : " + NUMBER_OF_SEGMENTS + " / " + NUMBER_OF_QUERIES);
+  }
+
+  @Benchmark
+  public double measureCostStrategy() throws InterruptedException
+  {
+    double cost = 0.0;
+    for (DataSegment segment : segmentQueries) {
+      cost += CostBalancerStrategy.computeJointSegmentsCost(segment, segments);
+    }
+    return cost;
+  }
+
+  @Benchmark
+  public double measureCachingCostStrategy() throws InterruptedException
+  {
+    double cost = 0.0;
+    for (DataSegment segment : segmentQueries) {
+      cost += segmentsCostCache.cost(segment);
+    }
+    return cost;
+  }
+
+  private DataSegment createSegment(int shift)
+  {
+    return new DataSegment(
+        "dataSource",
+        new Interval(referenceTime.plusHours(shift), referenceTime.plusHours(shift).plusHours(1)),
+        "version",
+        Collections.<String, Object>emptyMap(),
+        Collections.<String>emptyList(),
+        Collections.<String>emptyList(),
+        null,
+        0,
+        100
+    );
+  }
+
+}

indexing-service/src/test/java/io/druid/indexing/test/TestServerView.java

@@ -48,7 +48,7 @@ public void registerSegmentCallback(
   }
 
   @Override
-  public void registerServerCallback(Executor exec, ServerView.ServerCallback callback)
+  public void registerServerRemovedCallback(Executor exec, ServerView.ServerRemovedCallback callback)
   {
     // No-op
   }

server/src/main/java/io/druid/client/AbstractCuratorServerInventoryView.java

@@ -52,7 +52,7 @@ public abstract class AbstractCuratorServerInventoryView<InventoryType> implemen
   private final CuratorInventoryManager<DruidServer, InventoryType> inventoryManager;
   private final AtomicBoolean started = new AtomicBoolean(false);
 
-  private final ConcurrentMap<ServerCallback, Executor> serverCallbacks = new MapMaker().makeMap();
+  private final ConcurrentMap<ServerRemovedCallback, Executor> serverCallbacks = new MapMaker().makeMap();
   private final ConcurrentMap<SegmentCallback, Executor> segmentCallbacks = new MapMaker().makeMap();
 
   public AbstractCuratorServerInventoryView(
@@ -210,7 +210,7 @@ public Iterable<DruidServer> getInventory()
   }
 
   @Override
-  public void registerServerCallback(Executor exec, ServerCallback callback)
+  public void registerServerRemovedCallback(Executor exec, ServerRemovedCallback callback)
   {
     serverCallbacks.put(callback, exec);
   }
@@ -249,7 +249,7 @@ public void run()
 
   protected void runServerCallbacks(final DruidServer server)
   {
-    for (final Map.Entry<ServerCallback, Executor> entry : serverCallbacks.entrySet()) {
+    for (final Map.Entry<ServerRemovedCallback, Executor> entry : serverCallbacks.entrySet()) {
       entry.getValue().execute(
           new Runnable()
           {

server/src/main/java/io/druid/client/BrokerServerView.java

@@ -145,9 +145,9 @@ public CallbackAction segmentViewInitialized()
         segmentFilter
     );
 
-    baseView.registerServerCallback(
+    baseView.registerServerRemovedCallback(
         exec,
-        new ServerView.ServerCallback()
+        new ServerRemovedCallback()
         {
           @Override
           public ServerView.CallbackAction serverRemoved(DruidServer server)
@@ -307,9 +307,9 @@ public <T> QueryRunner<T> getQueryRunner(DruidServer server)
   }
 
   @Override
-  public void registerServerCallback(Executor exec, ServerCallback callback)
+  public void registerServerRemovedCallback(Executor exec, ServerRemovedCallback callback)
   {
-    baseView.registerServerCallback(exec, callback);
+    baseView.registerServerRemovedCallback(exec, callback);
   }
 
   @Override

server/src/main/java/io/druid/client/CoordinatorServerView.java

@@ -88,9 +88,9 @@ public ServerView.CallbackAction segmentViewInitialized()
         }
     );
 
-    baseView.registerServerCallback(
+    baseView.registerServerRemovedCallback(
         exec,
-        new ServerView.ServerCallback()
+        new ServerView.ServerRemovedCallback()
         {
           @Override
           public ServerView.CallbackAction serverRemoved(DruidServer server)

server/src/main/java/io/druid/client/FilteredServerInventoryView.java

@@ -33,5 +33,5 @@ public void registerSegmentCallback(
       Executor exec, ServerView.SegmentCallback callback, Predicate<Pair<DruidServerMetadata, DataSegment>> filter
   );
 
-  public void registerServerCallback(Executor exec, ServerView.ServerCallback callback);
+  public void registerServerRemovedCallback(Executor exec, ServerView.ServerRemovedCallback callback);
 }

server/src/main/java/io/druid/client/HttpServerInventoryView.java

@@ -86,7 +86,7 @@ public class HttpServerInventoryView implements ServerInventoryView, FilteredSer
 
   private final LifecycleLock lifecycleLock = new LifecycleLock();
 
-  private final ConcurrentMap<ServerCallback, Executor> serverCallbacks = new MapMaker().makeMap();
+  private final ConcurrentMap<ServerRemovedCallback, Executor> serverCallbacks = new MapMaker().makeMap();
   private final ConcurrentMap<SegmentCallback, Executor> segmentCallbacks = new MapMaker().makeMap();
 
   private final ConcurrentMap<SegmentCallback, Predicate<Pair<DruidServerMetadata, DataSegment>>> segmentPredicates = new MapMaker()
@@ -258,7 +258,7 @@ public void registerSegmentCallback(
   }
 
   @Override
-  public void registerServerCallback(Executor exec, ServerCallback callback)
+  public void registerServerRemovedCallback(Executor exec, ServerRemovedCallback callback)
   {
     serverCallbacks.put(callback, exec);
   }
@@ -327,7 +327,7 @@ public void run()
 
   private void runServerCallbacks(final DruidServer server)
   {
-    for (final Map.Entry<ServerCallback, Executor> entry : serverCallbacks.entrySet()) {
+    for (final Map.Entry<ServerRemovedCallback, Executor> entry : serverCallbacks.entrySet()) {
       entry.getValue().execute(
           new Runnable()
           {

server/src/main/java/io/druid/client/ServerView.java

@@ -28,7 +28,7 @@
  */
 public interface ServerView
 {
-  public void registerServerCallback(Executor exec, ServerCallback callback);
+  public void registerServerRemovedCallback(Executor exec, ServerRemovedCallback callback);
   public void registerSegmentCallback(Executor exec, SegmentCallback callback);
 
   public enum CallbackAction
@@ -37,7 +37,7 @@ public enum CallbackAction
     UNREGISTER,
   }
 
-  public static interface ServerCallback
+  public static interface ServerRemovedCallback
   {
     /**
      * Called when a server is removed.

server/src/main/java/io/druid/server/coordinator/CachingCostBalancerStrategy.java

@@ -0,0 +1,75 @@
+/*
+ * 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.coordinator;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.util.concurrent.ListeningExecutorService;
+import io.druid.server.coordinator.cost.ClusterCostCache;
+import io.druid.timeline.DataSegment;
+
+import java.util.Set;
+
+
+public class CachingCostBalancerStrategy extends CostBalancerStrategy
+{
+
+  private final ClusterCostCache clusterCostCache;
+
+  public CachingCostBalancerStrategy(ClusterCostCache clusterCostCache, ListeningExecutorService exec)
+  {
+    super(exec);
+    this.clusterCostCache = Preconditions.checkNotNull(clusterCostCache);
+  }
+
+  @Override
+  protected double computeCost(
+      DataSegment proposalSegment, ServerHolder server, boolean includeCurrentServer
+  )
+  {
+    final long proposalSegmentSize = proposalSegment.getSize();
+
+    // (optional) Don't include server if it is already serving segment
+    if (!includeCurrentServer && server.isServingSegment(proposalSegment)) {
+      return Double.POSITIVE_INFINITY;
+    }
+
+    // Don't calculate cost if the server doesn't have enough space or is loading the segment
+    if (proposalSegmentSize > server.getAvailableSize() || server.isLoadingSegment(proposalSegment)) {
+      return Double.POSITIVE_INFINITY;
+    }
+
+    final String serverName = server.getServer().getName();
+
+    double cost = clusterCostCache.computeCost(serverName, proposalSegment);
+
+    // add segments that will be loaded to the cost
+    cost += costCacheForLoadingSegments(server).computeCost(serverName, proposalSegment);
+
+    return cost;
+  }
+
+  private ClusterCostCache costCacheForLoadingSegments(ServerHolder server)
+  {
+    final Set<DataSegment> loadingSegments = server.getPeon().getSegmentsToLoad();
+    return ClusterCostCache.builder(ImmutableMap.of(server.getServer().getName(), loadingSegments)).build();
+  }
+
+}