address cr

indexing-service/src/main/java/io/druid/indexing/overlord/ImmutableZkWorker.java

@@ -0,0 +1,78 @@
+/*
+ * Druid - a distributed column store.
+ * Copyright (C) 2012, 2013  Metamarkets Group Inc.
+ *
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+ */
+
+package io.druid.indexing.overlord;
+
+import com.google.common.collect.ImmutableSet;
+import io.druid.indexing.common.task.Task;
+import io.druid.indexing.worker.Worker;
+
+import java.util.Set;
+
+/**
+ * A snapshot of a {@link io.druid.indexing.overlord.ZkWorker}
+ */
+public class ImmutableZkWorker
+{
+  private final ZkWorker mutableZkWorker;
+  private final Worker worker;
+  private final int currCapacityUsed;
+  private final Set<String> availabilityGroups;
+
+  public ImmutableZkWorker(
+      ZkWorker mutableZkWorker
+  )
+  {
+    this.mutableZkWorker = mutableZkWorker;
+    this.worker = mutableZkWorker.getWorker();
+    this.currCapacityUsed = mutableZkWorker.getCurrCapacityUsed();
+    this.availabilityGroups = ImmutableSet.copyOf(mutableZkWorker.getAvailabilityGroups());
+  }
+
+  public ZkWorker getMutableZkWorker()
+  {
+    return mutableZkWorker;
+  }
+
+  public Worker getWorker()
+  {
+    return worker;
+  }
+
+  public int getCurrCapacityUsed()
+  {
+    return currCapacityUsed;
+  }
+
+  public Set<String> getAvailabilityGroups()
+  {
+    return availabilityGroups;
+  }
+
+  public boolean isValidVersion(String minVersion)
+  {
+    return worker.getVersion().compareTo(minVersion) >= 0;
+  }
+
+  public boolean canRunTask(Task task)
+  {
+    return (worker.getCapacity() - getCurrCapacityUsed() >= task.getTaskResource().getRequiredCapacity()
+            && !getAvailabilityGroups().contains(task.getTaskResource().getAvailabilityGroup()));
+  }
+}

indexing-service/src/main/java/io/druid/indexing/overlord/RemoteTaskRunner.java

@@ -27,8 +27,10 @@ import com.google.common.base.Preconditions;
 import com.google.common.base.Stopwatch;
 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
+import com.google.common.collect.Maps;
 import com.google.common.io.InputSupplier;
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
@@ -62,6 +64,7 @@ import org.apache.zookeeper.KeeperException;
 import org.jboss.netty.handler.codec.http.HttpResponseStatus;
 import org.joda.time.DateTime;
 
+import javax.annotation.Nullable;
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.MalformedURLException;
@@ -83,13 +86,13 @@ import java.util.concurrent.TimeUnit;
  * creating ephemeral nodes in ZK that workers must remove. Workers announce the statuses of the tasks they are running.
  * Once a task completes, it is up to the RTR to remove the task status and run any necessary cleanup.
  * The RemoteTaskRunner is event driven and updates state according to ephemeral node changes in ZK.
- * 
+ * <p/>
  * The RemoteTaskRunner will assign tasks to a node until the node hits capacity. At that point, task assignment will
  * fail. The RemoteTaskRunner depends on another component to create additional worker resources.
  * For example, {@link io.druid.indexing.overlord.scaling.ResourceManagementScheduler} can take care of these duties.
- * 
+ * <p/>
  * If a worker node becomes inexplicably disconnected from Zk, the RemoteTaskRunner will fail any tasks associated with the worker.
- * 
+ * <p/>
  * The RemoteTaskRunner uses ZK for job management and assignment and http for IPC messages.
  */
 public class RemoteTaskRunner implements TaskRunner, TaskLogStreamer
