src/main/java/org/apache/aurora/scheduler/storage/durability/RecoveryTool.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.io.File;
import java.net.InetSocketAddress;
import java.util.Map;

import com.beust.jcommander.IStringConverter;
import com.beust.jcommander.IStringConverterFactory;
import com.beust.jcommander.JCommander;
import com.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.inject.AbstractModule;
import com.google.inject.Guice;
import com.google.inject.Injector;

import org.apache.aurora.common.util.BuildInfo;
import org.apache.aurora.common.util.Clock;
import org.apache.aurora.scheduler.app.LifecycleModule;
import org.apache.aurora.scheduler.config.converters.DataAmountConverter;
import org.apache.aurora.scheduler.config.converters.InetSocketAddressConverter;
import org.apache.aurora.scheduler.config.converters.TimeAmountConverter;
import org.apache.aurora.scheduler.config.types.DataAmount;
import org.apache.aurora.scheduler.config.types.TimeAmount;
import org.apache.aurora.scheduler.discovery.FlaggedZooKeeperConfig;
import org.apache.aurora.scheduler.discovery.ServiceDiscoveryBindings;
import org.apache.aurora.scheduler.log.mesos.MesosLogStreamModule;
import org.apache.aurora.scheduler.storage.Snapshotter;
import org.apache.aurora.scheduler.storage.backup.BackupReader;
import org.apache.aurora.scheduler.storage.log.LogPersistenceModule;
import org.apache.aurora.scheduler.storage.log.SnapshotterImpl;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A utility to recover the contents of one persistence into another.
 */
public final class RecoveryTool {

  private static final Logger LOG = LoggerFactory.getLogger(RecoveryTool.class);

  private RecoveryTool() {
    // Main-only class.
  }

  interface RecoveryEndpoint {
    Iterable<Object> getOptions();

    Persistence create();
  }

  private static class Log implements RecoveryEndpoint {
    private final FlaggedZooKeeperConfig.Options zkOptions = new FlaggedZooKeeperConfig.Options();
    private final MesosLogStreamModule.Options logOptions = new MesosLogStreamModule.Options();
    private final LogPersistenceModule.Options options = new LogPersistenceModule.Options();

    @Override
    public Iterable<Object> getOptions() {
      return ImmutableList.of(logOptions, options, zkOptions);
    }

    @Override
    public Persistence create() {
      Injector injector = Guice.createInjector(
          new MesosLogStreamModule(logOptions, FlaggedZooKeeperConfig.create(zkOptions)),
          new LogPersistenceModule(options),
          new LifecycleModule(),
          new AbstractModule() {
            @Override
            protected void configure() {
              bind(ServiceDiscoveryBindings.ZOO_KEEPER_CLUSTER_KEY)
                  .toInstance(zkOptions.zkEndpoints);
              bind(Snapshotter.class).to(SnapshotterImpl.class);
              bind(Clock.class).toInstance(Clock.SYSTEM_CLOCK);
              bind(BuildInfo.class).toInstance(new BuildInfo());
            }
          });
      return injector.getInstance(Persistence.class);
    }
  }

  private static class Backup implements RecoveryEndpoint {
    @Parameters(separators = "=")
    private static class Options {
      @Parameter(names = "-backup", description = "Backup file to load")
      File backup;
    }

    private final Options options = new Options();

    @Override
    public Iterable<Object> getOptions() {
      return ImmutableList.of(options);
    }

    @Override
    public Persistence create() {
      return new BackupReader(
          options.backup,
          new SnapshotterImpl(new BuildInfo(), Clock.SYSTEM_CLOCK));
    }
  }

  enum Endpoint {
    LOG(new Log()),
    BACKUP(new Backup());

    private final RecoveryEndpoint impl;

    Endpoint(RecoveryEndpoint impl) {
      this.impl = impl;
    }
  }

  @Parameters(separators = "=")
  private static class Options {
    @Parameter(names = "-from",
        required = true,
        description = "Persistence to read state from")
    Endpoint from;

    @Parameter(names = "-to",
        required = true,
        description = "Persistence to write recovered state into")
    Endpoint to;

    @Parameter(names = "-batch-size",
        description = "Write in batches of this may ops.")
    int batchSize = 50;

    @Parameter(names = "--help", description = "Print usage", help = true)
    boolean help;
  }

  private static JCommander configure(Options options, String... args) {
    JCommander.Builder builder = JCommander.newBuilder().programName(RecoveryTool.class.getName());
    builder.addConverterFactory(new IStringConverterFactory() {
      private Map<Class<?>, Class<? extends IStringConverter<?>>> classConverters =
          ImmutableMap.<Class<?>, Class<? extends IStringConverter<?>>>builder()
              .put(DataAmount.class, DataAmountConverter.class)
              .put(InetSocketAddress.class, InetSocketAddressConverter.class)
              .put(TimeAmount.class, TimeAmountConverter.class)
              .build();

      @SuppressWarnings("unchecked")
      @Override
      public <T> Class<? extends IStringConverter<T>> getConverter(Class<T> forType) {
        return (Class<IStringConverter<T>>) classConverters.get(forType);
      }
    });

    builder.addObject(options);
    for (Endpoint endpoint : Endpoint.values()) {
      endpoint.impl.getOptions().forEach(builder::addObject);
    }

    JCommander parser = builder.build();
    parser.parse(args);
    return parser;
  }

  public static void main(String[] args) {
    Options options = new Options();
    JCommander parser = configure(options, args);
    if (options.help) {
      parser.usage();
      System.exit(1);
    }

    LOG.info("Recovering from " + options.from + " to " + options.to);
    Persistence from = options.from.impl.create();
    Persistence to = options.to.impl.create();

    from.prepare();
    to.prepare();

    Recovery.copy(from, to, options.batchSize);
  }
}