aurora-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From ma...@apache.org
Subject [1/3] aurora git commit: Batching writes - Part 1 (of 3): Introducing BatchWorker and task event batching.
Date Fri, 16 Sep 2016 21:45:17 GMT
Repository: aurora
Updated Branches:
  refs/heads/master f1e09a9c7 -> 496397aa5


Batching writes - Part 1 (of 3): Introducing BatchWorker and task event batching.

Reviewed at https://reviews.apache.org/r/51759/


Project: http://git-wip-us.apache.org/repos/asf/aurora/repo
Commit: http://git-wip-us.apache.org/repos/asf/aurora/commit/ebfeb3e6
Tree: http://git-wip-us.apache.org/repos/asf/aurora/tree/ebfeb3e6
Diff: http://git-wip-us.apache.org/repos/asf/aurora/diff/ebfeb3e6

Branch: refs/heads/master
Commit: ebfeb3e602faa9281ff7ff50f42bd21885518953
Parents: f1e09a9
Author: Maxim Khutornenko <maxim@apache.org>
Authored: Fri Sep 16 14:17:04 2016 -0700
Committer: Maxim Khutornenko <maxim@apache.org>
Committed: Fri Sep 16 14:17:04 2016 -0700

----------------------------------------------------------------------
 .../apache/aurora/scheduler/BatchWorker.java    | 254 +++++++++++++++++++
 .../aurora/scheduler/SchedulerModule.java       |  25 ++
 .../scheduler/pruning/TaskHistoryPruner.java    |  14 +-
 .../scheduler/scheduling/TaskThrottler.java     |  28 +-
 .../scheduler/state/MaintenanceController.java  |  14 +-
 .../updater/JobUpdateControllerImpl.java        |  10 +-
 .../aurora/scheduler/BatchWorkerTest.java       |  96 +++++++
 .../pruning/TaskHistoryPrunerTest.java          |  10 +-
 .../scheduler/scheduling/TaskThrottlerTest.java |   9 +-
 .../state/MaintenanceControllerImplTest.java    |   5 +
 .../scheduler/testing/BatchWorkerUtil.java      |  59 +++++
 .../aurora/scheduler/updater/JobUpdaterIT.java  |   7 +-
 12 files changed, 505 insertions(+), 26 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/aurora/blob/ebfeb3e6/src/main/java/org/apache/aurora/scheduler/BatchWorker.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/BatchWorker.java b/src/main/java/org/apache/aurora/scheduler/BatchWorker.java