@@ -519,7 +522,24 @@ public class RemoteTaskRunner implements TaskRunner, TaskLogStreamer
       return true;
     } else {
       // Nothing running this task, announce it in ZK for a worker to run it
-      Optional<ZkWorker> zkWorker = strategy.findWorkerForTask(zkWorkers, task);
+      Optional<ZkWorker> zkWorker = strategy.findWorkerForTask(
+          ImmutableMap.copyOf(
+              Maps.transformEntries(
+                  zkWorkers,
+                  new Maps.EntryTransformer<String, ZkWorker, ImmutableZkWorker>()
+                  {
+                    @Override
+                    public ImmutableZkWorker transformEntry(
+                        String key, ZkWorker value
+                    )
+                    {
+                      return value.toImmutable();
+                    }
+                  }
+              )
+          ),
+          task
+      );
       if (zkWorker.isPresent()) {
         announceTask(task, zkWorker.get(), taskRunnerWorkItem);
         return true;

indexing-service/src/main/java/io/druid/indexing/overlord/ZkWorker.java

@@ -158,6 +158,11 @@ public class ZkWorker implements Closeable
     lastCompletedTaskTime.getAndSet(completedTaskTime);
   }
 
+  public ImmutableZkWorker toImmutable()
+  {
+    return new ImmutableZkWorker(this);
+  }
+
   @Override
   public void close() throws IOException
   {

indexing-service/src/main/java/io/druid/indexing/overlord/setup/FillCapacityWorkerSelectStrategy.java

@@ -20,16 +20,17 @@
 package io.druid.indexing.overlord.setup;
 
 import com.google.common.base.Optional;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Sets;
 import com.google.common.primitives.Ints;
 import com.google.inject.Inject;
 import com.metamx.emitter.EmittingLogger;
 import io.druid.indexing.common.task.Task;
+import io.druid.indexing.overlord.ImmutableZkWorker;
 import io.druid.indexing.overlord.ZkWorker;
 import io.druid.indexing.overlord.config.RemoteTaskRunnerConfig;
 
 import java.util.Comparator;
-import java.util.Map;
 import java.util.TreeSet;
 
 /**
@@ -46,14 +47,17 @@ public class FillCapacityWorkerSelectStrategy implements WorkerSelectStrategy
     this.config = config;
   }
 
-  public Optional<ZkWorker> findWorkerForTask(final Map<String, ZkWorker> zkWorkers, final Task task)
+  public Optional<ZkWorker> findWorkerForTask(
+      final ImmutableMap<String, ImmutableZkWorker> zkWorkers,
+      final Task task
+  )
   {
-    TreeSet<ZkWorker> sortedWorkers = Sets.newTreeSet(
-        new Comparator<ZkWorker>()
+    TreeSet<ImmutableZkWorker> sortedWorkers = Sets.newTreeSet(
+        new Comparator<ImmutableZkWorker>()
         {
           @Override
           public int compare(
-              ZkWorker zkWorker, ZkWorker zkWorker2
+              ImmutableZkWorker zkWorker, ImmutableZkWorker zkWorker2
           )
           {
             int retVal = Ints.compare(zkWorker2.getCurrCapacityUsed(), zkWorker.getCurrCapacityUsed());
@@ -68,9 +72,9 @@ public class FillCapacityWorkerSelectStrategy implements WorkerSelectStrategy
     sortedWorkers.addAll(zkWorkers.values());
     final String minWorkerVer = config.getMinWorkerVersion();
 
-    for (ZkWorker zkWorker : sortedWorkers) {
+    for (ImmutableZkWorker zkWorker : sortedWorkers) {
       if (zkWorker.canRunTask(task) && zkWorker.isValidVersion(minWorkerVer)) {
-        return Optional.of(zkWorker);
+        return Optional.of(zkWorker.getMutableZkWorker());
       }
     }
     log.debug("Worker nodes %s do not have capacity to run any more tasks!", zkWorkers.values());

indexing-service/src/main/java/io/druid/indexing/overlord/setup/WorkerSelectStrategy.java

@@ -22,12 +22,15 @@ package io.druid.indexing.overlord.setup;
 import com.fasterxml.jackson.annotation.JsonSubTypes;
 import com.fasterxml.jackson.annotation.JsonTypeInfo;
 import com.google.common.base.Optional;
+import com.google.common.collect.ImmutableMap;
 import io.druid.indexing.common.task.Task;
+import io.druid.indexing.overlord.ImmutableZkWorker;
 import io.druid.indexing.overlord.ZkWorker;
 
 import java.util.Map;
 
 /**
+ * The {@link io.druid.indexing.overlord.RemoteTaskRunner} uses this class to select a worker to assign tasks to.
  */
 @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "type")
 @JsonSubTypes(value = {
@@ -35,5 +38,16 @@ import java.util.Map;
 })
 public interface WorkerSelectStrategy
 {
-  public Optional<ZkWorker> findWorkerForTask(final Map<String, ZkWorker> zkWorkers, final Task task);
+  /**
+   * Customizable logic for selecting a worker to run a task.
+   *
+   * @param zkWorkers An immutable map of workers to choose from.
+   * @param task      The task to assign.
+   *
+   * @return A {@link io.druid.indexing.overlord.ZkWorker} to run the task if one is available.
+   */
+  public Optional<ZkWorker> findWorkerForTask(
+      final ImmutableMap<String, ImmutableZkWorker> zkWorkers,
+      final Task task
+  );
 }