/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.giraph.yarn;

import org.apache.giraph.conf.GiraphConfiguration;
import org.apache.giraph.conf.GiraphConstants;
import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;
import org.apache.giraph.graph.GraphTaskManager;

import org.apache.giraph.io.VertexOutputFormat;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.io.WritableComparable;
import org.apache.hadoop.mapreduce.MapContext;
import org.apache.hadoop.mapreduce.Mapper.Context;
import org.apache.hadoop.mapreduce.OutputCommitter;
import org.apache.hadoop.mapreduce.TaskID;
import org.apache.hadoop.mapreduce.lib.map.WrappedMapper;
import org.apache.hadoop.mapreduce.task.MapContextImpl;
import org.apache.hadoop.mapreduce.TaskAttemptID;
import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl;
import org.apache.log4j.Logger;

import java.io.IOException;

/**
 * This process will execute the BSP graph tasks alloted to this YARN
 * execution container. All tasks will be performed by calling the
 * GraphTaskManager object. Since this GiraphYarnTask will
 * not be passing data by key-value pairs through the MR framework, the
 * Mapper parameter types are irrelevant, and set to <code>Object</code> type.
 *
 * @param <I> Vertex id
 * @param <V> Vertex data
 * @param <E> Edge data
 */
public class GiraphYarnTask<I extends WritableComparable, V extends Writable,
    E extends Writable> {
  static {
    Configuration.addDefaultResource(GiraphConstants.GIRAPH_YARN_CONF_FILE);
  }
  /** Class logger */
  private static final Logger LOG = Logger.getLogger(GiraphYarnTask.class);
  /** Manage the framework-agnostic Giraph task for this job run */
  private GraphTaskManager<I, V, E> graphTaskManager;
  /** Giraph task ID number must start @ index 0. Used by ZK, BSP, etc. */
  private final int bspTaskId;
  /** A special "dummy" override of Mapper#Context, used to deliver MRv1 deps */
  private Context proxy;
  /** Configuration to hand off into Giraph, through wrapper Mapper#Context */
  private ImmutableClassesGiraphConfiguration conf;

  /**
   * Constructor. Build our DUMMY MRv1 data structures to pass to our
   * GiraphTaskManager. This allows us to continue to look the other way
   * while Giraph relies on MRv1 under the hood.
   * @param taskAttemptId the MRv1 TaskAttemptID we constructed from CLI args
   *                      supplied by GiraphApplicationMaster.
   */
  public GiraphYarnTask(final TaskAttemptID taskAttemptId) {
    conf = new ImmutableClassesGiraphConfiguration<I, V, E>(
      new GiraphConfiguration());
    bspTaskId = taskAttemptId.getTaskID().getId();
    conf.setInt("mapred.task.partition", bspTaskId);
    proxy = buildProxyMapperContext(taskAttemptId);
    graphTaskManager = new GraphTaskManager<I, V, E>(proxy);
  }

  /**
   * Run one Giraph worker (or master) task, hosted in this execution container.
   */
  public void run() {
    // Notify the master quicker if there is worker failure rather than
    // waiting for ZooKeeper to timeout and delete the ephemeral znodes
    try {
      graphTaskManager.setup(null); // defaults GTM to "assume fatjar mode"
      graphTaskManager.execute();
      graphTaskManager.cleanup();
      graphTaskManager.sendWorkerCountersAndFinishCleanup();
    } catch (InterruptedException ie) {
      LOG.error("run() caught an unrecoverable InterruptedException.", ie);
    } catch (IOException ioe) {
      throw new RuntimeException(
        "run() caught an unrecoverable IOException.", ioe);
      // CHECKSTYLE: stop IllegalCatch
    } catch (RuntimeException e) {
      // CHECKSTYLE: resume IllegalCatch
      graphTaskManager.zooKeeperCleanup();
      graphTaskManager.workerFailureCleanup();
      throw new RuntimeException(
        "run: Caught an unrecoverable exception " + e.getMessage(), e);
    } finally {
      // YARN: must complete the commit of the final output, Hadoop isn't there.
      finalizeYarnJob();
    }
  }

  /**
   * Without Hadoop MR to finish the consolidation of all the task output from
   * each HDFS task tmp dir, it won't get done. YARN has some job finalization
   * it must do "for us." -- AND must delete "jar cache" in HDFS too!
   */
  private void finalizeYarnJob() {
    if (conf.isPureYarnJob() && graphTaskManager.isMaster() &&
      conf.getVertexOutputFormatClass() != null) {
      try {
        LOG.info("Master is ready to commit final job output data.");
        VertexOutputFormat vertexOutputFormat =
          conf.createWrappedVertexOutputFormat();
        OutputCommitter outputCommitter =
          vertexOutputFormat.getOutputCommitter(proxy);
        // now we will have our output in OUTDIR if all went well...
        outputCommitter.commitJob(proxy);
        LOG.info("Master has committed the final job output data.");
      } catch (InterruptedException ie) {
        LOG.error("Interrupted while attempting to obtain " +
          "OutputCommitter.", ie);
      } catch (IOException ioe) {
        LOG.error("Master task's attempt to commit output has " +
          "FAILED.", ioe);
      }
    }
  }

  /**
   * Utility to generate dummy Mapper#Context for use in Giraph internals.
   * This is the "key hack" to inject MapReduce-related data structures
   * containing YARN cluster metadata (and our GiraphConf from the AppMaster)
   * into our Giraph BSP task code.
   * @param tid the TaskAttemptID to construct this Mapper#Context from.
   * @return sort of a Mapper#Context if you squint just right.
   */
  private Context buildProxyMapperContext(final TaskAttemptID tid) {
    MapContext mc = new MapContextImpl<Object, Object, Object, Object>(
      conf, // our Configuration, populated back at the GiraphYarnClient.
      tid,  // our TaskAttemptId, generated w/YARN app, container, attempt IDs
      null, // RecordReader here will never be used by Giraph
      null, // RecordWriter here will never be used by Giraph
      null, // OutputCommitter here will never be used by Giraph
      new TaskAttemptContextImpl.DummyReporter() { // goes in task logs for now
        @Override
        public void setStatus(String msg) {
          LOG.info("[STATUS: task-" + bspTaskId + "] " + msg);
        }
      },
      null); // Input split setting here will never be used by Giraph

    // now, we wrap our MapContext ref so we can produce a Mapper#Context
    WrappedMapper<Object, Object, Object, Object> wrappedMapper
      = new WrappedMapper<Object, Object, Object, Object>();
    return wrappedMapper.getMapContext(mc);
  }

  /**
   * Task entry point.
   * @param args CLI arguments injected by GiraphApplicationMaster to hand off
   *             job, task, and attempt ID's to this (and every) Giraph task.
   *             Args should be: <code>AppId ContainerId AppAttemptId</code>
   */
  @SuppressWarnings("rawtypes")
  public static void main(String[] args) {
    if (args.length != 4) {
      throw new IllegalStateException("GiraphYarnTask could not construct " +
        "a TaskAttemptID for the Giraph job from args: " + printArgs(args));
    }
    try {
      GiraphYarnTask<?, ?, ?> giraphYarnTask =
        new GiraphYarnTask(getTaskAttemptID(args));
      giraphYarnTask.run();
      // CHECKSTYLE: stop IllegalCatch
    } catch (Throwable t) {
      // CHECKSTYLE resume IllegalCatch
      LOG.error("GiraphYarnTask threw a top-level exception, failing task", t);
      System.exit(2);
    } // ALWAYS finish a YARN task or AppMaster with System#exit!!!
    System.exit(0);
  }

  /**
   * Utility to create a TaskAttemptId we can feed to our fake Mapper#Context.
   *
   * NOTE: ContainerId will serve as MR TaskID for Giraph tasks.
   * YARN container 1 is always AppMaster, so the least container id we will
   * ever get from YARN for a Giraph task is container id 2. Giraph on MapReduce
   * tasks must start at index 0. So we SUBTRACT TWO from each container id.
   *
   * @param args the command line args, fed to us by GiraphApplicationMaster.
   * @return the TaskAttemptId object, populated with YARN job data.
   */
  private static TaskAttemptID getTaskAttemptID(String[] args) {
    return new TaskAttemptID(
      args[0], // YARN ApplicationId Cluster Timestamp
      Integer.parseInt(args[1]), // YARN ApplicationId #
      TaskID.getTaskType('m'),  // Make Giraph think this is a Mapper task.
      Integer.parseInt(args[2]) - 2, // YARN ContainerId MINUS TWO (see above)
      Integer.parseInt(args[3])); // YARN AppAttemptId #
  }

  /**
   * Utility to help log command line args in the event of an error.
   * @param args the CLI args.
   * @return a pretty-print of the input args.
   */
  private static String printArgs(String[] args) {
    int count = 0;
    StringBuilder sb = new StringBuilder();
    for (String arg : args) {
      sb.append("arg[" + (count++) + "] == " + arg + ", ");
    }
    return sb.toString();
  }
}