new file mode 100644
index 0000000..e05d4b4
--- /dev/null
+++ b/src/main/java/org/apache/aurora/scheduler/BatchWorker.java
@@ -0,0 +1,254 @@
+/**
+ * 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 org.apache.aurora.scheduler;
+
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+
+import javax.inject.Inject;
+
+import com.google.common.util.concurrent.AbstractExecutionThreadService;
+
+import org.apache.aurora.common.stats.SlidingStats;
+import org.apache.aurora.common.stats.StatsProvider;
+import org.apache.aurora.common.util.BackoffStrategy;
+import org.apache.aurora.scheduler.base.AsyncUtil;
+import org.apache.aurora.scheduler.storage.Storage;
+import org.apache.aurora.scheduler.storage.Storage.MutableStoreProvider;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import static java.util.Objects.requireNonNull;
+
+import static com.google.common.base.Preconditions.checkState;
+
+/**
+ * Generic helper that allows bundling multiple work items into a single {@link Storage}
+ * transaction aiming to reduce the write lock contention.
+ *
+ * @param <T> Expected result type.
+ */
+public class BatchWorker<T> extends AbstractExecutionThreadService {
+  /**
+   * Empty result placeholder.
+   */
+  public interface NoResult { }
+
+  /**
+   * Convenience wrapper for a non-repeatable no value work {@link Result}.
+   */
+  public static final NoResult NO_RESULT = new NoResult() { };
+
+  private static final Logger LOG = LoggerFactory.getLogger(BatchWorker.class);
+  private final Storage storage;
+  private final int maxBatchSize;
+  private final SlidingStats batchUnlocked;
+  private final SlidingStats batchLocked;
+  private final BlockingQueue<WorkItem<T>> workQueue = new LinkedBlockingQueue<>();
+  private final ScheduledExecutorService scheduledExecutor;
+  private final AtomicInteger lastBatchSize = new AtomicInteger(0);
+  private final AtomicLong itemsProcessed;
+  private final AtomicLong batchesProcessed;
+
+  /**
+   * Wraps result returned by the {@link RepeatableWork} item.
+   *
+   * @param <T> Expected result type.
+   */
+  public static class Result<T> {
+    private final boolean isCompleted;
+    private final T value;
+
+    /**
+     * Initializes a {@link Result} instance with {@code isCompleted} and {@code value}.
+     * <p>
+     * The {@code isCompleted} may be set to {@code False} for a {@link RepeatableWork} that
has
+     * not finished yet. Otherwise, it must be set to {@code True}.
+     *
+     * @param isCompleted Flag indicating if the {@link RepeatableWork} has completed.
+     * @param value result value.
+     */
+    public Result(boolean isCompleted, T value) {
+      this.isCompleted = isCompleted;
+      this.value = value;
+    }
+  }
+
+  /**
+   * Encapsulates a potentially repeatable operation.
+   */
+  public interface RepeatableWork<T> {
+    /**
+     * Abstracts a unit of repeatable (i.e.: "repeat until completed") work.
+     * <p>
+     * The work unit may be repeated as instructed by the {@link Result}.
+     *
+     * @param storeProvider {@link MutableStoreProvider} instance.
+     * @return {@link Result}
+     */
+    Result<T> apply(MutableStoreProvider storeProvider);
+  }
+
+  /**
+   * Encapsulates a non-repeatable operation.
+   */
+  public interface Work<T> extends RepeatableWork<T> {
+    @Override
+    default Result<T> apply(MutableStoreProvider storeProvider) {
+      T value = execute(storeProvider);
+      return new Result<>(true, value);
+    }
+
+    /**
+     * Abstracts a unit of non-repeatable (i.e.: "run exactly once") work.
+     *
+     * @param storeProvider {@link MutableStoreProvider} instance.
+     * @return result value.
+     */
+    T execute(MutableStoreProvider storeProvider);
+  }
+
+  @Inject
+  protected BatchWorker(Storage storage, StatsProvider statsProvider, int maxBatchSize) {
+    this.storage = requireNonNull(storage);
+    this.maxBatchSize = maxBatchSize;
+
+    scheduledExecutor = AsyncUtil.singleThreadLoggingScheduledExecutor(serviceName() + "-%d",
LOG);
+    statsProvider.makeGauge(serviceName() + "_queue_size", () -> workQueue.size());
+    statsProvider.makeGauge(
+        serviceName() + "_last_processed_batch_size",
+        () -> lastBatchSize.intValue());
+    batchUnlocked = new SlidingStats(serviceName() + "_batch_unlocked", "nanos");
+    batchLocked = new SlidingStats(serviceName() + "_batch_locked", "nanos");
+    itemsProcessed = statsProvider.makeCounter(serviceName() + "_items_processed");
+    batchesProcessed = statsProvider.makeCounter(serviceName() + "_batches_processed");
+  }
+
+  /**
+   * Executes a non-repeatable {@link Work} and returns {@link CompletableFuture} to wait
on.
+   *
+   * @param work A non-repeatable {@link Work} to execute.
+   * @return {@link CompletableFuture} to wait on.
+   */
+  public CompletableFuture<T> execute(Work<T> work) {
+    CompletableFuture<T> result = new CompletableFuture<>();
+    workQueue.add(new WorkItem<>(
+        work,
+        result,
+        Optional.empty(),
+        Optional.empty()));
+
+    return result;
+  }
+
+  /**
+   * Executes a {@link RepeatableWork} until it completes and returns {@link CompletableFuture}
+   * to wait on.
+   *
+   * @param backoffStrategy A {@link BackoffStrategy} instance to backoff subsequent runs.
+   * @param work A {@link RepeatableWork} to execute.
+   */
+  public CompletableFuture<T> executeWithReplay(
+      BackoffStrategy backoffStrategy,
+      RepeatableWork<T> work) {
+
+    CompletableFuture<T> result = new CompletableFuture<>();
+    workQueue.add(new WorkItem<>(
+        work,
+        result,
+        Optional.of(backoffStrategy),
+        Optional.of(0L)));
+
+    return result;
+  }
+
+  @Override
+  protected void run() throws Exception {
+    while (isRunning()) {
+      List<WorkItem<T>> batch = new LinkedList<>();
+      batch.add(workQueue.take());
+      workQueue.drainTo(batch, maxBatchSize - batch.size());
+      processBatch(batch);
+    }
+  }
+
+  private void processBatch(List<WorkItem<T>> batch) {
+    if (!batch.isEmpty()) {
+      long unlockedStart = System.nanoTime();
+      storage.write((Storage.MutateWork.NoResult.Quiet) storeProvider -> {
+        long lockedStart = System.nanoTime();
+        for (WorkItem<T> item : batch) {
+          try {
+            Result<T> itemResult = item.work.apply(storeProvider);
+            if (itemResult.isCompleted) {
+              item.result.complete(itemResult.value);
+            } else {
+              // Work not finished yet - re-queue for a followup later.
+              long backoffMsec = backoffFor(item);
+              scheduledExecutor.schedule(
+                  () -> workQueue.add(new WorkItem<>(
+                      item.work,
+                      item.result,
+                      item.backoffStrategy,
+                      Optional.of(backoffMsec))),
+                  backoffMsec,
+                  TimeUnit.MILLISECONDS);
+            }
+          } catch (RuntimeException e) {
+            LOG.error("{}: Failed to process batch item. Error: {}", serviceName(), e);
+            item.result.completeExceptionally(e);
+          }
+        }
+        batchLocked.accumulate(System.nanoTime() - lockedStart);
+      });
+      batchUnlocked.accumulate(System.nanoTime() - unlockedStart);
+      batchesProcessed.incrementAndGet();
+      lastBatchSize.set(batch.size());
+      itemsProcessed.addAndGet(batch.size());
+    }
+  }
+
+  private long backoffFor(WorkItem<T> item) {
+    checkState(item.backoffStrategy.isPresent());
+    checkState(item.lastBackoffMsec.isPresent());
+    return item.backoffStrategy.get().calculateBackoffMs(item.lastBackoffMsec.get());
+  }
+
+  private class WorkItem<V> {
+    private final RepeatableWork<V> work;
+    private final CompletableFuture<T> result;
+    private final Optional<BackoffStrategy> backoffStrategy;
+    private final Optional<Long> lastBackoffMsec;
+
+    WorkItem(
+        RepeatableWork<V> work,
+        CompletableFuture<T> result,
+        Optional<BackoffStrategy> backoffStrategy,
+        Optional<Long> lastBackoffMsec) {
+
+      this.work = work;
+      this.result = result;
+      this.backoffStrategy = backoffStrategy;
+      this.lastBackoffMsec = lastBackoffMsec;
+    }
+  }
+}

