src/main/java/org/apache/aurora/scheduler/storage/Storage.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.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
import java.util.Optional;

import javax.inject.Qualifier;

import org.apache.aurora.scheduler.base.Query.Builder;
import org.apache.aurora.scheduler.base.SchedulerException;
import org.apache.aurora.scheduler.storage.entities.IJobConfiguration;
import org.apache.aurora.scheduler.storage.entities.IScheduledTask;

/**
 * Manages scheduler storage operations providing an interface to perform atomic changes.
 */
public interface Storage {

  /**
   * Provider for read-only stores.  Store implementations must be thread-safe, and should support
   * concurrent reads where appropriate.
   */
  interface StoreProvider {
    SchedulerStore getSchedulerStore();
    CronJobStore getCronJobStore();
    TaskStore getTaskStore();
    QuotaStore getQuotaStore();
    AttributeStore getAttributeStore();
    JobUpdateStore getJobUpdateStore();
    HostMaintenanceStore getHostMaintenanceStore();
  }

  /**
   * Provider for stores that permit mutations.  Store implementations need not support concurrent
   * writes, as a global reentrant write lock is used to serialize write operations.
   */
  interface MutableStoreProvider extends StoreProvider {
    SchedulerStore.Mutable getSchedulerStore();
    CronJobStore.Mutable getCronJobStore();

    /**
     * Gets access to the mutable task store.
     * <p>
     * This is labeled as unsafe, since it's rare that a caller should be using this.  In most
     * cases, mutations to the task store should be done through
     * {@link org.apache.aurora.scheduler.state.StateManager}.
     * <p>
     * TODO(William Farner): Come up with a way to restrict access to this interface.  As it stands,
     * it's trivial for an unsuspecting caller to modify the task store directly and subvert the
     * state machine and side effect systems.
     *
     * @return The mutable task store.
     */
    TaskStore.Mutable getUnsafeTaskStore();

    QuotaStore.Mutable getQuotaStore();
    AttributeStore.Mutable getAttributeStore();
    JobUpdateStore.Mutable getJobUpdateStore();
    HostMaintenanceStore.Mutable getHostMaintenanceStore();
  }

  /**
   * A unit of work to perform against the storage system.
   *
   * @param <S> Type of stores the work needs access to.
   * @param <T> Return type of the operation.
   * @param <E> Exception type thrown by the operation.
   */
  @FunctionalInterface
  interface StorageOperation<S extends StoreProvider, T, E extends Exception> {
    /**
     * Abstracts a unit of work that has a result, but may also throw a specific exception.
     *
     * @param storeProvider A provider to give access to different available stores.
     * @return the result of the successfully completed unit of work
     * @throws E if the unit of work could not be completed
     */
    T apply(S storeProvider) throws E;
  }

  /**
   * Encapsulates a read-only storage operation.
   *
   * @param <T> The type of result this unit of work produces.
   * @param <E> The type of exception this unit of work can throw.
   */
  @FunctionalInterface
  interface Work<T, E extends Exception> extends StorageOperation<StoreProvider, T, E> {

    /**
     * A convenient typedef for Work that throws no checked exceptions - it runs quietly.
     *
     * @param <T> The type of result this unit of work produces.
     */
    @FunctionalInterface
    interface Quiet<T> extends Work<T, RuntimeException> {
      // typedef
    }
  }

  /**
   * Encapsulates a storage operation, which has mutable storage access.
   *
   * @param <T> The type of result this unit of work produces.
   * @param <E> The type of exception this unit of work can throw.
   */
  @FunctionalInterface
  interface MutateWork<T, E extends Exception>
      extends StorageOperation<MutableStoreProvider, T, E> {

    /**
     * A convenient typedef for Work that throws no checked exceptions - it runs quietly.
     *
     * @param <T> The type of result this unit of work produces.
     */
    @FunctionalInterface
    interface Quiet<T> extends MutateWork<T, RuntimeException> {
      // typedef
    }

    /**
     * Encapsulates work that returns no result.
     *
     * @param <E> The type of exception this unit of work can throw.
     */
    @FunctionalInterface
    interface NoResult<E extends Exception> extends MutateWork<Void, E> {

      @Override
      default Void apply(MutableStoreProvider storeProvider) throws E {
        execute(storeProvider);
        return null;
      }

      /**
       * Similar to {@link #apply(MutableStoreProvider)} except that no result is
       * returned.
       *
       * @param storeProvider A provider to give access to different available stores.
       * @throws E if the unit of work could not be completed
       */
      void execute(MutableStoreProvider storeProvider) throws E;

      /**
       * A convenient typedef for Work with no result that throws no checked exceptions - it runs
       * quitely.
       */
      @FunctionalInterface
      interface Quiet extends NoResult<RuntimeException> {
        // typedef
      }
    }
  }

