io.cdap.cdap.internal.app.runtime.distributed.ProgramLaunchConfig Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of cdap-app-fabric Show documentation
There is a newer version: 6.10.1
/*
 * Copyright © 2018 Cask Data, Inc.
 *
 * 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 io.cdap.cdap.internal.app.runtime.distributed;

import ch.qos.logback.classic.Level;
import com.google.common.collect.Iterables;
import io.cdap.cdap.api.Resources;
import io.cdap.cdap.app.runtime.ProgramOptions;
import io.cdap.cdap.common.twill.ProgramRuntimeClassAcceptor;
import io.cdap.cdap.internal.app.runtime.SystemArguments;
import org.apache.twill.api.ClassAcceptor;
import org.apache.twill.api.ResourceSpecification;
import org.apache.twill.api.TwillRunnable;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import javax.annotation.Nullable;

/**
 * Configuration for launching Twill container for a program.
 */
@SuppressWarnings({"WeakerAccess", "UnusedReturnValue"})
public final class ProgramLaunchConfig {

  private final Map extraResources = new HashMap<>();
  private final List extraClasspath = new ArrayList<>();
  private final Map extraEnv = new HashMap<>();
  private final Map runnables = new HashMap<>();
  private final List> launchOrder = new ArrayList<>();
  private final Set> extraDependencies = new HashSet<>();
  private final Map extraSystemArguments = new HashMap<>();
  private ClassAcceptor classAcceptor = new ProgramRuntimeClassAcceptor();

  /**
   * Adds extra system arguments that will be available through the {@link ProgramOptions#getArguments()}
   * in the program container.
   */
  public ProgramLaunchConfig addExtraSystemArguments(Map args) {
    extraSystemArguments.putAll(args);
    return this;
  }

  public ProgramLaunchConfig addExtraSystemArgument(String key, String value) {
    extraSystemArguments.put(key, value);
    return this;
  }

  public ProgramLaunchConfig addExtraResources(Map resources) {
    extraResources.putAll(resources);
    return this;
  }

  public ProgramLaunchConfig addExtraClasspath(Iterable classpath) {
    Iterables.addAll(extraClasspath, classpath);
    return this;
  }

  public ProgramLaunchConfig addExtraEnv(Map env) {
    extraEnv.putAll(env);
    return this;
  }

  public ProgramLaunchConfig addRunnable(String name, TwillRunnable runnable, int instances,
                                         Map args, Resources defaultResource) {
    return addRunnable(name, runnable, instances, args, defaultResource, null);
  }

  public ProgramLaunchConfig addRunnable(String name, TwillRunnable runnable, int instances,
                                         Map args, Resources defaultResources,
                                         @Nullable Integer maxRetries) {
    ResourceSpecification resourceSpec = createResourceSpec(SystemArguments.getResources(args, defaultResources),
                                                            instances);

    Map configs = SystemArguments.getTwillContainerConfigs(args, resourceSpec.getMemorySize());
    Map logLevels = SystemArguments.getLogLevels(args);

    runnables.put(name, new RunnableDefinition(runnable, resourceSpec, configs, logLevels, maxRetries));
    return this;
  }

  public ProgramLaunchConfig setLaunchOrder(Iterable> order) {
    launchOrder.clear();
    Iterables.addAll(launchOrder, order);
    return this;
  }

  public ProgramLaunchConfig setClassAcceptor(ClassAcceptor classAcceptor) {
    this.classAcceptor = classAcceptor;
    return this;
  }

  public ProgramLaunchConfig addExtraDependencies(Class...classes) {
    return addExtraDependencies(Arrays.asList(classes));
  }

  public ProgramLaunchConfig addExtraDependencies(Iterable> classes) {
    Iterables.addAll(extraDependencies, classes);
    return this;
  }

  /**
   * Returns the set of extra system arguments.
   */
  public Map getExtraSystemArguments() {
    return extraSystemArguments;
  }

  public Map getExtraResources() {
    return extraResources;
  }

  public List getExtraClasspath() {
    return extraClasspath;
  }

  public Map getExtraEnv() {
    return extraEnv;
  }

  public ClassAcceptor getClassAcceptor() {
    return classAcceptor;
  }

  public Map getRunnables() {
    return runnables;
  }

  public List> getLaunchOrder() {
    return launchOrder;
  }

  public Set> getExtraDependencies() {
    return extraDependencies;
  }

  public ProgramLaunchConfig clearRunnables() {
    runnables.clear();
    return this;
  }

  /**
   * Returns a {@link ResourceSpecification} created from the given {@link Resources} and number of instances.
   */
  private ResourceSpecification createResourceSpec(Resources resources, int instances) {
    return ResourceSpecification.Builder.with()
      .setVirtualCores(resources.getVirtualCores())
      .setMemory(resources.getMemoryMB(), ResourceSpecification.SizeUnit.MEGA)
      .setInstances(instances)
      .build();
  }
}