/*

  * 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.examples;

 import org.apache.commons.cli.CommandLine;

 import org.apache.commons.cli.CommandLineParser;

 import org.apache.commons.cli.HelpFormatter;

 import org.apache.commons.cli.Options;

 import org.apache.commons.cli.PosixParser;

 import org.apache.giraph.aggregators.LongSumAggregator;

 import org.apache.giraph.graph.BasicComputation;

 import org.apache.giraph.edge.Edge;

 import org.apache.giraph.edge.EdgeFactory;

 import org.apache.giraph.io.formats.FileOutputFormatUtil;

 import org.apache.giraph.io.formats.GeneratedVertexInputFormat;

 import org.apache.giraph.io.formats.IdWithValueTextOutputFormat;

 import org.apache.giraph.job.GiraphJob;

 import org.apache.giraph.master.DefaultMasterCompute;

 import org.apache.giraph.graph.Vertex;

 import org.apache.giraph.worker.WorkerContext;

 import org.apache.hadoop.conf.Configuration;

 import org.apache.hadoop.fs.Path;

 import org.apache.hadoop.io.FloatWritable;

 import org.apache.hadoop.io.IntWritable;

 import org.apache.hadoop.io.LongWritable;

 import org.apache.hadoop.util.Tool;

 import org.apache.hadoop.util.ToolRunner;

 import org.apache.log4j.Logger;

 import java.io.IOException;

 /**

  * An example that simply uses its id, value, and edges to compute new data

  * every iteration to verify that checkpoint restarting works.  Fault injection

  * can also test automated checkpoint restarts.

  */

 public class SimpleCheckpoint implements Tool {

   /** Which superstep to cause the worker to fail */

   public static final int FAULTING_SUPERSTEP = 4;

   /** Vertex id to fault on */

   public static final long FAULTING_VERTEX_ID = 1;

   /** Dynamically set number of supersteps */

   public static final String SUPERSTEP_COUNT =

       "simpleCheckpointVertex.superstepCount";

   /** Should fault? */

   public static final String ENABLE_FAULT =

       "simpleCheckpointVertex.enableFault";

   /** Class logger */

   private static final Logger LOG =

       Logger.getLogger(SimpleCheckpoint.class);

   /** Configuration */

   private Configuration conf;

   /**

    * Actual computation.

    */

   public static class SimpleCheckpointComputation extends

       BasicComputation<LongWritable, IntWritable, FloatWritable,

           FloatWritable> {

     @Override

     public void compute(

         Vertex<LongWritable, IntWritable, FloatWritable> vertex,

         Iterable<FloatWritable> messages) throws IOException {

       SimpleCheckpointVertexWorkerContext workerContext = getWorkerContext();

       boolean enableFault = workerContext.getEnableFault();

       int supersteps = workerContext.getSupersteps();

       if (enableFault && (getSuperstep() == FAULTING_SUPERSTEP) &&

           (getContext().getTaskAttemptID().getId() == 0) &&

           (vertex.getId().get() == FAULTING_VERTEX_ID)) {

         LOG.info("compute: Forced a fault on the first " +

             "attempt of superstep " +

             FAULTING_SUPERSTEP + " and vertex id " +

             FAULTING_VERTEX_ID);

         System.exit(-1);

       }

       if (getSuperstep() > supersteps) {

         vertex.voteToHalt();

         return;

       }

       long sumAgg = this.<LongWritable>getAggregatedValue(

           LongSumAggregator.class.getName()).get();

       LOG.info("compute: " + sumAgg);

       aggregate(LongSumAggregator.class.getName(),

           new LongWritable(vertex.getId().get()));

       LOG.info("compute: sum = " + sumAgg +

           " for vertex " + vertex.getId());

       float msgValue = 0.0f;

       for (FloatWritable message : messages) {

         float curMsgValue = message.get();

         msgValue += curMsgValue;

         LOG.info("compute: got msgValue = " + curMsgValue +

             " for vertex " + vertex.getId() +

             " on superstep " + getSuperstep());

       }

       int vertexValue = vertex.getValue().get();

       vertex.setValue(new IntWritable(vertexValue + (int) msgValue));

       LOG.info("compute: vertex " + vertex.getId() +

           " has value " + vertex.getValue() +

           " on superstep " + getSuperstep());

       for (Edge<LongWritable, FloatWritable> edge : vertex.getEdges()) {

         FloatWritable newEdgeValue = new FloatWritable(edge.getValue().get() +

             (float) vertexValue);

         Edge<LongWritable, FloatWritable> newEdge =

             EdgeFactory.create(edge.getTargetVertexId(), newEdgeValue);

         LOG.info("compute: vertex " + vertex.getId() +

             " sending edgeValue " + edge.getValue() +

             " vertexValue " + vertexValue +

             " total " + newEdgeValue +

             " to vertex " + edge.getTargetVertexId() +

             " on superstep " + getSuperstep());

         vertex.addEdge(newEdge);

         sendMessage(edge.getTargetVertexId(), newEdgeValue);

       }

     }

   }

   /**

    * Worker context associated with {@link SimpleCheckpoint}.

    */

   public static class SimpleCheckpointVertexWorkerContext

       extends WorkerContext {

     /** Filename to indicate whether a fault was found */

     public static final String FAULT_FILE = "/tmp/faultFile";

     /** User can access this after the application finishes if local */

     private static long FINAL_SUM;

     /** Number of supersteps to run (6 by default) */

     private int supersteps = 6;

     /** Enable the fault at the particular vertex id and superstep? */

     private boolean enableFault = false;

     public static long getFinalSum() {

       return FINAL_SUM;

     }

     @Override

     public void preApplication()

       throws InstantiationException, IllegalAccessException {

       supersteps = getContext().getConfiguration()

           .getInt(SUPERSTEP_COUNT, supersteps);

       enableFault = getContext().getConfiguration()

           .getBoolean(ENABLE_FAULT, false);

     }

     @Override

     public void postApplication() {

       setFinalSum(this.<LongWritable>getAggregatedValue(

           LongSumAggregator.class.getName()).get());

       LOG.info("FINAL_SUM=" + FINAL_SUM);

     }

     /**

      * Set the final sum

      *

      * @param value sum

      */

     private static void setFinalSum(long value) {

       FINAL_SUM = value;

     }

     @Override

     public void preSuperstep() {

     }

     @Override

     public void postSuperstep() { }

     public int getSupersteps() {

       return this.supersteps;

     }

     public boolean getEnableFault() {

       return this.enableFault;

     }

   }

   @Override

   public int run(String[] args) throws Exception {

     Options options = new Options();

     options.addOption("h", "help", false, "Help");

     options.addOption("v", "verbose", false, "Verbose");

     options.addOption("w",

         "workers",

         true,

         "Number of workers");

     options.addOption("s",

         "supersteps",

         true,

         "Supersteps to execute before finishing");

     options.addOption("w",

         "workers",

         true,

         "Minimum number of workers");

     options.addOption("o",

         "outputDirectory",

         true,

         "Output directory");

     HelpFormatter formatter = new HelpFormatter();

     if (args.length == 0) {

       formatter.printHelp(getClass().getName(), options, true);

       return 0;

     }

     CommandLineParser parser = new PosixParser();

     CommandLine cmd = parser.parse(options, args);

     if (cmd.hasOption('h')) {

       formatter.printHelp(getClass().getName(), options, true);

       return 0;

     }

     if (!cmd.hasOption('w')) {

       LOG.info("Need to choose the number of workers (-w)");

       return -1;

     }

     if (!cmd.hasOption('o')) {

       LOG.info("Need to set the output directory (-o)");

       return -1;

     }

     GiraphJob bspJob = new GiraphJob(getConf(), getClass().getName());

     bspJob.getConfiguration().setComputationClass(

         SimpleCheckpointComputation.class);

     bspJob.getConfiguration().setVertexInputFormatClass(

         GeneratedVertexInputFormat.class);

     bspJob.getConfiguration().setVertexOutputFormatClass(

         IdWithValueTextOutputFormat.class);

     bspJob.getConfiguration().setWorkerContextClass(

         SimpleCheckpointVertexWorkerContext.class);

     bspJob.getConfiguration().setMasterComputeClass(

         SimpleCheckpointVertexMasterCompute.class);

     int minWorkers = Integer.parseInt(cmd.getOptionValue('w'));

     int maxWorkers = Integer.parseInt(cmd.getOptionValue('w'));

     bspJob.getConfiguration().setWorkerConfiguration(

         minWorkers, maxWorkers, 100.0f);

     FileOutputFormatUtil.setOutputPath(bspJob.getInternalJob(),

                                    new Path(cmd.getOptionValue('o')));

     boolean verbose = false;

     if (cmd.hasOption('v')) {

       verbose = true;

     }

     if (cmd.hasOption('s')) {

       getConf().setInt(SUPERSTEP_COUNT,

           Integer.parseInt(cmd.getOptionValue('s')));

     }

     if (bspJob.run(verbose)) {

       return 0;

     } else {

       return -1;

     }

   }

   /**

    * Master compute associated with {@link SimpleCheckpoint}.

    * It registers required aggregators.

    */

   public static class SimpleCheckpointVertexMasterCompute extends

       DefaultMasterCompute {

     @Override

     public void initialize() throws InstantiationException,

         IllegalAccessException {

       registerAggregator(LongSumAggregator.class.getName(),

           LongSumAggregator.class);

     }

   }

   /**

    * Executable from the command line.

    *

    * @param args Command line args.

    * @throws Exception

    */

   public static void main(String[] args) throws Exception {

     System.exit(ToolRunner.run(new SimpleCheckpoint(), args));

   }

   @Override

   public Configuration getConf() {

     return conf;

   }

   @Override

   public void setConf(Configuration conf) {

     this.conf = conf;

   }

 }