Recovery.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.backup;

import java.io.File;
import java.io.IOException;
import java.util.Set;
import java.util.concurrent.atomic.AtomicReference;

import javax.annotation.Nullable;
import javax.inject.Inject;

import com.google.common.base.Function;
import com.google.common.collect.ImmutableSet;
import com.google.common.io.Files;
import com.google.common.util.concurrent.Atomics;

import org.apache.aurora.codec.ThriftBinaryCodec;
import org.apache.aurora.codec.ThriftBinaryCodec.CodingException;
import org.apache.aurora.common.base.Command;
import org.apache.aurora.gen.storage.Snapshot;
import org.apache.aurora.scheduler.base.Query;
import org.apache.aurora.scheduler.storage.DistributedSnapshotStore;
import org.apache.aurora.scheduler.storage.Storage;
import org.apache.aurora.scheduler.storage.Storage.MutateWork.NoResult;
import org.apache.aurora.scheduler.storage.entities.IScheduledTask;

import static java.util.Objects.requireNonNull;

/**
 * A recovery mechanism that works with {@link StorageBackup} to provide a two-step storage
 * recovery process.
 */
public interface Recovery {

  /**
   * List backups available for recovery.
   *
   * @return Available backup IDs.
   */
  Set<String> listBackups();

  /**
   * Loads a backup in 'staging' so that it may be queried and modified prior to committing.
   *
   * @param backupName Name of the backup to load.
   * @throws RecoveryException If the backup could not be found or loaded.
   */
  void stage(String backupName) throws RecoveryException;

  /**
   * Queries a staged backup.
   *
   * @param query Builder of query to perform.
   * @return Tasks matching the query.
   * @throws RecoveryException If a backup is not staged, or could not be queried.
   */
  Iterable<IScheduledTask> query(Query.Builder query) throws RecoveryException;

  /**
   * Deletes tasks from a staged backup.
   *
   * @param query Query builder for tasks to delete.
   * @throws RecoveryException If a backup is not staged, or tasks could not be deleted.
   */
  void deleteTasks(Query.Builder query) throws RecoveryException;

  /**
   * Unloads a staged backup.
   */
  void unload();

  /**
   * Commits a staged backup the main storage system.
   *
   * @throws RecoveryException If a backup is not staged, or the commit failed.
   */
  void commit() throws RecoveryException;

  /**
   * Thrown when a recovery operation could not be completed due to internal errors or improper
   * invocation order.
   */
  class RecoveryException extends RuntimeException {
    public RecoveryException(String message) {
      super(message);
    }

    public RecoveryException(String message, Throwable cause) {
      super(message, cause);
    }
  }

  class RecoveryImpl implements Recovery {
    private final File backupDir;
    private final Function<Snapshot, TemporaryStorage> tempStorageFactory;
    private final AtomicReference<PendingRecovery> recovery;
    private final Storage primaryStorage;
    private final DistributedSnapshotStore distributedStore;
    private final Command shutDownNow;

    @Inject
    RecoveryImpl(
        File backupDir,
        Function<Snapshot, TemporaryStorage> tempStorageFactory,
        Storage primaryStorage,
        DistributedSnapshotStore distributedStore,
        Command shutDownNow) {

      this.backupDir = requireNonNull(backupDir);
      this.tempStorageFactory = requireNonNull(tempStorageFactory);
      this.recovery = Atomics.newReference();
      this.primaryStorage = requireNonNull(primaryStorage);
      this.distributedStore = requireNonNull(distributedStore);
      this.shutDownNow = requireNonNull(shutDownNow);
    }

    @Override
    public Set<String> listBackups() {
      return ImmutableSet.<String>builder().add(backupDir.list()).build();
    }

    @Override
    public void stage(String backupName) throws RecoveryException {
      File backupFile = new File(backupDir, backupName);
      if (!backupFile.exists()) {
        throw new RecoveryException("Backup " + backupName + " does not exist.");
      }

      Snapshot snapshot;
      try {
        snapshot = ThriftBinaryCodec.decode(Snapshot.class, Files.toByteArray(backupFile));
      } catch (CodingException e) {
        throw new RecoveryException("Failed to decode backup " + e, e);
      } catch (IOException e) {
        throw new RecoveryException("Failed to read backup " + e, e);
      }
      boolean applied =
          recovery.compareAndSet(null, new PendingRecovery(tempStorageFactory.apply(snapshot)));
      if (!applied) {
        throw new RecoveryException("Another backup is already loaded.");
      }
    }

    private PendingRecovery getLoadedRecovery() throws RecoveryException {
      @Nullable PendingRecovery loaded = this.recovery.get();
      if (loaded == null) {
        throw new RecoveryException("No backup loaded.");
      }
      return loaded;
    }

    @Override
    public Iterable<IScheduledTask> query(Query.Builder query) throws RecoveryException {
      return getLoadedRecovery().query(query);
    }

    @Override
    public void deleteTasks(Query.Builder query) throws RecoveryException {
      getLoadedRecovery().delete(query);
    }

    @Override
    public void unload() {
      recovery.set(null);
    }

    @Override
    public void commit() throws RecoveryException {
      getLoadedRecovery().commit();
    }

    private class PendingRecovery {
      private final TemporaryStorage tempStorage;

      PendingRecovery(TemporaryStorage tempStorage) {
        this.tempStorage = tempStorage;
      }

      void commit() {
        primaryStorage.write((NoResult.Quiet) storeProvider -> {
          try {
            distributedStore.persist(tempStorage.toSnapshot());
            shutDownNow.execute();
          } catch (CodingException e) {
            throw new IllegalStateException("Failed to encode snapshot.", e);
          }
        });
      }

      Iterable<IScheduledTask> query(final Query.Builder query) {
        return tempStorage.fetchTasks(query);
      }

      void delete(final Query.Builder query) {
        tempStorage.deleteTasks(query);
      }
    }
  }
}