http://git-wip-us.apache.org/repos/asf/aurora/blob/ebfeb3e6/src/main/java/org/apache/aurora/scheduler/SchedulerModule.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/SchedulerModule.java b/src/main/java/org/apache/aurora/scheduler/SchedulerModule.java
index 4a7ef0b..2ec3967 100644
--- a/src/main/java/org/apache/aurora/scheduler/SchedulerModule.java
+++ b/src/main/java/org/apache/aurora/scheduler/SchedulerModule.java
@@ -16,6 +16,8 @@ package org.apache.aurora.scheduler;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ScheduledExecutorService;
+
+import javax.inject.Inject;
 import javax.inject.Singleton;
 
 import com.google.inject.AbstractModule;
@@ -27,10 +29,13 @@ import org.apache.aurora.common.args.CmdLine;
 import org.apache.aurora.common.args.constraints.Positive;
 import org.apache.aurora.common.quantity.Amount;
 import org.apache.aurora.common.quantity.Time;
+import org.apache.aurora.common.stats.StatsProvider;
+import org.apache.aurora.scheduler.BatchWorker.NoResult;
 import org.apache.aurora.scheduler.SchedulerLifecycle.LeadingOptions;
 import org.apache.aurora.scheduler.TaskIdGenerator.TaskIdGeneratorImpl;
 import org.apache.aurora.scheduler.base.AsyncUtil;
 import org.apache.aurora.scheduler.events.PubsubEventModule;
+import org.apache.aurora.scheduler.storage.Storage;
 import org.apache.mesos.Protos;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -59,6 +64,11 @@ public class SchedulerModule extends AbstractModule {
       help = "The maximum number of status updates that can be processed in a batch.")
   private static final Arg<Integer> MAX_STATUS_UPDATE_BATCH_SIZE = Arg.create(1000);
 
+  @Positive
+  @CmdLine(name = "max_task_event_batch_size",
+      help = "The maximum number of task state change events that can be processed in a batch.")
+  private static final Arg<Integer> MAX_TASK_EVENT_BATCH_SIZE = Arg.create(300);
+
   @Override
   protected void configure() {
     bind(TaskIdGenerator.class).to(TaskIdGeneratorImpl.class);
@@ -93,6 +103,21 @@ public class SchedulerModule extends AbstractModule {
     bind(TaskStatusHandler.class).to(TaskStatusHandlerImpl.class);
     bind(TaskStatusHandlerImpl.class).in(Singleton.class);
     addSchedulerActiveServiceBinding(binder()).to(TaskStatusHandlerImpl.class);
+
+    bind(TaskEventBatchWorker.class).in(Singleton.class);
+    addSchedulerActiveServiceBinding(binder()).to(TaskEventBatchWorker.class);
   }
 
+  public static class TaskEventBatchWorker extends BatchWorker<NoResult> {
+    @Inject
+    TaskEventBatchWorker(Storage storage, StatsProvider statsProvider) {
+
+      super(storage, statsProvider, MAX_TASK_EVENT_BATCH_SIZE.get());
+    }
+
+    @Override
+    protected String serviceName() {
+      return "TaskEventBatchWorker";
+    }
+  }
 }

http://git-wip-us.apache.org/repos/asf/aurora/blob/ebfeb3e6/src/main/java/org/apache/aurora/scheduler/pruning/TaskHistoryPruner.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/pruning/TaskHistoryPruner.java b/src/main/java/org/apache/aurora/scheduler/pruning/TaskHistoryPruner.java
index f07746c..c672826 100644
--- a/src/main/java/org/apache/aurora/scheduler/pruning/TaskHistoryPruner.java
+++ b/src/main/java/org/apache/aurora/scheduler/pruning/TaskHistoryPruner.java
@@ -29,13 +29,14 @@ import org.apache.aurora.common.quantity.Amount;
 import org.apache.aurora.common.quantity.Time;
 import org.apache.aurora.common.util.Clock;
 import org.apache.aurora.gen.apiConstants;
+import org.apache.aurora.scheduler.BatchWorker;
+import org.apache.aurora.scheduler.SchedulerModule.TaskEventBatchWorker;
 import org.apache.aurora.scheduler.async.AsyncModule.AsyncExecutor;
 import org.apache.aurora.scheduler.async.DelayExecutor;
 import org.apache.aurora.scheduler.base.Query;
 import org.apache.aurora.scheduler.base.Tasks;
 import org.apache.aurora.scheduler.state.StateManager;
 import org.apache.aurora.scheduler.storage.Storage;
-import org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult;
 import org.apache.aurora.scheduler.storage.entities.IJobKey;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
 import org.slf4j.Logger;