  /**
   * Indicates a problem reading from or writing to stable storage.
   */
  class StorageException extends SchedulerException {
    public StorageException(String message, Throwable cause) {
      super(message, cause);
    }

    public StorageException(String message) {
      super(message);
    }
  }

  /**
   * Indicates that stable storage is temporarily unavailable.
   */
  class TransientStorageException extends StorageException {
    public TransientStorageException(String message, Throwable cause) {
      super(message, cause);
    }

    public TransientStorageException(String message) {
      super(message);
    }
  }

  /**
   * Executes the unit of read-only {@code work}.  The consistency model creates the possibility
   * for a reader to read uncommitted state from a concurrent writer.
   * <p>
   * TODO(wfarner): This method no longer needs to exist now that there is no global locking for
   * reads.  We could instead directly inject the individual stores where they are used, as long
   * as the stores have a layer to replicate what is currently done by
   * {@link CallOrderEnforcingStorage}.
   *
   * @param work The unit of work to execute.
   * @param <T> The type of result this unit of work produces.
   * @param <E> The type of exception this unit of work can throw.
   * @return the result when the unit of work completes successfully
   * @throws StorageException if there was a problem reading from stable storage.
   * @throws E bubbled transparently when the unit of work throws
   */
  <T, E extends Exception> T read(Work<T, E> work) throws StorageException, E;

  /**
   * Executes the unit of mutating {@code work}.
   * TODO(wfarner): Add a mechanism by which mutating work can add side-effect operations to be
   * performed after completion of the outer-most transaction.  As it stands, it's somewhat
   * futile to try to achieve this within a transaction, since the local code does not know
   * if the current transaction is nested.
   *
   * @param work The unit of work to execute.
   * @param <T> The type of result this unit of work produces.
   * @param <E> The type of exception this unit of work can throw.
   * @return the result when the unit of work completes successfully
   * @throws StorageException if there was a problem reading from or writing to stable storage.
   * @throws E bubbled transparently when the unit of work throws
   */
  <T, E extends Exception> T write(MutateWork<T, E> work) throws StorageException, E;

  /**
   * Requests the underlying storage prepare its data set; ie: initialize schemas, begin syncing
   * out of date data, etc.  This method should not block.
   *
   * @throws StorageException if there was a problem preparing storage.
   */
  void prepare() throws StorageException;

  /**
   * A non-volatile storage that has additional methods to control its lifecycle.
   */
  interface NonVolatileStorage extends Storage {
    /**
     * Prepares the underlying storage for serving traffic.
     *
     * @param initializationLogic work to perform after this storage system is ready but before
     *     allowing general use of
     *     {@link #read}.
     * @throws StorageException if there was a starting storage.
     */
    void start(MutateWork.NoResult.Quiet initializationLogic) throws StorageException;

    /**
     * Prepares the underlying storage system for clean shutdown.
     */
    void stop();
  }

  /**
   * Identifies a storage layer that is in-memory only.
   * This generally should only be used when the storage is first starting up, to perform queries
   * related to initially load the storage.
   */
  @Retention(RetentionPolicy.RUNTIME)
  @Target({ ElementType.PARAMETER, ElementType.METHOD })
  @Qualifier
  @interface Volatile { }

  /**
   * Utility functions for interacting with a Storage instance.
   */
  final class Util {

    private Util() {
      // Utility class.
    }

    /**
     * Fetch tasks matching the query returned by {@code query} from {@code storage} in a
     * read operation.
     *
     * @see #fetchTasks
     * @param storage Storage instance to query from.
     * @param query Builder of the query to perform.
     * @return Tasks returned from the query.
     */
    public static Iterable<IScheduledTask> fetchTasks(Storage storage, Builder query) {
      return storage.read(storeProvider -> storeProvider.getTaskStore().fetchTasks(query));
    }

    public static Optional<IScheduledTask> fetchTask(Storage storage, String taskId) {
      return storage.read(storeProvider -> storeProvider.getTaskStore().fetchTask(taskId));
    }

    public static Iterable<IJobConfiguration> fetchCronJobs(Storage storage) {
      return storage.read(storeProvider -> storeProvider.getCronJobStore().fetchJobs());
    }
  }
}