LauncherUtils.java

// Copyright 2016 Twitter. All rights reserved.
//
// 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 com.twitter.heron.spi.utils;

import java.util.logging.Level;
import java.util.logging.Logger;

import com.twitter.heron.api.generated.TopologyAPI;
import com.twitter.heron.common.basics.SysUtils;
import com.twitter.heron.spi.common.Config;
import com.twitter.heron.spi.common.Context;
import com.twitter.heron.spi.common.Keys;
import com.twitter.heron.spi.packing.IPacking;
import com.twitter.heron.spi.packing.PackingPlan;
import com.twitter.heron.spi.scheduler.IScheduler;
import com.twitter.heron.spi.statemgr.SchedulerStateManagerAdaptor;

/**
 * {@link LauncherUtils} contains helper methods used by the server and client side launch
 * controllers; {@link com.twitter.heron.scheduler.LaunchRunner} and
 * {@link com.twitter.heron.scheduler.SchedulerMain}
 */
public class LauncherUtils {
  private static final Logger LOG = Logger.getLogger(LauncherUtils.class.getName());

  private static LauncherUtils instance = new LauncherUtils();

  public static LauncherUtils getInstance() {
    return instance;
  }

  /**
   * Returns a packing plan generated by configured packing class
   */
  public PackingPlan createPackingPlan(final Config config, final Config runtime) {
    // Create an instance of the packing class
    String packingClass = Context.packingClass(config);
    IPacking packing;
    try {
      // create an instance of the packing class
      packing = ReflectionUtils.newInstance(packingClass);
    } catch (IllegalAccessException | InstantiationException | ClassNotFoundException e) {
      LOG.log(Level.SEVERE, "Failed to instantiate packing instance", e);
      return null;
    }

    try {
      TopologyAPI.Topology topology = com.twitter.heron.spi.utils.Runtime.topology(runtime);
      packing.initialize(config, topology);
      PackingPlan packedPlan = packing.pack();
      return packedPlan;
    } finally {
      SysUtils.closeIgnoringExceptions(packing);
    }
  }

  /**
   * Invoke the onScheduler() in IScheduler directly as a library
   *
   * @param config The Config to initialize IScheduler
   * @param runtime The runtime Config to initialize IScheduler
   * @param scheduler the IScheduler to invoke
   * @param packing The PackingPlan to scheduler for OnSchedule()
   * @return true if scheduling successfully
   */
  public boolean onScheduleAsLibrary(
      Config config,
      Config runtime,
      IScheduler scheduler,
      PackingPlan packing) {
    boolean ret = false;

    try {
      scheduler.initialize(config, runtime);
      ret = scheduler.onSchedule(packing);

      if (ret) {
        // Set the SchedulerLocation at last step,
        // since some methods in IScheduler will provide correct values
        // only after IScheduler.onSchedule is invoked correctly
        ret = SchedulerUtils.setLibSchedulerLocation(runtime, scheduler, false);
      } else {
        LOG.severe("Failed to invoke IScheduler as library");
      }
    } finally {
      scheduler.close();
    }

    return ret;
  }

  /**
   * Creates and initializes scheduler instance
   *
   * @return initialized scheduler instances
   */
  public IScheduler getSchedulerInstance(Config config, Config runtime) {
    String schedulerClass = Context.schedulerClass(config);
    IScheduler scheduler;
    try {
      // create an instance of scheduler
      scheduler = ReflectionUtils.newInstance(schedulerClass);
    } catch (IllegalAccessException | InstantiationException | ClassNotFoundException e) {
      LOG.log(Level.SEVERE, "Failed to instantiate scheduler", e);
      return null;
    }

    scheduler.initialize(config, runtime);
    return scheduler;
  }

  /**
   * Builds initial runtime config instance using topology information.
   *
   * @return initial runtime config instance
   */
  public Config getPrimaryRuntime(TopologyAPI.Topology topology,
                                  SchedulerStateManagerAdaptor adaptor) {
    return Config.newBuilder()
        .put(Keys.topologyId(), topology.getId())
        .put(Keys.topologyName(), topology.getName())
        .put(Keys.topologyDefinition(), topology)
        .put(Keys.schedulerStateManagerAdaptor(), adaptor)
        .put(Keys.numContainers(), 1 + TopologyUtils.getNumContainers(topology))
        .build();
  }

  /**
   * Creates a config instance with packing plan info added to runtime config
   */
  public Config createConfigWithPackingDetails(Config runtime, PackingPlan packing) {
    Config ytruntime;
    ytruntime = Config.newBuilder()
        .putAll(runtime)
        .put(Keys.componentRamMap(), packing.getComponentRamDistribution())
        .put(Keys.numContainers(), 1 + packing.getContainers().size())
        .build();
    return ytruntime;
  }
}