/*
 * Copyright 2012 LinkedIn Corp.
 *
 * 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 azkaban.jobExecutor;

import java.io.File;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

import org.apache.log4j.Logger;

import azkaban.jobExecutor.utils.process.AzkabanProcess;
import azkaban.jobExecutor.utils.process.AzkabanProcessBuilder;
import azkaban.utils.Pair;
import azkaban.utils.Props;
import azkaban.utils.SystemMemoryInfo;

/**
 * A job that runs a simple unix command
 */
public class ProcessJob extends AbstractProcessJob {

  public static final String COMMAND = "command";
  private static final long KILL_TIME_MS = 5000;
  private volatile AzkabanProcess process;
  private static final String MEMCHECK_ENABLED = "memCheck.enabled";
  private static final String MEMCHECK_FREEMEMDECRAMT = "memCheck.freeMemDecrAmt";

  public ProcessJob(final String jobId, final Props sysProps,
      final Props jobProps, final Logger log) {
    super(jobId, sysProps, jobProps, log);
  }

  @Override
  public void run() throws Exception {
    try {
      resolveProps();
    } catch (Exception e) {
      handleError("Bad property definition! " + e.getMessage(), e);
    }

    if (sysProps.getBoolean(MEMCHECK_ENABLED, true)) {
      long freeMemDecrAmt = sysProps.getLong(MEMCHECK_FREEMEMDECRAMT, 0);
      Pair<Long, Long> memPair = getProcMemoryRequirement();
      boolean isMemGranted = SystemMemoryInfo.canSystemGrantMemory(memPair.getFirst(), memPair.getSecond(), freeMemDecrAmt);
      if (!isMemGranted) {
        throw new Exception(String.format("Cannot request memory (Xms %d kb, Xmx %d kb) from system for job %s",
                memPair.getFirst(), memPair.getSecond(), getId()));
      }
    }

    List<String> commands = null;
    try {
      commands = getCommandList();
    } catch (Exception e) {
      handleError("Job set up failed " + e.getCause(), e);
    }

    long startMs = System.currentTimeMillis();

    if (commands == null) {
      handleError("There are no commands to execute", null);
    }

    info(commands.size() + " commands to execute.");
    File[] propFiles = initPropsFiles();
    Map<String, String> envVars = getEnvironmentVariables();

    for (String command : commands) {
      info("Command: " + command);
      AzkabanProcessBuilder builder =
          new AzkabanProcessBuilder(partitionCommandLine(command))
              .setEnv(envVars).setWorkingDir(getCwd()).setLogger(getLog());

      if (builder.getEnv().size() > 0) {
        info("Environment variables: " + builder.getEnv());
      }
      info("Working directory: " + builder.getWorkingDir());

      boolean success = false;
      this.process = builder.build();

      try {
        this.process.run();
        success = true;
      } catch (Throwable e) {
        for (File file : propFiles)
          if (file != null && file.exists())
            file.delete();
        throw new RuntimeException(e);
      } finally {
        this.process = null;
        info("Process completed "
            + (success ? "successfully" : "unsuccessfully") + " in "
            + ((System.currentTimeMillis() - startMs) / 1000) + " seconds.");
      }
    }

    // Get the output properties from this job.
    generateProperties(propFiles[1]);
  }

  protected Pair<Long, Long> getProcMemoryRequirement() {
    return new Pair<Long, Long>(0L, 0L);
  }

  protected void handleError(String errorMsg, Exception e) throws Exception {
    error(errorMsg);
    if (e != null) {
      throw new Exception(errorMsg, e);
    } else {
      throw new Exception(errorMsg);
    }
  }

  protected List<String> getCommandList() {
    List<String> commands = new ArrayList<String>();
    commands.add(jobProps.getString(COMMAND));
    for (int i = 1; jobProps.containsKey(COMMAND + "." + i); i++) {
      commands.add(jobProps.getString(COMMAND + "." + i));
    }

    return commands;
  }

  @Override
  public void cancel() throws InterruptedException {
    if (process == null)
      throw new IllegalStateException("Not started.");
    boolean killed = process.softKill(KILL_TIME_MS, TimeUnit.MILLISECONDS);
    if (!killed) {
      warn("Kill with signal TERM failed. Killing with KILL signal.");
      process.hardKill();
    }
  }

  @Override
  public double getProgress() {
    return process != null && process.isComplete() ? 1.0 : 0.0;
  }

  public int getProcessId() {
    return process.getProcessId();
  }

  public String getPath() {
    return _jobPath == null ? "" : _jobPath;
  }

  /**
   * Splits the command into a unix like command line structure. Quotes and
   * single quotes are treated as nested strings.
   *
   * @param command
   * @return
   */
  public static String[] partitionCommandLine(final String command) {
    ArrayList<String> commands = new ArrayList<String>();

    int index = 0;

    StringBuffer buffer = new StringBuffer(command.length());

    boolean isApos = false;
    boolean isQuote = false;
    while (index < command.length()) {
      char c = command.charAt(index);

      switch (c) {
      case ' ':
        if (!isQuote && !isApos) {
          String arg = buffer.toString();
          buffer = new StringBuffer(command.length() - index);
          if (arg.length() > 0) {
            commands.add(arg);
          }
        } else {
          buffer.append(c);
        }
        break;
      case '\'':
        if (!isQuote) {
          isApos = !isApos;
        } else {
          buffer.append(c);
        }
        break;
      case '"':
        if (!isApos) {
          isQuote = !isQuote;
        } else {
          buffer.append(c);
        }
        break;
      default:
        buffer.append(c);
      }

      index++;
    }

    if (buffer.length() > 0) {
      String arg = buffer.toString();
      commands.add(arg);
    }

    return commands.toArray(new String[commands.size()]);
  }
}

@@ -3816,32 +3816,314 @@
 Files%5B1%5D);%0A  %7D%0A%0A
+  /**%0A   * This is used to get the min/max memory size requirement by processes.%0A   * SystemMemoryInfo can use the info to determine if the memory request%0A   * can be fulfilled. For Java process, this should be Xms/Xmx setting.%0A   *  %0A   * @return pair of min/max memory size%0A   */%0A
   protected Pair

add system memory awareness to Azkaban