@@ -62,6 +63,7 @@ public class TaskHistoryPruner implements EventSubscriber {
   private final HistoryPrunnerSettings settings;
   private final Storage storage;
   private final Lifecycle lifecycle;
+  private final TaskEventBatchWorker batchWorker;
 
   private final Predicate<IScheduledTask> safeToDelete = new Predicate<IScheduledTask>()
{
     @Override
@@ -94,7 +96,8 @@ public class TaskHistoryPruner implements EventSubscriber {
       Clock clock,
       HistoryPrunnerSettings settings,
       Storage storage,
-      Lifecycle lifecycle) {
+      Lifecycle lifecycle,
+      TaskEventBatchWorker batchWorker) {
 
     this.executor = requireNonNull(executor);
     this.stateManager = requireNonNull(stateManager);
@@ -102,6 +105,7 @@ public class TaskHistoryPruner implements EventSubscriber {
     this.settings = requireNonNull(settings);
     this.storage = requireNonNull(storage);
     this.lifecycle = requireNonNull(lifecycle);
+    this.batchWorker = requireNonNull(batchWorker);
   }
 
   @VisibleForTesting
@@ -131,8 +135,10 @@ public class TaskHistoryPruner implements EventSubscriber {
 
   private void deleteTasks(final Set<String> taskIds) {
     LOG.info("Pruning inactive tasks " + taskIds);
-    storage.write(
-        (NoResult.Quiet) storeProvider -> stateManager.deleteTasks(storeProvider, taskIds));
+    batchWorker.execute(storeProvider -> {
+      stateManager.deleteTasks(storeProvider, taskIds);
+      return BatchWorker.NO_RESULT;
+    });
   }
 
   @VisibleForTesting

http://git-wip-us.apache.org/repos/asf/aurora/blob/ebfeb3e6/src/main/java/org/apache/aurora/scheduler/scheduling/TaskThrottler.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/scheduling/TaskThrottler.java b/src/main/java/org/apache/aurora/scheduler/scheduling/TaskThrottler.java
index bbd971a..867c9bd 100644
--- a/src/main/java/org/apache/aurora/scheduler/scheduling/TaskThrottler.java
+++ b/src/main/java/org/apache/aurora/scheduler/scheduling/TaskThrottler.java
@@ -22,13 +22,14 @@ import org.apache.aurora.common.quantity.Amount;
 import org.apache.aurora.common.quantity.Time;
 import org.apache.aurora.common.stats.SlidingStats;
 import org.apache.aurora.common.util.Clock;
+import org.apache.aurora.scheduler.BatchWorker;
+import org.apache.aurora.scheduler.SchedulerModule.TaskEventBatchWorker;
 import org.apache.aurora.scheduler.async.AsyncModule.AsyncExecutor;
 import org.apache.aurora.scheduler.async.DelayExecutor;
 import org.apache.aurora.scheduler.base.Tasks;
 import org.apache.aurora.scheduler.events.PubsubEvent.EventSubscriber;
 import org.apache.aurora.scheduler.events.PubsubEvent.TaskStateChange;
 import org.apache.aurora.scheduler.state.StateManager;
-import org.apache.aurora.scheduler.storage.Storage;
 
 import static java.util.Objects.requireNonNull;
 
@@ -46,8 +47,8 @@ class TaskThrottler implements EventSubscriber {
   private final RescheduleCalculator rescheduleCalculator;
   private final Clock clock;
   private final DelayExecutor executor;
-  private final Storage storage;
   private final StateManager stateManager;
+  private final TaskEventBatchWorker batchWorker;
 
   private final SlidingStats throttleStats = new SlidingStats("task_throttle", "ms");
 
@@ -56,14 +57,14 @@ class TaskThrottler implements EventSubscriber {
       RescheduleCalculator rescheduleCalculator,
       Clock clock,
       @AsyncExecutor DelayExecutor executor,
-      Storage storage,
-      StateManager stateManager) {
+      StateManager stateManager,
+      TaskEventBatchWorker batchWorker) {
 
     this.rescheduleCalculator = requireNonNull(rescheduleCalculator);
     this.clock = requireNonNull(clock);
     this.executor = requireNonNull(executor);
-    this.storage = requireNonNull(storage);
     this.stateManager = requireNonNull(stateManager);
+    this.batchWorker = requireNonNull(batchWorker);
   }
 
   @Subscribe
@@ -73,13 +74,16 @@ class TaskThrottler implements EventSubscriber {
           + rescheduleCalculator.getFlappingPenaltyMs(stateChange.getTask());
       long delayMs = Math.max(0, readyAtMs - clock.nowMillis());
       throttleStats.accumulate(delayMs);
-      executor.execute(
-          () -> storage.write(storeProvider -> stateManager.changeState(
-              storeProvider,
-              stateChange.getTaskId(),
-              Optional.of(THROTTLED),
-              PENDING,
-              Optional.absent())),
+      executor.execute(() ->
+              batchWorker.execute(storeProvider -> {
+                stateManager.changeState(
+                    storeProvider,
+                    stateChange.getTaskId(),
+                    Optional.of(THROTTLED),
+                    PENDING,
+                    Optional.absent());
+                return BatchWorker.NO_RESULT;
+              }),
           Amount.of(delayMs, Time.MILLISECONDS));
     }
   }

http://git-wip-us.apache.org/repos/asf/aurora/blob/ebfeb3e6/src/main/java/org/apache/aurora/scheduler/state/MaintenanceController.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/state/MaintenanceController.java b/src/main/java/org/apache/aurora/scheduler/state/MaintenanceController.java
index 3c7cda0..574efc9 100644
--- a/src/main/java/org/apache/aurora/scheduler/state/MaintenanceController.java
+++ b/src/main/java/org/apache/aurora/scheduler/state/MaintenanceController.java
@@ -30,6 +30,8 @@ import com.google.common.eventbus.Subscribe;
 import org.apache.aurora.gen.HostStatus;
 import org.apache.aurora.gen.MaintenanceMode;
 import org.apache.aurora.gen.ScheduleStatus;
+import org.apache.aurora.scheduler.BatchWorker;
+import org.apache.aurora.scheduler.SchedulerModule.TaskEventBatchWorker;
 import org.apache.aurora.scheduler.base.Query;
 import org.apache.aurora.scheduler.base.Tasks;
 import org.apache.aurora.scheduler.events.PubsubEvent.EventSubscriber;
@@ -37,7 +39,6 @@ import org.apache.aurora.scheduler.events.PubsubEvent.TaskStateChange;
 import org.apache.aurora.scheduler.storage.AttributeStore;
 import org.apache.aurora.scheduler.storage.Storage;
 import org.apache.aurora.scheduler.storage.Storage.MutableStoreProvider;
-import org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult;
 import org.apache.aurora.scheduler.storage.entities.IHostAttributes;
 import org.apache.aurora.scheduler.storage.entities.IHostStatus;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
@@ -106,11 +107,17 @@ public interface MaintenanceController {
     private static final Logger LOG = LoggerFactory.getLogger(MaintenanceControllerImpl.class);
     private final Storage storage;
     private final StateManager stateManager;
+    private final TaskEventBatchWorker batchWorker;
 
     @Inject
-    public MaintenanceControllerImpl(Storage storage, StateManager stateManager) {
+    public MaintenanceControllerImpl(
+        Storage storage,
+        StateManager stateManager,
+        TaskEventBatchWorker batchWorker) {
+
       this.storage = requireNonNull(storage);
       this.stateManager = requireNonNull(stateManager);
+      this.batchWorker = requireNonNull(batchWorker);
     }
 
     private Set<IHostStatus> watchDrainingTasks(MutableStoreProvider store, Set<String>
hosts) {
@@ -153,7 +160,7 @@ public interface MaintenanceController {
     public void taskChangedState(final TaskStateChange change) {
       if (Tasks.isTerminated(change.getNewState())) {
         final String host = change.getTask().getAssignedTask().getSlaveHost();
-        storage.write((NoResult.Quiet) (MutableStoreProvider store) -> {
+        batchWorker.execute(store -> {
           // If the task _was_ associated with a draining host, and it was the last task
on the
           // host.
           Optional<IHostAttributes> attributes =
@@ -168,6 +175,7 @@ public interface MaintenanceController {
               LOG.info("Host {} is DRAINING with active tasks: {}", host, Tasks.ids(activeTasks));
             }
           }
+          return BatchWorker.NO_RESULT;
         });
       }
     }

http://git-wip-us.apache.org/repos/asf/aurora/blob/ebfeb3e6/src/main/java/org/apache/aurora/scheduler/updater/JobUpdateControllerImpl.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/updater/JobUpdateControllerImpl.java
b/src/main/java/org/apache/aurora/scheduler/updater/JobUpdateControllerImpl.java
index ef6253e..25b3f37 100644
--- a/src/main/java/org/apache/aurora/scheduler/updater/JobUpdateControllerImpl.java
+++ b/src/main/java/org/apache/aurora/scheduler/updater/JobUpdateControllerImpl.java
@@ -44,6 +44,8 @@ import org.apache.aurora.gen.JobUpdateQuery;
 import org.apache.aurora.gen.JobUpdateStatus;
 import org.apache.aurora.gen.Lock;
 import org.apache.aurora.gen.LockKey;
+import org.apache.aurora.scheduler.BatchWorker;
+import org.apache.aurora.scheduler.SchedulerModule.TaskEventBatchWorker;
 import org.apache.aurora.scheduler.base.InstanceKeys;
 import org.apache.aurora.scheduler.base.JobKeys;
 import org.apache.aurora.scheduler.base.Query;
@@ -120,6 +122,7 @@ class JobUpdateControllerImpl implements JobUpdateController {
   private final Clock clock;
   private final PulseHandler pulseHandler;
   private final Lifecycle lifecycle;
+  private final TaskEventBatchWorker batchWorker;
 
   // Currently-active updaters. An active updater is one that is rolling forward or back.
Paused
   // and completed updates are represented only in storage, not here.
@@ -134,7 +137,8 @@ class JobUpdateControllerImpl implements JobUpdateController {
       ScheduledExecutorService executor,
       StateManager stateManager,
       Clock clock,
-      Lifecycle lifecycle) {
+      Lifecycle lifecycle,
+      TaskEventBatchWorker batchWorker) {
 
     this.updateFactory = requireNonNull(updateFactory);
     this.lockManager = requireNonNull(lockManager);
@@ -143,6 +147,7 @@ class JobUpdateControllerImpl implements JobUpdateController {
     this.stateManager = requireNonNull(stateManager);
     this.clock = requireNonNull(clock);
     this.lifecycle = requireNonNull(lifecycle);
+    this.batchWorker = requireNonNull(batchWorker);
     this.pulseHandler = new PulseHandler(clock);
   }
 
@@ -346,7 +351,7 @@ class JobUpdateControllerImpl implements JobUpdateController {
   }
 
   private void instanceChanged(final IInstanceKey instance, final Optional<IScheduledTask>
state) {
-    storage.write((NoResult.Quiet) storeProvider -> {
+    batchWorker.execute(storeProvider -> {
       IJobKey job = instance.getJobKey();
       UpdateFactory.Update update = updates.get(job);
       if (update != null) {
@@ -366,6 +371,7 @@ class JobUpdateControllerImpl implements JobUpdateController {
               + JobKeys.canonicalString(job));
         }
       }
+      return BatchWorker.NO_RESULT;
     });
   }
 

http://git-wip-us.apache.org/repos/asf/aurora/blob/ebfeb3e6/src/test/java/org/apache/aurora/scheduler/BatchWorkerTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/BatchWorkerTest.java b/src/test/java/org/apache/aurora/scheduler/BatchWorkerTest.java
new file mode 100644
index 0000000..a86dc82
--- /dev/null
+++ b/src/test/java/org/apache/aurora/scheduler/BatchWorkerTest.java
@@ -0,0 +1,96 @@
+/**
+ * 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 org.apache.aurora.scheduler;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+
+import org.apache.aurora.common.testing.easymock.EasyMockTest;
+import org.apache.aurora.common.util.BackoffStrategy;
+import org.apache.aurora.scheduler.BatchWorker.Result;
+import org.apache.aurora.scheduler.storage.testing.StorageTestUtil;
+import org.apache.aurora.scheduler.testing.FakeStatsProvider;
+import org.easymock.EasyMock;
+import org.junit.Before;
+import org.junit.Test;
+
+import static org.easymock.EasyMock.expect;
+import static org.junit.Assert.assertTrue;
+
+public class BatchWorkerTest extends EasyMockTest {
+  private static final String SERVICE_NAME = "TestWorker";
+  private static final String BATCH_STAT = SERVICE_NAME + "_batches_processed";
+  private FakeStatsProvider statsProvider;
+  private BatchWorker<Boolean> batchWorker;
+
+  @Before
+  public void setUp() {
+    StorageTestUtil storageUtil = new StorageTestUtil(this);
+    storageUtil.expectOperations();
+    statsProvider = new FakeStatsProvider();
+    batchWorker = new BatchWorker<Boolean>(storageUtil.storage, statsProvider, 2) {
+      @Override
+      protected String serviceName() {
+        return SERVICE_NAME;
+      }
+    };
+  }
+
+  @Test
+  public void testExecute() throws Exception {
+    control.replay();
+
+    CompletableFuture<Boolean> result1 = batchWorker.execute(store -> true);
+    CompletableFuture<Boolean> result2 = batchWorker.execute(store -> true);
+    CompletableFuture<Boolean> result3 = batchWorker.execute(store -> true);
+    batchWorker.startAsync().awaitRunning();
+
+    assertTrue(result1.get());
+    assertTrue(result2.get());
+    assertTrue(result3.get());
+  }
+
+  @Test(expected = ExecutionException.class)
+  public void testExecuteThrows() throws Exception {
+    control.replay();
+
+    CompletableFuture<Boolean> result =
+        batchWorker.execute(store -> { throw new IllegalArgumentException(); });
+    batchWorker.startAsync().awaitRunning();
+
+    result.get();
+  }
+
+  @Test
+  public void testExecuteWithReplay() throws Exception {
+    BackoffStrategy backoff = createMock(BackoffStrategy.class);
+    final CountDownLatch complete = new CountDownLatch(1);
+
+    expect(backoff.calculateBackoffMs(EasyMock.anyLong())).andReturn(0L).anyTimes();
+
+    control.replay();
+
+    batchWorker.startAsync().awaitRunning();
+    batchWorker.executeWithReplay(
+        backoff,
+        store -> statsProvider.getValue(BATCH_STAT).longValue() > 1L
+            ? new Result<>(true, true)
+            : new Result<>(false, false))
+        .thenAccept(result -> complete.countDown());
+
+    assertTrue(complete.await(10L, TimeUnit.SECONDS));
+  }
+}

http://git-wip-us.apache.org/repos/asf/aurora/blob/ebfeb3e6/src/test/java/org/apache/aurora/scheduler/pruning/TaskHistoryPrunerTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/pruning/TaskHistoryPrunerTest.java
b/src/test/java/org/apache/aurora/scheduler/pruning/TaskHistoryPrunerTest.java
index 99c27e8..8469596 100644
--- a/src/test/java/org/apache/aurora/scheduler/pruning/TaskHistoryPrunerTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/pruning/TaskHistoryPrunerTest.java
@@ -27,6 +27,7 @@ import org.apache.aurora.common.testing.easymock.EasyMockTest;
 import org.apache.aurora.common.util.testing.FakeClock;
 import org.apache.aurora.gen.ScheduleStatus;
 import org.apache.aurora.gen.ScheduledTask;
+import org.apache.aurora.scheduler.SchedulerModule.TaskEventBatchWorker;
 import org.apache.aurora.scheduler.async.DelayExecutor;
 import org.apache.aurora.scheduler.base.JobKeys;
 import org.apache.aurora.scheduler.base.TaskTestUtil;
@@ -48,6 +49,7 @@ import static org.apache.aurora.gen.ScheduleStatus.KILLED;
 import static org.apache.aurora.gen.ScheduleStatus.LOST;
 import static org.apache.aurora.gen.ScheduleStatus.RUNNING;
 import static org.apache.aurora.gen.ScheduleStatus.STARTING;
+import static org.apache.aurora.scheduler.testing.BatchWorkerUtil.expectBatchExecute;
 import static org.easymock.EasyMock.eq;
 import static org.easymock.EasyMock.expectLastCall;
 
@@ -68,20 +70,24 @@ public class TaskHistoryPrunerTest extends EasyMockTest {
   private Command shutdownCommand;
 
   @Before
-  public void setUp() {
+  public void setUp() throws Exception {
     executor = createMock(DelayExecutor.class);
     clock = new FakeClock();
     stateManager = createMock(StateManager.class);
     storageUtil = new StorageTestUtil(this);
     storageUtil.expectOperations();
     shutdownCommand = createMock(Command.class);
+    TaskEventBatchWorker batchWorker = createMock(TaskEventBatchWorker.class);
+    expectBatchExecute(batchWorker, storageUtil.storage, control).anyTimes();
+
     pruner = new TaskHistoryPruner(
         executor,
         stateManager,
         clock,
         new HistoryPrunnerSettings(ONE_DAY, ONE_MINUTE, PER_JOB_HISTORY),
         storageUtil.storage,
-        new Lifecycle(shutdownCommand));
+        new Lifecycle(shutdownCommand),
+        batchWorker);
     closer = Closer.create();
   }
 

http://git-wip-us.apache.org/repos/asf/aurora/blob/ebfeb3e6/src/test/java/org/apache/aurora/scheduler/scheduling/TaskThrottlerTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/scheduling/TaskThrottlerTest.java b/src/test/java/org/apache/aurora/scheduler/scheduling/TaskThrottlerTest.java
index 7d104aa..433f791 100644
--- a/src/test/java/org/apache/aurora/scheduler/scheduling/TaskThrottlerTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/scheduling/TaskThrottlerTest.java
@@ -24,6 +24,7 @@ import org.apache.aurora.gen.AssignedTask;
 import org.apache.aurora.gen.ScheduleStatus;
 import org.apache.aurora.gen.ScheduledTask;
 import org.apache.aurora.gen.TaskEvent;
+import org.apache.aurora.scheduler.SchedulerModule.TaskEventBatchWorker;
 import org.apache.aurora.scheduler.async.DelayExecutor;
 import org.apache.aurora.scheduler.base.Tasks;
 import org.apache.aurora.scheduler.events.PubsubEvent.TaskStateChange;
@@ -39,6 +40,7 @@ import static org.apache.aurora.gen.ScheduleStatus.INIT;
 import static org.apache.aurora.gen.ScheduleStatus.PENDING;
 import static org.apache.aurora.gen.ScheduleStatus.RUNNING;
 import static org.apache.aurora.gen.ScheduleStatus.THROTTLED;
+import static org.apache.aurora.scheduler.testing.BatchWorkerUtil.expectBatchExecute;
 import static org.easymock.EasyMock.capture;
 import static org.easymock.EasyMock.eq;
 import static org.easymock.EasyMock.expect;
@@ -60,12 +62,15 @@ public class TaskThrottlerTest extends EasyMockTest {
     storageUtil = new StorageTestUtil(this);
     storageUtil.expectOperations();
     stateManager = createMock(StateManager.class);
+    TaskEventBatchWorker batchWorker = createMock(TaskEventBatchWorker.class);
+    expectBatchExecute(batchWorker, storageUtil.storage, control).anyTimes();
+
     throttler = new TaskThrottler(
         rescheduleCalculator,
         clock,
         executor,
-        storageUtil.storage,
-        stateManager);
+        stateManager,
+        batchWorker);
   }
 
   @Test

http://git-wip-us.apache.org/repos/asf/aurora/blob/ebfeb3e6/src/test/java/org/apache/aurora/scheduler/state/MaintenanceControllerImplTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/state/MaintenanceControllerImplTest.java
b/src/test/java/org/apache/aurora/scheduler/state/MaintenanceControllerImplTest.java
index 94f5ca5..ae83dea 100644
--- a/src/test/java/org/apache/aurora/scheduler/state/MaintenanceControllerImplTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/state/MaintenanceControllerImplTest.java
@@ -30,6 +30,7 @@ import org.apache.aurora.gen.HostStatus;
 import org.apache.aurora.gen.MaintenanceMode;
 import org.apache.aurora.gen.ScheduleStatus;
 import org.apache.aurora.gen.ScheduledTask;
+import org.apache.aurora.scheduler.SchedulerModule.TaskEventBatchWorker;
 import org.apache.aurora.scheduler.async.AsyncModule.AsyncExecutor;
 import org.apache.aurora.scheduler.base.Query;
 import org.apache.aurora.scheduler.base.TaskTestUtil;
@@ -53,6 +54,7 @@ import static org.apache.aurora.gen.MaintenanceMode.SCHEDULED;
 import static org.apache.aurora.gen.ScheduleStatus.KILLED;
 import static org.apache.aurora.gen.ScheduleStatus.RUNNING;
 import static org.apache.aurora.scheduler.state.MaintenanceController.MaintenanceControllerImpl;
+import static org.apache.aurora.scheduler.testing.BatchWorkerUtil.expectBatchExecute;
 import static org.easymock.EasyMock.expect;
 import static org.junit.Assert.assertEquals;
 
@@ -71,6 +73,8 @@ public class MaintenanceControllerImplTest extends EasyMockTest {
     storageUtil = new StorageTestUtil(this);
     storageUtil.expectOperations();
     stateManager = createMock(StateManager.class);
+    TaskEventBatchWorker batchWorker = createMock(TaskEventBatchWorker.class);
+    expectBatchExecute(batchWorker, storageUtil.storage, control).anyTimes();
 
     Injector injector = Guice.createInjector(
         new PubsubEventModule(),
@@ -83,6 +87,7 @@ public class MaintenanceControllerImplTest extends EasyMockTest {
             bind(StatsProvider.class).toInstance(new FakeStatsProvider());
             bind(Executor.class).annotatedWith(AsyncExecutor.class)
                 .toInstance(MoreExecutors.directExecutor());
+            bind(TaskEventBatchWorker.class).toInstance(batchWorker);
           }
         });
     maintenance = injector.getInstance(MaintenanceController.class);

http://git-wip-us.apache.org/repos/asf/aurora/blob/ebfeb3e6/src/test/java/org/apache/aurora/scheduler/testing/BatchWorkerUtil.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/testing/BatchWorkerUtil.java b/src/test/java/org/apache/aurora/scheduler/testing/BatchWorkerUtil.java
new file mode 100644
index 0000000..46b2e36
--- /dev/null
+++ b/src/test/java/org/apache/aurora/scheduler/testing/BatchWorkerUtil.java
@@ -0,0 +1,59 @@
+/**
+ * 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 org.apache.aurora.scheduler.testing;
+
+import java.util.concurrent.CompletableFuture;
+
+import org.apache.aurora.common.testing.easymock.EasyMockTest;
+import org.apache.aurora.scheduler.BatchWorker;
+import org.apache.aurora.scheduler.BatchWorker.Work;
+import org.apache.aurora.scheduler.storage.Storage;
+import org.easymock.Capture;
+import org.easymock.IExpectationSetters;
+import org.easymock.IMocksControl;
+
+import static org.apache.aurora.common.testing.easymock.EasyMockTest.createCapture;
+import static org.easymock.EasyMock.capture;
+import static org.easymock.EasyMock.expect;
+
+public final class BatchWorkerUtil {
+  private BatchWorkerUtil() {
+    // Utility class.
+  }
+
+  public static <T> IExpectationSetters<CompletableFuture<T>> expectBatchExecute(
+      BatchWorker<T> batchWorker,
+      Storage storage,
+      IMocksControl control,
+      T resultValue) throws Exception {
+
+    final CompletableFuture<T> result = new EasyMockTest.Clazz<CompletableFuture<T>>()
{ }
+        .createMock(control);
+    expect(result.get()).andReturn(resultValue).anyTimes();
+
+    final Capture<Work<T>> capture = createCapture();
+    return expect(batchWorker.execute(capture(capture))).andAnswer(() -> {
+      storage.write((Storage.MutateWork.NoResult.Quiet) store -> capture.getValue().apply(store));
+      return result;
+    });
+  }
+
+  public static <T> IExpectationSetters<CompletableFuture<T>> expectBatchExecute(
+      BatchWorker<T> batchWorker,
+      Storage storage,
+      IMocksControl control) throws Exception {
+
+    return expectBatchExecute(batchWorker, storage, control, null);
+  }
+}

http://git-wip-us.apache.org/repos/asf/aurora/blob/ebfeb3e6/src/test/java/org/apache/aurora/scheduler/updater/JobUpdaterIT.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/updater/JobUpdaterIT.java b/src/test/java/org/apache/aurora/scheduler/updater/JobUpdaterIT.java
index f879827..ea0b89a 100644
--- a/src/test/java/org/apache/aurora/scheduler/updater/JobUpdaterIT.java
+++ b/src/test/java/org/apache/aurora/scheduler/updater/JobUpdaterIT.java
@@ -60,6 +60,7 @@ import org.apache.aurora.gen.Range;
 import org.apache.aurora.gen.ScheduleStatus;
 import org.apache.aurora.gen.ScheduledTask;
 import org.apache.aurora.gen.TaskConfig;
+import org.apache.aurora.scheduler.SchedulerModule.TaskEventBatchWorker;
 import org.apache.aurora.scheduler.TaskIdGenerator;
 import org.apache.aurora.scheduler.TaskIdGenerator.TaskIdGeneratorImpl;
 import org.apache.aurora.scheduler.base.JobKeys;
@@ -125,6 +126,7 @@ import static org.apache.aurora.gen.ScheduleStatus.KILLED;
 import static org.apache.aurora.gen.ScheduleStatus.RUNNING;
 import static org.apache.aurora.gen.ScheduleStatus.STARTING;
 import static org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult;
+import static org.apache.aurora.scheduler.testing.BatchWorkerUtil.expectBatchExecute;
 import static org.apache.aurora.scheduler.updater.UpdateFactory.UpdateFactoryImpl.expandInstanceIds;
 import static org.easymock.EasyMock.expectLastCall;
 import static org.junit.Assert.assertEquals;
@@ -164,7 +166,7 @@ public class JobUpdaterIT extends EasyMockTest {
   }
 
   @Before
-  public void setUp() {
+  public void setUp() throws Exception {
     // Avoid console spam due to stats registered multiple times.
     Stats.flush();
     ScheduledExecutorService executor = createMock(ScheduledExecutorService.class);
@@ -172,6 +174,7 @@ public class JobUpdaterIT extends EasyMockTest {
     driver = createMock(Driver.class);
     shutdownCommand = createMock(Command.class);
     eventBus = new EventBus();
+    TaskEventBatchWorker batchWorker = createMock(TaskEventBatchWorker.class);
 
     Injector injector = Guice.createInjector(
         new UpdaterModule(executor),
@@ -195,6 +198,7 @@ public class JobUpdaterIT extends EasyMockTest {
             bind(LockManager.class).to(LockManagerImpl.class);
             bind(UUIDGenerator.class).to(UUIDGeneratorImpl.class);
             bind(Lifecycle.class).toInstance(new Lifecycle(shutdownCommand));
+            bind(TaskEventBatchWorker.class).toInstance(batchWorker);
           }
         });
     updater = injector.getInstance(JobUpdateController.class);
@@ -204,6 +208,7 @@ public class JobUpdaterIT extends EasyMockTest {
     stateManager = injector.getInstance(StateManager.class);
     eventBus.register(injector.getInstance(JobUpdateEventSubscriber.class));
     subscriber = injector.getInstance(JobUpdateEventSubscriber.class);
+    expectBatchExecute(batchWorker, storage, control).anyTimes();
   }
 
   @After


Mime
View raw message