private IScheduledTask makeFlappyTaskWithStates(
String taskId, Iterable<ScheduleStatus> states, @Nullable String ancestorId) {
Amount<Long, Time> timeInState = Amount.of(10L, Time.SECONDS);
ScheduledTask base = makeTask(taskId, INIT).newBuilder();
for (ScheduleStatus status : states) {
base.addToTaskEvents(new TaskEvent(clock.nowMillis(), status));
clock.advance(timeInState);
}
base.setAncestorId(ancestorId);
final IScheduledTask result = IScheduledTask.build(base);
// Insert the task if it doesn't already exist.
storage.write(
new MutateWork.NoResult.Quiet() {
@Override
protected void execute(MutableStoreProvider storeProvider) {
TaskStore.Mutable taskStore = storeProvider.getUnsafeTaskStore();
if (taskStore.fetchTasks(Query.taskScoped(Tasks.id(result))).isEmpty()) {
taskStore.saveTasks(ImmutableSet.of(result));
}
}
});
return result;
}
private Capture<Runnable> expectOfferDeclineIn(int delayMillis) {
expect(returnDelay.get()).andReturn(Amount.of(delayMillis, Time.MILLISECONDS));
Capture<Runnable> runnable = createCapture();
executor.schedule(capture(runnable), eq((long) delayMillis), eq(TimeUnit.MILLISECONDS));
expectLastCall().andReturn(createMock(ScheduledFuture.class));
return runnable;
}
@Test
public void testCalculateBackoffMs() {
TruncatedBinaryBackoff backoff =
new TruncatedBinaryBackoff(
Amount.of(1L, Time.MILLISECONDS), Amount.of(6L, Time.MILLISECONDS));
try {
backoff.calculateBackoffMs(-1L);
} catch (IllegalArgumentException e) {
// expected
}
assertEquals(1, backoff.calculateBackoffMs(0));
assertEquals(2, backoff.calculateBackoffMs(1));
assertEquals(4, backoff.calculateBackoffMs(2));
assertEquals(6, backoff.calculateBackoffMs(4));
assertEquals(6, backoff.calculateBackoffMs(8));
}
@Override
public Amount<Long, Time> getReevaluationDelay(
IInstanceKey instance,
IJobUpdateInstructions instructions,
MutableStoreProvider storeProvider,
StateManager stateManager,
JobUpdateStatus status) {
return Amount.of(
(long) instructions.getSettings().getMinWaitInInstanceRunningMs(), Time.MILLISECONDS);
}
/** Returns offers after a random duration within a fixed window. */
private static class RandomJitterReturnDelay implements OfferReturnDelay {
private static final int JITTER_WINDOW_MS = Amount.of(1, Time.MINUTES).as(Time.MILLISECONDS);
private final int minHoldTimeMs = MIN_OFFER_HOLD_TIME.get().as(Time.MILLISECONDS);
private final Random random = new Random.SystemRandom(new java.util.Random());
@Override
public Amount<Integer, Time> get() {
return Amount.of(minHoldTimeMs + random.nextInt(JITTER_WINDOW_MS), Time.MILLISECONDS);
}
}
private void replayAndCreateScheduler() {
control.replay();
offerQueue = new OfferQueueImpl(driver, returnDelay, executor, maintenance);
RateLimiter rateLimiter = RateLimiter.create(1);
Amount<Long, Time> flappingThreshold = Amount.of(5L, Time.MINUTES);
SchedulingAction scheduler = new TaskScheduler(storage, stateManager, assigner, offerQueue);
taskGroups =
new TaskGroups(
executor,
storage,
retryStrategy,
rateLimiter,
scheduler,
clock,
new RescheduleCalculatorImpl(
storage,
new RescheduleCalculatorSettings(
flappingStrategy, flappingThreshold, Amount.of(5, Time.SECONDS)),
clock),
preemptor);
}
@Before
@SuppressWarnings("unchecked") // Due to createMock.
public void setUp() {
wrappedLog = createMock(Log.class);
bufferedLog =
BufferedLog.<String, Boolean>builder()
.buffer(wrappedLog)
.withRetryFilter(RETRY_FILTER)
.withChunkLength(BUFFER_SIZE)
.withMaxBuffer(MAX_BUFFER_SIZE)
.withFlushInterval(Amount.of(10000, Time.SECONDS))
.withExecutorService(MoreExecutors.sameThreadExecutor())
.build();
}
@Override
public Amount<Long, Time> getReevaluationDelay(
IInstanceKey instance,
IJobUpdateInstructions instructions,
MutableStoreProvider storeProvider,
StateManager stateManager,
JobUpdateStatus status) {
LOG.info("Adding instance " + instance + " while " + status);
ITaskConfig replacement =
getTargetConfig(instructions, status == ROLLING_FORWARD, instance.getInstanceId());
stateManager.insertPendingTasks(
storeProvider, replacement, ImmutableSet.of(instance.getInstanceId()));
return Amount.of(
(long) instructions.getSettings().getMaxWaitToInstanceRunningMs(), Time.MILLISECONDS);
}
@Override
public Amount<Long, Time> getReevaluationDelay(
IInstanceKey instance,
IJobUpdateInstructions instructions,
MutableStoreProvider storeProvider,
StateManager stateManager,
JobUpdateStatus status) {
String taskId =
Tasks.id(
Iterables.getOnlyElement(
storeProvider
.getTaskStore()
.fetchTasks(Query.instanceScoped(instance).active())));
LOG.info("Killing " + instance + " while " + status);
stateManager.changeState(
storeProvider,
taskId,
Optional.absent(),
ScheduleStatus.KILLING,
Optional.of("Killed for job update."));
return Amount.of(
(long) instructions.getSettings().getMaxWaitToInstanceRunningMs(), Time.MILLISECONDS);
}
/** Binding module for async task management. */
public class AsyncModule extends AbstractModule {
private static final Logger LOG = Logger.getLogger(AsyncModule.class.getName());
@CmdLine(
name = "async_worker_threads",
help = "The number of worker threads to process async task operations with.")
private static final Arg<Integer> ASYNC_WORKER_THREADS = Arg.create(1);
@CmdLine(
name = "transient_task_state_timeout",
help = "The amount of time after which to treat a task stuck in a transient state as LOST.")
private static final Arg<Amount<Long, Time>> TRANSIENT_TASK_STATE_TIMEOUT =
Arg.create(Amount.of(5L, Time.MINUTES));
@CmdLine(
name = "initial_schedule_delay",
help = "Initial amount of time to wait before attempting to schedule a PENDING task.")
private static final Arg<Amount<Long, Time>> INITIAL_SCHEDULE_DELAY =
Arg.create(Amount.of(1L, Time.SECONDS));
@CmdLine(
name = "max_schedule_delay",
help = "Maximum delay between attempts to schedule a PENDING tasks.")
private static final Arg<Amount<Long, Time>> MAX_SCHEDULE_DELAY =
Arg.create(Amount.of(30L, Time.SECONDS));
@CmdLine(
name = "min_offer_hold_time",
help = "Minimum amount of time to hold a resource offer before declining.")
private static final Arg<Amount<Integer, Time>> MIN_OFFER_HOLD_TIME =
Arg.create(Amount.of(5, Time.MINUTES));
@CmdLine(
name = "history_prune_threshold",
help = "Time after which the scheduler will prune terminated task history.")
private static final Arg<Amount<Long, Time>> HISTORY_PRUNE_THRESHOLD =
Arg.create(Amount.of(2L, Time.DAYS));
@CmdLine(
name = "max_schedule_attempts_per_sec",
help = "Maximum number of scheduling attempts to make per second.")
private static final Arg<Double> MAX_SCHEDULE_ATTEMPTS_PER_SEC = Arg.create(10D);
@CmdLine(
name = "flapping_task_threshold",
help = "A task that repeatedly runs for less than this time is considered to be flapping.")
private static final Arg<Amount<Long, Time>> FLAPPING_THRESHOLD =
Arg.create(Amount.of(5L, Time.MINUTES));
@CmdLine(
name = "initial_flapping_task_delay",
help = "Initial amount of time to wait before attempting to schedule a flapping task.")
private static final Arg<Amount<Long, Time>> INITIAL_FLAPPING_DELAY =
Arg.create(Amount.of(30L, Time.SECONDS));
@CmdLine(
name = "max_flapping_task_delay",
help = "Maximum delay between attempts to schedule a flapping task.")
private static final Arg<Amount<Long, Time>> MAX_FLAPPING_DELAY =
Arg.create(Amount.of(5L, Time.MINUTES));
@CmdLine(
name = "max_reschedule_task_delay_on_startup",
help = "Upper bound of random delay for pending task rescheduling on scheduler startup.")
private static final Arg<Amount<Integer, Time>> MAX_RESCHEDULING_DELAY =
Arg.create(Amount.of(30, Time.SECONDS));
@CmdLine(
name = "preemption_delay",
help = "Time interval after which a pending task becomes eligible to preempt other tasks")
private static final Arg<Amount<Long, Time>> PREEMPTION_DELAY =
Arg.create(Amount.of(10L, Time.MINUTES));
@CmdLine(name = "enable_preemptor", help = "Enable the preemptor and preemption")
private static final Arg<Boolean> ENABLE_PREEMPTOR = Arg.create(true);
private static final Preemptor NULL_PREEMPTOR =
new Preemptor() {
@Override
public Optional<String> findPreemptionSlotFor(String taskId) {
return Optional.absent();
}
};
@CmdLine(
name = "offer_reservation_duration",
help =
"Time to reserve a slave's offers while "
+ "trying to satisfy a task preempting another.")
private static final Arg<Amount<Long, Time>> RESERVATION_DURATION =
Arg.create(Amount.of(3L, Time.MINUTES));
@BindingAnnotation
@Target({FIELD, PARAMETER, METHOD})
@Retention(RUNTIME)
private @interface PreemptionBinding {}
@VisibleForTesting
static final Key<Preemptor> PREEMPTOR_KEY = Key.get(Preemptor.class, PreemptionBinding.class);
@Override
protected void configure() {
// Don't worry about clean shutdown, these can be daemon and cleanup-free.
final ScheduledThreadPoolExecutor executor =
new ScheduledThreadPoolExecutor(
ASYNC_WORKER_THREADS.get(),
new ThreadFactoryBuilder().setNameFormat("AsyncProcessor-%d").setDaemon(true).build());
Stats.exportSize("timeout_queue_size", executor.getQueue());
Stats.export(
new StatImpl<Long>("async_tasks_completed") {
@Override
public Long read() {
return executor.getCompletedTaskCount();
}
});
// AsyncModule itself is not a subclass of PrivateModule because TaskEventModule internally uses
// a MultiBinder, which cannot span multiple injectors.
binder()
.install(
new PrivateModule() {
@Override
protected void configure() {
bind(new TypeLiteral<Amount<Long, Time>>() {})
.toInstance(TRANSIENT_TASK_STATE_TIMEOUT.get());
bind(ScheduledExecutorService.class).toInstance(executor);
bind(TaskTimeout.class).in(Singleton.class);
requireBinding(StatsProvider.class);
expose(TaskTimeout.class);
}
});
PubsubEventModule.bindSubscriber(binder(), TaskTimeout.class);
binder()
.install(
new PrivateModule() {
@Override
protected void configure() {
bind(TaskGroupsSettings.class)
.toInstance(
new TaskGroupsSettings(
new TruncatedBinaryBackoff(
INITIAL_SCHEDULE_DELAY.get(), MAX_SCHEDULE_DELAY.get()),
RateLimiter.create(MAX_SCHEDULE_ATTEMPTS_PER_SEC.get())));
bind(RescheduleCalculatorImpl.RescheduleCalculatorSettings.class)
.toInstance(
new RescheduleCalculatorImpl.RescheduleCalculatorSettings(
new TruncatedBinaryBackoff(
INITIAL_FLAPPING_DELAY.get(), MAX_FLAPPING_DELAY.get()),
FLAPPING_THRESHOLD.get(),
MAX_RESCHEDULING_DELAY.get()));
bind(RescheduleCalculator.class)
.to(RescheduleCalculatorImpl.class)
.in(Singleton.class);
if (ENABLE_PREEMPTOR.get()) {
bind(PREEMPTOR_KEY).to(PreemptorImpl.class);
bind(PreemptorImpl.class).in(Singleton.class);
LOG.info("Preemptor Enabled.");
} else {
bind(PREEMPTOR_KEY).toInstance(NULL_PREEMPTOR);
LOG.warning("Preemptor Disabled.");
}
expose(PREEMPTOR_KEY);
bind(new TypeLiteral<Amount<Long, Time>>() {})
.annotatedWith(PreemptionDelay.class)
.toInstance(PREEMPTION_DELAY.get());
bind(TaskGroups.class).in(Singleton.class);
expose(TaskGroups.class);
}
});
bindTaskScheduler(binder(), PREEMPTOR_KEY, RESERVATION_DURATION.get());
PubsubEventModule.bindSubscriber(binder(), TaskGroups.class);
binder()
.install(
new PrivateModule() {
@Override
protected void configure() {
bind(OfferReturnDelay.class).to(RandomJitterReturnDelay.class);
bind(ScheduledExecutorService.class).toInstance(executor);
bind(OfferQueue.class).to(OfferQueueImpl.class);
bind(OfferQueueImpl.class).in(Singleton.class);
expose(OfferQueue.class);
}
});
PubsubEventModule.bindSubscriber(binder(), OfferQueue.class);
binder()
.install(
new PrivateModule() {
@Override
protected void configure() {
// TODO(ksweeney): Create a configuration validator module so this can be injected.
// TODO(William Farner): Revert this once large task counts is cheap ala
// hierarchichal store
bind(Integer.class).annotatedWith(PruneThreshold.class).toInstance(100);
bind(new TypeLiteral<Amount<Long, Time>>() {})
.annotatedWith(PruneThreshold.class)
.toInstance(HISTORY_PRUNE_THRESHOLD.get());
bind(ScheduledExecutorService.class).toInstance(executor);
bind(HistoryPruner.class).in(Singleton.class);
expose(HistoryPruner.class);
}
});
PubsubEventModule.bindSubscriber(binder(), HistoryPruner.class);
}
/**
* This method exists because we want to test the wiring up of TaskSchedulerImpl class to the
* PubSub system in the TaskSchedulerImplTest class. The method has a complex signature because
* the binding of the TaskScheduler and friends occurs in a PrivateModule which does not interact
* well with the MultiBinder that backs the PubSub system.
*/
@VisibleForTesting
static void bindTaskScheduler(
Binder binder,
final Key<Preemptor> preemptorKey,
final Amount<Long, Time> reservationDuration) {
binder.install(
new PrivateModule() {
@Override
protected void configure() {
bind(Preemptor.class).to(preemptorKey);
bind(new TypeLiteral<Amount<Long, Time>>() {})
.annotatedWith(ReservationDuration.class)
.toInstance(reservationDuration);
bind(TaskScheduler.class).to(TaskSchedulerImpl.class);
bind(TaskSchedulerImpl.class).in(Singleton.class);
expose(TaskScheduler.class);
}
});
PubsubEventModule.bindSubscriber(binder, TaskScheduler.class);
}
/** Returns offers after a random duration within a fixed window. */
private static class RandomJitterReturnDelay implements OfferReturnDelay {
private static final int JITTER_WINDOW_MS = Amount.of(1, Time.MINUTES).as(Time.MILLISECONDS);
private final int minHoldTimeMs = MIN_OFFER_HOLD_TIME.get().as(Time.MILLISECONDS);
private final Random random = new Random.SystemRandom(new java.util.Random());
@Override
public Amount<Integer, Time> get() {
return Amount.of(minHoldTimeMs + random.nextInt(JITTER_WINDOW_MS), Time.MILLISECONDS);
}
}
}
/**
* Convenience class to export statistics about the JVM.
*
* <p>TODO(William Farner): Some of these are fixed values, make sure to export them so that they
* are not collected as time series.
*
* @author William Farner
*/
public class JvmStats {
private static final long BYTES_PER_MB = Amount.of(1L, Data.MB).as(Data.BYTES);
private static final double SECS_PER_NANO =
((double) 1) / Amount.of(1L, Time.SECONDS).as(Time.NANOSECONDS);
private JvmStats() {
// Utility class.
}
/** Exports stats related to the JVM and runtime environment. */
public static void export() {
final OperatingSystemMXBean osMbean = ManagementFactory.getOperatingSystemMXBean();
if (osMbean instanceof com.sun.management.OperatingSystemMXBean) {
final com.sun.management.OperatingSystemMXBean sunOsMbean =
(com.sun.management.OperatingSystemMXBean) osMbean;
Stats.exportAll(
ImmutableList.<Stat<? extends Number>>builder()
.add(
new StatImpl<Long>("system_free_physical_memory_mb") {
@Override
public Long read() {
return sunOsMbean.getFreePhysicalMemorySize() / BYTES_PER_MB;
}
})
.add(
new StatImpl<Long>("system_free_swap_mb") {
@Override
public Long read() {
return sunOsMbean.getFreeSwapSpaceSize() / BYTES_PER_MB;
}
})
.add(
Rate.of(
new StatImpl<Long>("process_cpu_time_nanos") {
@Override
public Long read() {
return sunOsMbean.getProcessCpuTime();
}
})
.withName("process_cpu_cores_utilized")
.withScaleFactor(SECS_PER_NANO)
.build())
.build());
}
if (osMbean instanceof com.sun.management.UnixOperatingSystemMXBean) {
final com.sun.management.UnixOperatingSystemMXBean unixOsMbean =
(com.sun.management.UnixOperatingSystemMXBean) osMbean;
Stats.exportAll(
ImmutableList.<Stat<? extends Number>>builder()
.add(
new StatImpl<Long>("process_max_fd_count") {
@Override
public Long read() {
return unixOsMbean.getMaxFileDescriptorCount();
}
})
.add(
new StatImpl<Long>("process_open_fd_count") {
@Override
public Long read() {
return unixOsMbean.getOpenFileDescriptorCount();
}
})
.build());
}
final Runtime runtime = Runtime.getRuntime();
final ClassLoadingMXBean classLoadingBean = ManagementFactory.getClassLoadingMXBean();
MemoryMXBean memoryBean = ManagementFactory.getMemoryMXBean();
final MemoryUsage heapUsage = memoryBean.getHeapMemoryUsage();
final MemoryUsage nonHeapUsage = memoryBean.getNonHeapMemoryUsage();
final ThreadMXBean threads = ManagementFactory.getThreadMXBean();
final RuntimeMXBean runtimeMXBean = ManagementFactory.getRuntimeMXBean();
Stats.exportAll(
ImmutableList.<Stat<? extends Number>>builder()
.add(
new StatImpl<Long>("jvm_time_ms") {
@Override
public Long read() {
return System.currentTimeMillis();
}
})
.add(
new StatImpl<Integer>("jvm_available_processors") {
@Override
public Integer read() {
return runtime.availableProcessors();
}
})
.add(
new StatImpl<Long>("jvm_memory_free_mb") {
@Override
public Long read() {
return runtime.freeMemory() / BYTES_PER_MB;
}
})
.add(
new StatImpl<Long>("jvm_memory_max_mb") {
@Override
public Long read() {
return runtime.maxMemory() / BYTES_PER_MB;
}
})
.add(
new StatImpl<Long>("jvm_memory_mb_total") {
@Override
public Long read() {
return runtime.totalMemory() / BYTES_PER_MB;
}
})
.add(
new StatImpl<Integer>("jvm_class_loaded_count") {
@Override
public Integer read() {
return classLoadingBean.getLoadedClassCount();
}
})
.add(
new StatImpl<Long>("jvm_class_total_loaded_count") {
@Override
public Long read() {
return classLoadingBean.getTotalLoadedClassCount();
}
})
.add(
new StatImpl<Long>("jvm_class_unloaded_count") {
@Override
public Long read() {
return classLoadingBean.getUnloadedClassCount();
}
})
.add(
new StatImpl<Long>("jvm_gc_collection_time_ms") {
@Override
public Long read() {
long collectionTimeMs = 0;
for (GarbageCollectorMXBean bean :
ManagementFactory.getGarbageCollectorMXBeans()) {
collectionTimeMs += bean.getCollectionTime();
}
return collectionTimeMs;
}
})
.add(
new StatImpl<Long>("jvm_gc_collection_count") {
@Override
public Long read() {
long collections = 0;
for (GarbageCollectorMXBean bean :
ManagementFactory.getGarbageCollectorMXBeans()) {
collections += bean.getCollectionCount();
}
return collections;
}
})
.add(
new StatImpl<Long>("jvm_memory_heap_mb_used") {
@Override
public Long read() {
return heapUsage.getUsed() / BYTES_PER_MB;
}
})
.add(
new StatImpl<Long>("jvm_memory_heap_mb_committed") {
@Override
public Long read() {
return heapUsage.getCommitted() / BYTES_PER_MB;
}
})
.add(
new StatImpl<Long>("jvm_memory_heap_mb_max") {
@Override
public Long read() {
return heapUsage.getMax() / BYTES_PER_MB;
}
})
.add(
new StatImpl<Long>("jvm_memory_non_heap_mb_used") {
@Override
public Long read() {
return nonHeapUsage.getUsed() / BYTES_PER_MB;
}
})
.add(
new StatImpl<Long>("jvm_memory_non_heap_mb_committed") {
@Override
public Long read() {
return nonHeapUsage.getCommitted() / BYTES_PER_MB;
}
})
.add(
new StatImpl<Long>("jvm_memory_non_heap_mb_max") {
@Override
public Long read() {
return nonHeapUsage.getMax() / BYTES_PER_MB;
}
})
.add(
new StatImpl<Long>("jvm_uptime_secs") {
@Override
public Long read() {
return runtimeMXBean.getUptime() / 1000;
}
})
.add(
new StatImpl<Double>("system_load_avg") {
@Override
public Double read() {
return osMbean.getSystemLoadAverage();
}
})
.add(
new StatImpl<Integer>("jvm_threads_peak") {
@Override
public Integer read() {
return threads.getPeakThreadCount();
}
})
.add(
new StatImpl<Long>("jvm_threads_started") {
@Override
public Long read() {
return threads.getTotalStartedThreadCount();
}
})
.add(
new StatImpl<Integer>("jvm_threads_daemon") {
@Override
public Integer read() {
return threads.getDaemonThreadCount();
}
})
.add(
new StatImpl<Integer>("jvm_threads_active") {
@Override
public Integer read() {
return threads.getThreadCount();
}
})
.build());
// Export per memory pool gc time and cycle count like Ostrich
// This is based on code in Bridcage: https://cgit.twitter.biz/birdcage/tree/ \
// ostrich/src/main/scala/com/twitter/ostrich/stats/StatsCollection.scala
Stats.exportAll(
Iterables.transform(
ManagementFactory.getGarbageCollectorMXBeans(),
new Function<GarbageCollectorMXBean, Stat<? extends Number>>() {
@Override
public Stat<? extends Number> apply(final GarbageCollectorMXBean gcMXBean) {
return new StatImpl<Long>(
"jvm_gc_" + Stats.normalizeName(gcMXBean.getName()) + "_collection_count") {
@Override
public Long read() {
return gcMXBean.getCollectionCount();
}
};
}
}));
Stats.exportAll(
Iterables.transform(
ManagementFactory.getGarbageCollectorMXBeans(),
new Function<GarbageCollectorMXBean, Stat<? extends Number>>() {
@Override
public Stat<? extends Number> apply(final GarbageCollectorMXBean gcMXBean) {
return new StatImpl<Long>(
"jvm_gc_" + Stats.normalizeName(gcMXBean.getName()) + "_collection_time_ms") {
@Override
public Long read() {
return gcMXBean.getCollectionTime();
}
};
}
}));
Stats.exportString(
new StatImpl<String>("jvm_input_arguments") {
@Override
public String read() {
return runtimeMXBean.getInputArguments().toString();
}
});
for (final String property : System.getProperties().stringPropertyNames()) {
Stats.exportString(
new StatImpl<String>("jvm_prop_" + Stats.normalizeName(property)) {
@Override
public String read() {
return System.getProperty(property);
}
});
}
}
}
@Test(expected = IllegalArgumentException.class)
public void testMaximumBackoffLessThanInitialBackoffRejected() {
new TruncatedBinaryBackoff(Amount.of(2L, Time.SECONDS), Amount.of(1L, Time.SECONDS));
}
@Test(expected = NullPointerException.class)
public void testNullMaximumBackoffRejected() {
new TruncatedBinaryBackoff(Amount.of(1L, Time.SECONDS), null);
}
@Test(expected = IllegalArgumentException.class)
public void testNegativeInitialBackoffRejected() {
new TruncatedBinaryBackoff(Amount.of(-1L, Time.SECONDS), Amount.of(1L, Time.SECONDS));
}
@Test(expected = NullPointerException.class)
public void testNullInitialBackoffRejected() {
new TruncatedBinaryBackoff(null, Amount.of(1L, Time.SECONDS));
}
public GroupTest() {
super(Amount.of(1, Time.DAYS));
}
public class RecoveryTest extends EasyMockTest {
private static final Amount<Long, Time> INTERVAL = Amount.of(1L, Time.HOURS);
private static final ScheduledTask TASK1 = makeTask("task1");
private static final ScheduledTask TASK2 = makeTask("task2");
private static final Snapshot SNAPSHOT1 = makeSnapshot(TASK1, TASK2);
private SnapshotStore<Snapshot> snapshotStore;
private DistributedSnapshotStore distributedStore;
private Storage primaryStorage;
private MutableStoreProvider storeProvider;
private Command shutDownNow;
private FakeClock clock;
private StorageBackupImpl storageBackup;
private RecoveryImpl recovery;
@Before
public void setUp() {
final File backupDir = FileUtils.createTempDir();
addTearDown(
new TearDown() {
@Override
public void tearDown() throws Exception {
org.apache.commons.io.FileUtils.deleteDirectory(backupDir);
}
});
snapshotStore = createMock(new Clazz<SnapshotStore<Snapshot>>() {});
distributedStore = createMock(DistributedSnapshotStore.class);
primaryStorage = createMock(Storage.class);
storeProvider = createMock(MutableStoreProvider.class);
shutDownNow = createMock(Command.class);
clock = new FakeClock();
TemporaryStorageFactory factory = new TemporaryStorageFactory();
storageBackup =
new StorageBackupImpl(snapshotStore, clock, new BackupConfig(backupDir, 5, INTERVAL));
recovery = new RecoveryImpl(backupDir, factory, primaryStorage, distributedStore, shutDownNow);
}
@Test
public void testRecover() throws Exception {
expect(snapshotStore.createSnapshot()).andReturn(SNAPSHOT1);
Capture<MutateWork<Object, Exception>> transaction = createCapture();
expect(primaryStorage.write(capture(transaction))).andReturn(null);
distributedStore.persist(SNAPSHOT1);
shutDownNow.execute();
control.replay();
assertEquals(ImmutableSet.<String>of(), recovery.listBackups());
clock.advance(INTERVAL);
storageBackup.createSnapshot();
String backup1 = storageBackup.createBackupName();
assertEquals(ImmutableSet.of(backup1), recovery.listBackups());
recovery.stage(backup1);
assertEquals(
IScheduledTask.setFromBuilders(SNAPSHOT1.getTasks()), recovery.query(Query.unscoped()));
recovery.commit();
transaction.getValue().apply(storeProvider);
}
@Test
public void testModifySnapshotBeforeCommit() throws Exception {
expect(snapshotStore.createSnapshot()).andReturn(SNAPSHOT1);
Snapshot modified = SNAPSHOT1.deepCopy().setTasks(ImmutableSet.of(TASK1));
Capture<MutateWork<Object, Exception>> transaction = createCapture();
expect(primaryStorage.write(capture(transaction))).andReturn(null);
distributedStore.persist(modified);
shutDownNow.execute();
control.replay();
clock.advance(INTERVAL);
storageBackup.createSnapshot();
String backup1 = storageBackup.createBackupName();
recovery.stage(backup1);
assertEquals(
IScheduledTask.setFromBuilders(SNAPSHOT1.getTasks()), recovery.query(Query.unscoped()));
recovery.deleteTasks(Query.taskScoped(Tasks.id(TASK2)));
assertEquals(
IScheduledTask.setFromBuilders(modified.getTasks()), recovery.query(Query.unscoped()));
recovery.commit();
transaction.getValue().apply(storeProvider);
}
@Test(expected = RecoveryException.class)
public void testLoadUnknownBackup() throws Exception {
control.replay();
recovery.stage("foo");
}
@Test(expected = RecoveryException.class)
public void backupNotLoaded() throws Exception {
control.replay();
recovery.commit();
}
private static Snapshot makeSnapshot(ScheduledTask... tasks) {
return new Snapshot()
.setHostAttributes(ImmutableSet.<HostAttributes>of())
.setJobs(ImmutableSet.<StoredJob>of())
.setSchedulerMetadata(new SchedulerMetadata().setVersion(CURRENT_API_VERSION))
.setQuotaConfigurations(ImmutableSet.<QuotaConfiguration>of())
.setTasks(ImmutableSet.<ScheduledTask>builder().add(tasks).build())
.setLocks(ImmutableSet.<Lock>of())
.setJobUpdateDetails(ImmutableSet.<JobUpdateDetails>of());
}
private static ScheduledTask makeTask(String taskId) {
return new ScheduledTask()
.setAssignedTask(
new AssignedTask()
.setTaskId(taskId)
.setTask(
new TaskConfig()
.setJobName("job-" + taskId)
.setEnvironment("test")
.setOwner(
new Identity().setRole("role-" + taskId).setUser("user-" + taskId))));
}
}
@Override
public Amount<Integer, Time> get() {
return Amount.of(minHoldTimeMs + random.nextInt(JITTER_WINDOW_MS), Time.MILLISECONDS);
}