src/test/java/org/apache/aurora/scheduler/storage/AbstractTaskStoreTest.java

/**
 * 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.storage;

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

import com.google.common.collect.FluentIterable;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import com.google.inject.Guice;
import com.google.inject.Injector;
import com.google.inject.Module;

import org.apache.aurora.common.testing.TearDownTestCase;
import org.apache.aurora.gen.AppcImage;
import org.apache.aurora.gen.Attribute;
import org.apache.aurora.gen.Container;
import org.apache.aurora.gen.DockerImage;
import org.apache.aurora.gen.HostAttributes;
import org.apache.aurora.gen.Image;
import org.apache.aurora.gen.MaintenanceMode;
import org.apache.aurora.gen.MesosContainer;
import org.apache.aurora.gen.MesosFetcherURI;
import org.apache.aurora.gen.Metadata;
import org.apache.aurora.gen.Mode;
import org.apache.aurora.gen.ScheduledTask;
import org.apache.aurora.gen.TaskQuery;
import org.apache.aurora.gen.Volume;
import org.apache.aurora.scheduler.base.JobKeys;
import org.apache.aurora.scheduler.base.Query;
import org.apache.aurora.scheduler.base.TaskTestUtil;
import org.apache.aurora.scheduler.base.Tasks;
import org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult;
import org.apache.aurora.scheduler.storage.TaskStore.Mutable.TaskMutation;
import org.apache.aurora.scheduler.storage.entities.IHostAttributes;
import org.apache.aurora.scheduler.storage.entities.IJobKey;
import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
import org.apache.aurora.scheduler.storage.testing.StorageEntityUtil;
import org.junit.Before;
import org.junit.Ignore;
import org.junit.Test;

import static org.apache.aurora.gen.ScheduleStatus.ASSIGNED;
import static org.apache.aurora.gen.ScheduleStatus.RUNNING;
import static org.apache.aurora.scheduler.base.TaskTestUtil.makeTask;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;

public abstract class AbstractTaskStoreTest extends TearDownTestCase {
  protected static final IHostAttributes HOST_A = IHostAttributes.build(
      new HostAttributes(
          "hostA",
          ImmutableSet.of(new Attribute("zone", ImmutableSet.of("1a"))))
          .setSlaveId("slaveIdA")
          .setMode(MaintenanceMode.NONE));
  protected static final IHostAttributes HOST_B = IHostAttributes.build(
      new HostAttributes(
          "hostB",
          ImmutableSet.of(new Attribute("zone", ImmutableSet.of("1a"))))
          .setSlaveId("slaveIdB")
          .setMode(MaintenanceMode.NONE));
  protected static final IScheduledTask TASK_A = createTask("a");
  protected static final IScheduledTask TASK_B =
      setContainer(createTask("b"), Container.mesos(new MesosContainer()));
  protected static final IScheduledTask TASK_C = createTask("c");
  protected static final IScheduledTask TASK_D = createTask("d");

  protected Injector injector;
  protected Storage storage;

  protected abstract Module getStorageModule();

  @Before
  public void setUp() {
    injector = Guice.createInjector(getStorageModule());
    storage = injector.getInstance(Storage.class);
    storage.prepare();

    storage.write((NoResult.Quiet) storeProvider -> {
      AttributeStore.Mutable attributeStore = storeProvider.getAttributeStore();
      attributeStore.saveHostAttributes(HOST_A);
      attributeStore.saveHostAttributes(HOST_B);
    });
  }

  private Optional<IScheduledTask> fetchTask(String taskId) {
    return storage.read(storeProvider -> storeProvider.getTaskStore().fetchTask(taskId));
  }

  private Iterable<IScheduledTask> fetchTasks(Query.Builder query) {
    return storage.read(storeProvider -> storeProvider.getTaskStore().fetchTasks(query));
  }

  protected void saveTasks(IScheduledTask... tasks) {
    saveTasks(ImmutableSet.copyOf(tasks));
  }

  private void saveTasks(Set<IScheduledTask> tasks) {
    storage.write((NoResult.Quiet)
        storeProvider -> storeProvider.getUnsafeTaskStore().saveTasks(ImmutableSet.copyOf(tasks)));
  }

  private Optional<IScheduledTask> mutateTask(String taskId, TaskMutation mutation) {
    return storage.write(
        storeProvider -> storeProvider.getUnsafeTaskStore().mutateTask(taskId, mutation));
  }

  protected void deleteTasks(String... taskIds) {
    storage.write((NoResult.Quiet) storeProvider ->
        storeProvider.getUnsafeTaskStore().deleteTasks(ImmutableSet.copyOf(taskIds)));
  }

  protected void deleteAllTasks() {
    storage.write(
        (NoResult.Quiet) storeProvider -> storeProvider.getUnsafeTaskStore().deleteAllTasks());
  }

  @Test
  public void testSave() {
    IScheduledTask aWithHost = setHost(TASK_A, HOST_A);
    StorageEntityUtil.assertFullyPopulated(aWithHost.newBuilder());

    saveTasks(aWithHost, TASK_B);
    assertStoreContents(aWithHost, TASK_B);

    saveTasks(TASK_C, TASK_D);
    assertStoreContents(aWithHost, TASK_B, TASK_C, TASK_D);

    // Saving the same task should overwrite.
    IScheduledTask taskAModified = IScheduledTask.build(aWithHost.newBuilder().setStatus(RUNNING));
    saveTasks(taskAModified);
    assertStoreContents(taskAModified, TASK_B, TASK_C, TASK_D);
  }

  @Test
  public void testSaveWithMetadata() {
    ScheduledTask builder = TASK_A.newBuilder();
    builder.getAssignedTask().getTask().setMetadata(
        ImmutableSet.of(
            new Metadata("package", "a"),
            new Metadata("package", "b")));
    IScheduledTask task = IScheduledTask.build(builder);
    saveTasks(task);
    assertStoreContents(task);
  }

  @Test
  public void testSaveWithMesosFetcherUris() {
    ScheduledTask builder = TASK_A.newBuilder();
    builder.getAssignedTask().getTask().setMesosFetcherUris(
        ImmutableSet.of(
            new MesosFetcherURI("pathA").setExtract(true).setCache(true),
            new MesosFetcherURI("pathB").setExtract(true).setCache(true)));
    IScheduledTask task = IScheduledTask.build(builder);
    saveTasks(task);
    assertStoreContents(task);
  }

  @Test
  public void testSaveWithContainerVolumes() {
    ScheduledTask builder = TASK_B.newBuilder();
    Image image = Image.docker(new DockerImage().setName("some-name").setTag("some-tag"));
    List<Volume> volumes = ImmutableList.of(
        new Volume("container2", "host2", Mode.RW),
        new Volume("container", "host", Mode.RO));
    builder.getAssignedTask().getTask().getContainer().getMesos().setImage(image)
        .setVolumes(volumes);

    IScheduledTask task = IScheduledTask.build(builder);

    saveTasks(task);
    assertStoreContents(task);
  }

  @Test
  public void testSaveWithDockerImage() {
    ScheduledTask builder = TASK_A.newBuilder();

    Image image = new Image();
    image.setDocker(new DockerImage().setName("some-name").setTag("some-tag"));
    builder.getAssignedTask().getTask().setContainer(
        Container.mesos(new MesosContainer().setImage(image)));

    IScheduledTask task = IScheduledTask.build(builder);
    saveTasks(task);
    assertStoreContents(task);
  }

  @Test
  public void testSaveWithAppcImage() {
    ScheduledTask builder = TASK_A.newBuilder();

    Image image = new Image();
    image.setAppc(new AppcImage().setName("some-name").setImageId("some-tag"));
    builder.getAssignedTask().getTask().setContainer(
        Container.mesos(new MesosContainer().setImage(image)));

    IScheduledTask task = IScheduledTask.build(builder);
    saveTasks(task);
    assertStoreContents(task);
  }

  @Test
  public void testSaveWithMultiplePorts() {
    ScheduledTask builder = TASK_A.newBuilder();
    builder.getAssignedTask().setAssignedPorts(ImmutableMap.of("http", 1000, "tcp", 2000));
    IScheduledTask task = IScheduledTask.build(builder);
    saveTasks(task);

    Optional<IScheduledTask> maybeTask = fetchTask(task.getAssignedTask().getTaskId());
    saveTasks(maybeTask.get());

    assertStoreContents(task);
  }

  @Test
  public void testQuery() {
    assertStoreContents();
    saveTasks(TASK_A, TASK_B, TASK_C, TASK_D);

    assertQueryResults(Query.taskScoped("b"), TASK_B);
    assertQueryResults(Query.taskScoped("a", "d"), TASK_A, TASK_D);
    assertQueryResults(Query.roleScoped("role-c"), TASK_C);
    assertQueryResults(Query.envScoped("role-c", "env-c"), TASK_C);
    assertQueryResults(Query.envScoped("role-c", "devel"));
    assertQueryResults(
        Query.unscoped().byStatus(ASSIGNED),
        TASK_A, TASK_B, TASK_C, TASK_D);
    assertQueryResults(
        Query.instanceScoped(JobKeys.from("role-a", "env-a", "job-a"), 2).active(), TASK_A);
    assertQueryResults(Query.jobScoped(JobKeys.from("role-b", "env-b", "job-b")).active(), TASK_B);
    assertQueryResults(Query.jobScoped(JobKeys.from("role-b", "devel", "job-b")).active());

    // Explicitly call out the matching behaviors for different means of expressing empty query
    // conditions. Specifically - neither null collections nor empty collections should trigger
    // query restrictions and thus they should yield the same results.
    assertQueryResults(new TaskQuery().setTaskIds(null), TASK_A, TASK_B, TASK_C, TASK_D);
    assertQueryResults(new TaskQuery().setTaskIds(ImmutableSet.of()),
        TASK_A, TASK_B, TASK_C, TASK_D);

    assertQueryResults(new TaskQuery().setInstanceIds(null),  TASK_A, TASK_B, TASK_C, TASK_D);
    assertQueryResults(
        new TaskQuery().setInstanceIds(ImmutableSet.of()),
        TASK_A, TASK_B, TASK_C, TASK_D);

    assertQueryResults(new TaskQuery().setStatuses(null), TASK_A, TASK_B, TASK_C, TASK_D);
    assertQueryResults(
        new TaskQuery().setStatuses(ImmutableSet.of()),
        TASK_A, TASK_B, TASK_C, TASK_D);

    assertQueryResults(new TaskQuery().setSlaveHosts(null), TASK_A, TASK_B, TASK_C, TASK_D);
    assertQueryResults(new TaskQuery().setSlaveHosts(ImmutableSet.of()),
        TASK_A, TASK_B, TASK_C, TASK_D);

    assertQueryResults(new TaskQuery().setJobKeys(null), TASK_A, TASK_B, TASK_C, TASK_D);
    assertQueryResults(new TaskQuery().setJobKeys(ImmutableSet.of()),
        TASK_A, TASK_B, TASK_C, TASK_D);
  }

  @Test
  public void testQueryMultipleInstances() {
    ImmutableSet.Builder<IScheduledTask> tasksBuilder = ImmutableSet.builder();
    for (int i = 0; i < 100; i++) {
      ScheduledTask builder = TASK_A.newBuilder();
      builder.getAssignedTask()
          .setTaskId("id" + i)
          .setInstanceId(i);
      tasksBuilder.add(IScheduledTask.build(builder));
    }
    Set<IScheduledTask> tasks = tasksBuilder.build();
    saveTasks(tasks);
    assertQueryResults(Query.unscoped(), tasks);
  }

  @Test
  public void testQueryBySlaveHost() {
    IScheduledTask a = setHost(makeTask("a", JobKeys.from("role", "env", "job")), HOST_A);
    IScheduledTask b = setHost(makeTask("b", JobKeys.from("role", "env", "job")), HOST_B);
    saveTasks(a, b);

    assertQueryResults(Query.slaveScoped(HOST_A.getHost()), a);
    assertQueryResults(Query.slaveScoped(HOST_A.getHost(), HOST_B.getHost()), a, b);
  }

  @Test
  public void testQueryByJobKeys() {
    assertStoreContents();
    saveTasks(TASK_A, TASK_B, TASK_C, TASK_D);

    assertQueryResults(
        Query.jobScoped(ImmutableSet.of(
            JobKeys.from("role-a", "env-a", "job-a"),
            JobKeys.from("role-b", "env-b", "job-b"),
            JobKeys.from("role-c", "env-c", "job-c"))),
        TASK_A, TASK_B, TASK_C);

    // Conflicting jobs will produce the result from the last added JobKey
    assertQueryResults(
        Query.jobScoped(JobKeys.from("role-a", "env-a", "job-a"))
            .byJobKeys(ImmutableSet.of(JobKeys.from("role-b", "env-b", "job-b"))),
        TASK_B);

    // The .byJobKeys will override the previous scoping and OR all of the keys.
    assertQueryResults(
        Query.jobScoped(JobKeys.from("role-a", "env-a", "job-a"))
            .byJobKeys(ImmutableSet.of(
                JobKeys.from("role-b", "env-b", "job-b"),
                JobKeys.from("role-a", "env-a", "job-a"))),
        TASK_A, TASK_B);

    // Combination of individual field and jobKeys is allowed.
    assertQueryResults(
        Query.roleScoped("role-b")
            .byJobKeys(ImmutableSet.of(
                JobKeys.from("role-b", "env-b", "job-b"),
                JobKeys.from("role-a", "env-a", "job-a"))),
        TASK_B);
  }

  @Test
  public void testMutate() {
    saveTasks(TASK_A, TASK_B, TASK_C, TASK_D);
    assertQueryResults(Query.statusScoped(RUNNING));

    mutateTask("a", task -> IScheduledTask.build(task.newBuilder().setStatus(RUNNING)));

    assertQueryResults(
        Query.statusScoped(RUNNING),
        IScheduledTask.build(TASK_A.newBuilder().setStatus(RUNNING)));

    assertEquals(
        Optional.empty(),
        mutateTask(
            "nonexistent",
            task -> IScheduledTask.build(task.newBuilder().setStatus(RUNNING))));

    assertStoreContents(
        IScheduledTask.build(TASK_A.newBuilder().setStatus(RUNNING)),
        IScheduledTask.build(TASK_B.newBuilder().setStatus(ASSIGNED)),
        IScheduledTask.build(TASK_C.newBuilder().setStatus(ASSIGNED)),
        IScheduledTask.build(TASK_D.newBuilder().setStatus(ASSIGNED)));
  }

  @Test
  public void testDelete() {
    saveTasks(TASK_A, TASK_B, TASK_C, TASK_D);
    deleteTasks("a");
    assertStoreContents(TASK_B, TASK_C, TASK_D);
    deleteTasks("c");
    assertStoreContents(TASK_B, TASK_D);
    deleteTasks("b", "d");
    assertStoreContents();
  }

  @Test
  public void testDeleteAll() {
    saveTasks(TASK_A, TASK_B, TASK_C, TASK_D);
    deleteAllTasks();
    assertStoreContents();
  }

  @Test
  public void testConsistentJobIndex() {
    IScheduledTask a = makeTask("a", JobKeys.from("jim", "test", "job"));
    IScheduledTask b = makeTask("b", JobKeys.from("jim", "test", "job"));
    IScheduledTask c = makeTask("c", JobKeys.from("jim", "test", "job2"));
    IScheduledTask d = makeTask("d", JobKeys.from("joe", "test", "job"));
    IScheduledTask e = makeTask("e", JobKeys.from("jim", "prod", "job"));
    Query.Builder jimsJob = Query.jobScoped(JobKeys.from("jim", "test", "job"));
    Query.Builder jimsJob2 = Query.jobScoped(JobKeys.from("jim", "test", "job2"));
    Query.Builder joesJob = Query.jobScoped(JobKeys.from("joe", "test", "job"));

    saveTasks(a, b, c, d, e);
    assertQueryResults(jimsJob, a, b);
    assertQueryResults(jimsJob2, c);
    assertQueryResults(joesJob, d);

    deleteTasks(Tasks.id(b));
    assertQueryResults(jimsJob, a);
    assertQueryResults(jimsJob2, c);
    assertQueryResults(joesJob, d);

    mutateTask(Tasks.id(a), task -> IScheduledTask.build(task.newBuilder().setStatus(RUNNING)));
    IScheduledTask aRunning = IScheduledTask.build(a.newBuilder().setStatus(RUNNING));
    assertQueryResults(jimsJob, aRunning);
    assertQueryResults(jimsJob2, c);
    assertQueryResults(joesJob, d);

    deleteTasks(Tasks.id(d));
    assertQueryResults(joesJob);

    deleteTasks(Tasks.id(d));
    assertQueryResults(jimsJob, aRunning);
    assertQueryResults(jimsJob2, c);
    assertQueryResults(joesJob);

    saveTasks(b);
    assertQueryResults(jimsJob, aRunning, b);
    assertQueryResults(jimsJob2, c);
    assertQueryResults(joesJob);
  }

  @Test
  public void testCanonicalTaskConfigs() {
    IScheduledTask a = createTask("a");
    IScheduledTask b = createTask("a");
    IScheduledTask c = createTask("a");
    saveTasks(a, b, c);
    Set<IScheduledTask> inserted = ImmutableSet.of(a, b, c);

    Set<ITaskConfig> storedConfigs = FluentIterable.from(fetchTasks(Query.unscoped()))
        .transform(Tasks::getConfig)
        .toSet();
    assertEquals(
        FluentIterable.from(inserted).transform(Tasks::getConfig).toSet(),
        storedConfigs);
    Map<ITaskConfig, ITaskConfig> identityMap = Maps.newIdentityHashMap();
    for (ITaskConfig stored : storedConfigs) {
      identityMap.put(stored, stored);
    }
    assertEquals(
        ImmutableMap.of(Tasks.getConfig(a), Tasks.getConfig(a)),
        identityMap);
  }

  private static IScheduledTask setHost(IScheduledTask task, IHostAttributes host) {
    ScheduledTask builder = task.newBuilder();
    builder.getAssignedTask()
        .setSlaveHost(host.getHost())
        .setSlaveId(host.getSlaveId());
    return IScheduledTask.build(builder);
  }

  private static IScheduledTask unsetHost(IScheduledTask task) {
    ScheduledTask builder = task.newBuilder();
    builder.getAssignedTask()
        .setSlaveHost(null)
        .setSlaveId(null);
    return IScheduledTask.build(builder);
  }

  private static IScheduledTask setConfigData(IScheduledTask task, String configData) {
    ScheduledTask builder = task.newBuilder();
    builder.getAssignedTask().getTask().getExecutorConfig().setData(configData);
    return IScheduledTask.build(builder);
  }

  @Test
  public void testAddSlaveHost() {
    final IScheduledTask a = createTask("a");
    saveTasks(a);
    assertQueryResults(Query.slaveScoped(HOST_A.getHost()));

    final IScheduledTask b = setHost(a, HOST_A);
    Optional<IScheduledTask> result = mutateTask(Tasks.id(a),
        task -> {
          assertEquals(a, task);
          return b;
        });
    assertEquals(Optional.of(b), result);
    assertQueryResults(Query.slaveScoped(HOST_A.getHost()), b);

    // Unrealistic behavior, but proving that the secondary index can handle key mutations.
    final IScheduledTask c = setHost(b, HOST_B);
    Optional<IScheduledTask> result2 = mutateTask(Tasks.id(a),
        task -> {
          assertEquals(b, task);
          return c;
        });
    assertEquals(Optional.of(c), result2);
    assertQueryResults(Query.slaveScoped(HOST_B.getHost()), c);

    deleteTasks(Tasks.id(a));
    assertQueryResults(Query.slaveScoped(HOST_B.getHost()));
  }

  @Test
  public void testUnsetSlaveHost() {
    // Unrealistic behavior, but proving that the secondary index does not become stale.

    final IScheduledTask a = setHost(createTask("a"), HOST_A);
    saveTasks(a);
    assertQueryResults(Query.slaveScoped(HOST_A.getHost()), a);

    final IScheduledTask b = unsetHost(a);
    Optional<IScheduledTask> result = mutateTask(Tasks.id(a),
        task -> {
          assertEquals(a, task);
          return b;
        });
    assertEquals(Optional.of(b), result);
    assertQueryResults(Query.slaveScoped(HOST_A.getHost()));
    assertQueryResults(Query.taskScoped(Tasks.id(b)), b);
  }

  @Test
  public void testTasksOnSameHost() {
    final IScheduledTask a = setHost(createTask("a"), HOST_A);
    final IScheduledTask b = setHost(createTask("b"), HOST_A);
    saveTasks(a, b);
    assertQueryResults(Query.slaveScoped(HOST_A.getHost()), a, b);
  }

  @Test
  public void testSaveOverwrites() {
    // Ensures that saving a task with an existing task ID is effectively the same as a mutate,
    // and does not result in a duplicate object in the primary or secondary index.

    final IScheduledTask a = setHost(createTask("a"), HOST_A);
    saveTasks(a);

    final IScheduledTask updated = setConfigData(a, "new config data");
    saveTasks(updated);
    assertQueryResults(Query.taskScoped(Tasks.id(a)), updated);
    assertQueryResults(Query.slaveScoped(HOST_A.getHost()), updated);
  }

  private Set<IJobKey> getJobKeys() {
    return storage.read(storeProvider -> storeProvider.getTaskStore().getJobKeys());
  }

  private Set<IJobKey> toJobKeys(IScheduledTask... tasks) {
    return FluentIterable.from(ImmutableSet.copyOf(tasks))
        .transform(Tasks::getJob)
        .toSet();
  }

  @Test
  public void testGetsJobKeys() {
    assertEquals(ImmutableSet.of(), getJobKeys());
    saveTasks(TASK_A);
    assertEquals(toJobKeys(TASK_A), getJobKeys());
    saveTasks(TASK_B, TASK_C);
    assertEquals(toJobKeys(TASK_A, TASK_B, TASK_C), getJobKeys());
    deleteTasks(Tasks.id(TASK_B));
    assertEquals(toJobKeys(TASK_A, TASK_C), getJobKeys());
    IJobKey multiInstanceJob = JobKeys.from("role", "env", "instances");
    saveTasks(
        makeTask("instance1", multiInstanceJob),
        makeTask("instance2", multiInstanceJob),
        makeTask("instance3", multiInstanceJob));
    assertEquals(
        ImmutableSet.builder().addAll(toJobKeys(TASK_A, TASK_C)).add(multiInstanceJob).build(),
        getJobKeys());
  }

  @Ignore
  @Test
  public void testReadSecondaryIndexMultipleThreads() throws Exception {
    ExecutorService executor = Executors.newFixedThreadPool(4,
        new ThreadFactoryBuilder().setNameFormat("SlowRead-%d").setDaemon(true).build());

    try {
      ImmutableSet.Builder<IScheduledTask> builder = ImmutableSet.builder();
      final int numTasks = 100;
      final int numJobs = 100;
      for (int j = 0; j < numJobs; j++) {
        for (int t = 0; t < numTasks; t++) {
          builder.add(makeTask("" + j + "-" + t, JobKeys.from("role", "env", "name" + j)));
        }
      }
      saveTasks(builder.build());

      final CountDownLatch read = new CountDownLatch(numJobs);
      for (int j = 0; j < numJobs; j++) {
        final int id = j;
        executor.submit(() -> {
          assertNotNull(fetchTasks(Query.jobScoped(JobKeys.from("role", "env", "name" + id))));
          read.countDown();
        });
        executor.submit(() -> saveTasks(createTask("TaskNew1" + id)));
      }

      read.await();
    } finally {
      MoreExecutors.shutdownAndAwaitTermination(executor, 1, TimeUnit.SECONDS);
    }
  }

  @Test
  public void testNullVsEmptyRelations() throws Exception {
    // Test for regression of AURORA-1476.

    ITaskConfig nullMetadata =
        ITaskConfig.build(TaskTestUtil.makeConfig(TaskTestUtil.JOB).newBuilder().setMetadata(null));

    IScheduledTask a = makeTask("a", nullMetadata);
    IScheduledTask b = makeTask("a", nullMetadata);
    IScheduledTask c = makeTask("a", nullMetadata);
    saveTasks(a);
    saveTasks(b);
    saveTasks(c);
  }

  @Test
  public void testConcurrentFetchTasks() throws Exception {
    // Test for regression of AURORA-1625
    ListeningExecutorService executor =
        MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(5));
    assertStoreContents();
    saveTasks(TASK_A, TASK_B, TASK_C, TASK_D);

    List<ListenableFuture<Integer>> futures = Lists.newArrayList();

    for (int i = 0; i < 100; i++) {
      futures.add(executor.submit(() -> Iterables.size(fetchTasks(Query.unscoped()))));
    }

    Future<List<Integer>> f = Futures.allAsList(futures);

    executor.shutdown();
    executor.awaitTermination(1, TimeUnit.MINUTES);

    assertEquals(Iterables.getOnlyElement(ImmutableSet.copyOf(f.get())), (Integer) 4);
  }

  private void assertStoreContents(IScheduledTask... tasks) {
    assertQueryResults(Query.unscoped(), tasks);
  }

  private void assertQueryResults(TaskQuery query, IScheduledTask... tasks) {
    assertQueryResults(Query.arbitrary(query), tasks);
  }

  private void assertQueryResults(Query.Builder query, IScheduledTask... tasks) {
    assertQueryResults(query, ImmutableSet.copyOf(tasks));
  }

  private void assertQueryResults(Query.Builder query, Set<IScheduledTask> tasks) {
    Iterable<IScheduledTask> result = fetchTasks(query);
    assertQueryHasNoDupes(result);
    Set<IScheduledTask> set = ImmutableSet.copyOf(result);
    assertEquals(tasks, set);
  }

  private void assertQueryHasNoDupes(Iterable<IScheduledTask> result) {
    assertEquals(ImmutableList.copyOf(result).size(), ImmutableSet.copyOf(result).size());
  }

  private static IScheduledTask createTask(String id) {
    return makeTask(id, JobKeys.from("role-" + id, "env-" + id, "job-" + id));
  }

  private static IScheduledTask setContainer(IScheduledTask task, Container container) {
    ScheduledTask builder = task.newBuilder();
    builder.getAssignedTask().getTask().setContainer(container);
    return IScheduledTask.build(builder);
  }
}