src/main/java/org/apache/aurora/scheduler/storage/durability/DurableStorage.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.durability;

import java.util.List;
import java.util.concurrent.locks.ReentrantLock;

import javax.inject.Inject;

import org.apache.aurora.common.inject.TimedInterceptor.Timed;
import org.apache.aurora.common.stats.SlidingStats;
import org.apache.aurora.gen.storage.Op;
import org.apache.aurora.scheduler.base.SchedulerException;
import org.apache.aurora.scheduler.events.EventSink;
import org.apache.aurora.scheduler.storage.AttributeStore;
import org.apache.aurora.scheduler.storage.CronJobStore;
import org.apache.aurora.scheduler.storage.HostMaintenanceStore;
import org.apache.aurora.scheduler.storage.JobUpdateStore;
import org.apache.aurora.scheduler.storage.QuotaStore;
import org.apache.aurora.scheduler.storage.SchedulerStore;
import org.apache.aurora.scheduler.storage.Storage;
import org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult;
import org.apache.aurora.scheduler.storage.Storage.NonVolatileStorage;
import org.apache.aurora.scheduler.storage.TaskStore;
import org.apache.aurora.scheduler.storage.durability.Persistence.PersistenceException;
import org.slf4j.LoggerFactory;

import static java.util.Objects.requireNonNull;

/**
 * A storage implementation that ensures storage mutations are written to a persistence layer.
 *
 * <p>In the classic write-ahead log usage we'd perform mutations as follows:
 * <ol>
 *   <li>record op</li>
 *   <li>perform op locally</li>
 *   <li>persist ops</li>
 * </ol>
 *
 * <p>Writing the operation to persistences ensures we have a record of our mutation in case we
 * should need to recover state later after a crash or on a new host (assuming the scheduler is
 * distributed).  We then apply the mutation to a local (in-memory) data structure for serving fast
 * read requests.
 *
 * <p>This implementation leverages a local transaction to handle this:
 * <ol>
 *   <li>start local transaction</li>
 *   <li>perform op locally (uncommitted!)</li>
 *   <li>write op to persistence</li>
 * </ol>
 *
 * <p>If the op fails to apply to local storage we will never persist the op, and if the op
 * fails to persist, it'll throw and abort the local storage operation as well.
 */
public class DurableStorage implements NonVolatileStorage {

  /**
   * A maintainer for context about open transactions. Assumes that an external entity is
   * responsible for opening and closing transactions.
   */
  interface TransactionManager {

    /**
     * Checks whether there is an open transaction.
     *
     * @return {@code true} if there is an open transaction, {@code false} otherwise.
     */
    boolean hasActiveTransaction();

    /**
     * Adds an operation to the existing transaction.
     *
     * @param op Operation to include in the existing transaction.
     */
    void log(Op op);
  }

  private final Persistence persistence;
  private final Storage writeBehindStorage;
  private final ReentrantLock writeLock;
  private final ThriftBackfill thriftBackfill;

  private final WriteRecorder writeRecorder;

  private TransactionRecorder transaction = null;

  private final SlidingStats writerWaitStats = new SlidingStats("storage_write_lock_wait", "ns");

  @Inject
  DurableStorage(
      Persistence persistence,
      @Volatile Storage delegateStorage,
      @Volatile SchedulerStore.Mutable schedulerStore,
      @Volatile CronJobStore.Mutable jobStore,
      @Volatile TaskStore.Mutable taskStore,
      @Volatile QuotaStore.Mutable quotaStore,
      @Volatile AttributeStore.Mutable attributeStore,
      @Volatile JobUpdateStore.Mutable jobUpdateStore,
      @Volatile HostMaintenanceStore.Mutable hostMaintenanceStore,
      EventSink eventSink,
      ReentrantLock writeLock,
      ThriftBackfill thriftBackfill) {

    this.persistence = requireNonNull(persistence);

    // DurableStorage has two distinct operating modes: pre- and post-recovery.  When recovering,
    // we write directly to the writeBehind stores since we are replaying what's already persisted.
    // After that, all writes must succeed in Persistence before they may be considered successful.
    this.writeBehindStorage = requireNonNull(delegateStorage);
    this.writeLock = requireNonNull(writeLock);
    this.thriftBackfill = requireNonNull(thriftBackfill);
    TransactionManager transactionManager = new TransactionManager() {
      @Override
      public boolean hasActiveTransaction() {
        return transaction != null;
      }

      @Override
      public void log(Op op) {
        transaction.add(op);
      }
    };
    this.writeRecorder = new WriteRecorder(
        transactionManager,
        schedulerStore,
        jobStore,
        taskStore,
        quotaStore,
        attributeStore,
        jobUpdateStore,
        hostMaintenanceStore,
        LoggerFactory.getLogger(WriteRecorder.class),
        eventSink);
  }

  @Override
  @Timed("scheduler_storage_prepare")
  public synchronized void prepare() {
    writeBehindStorage.prepare();
    persistence.prepare();
  }

  @Override
  @Timed("scheduler_storage_start")
  public void start(final MutateWork.NoResult.Quiet initializationLogic) {
    writeLock.lock();
    try {
      // We recover directly into the forwarded system to avoid persisting replayed operations.
      writeBehindStorage.write((NoResult.Quiet) this::recover);

      // Now that we're recovered we should persist any mutations done in initializationLogic, so
      // run it in one of our transactions.
      write(initializationLogic);
    } finally {
      writeLock.unlock();
    }
  }

  @Override
  public void stop() {
    // No-op.
  }

  @Timed("scheduler_storage_recover")
  void recover(MutableStoreProvider stores) throws RecoveryFailedException {
    try {
      Loader.load(stores, thriftBackfill, persistence.recover());
    } catch (PersistenceException e) {
      throw new RecoveryFailedException(e);
    }
  }

  private static final class RecoveryFailedException extends SchedulerException {
    RecoveryFailedException(Throwable cause) {
      super(cause);
    }
  }

  private <T, E extends Exception> T doInTransaction(final MutateWork<T, E> work)
      throws StorageException, E {

    // The transaction has already been set up so we just need to delegate with our store provider
    // so any mutations may be persisted.
    if (transaction != null) {
      return work.apply(writeRecorder);
    }

    transaction = new TransactionRecorder();
    try {
      return writeBehindStorage.write(unused -> {
        T result = work.apply(writeRecorder);
        List<Op> ops = transaction.getOps();
        if (!ops.isEmpty()) {
          try {
            persistence.persist(ops.stream());
          } catch (PersistenceException e) {
            throw new StorageException("Failed to persist storage changes", e);
          }
        }
        return result;
      });
    } finally {
      transaction = null;
    }
  }

  @Override
  public <T, E extends Exception> T write(final MutateWork<T, E> work) throws StorageException, E {
    long waitStart = System.nanoTime();
    writeLock.lock();
    try {
      writerWaitStats.accumulate(System.nanoTime() - waitStart);
      return doInTransaction(work);
    } finally {
      writeLock.unlock();
    }
  }

  @Override
  public <T, E extends Exception> T read(Work<T, E> work) throws StorageException, E {
    return writeBehindStorage.read(work);
  }
}