Improve performance for Probabilistic priority queueing

src/main/java/io/github/aurora/scheduler/scheduling/ProbabilisticPriorityAssigner.java

@@ -16,6 +16,7 @@
 import java.lang.annotation.Retention;
 import java.lang.annotation.Target;
 import java.util.HashSet;
+import java.util.LinkedList;
 import java.util.Map;
 import java.util.Random;
 import java.util.Set;
@@ -54,7 +55,8 @@ public class ProbabilisticPriorityAssigner extends TaskAssignerImpl {
   private static final Logger LOG = LoggerFactory.
       getLogger(ProbabilisticPriorityAssigner.class);
 
-  private final Storage storage;
+  private static Iterable<IScheduledTask> pendindTasks = new LinkedList<>();
+
   private Double exponent;
 
   @VisibleForTesting
@@ -69,10 +71,8 @@ public ProbabilisticPriorityAssigner(
       OfferManager offerManager,
       UpdateAgentReserver updateAgentReserver,
       StatsProvider statsProvider,
-      Storage storage,
       @Exponent Double exponent) {
     super(stateManager, taskFactory, offerManager, updateAgentReserver, statsProvider);
-    this.storage = requireNonNull(storage);
     this.exponent = requireNonNull(exponent);
   }
 
@@ -87,12 +87,13 @@ public Set<String> maybeAssign(
 
     // probabilistic priority queueing: may not schedule these tasks if
     // there are pending tasks with higher priority.
-    Iterable<IScheduledTask> pendindTasks = getPendingTasks();
     Set<Integer> prioritySet = new HashSet<>();
-    for (IScheduledTask t: pendindTasks) {
-      prioritySet.add(t.getAssignedTask().getTask().getPriority());
+    synchronized (pendindTasks) {
+      for (IScheduledTask t: pendindTasks) {
+        prioritySet.add(t.getAssignedTask().getTask().getPriority());
+      }
     }
-    //TODO(lenhattan86): Is the group is always included in the pending task set?
+    //this group is not always included in the pending task set
     prioritySet.add(groupKey.getTask().getPriority());
 
     if (!isScheduled(prioritySet, groupKey.getTask().getPriority())) {
@@ -131,8 +132,9 @@ boolean isScheduled(Set<Integer> prioritySet, int priority) {
   }
 
   @VisibleForTesting
-  Iterable<IScheduledTask> getPendingTasks() {
-    return Storage.Util.fetchTasks(storage, Query.unscoped().byStatus(ScheduleStatus.PENDING));
+  public static synchronized void fetchPendingTasks(Storage storage) {
+    pendindTasks =
+        Storage.Util.fetchTasks(storage, Query.unscoped().byStatus(ScheduleStatus.PENDING));
   }
 
   @VisibleForTesting

src/main/java/io/github/aurora/scheduler/scheduling/ProbabilisticPriorityAssignerModule.java

@@ -13,26 +13,53 @@
  */
 package io.github.aurora.scheduler.scheduling;
 
+import java.lang.annotation.Retention;
+import java.lang.annotation.Target;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+
+import javax.inject.Inject;
+import javax.inject.Qualifier;
 import javax.inject.Singleton;
 
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.util.concurrent.AbstractIdleService;
 import com.google.inject.AbstractModule;
 
+import org.apache.aurora.common.quantity.Time;
+import org.apache.aurora.scheduler.SchedulerServicesModule;
+import org.apache.aurora.scheduler.base.AsyncUtil;
 import org.apache.aurora.scheduler.config.CliOptions;
 import org.apache.aurora.scheduler.config.CommandLine;
+import org.apache.aurora.scheduler.config.types.TimeAmount;
 import org.apache.aurora.scheduler.scheduling.TaskAssigner;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import static java.lang.annotation.ElementType.FIELD;
+import static java.lang.annotation.ElementType.METHOD;
+import static java.lang.annotation.ElementType.PARAMETER;
+import static java.lang.annotation.RetentionPolicy.RUNTIME;
+import static java.util.Objects.requireNonNull;
 
 /**
  * The default TaskAssigner implementation.
  */
 public class ProbabilisticPriorityAssignerModule extends AbstractModule {
+  private static final Logger LOG =
+      LoggerFactory.getLogger(ProbabilisticPriorityAssignerModule.class);
+
   private final Options options;
 
   @Parameters(separators = "=")
   public static class Options {
     @Parameter(names = "-probabilistic_priority_assigner_exponent")
     Double probabilisticPriorityAssignerExponent = 0.0;
+
+    @Parameter(names = "-probabilistic_priority_assigner_task_fetch_interval")
+    TimeAmount probabilisticPriorityAssignerTaskFetchInterval = new TimeAmount(1, Time.SECONDS);
   }
 
   public ProbabilisticPriorityAssignerModule(CliOptions mOptions) {
@@ -50,5 +77,53 @@ protected void configure() {
         .annotatedWith(ProbabilisticPriorityAssigner.Exponent.class)
         .toInstance(options.probabilisticPriorityAssignerExponent);
     bind(TaskAssigner.class).to(ProbabilisticPriorityAssigner.class).in(Singleton.class);
+    bind(TaskFetcher.class).in(com.google.inject.Singleton.class);
+    bind(ScheduledExecutorService.class)
+        .annotatedWith(Executor.class)
+        .toInstance(AsyncUtil.singleThreadLoggingScheduledExecutor("HttpOfferSet-%d", LOG));
+    bind(Integer.class)
+        .annotatedWith(TaskFetchInvervalMs.class)
+        .toInstance(options.probabilisticPriorityAssignerTaskFetchInterval.
+            as(Time.MILLISECONDS).intValue());
+    bind(StatUpdater.class).in(com.google.inject.Singleton.class);
+    SchedulerServicesModule.addSchedulerActiveServiceBinding(binder()).to(StatUpdater.class);
+  }
+
+  // to bind an executor object to StatUpdater using @Executor
+  @VisibleForTesting
+  @Qualifier
+  @Target({ FIELD, PARAMETER, METHOD }) @Retention(RUNTIME)
+  @interface Executor { }
+
+  // to bind probabilisticPriorityAssignerTaskFetchInterval value to StatUpdater
+  @VisibleForTesting
+  @Qualifier
+  @Target({ FIELD, PARAMETER, METHOD }) @Retention(RUNTIME)
+  @interface TaskFetchInvervalMs { }
+
+  static class StatUpdater extends AbstractIdleService {
+    private final ScheduledExecutorService executor;
+    private final TaskFetcher taskFetcher;
+    private final Integer taskFetchIntervalMs;
+
+    @Inject
+    StatUpdater(
+        @Executor ScheduledExecutorService mExecutor,
+        TaskFetcher mTaskFetcher,
+        @TaskFetchInvervalMs Integer mTaskFetchIntervalMs) {
+      executor = requireNonNull(mExecutor);
+      taskFetcher = requireNonNull(mTaskFetcher);
+      taskFetchIntervalMs = mTaskFetchIntervalMs;
+    }
+
+    @Override
+    protected void startUp() {
+      executor.scheduleAtFixedRate(taskFetcher, 0, taskFetchIntervalMs, TimeUnit.MILLISECONDS);
+    }
+
+    @Override
+    protected void shutDown() {
+      // Ignored.
+    }
   }
 }

src/main/java/io/github/aurora/scheduler/scheduling/TaskFetcher.java

@@ -0,0 +1,32 @@
+/**
+ * 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.github.aurora.scheduler.scheduling;
+
+import javax.inject.Inject;
+
+import org.apache.aurora.scheduler.storage.Storage;
+
+public class TaskFetcher implements Runnable {
+  private final Storage storage;
+
+  @Inject
+  TaskFetcher(Storage mStorage) {
+    storage = mStorage;
+  }
+
+  @Override
+  public void run() {
+    ProbabilisticPriorityAssigner.fetchPendingTasks(storage);
+  }
+}

src/test/java/io/github/aurora/scheduler/scheduling/ProbabilisticPriorityAssignerTest.java

@@ -13,47 +13,33 @@
  */
 package io.github.aurora.scheduler.scheduling;
 
-import java.util.Collection;
 import java.util.HashSet;
 import java.util.Set;
 
-import com.google.common.collect.ImmutableList;
-
 import org.apache.aurora.common.testing.easymock.EasyMockTest;
-import org.apache.aurora.gen.ScheduleStatus;
-import org.apache.aurora.gen.ScheduledTask;
 import org.apache.aurora.scheduler.mesos.MesosTaskFactory;
 import org.apache.aurora.scheduler.offers.OfferManager;
 import org.apache.aurora.scheduler.state.StateManager;
-import org.apache.aurora.scheduler.storage.Storage;
-import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
-import org.apache.aurora.scheduler.storage.mem.MemStorageModule;
 import org.apache.aurora.scheduler.testing.FakeStatsProvider;
 import org.apache.aurora.scheduler.updater.UpdateAgentReserver;
 import org.junit.Before;
 import org.junit.Test;
 
-import static org.apache.aurora.scheduler.base.TaskTestUtil.JOB;
-import static org.apache.aurora.scheduler.base.TaskTestUtil.makeTask;
-import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 
 public class ProbabilisticPriorityAssignerTest extends EasyMockTest {
 
   private ProbabilisticPriorityAssigner assigner;
-  private Storage storage;
 
   @Before
   public void setUp() {
-    storage = MemStorageModule.newEmptyStorage();
     assigner = new ProbabilisticPriorityAssigner(
         createMock(StateManager.class),
         createMock(MesosTaskFactory.class),
         createMock(OfferManager.class),
         createMock(UpdateAgentReserver.class),
         new FakeStatsProvider(),
-        storage,
         0.0);
   }
 
@@ -117,26 +103,4 @@ public void testIsScheduled() {
     }
     assertTrue(res);
   }
-
-  @Test
-  public void testGetPendingTasks() {
-    control.replay();
-
-    storage.write((Storage.MutateWork.NoResult.Quiet) sp -> {
-      ScheduledTask t0 = makeTask("t0", JOB)
-          .newBuilder()
-          .setStatus(ScheduleStatus.PENDING);
-      ScheduledTask t1 = makeTask("t1", JOB)
-          .newBuilder()
-          .setStatus(ScheduleStatus.PENDING);
-      ScheduledTask t2 = makeTask("t2", JOB)
-          .newBuilder()
-          .setStatus(ScheduleStatus.RUNNING);
-      sp.getUnsafeTaskStore().saveTasks(
-          IScheduledTask.setFromBuilders(ImmutableList.of(t0, t1, t2)));
-    });
-
-    Iterable<IScheduledTask> tasks = assigner.getPendingTasks();
-    assertEquals(((Collection<?>) tasks).size(), 2);
-  }
 }