/*

  * 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.block_app.framework.api.local;

 import java.io.IOException;

 import java.util.List;

 import java.util.concurrent.CountDownLatch;

 import java.util.concurrent.ExecutorService;

 import java.util.concurrent.Executors;

 import java.util.concurrent.atomic.AtomicBoolean;

 import java.util.concurrent.atomic.AtomicReference;

 import org.apache.giraph.block_app.framework.BlockFactory;

 import org.apache.giraph.block_app.framework.BlockUtils;

 import org.apache.giraph.block_app.framework.api.local.InternalApi.InternalWorkerApi;

 import org.apache.giraph.block_app.framework.block.Block;

 import org.apache.giraph.block_app.framework.internal.BlockMasterLogic;

 import org.apache.giraph.block_app.framework.internal.BlockWorkerContextLogic;

 import org.apache.giraph.block_app.framework.internal.BlockWorkerLogic;

 import org.apache.giraph.block_app.framework.internal.BlockWorkerPieces;

 import org.apache.giraph.block_app.framework.output.BlockOutputHandle;

 import org.apache.giraph.conf.BooleanConfOption;

 import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;

 import org.apache.giraph.conf.IntConfOption;

 import org.apache.giraph.graph.OnlyIdVertex;

 import org.apache.giraph.graph.Vertex;

 import org.apache.giraph.io.SimpleVertexWriter;

 import org.apache.giraph.partition.Partition;

 import org.apache.giraph.utils.InternalVertexRunner;

 import org.apache.giraph.utils.TestGraph;

 import org.apache.giraph.utils.Trimmable;

 import org.apache.giraph.utils.WritableUtils;

 import org.apache.giraph.writable.kryo.KryoWritableWrapper;

 import org.apache.hadoop.io.Writable;

 import org.apache.hadoop.io.WritableComparable;

 import org.apache.hadoop.util.Progressable;

 import com.google.common.base.Preconditions;

 import com.google.common.collect.Iterables;

 /**

  * Local in-memory Block application job runner.

  * Implementation should be faster then using InternalVertexRunner.

  *

  * Useful for fast testing.

  */

 @SuppressWarnings({ "rawtypes", "unchecked" })

 public class LocalBlockRunner {

   /** Number of threads to use */

   public static final IntConfOption NUM_THREADS = new IntConfOption(

       "test.LocalBlockRunner.NUM_THREADS", 3, "");

   /** Number of vertex partitions */

   public static final IntConfOption NUM_PARTITIONS = new IntConfOption(

       "test.LocalBlockRunner.NUM_PARTITIONS", 16, "");

   /**

    * Whether to run all supported checks. Disable if you are running this

    * not within a unit test, and on a large graph, where performance matters.

    */

   public static final BooleanConfOption RUN_ALL_CHECKS = new BooleanConfOption(

       "test.LocalBlockRunner.RUN_ALL_CHECKS", true, "");

   // merge into RUN_ALL_CHECKS, after SERIALIZE_MASTER starts working

   public static final BooleanConfOption SERIALIZE_MASTER =

       new BooleanConfOption(

           "test.LocalBlockRunner.SERIALIZE_MASTER", false, "");

   private LocalBlockRunner() { }

   /**

    * Run Block Application specified within the conf, on a given graph,

    * locally, in-memory.

    *

    * With a boolean flag, you can switch between LocalBlockRunner and

    * InternalVertexRunner implementations of local in-memory computation.

    */

   public static

   <I extends WritableComparable, V extends Writable, E extends Writable>

   TestGraph<I, V, E> runApp(

       TestGraph<I, V, E> graph, boolean useFullDigraphTests) throws Exception {

     if (useFullDigraphTests) {

       return InternalVertexRunner.runWithInMemoryOutput(graph.getConf(), graph);

     } else {

       runApp(graph);

       return graph;

     }

   }

   /**

    * Run Block Application specified within the conf, on a given graph,

    * locally, in-memory.

    */

   public static

   <I extends WritableComparable, V extends Writable, E extends Writable>

   void runApp(TestGraph<I, V, E> graph) {

     SimpleVertexWriter<I, V, E> noOpVertexSaver = noOpVertexSaver();

     runAppWithVertexOutput(graph, noOpVertexSaver);

   }

   /**

    * Run Block from a specified execution stage on a given graph,

    * locally, in-memory.

    */

   public static

   <I extends WritableComparable, V extends Writable, E extends Writable>

   void runBlock(

       TestGraph<I, V, E> graph, Block block, Object executionStage) {

     SimpleVertexWriter<I, V, E> noOpVertexSaver = noOpVertexSaver();

     runBlockWithVertexOutput(

         block, executionStage, graph, noOpVertexSaver);

   }

   /**

    * Run Block Application specified within the conf, on a given graph,

    * locally, in-memory, with a given vertexSaver.

    */

   public static

   <I extends WritableComparable, V extends Writable, E extends Writable>

   void runAppWithVertexOutput(

       TestGraph<I, V, E> graph, final SimpleVertexWriter<I, V, E> vertexSaver) {

     BlockFactory<?> factory = BlockUtils.createBlockFactory(graph.getConf());

     runBlockWithVertexOutput(

         factory.createBlock(graph.getConf()),

         factory.createExecutionStage(graph.getConf()),

         graph, vertexSaver);

   }

   /**

    * Run Block from a specified execution stage on a given graph,

    * locally, in-memory, with a given vertexSaver.

    */

   public static

   <I extends WritableComparable, V extends Writable, E extends Writable>

   void runBlockWithVertexOutput(

       Block block, Object executionStage, TestGraph<I, V, E> graph,

       final SimpleVertexWriter<I, V, E> vertexSaver

   ) {

     Preconditions.checkNotNull(block);

     Preconditions.checkNotNull(graph);

     ImmutableClassesGiraphConfiguration<I, V, E> conf = graph.getConf();

     int numThreads = NUM_THREADS.get(conf);

     int numPartitions = NUM_PARTITIONS.get(conf);

     boolean runAllChecks = RUN_ALL_CHECKS.get(conf);

     boolean serializeMaster = SERIALIZE_MASTER.get(conf);

     final boolean doOutputDuringComputation = conf.doOutputDuringComputation();

     final InternalApi internalApi =

         new InternalApi(graph, conf, numPartitions, runAllChecks);

     final InternalWorkerApi internalWorkerApi = internalApi.getWorkerApi();

     BlockUtils.checkBlockTypes(block, executionStage, conf);

     BlockMasterLogic<Object> blockMasterLogic = new BlockMasterLogic<>();

     blockMasterLogic.initialize(block, executionStage, internalApi);

     BlockWorkerContextLogic workerContextLogic =

         internalApi.getWorkerContextLogic();

     workerContextLogic.preApplication(internalWorkerApi,

         new BlockOutputHandle("", conf, new Progressable() {

           @Override

           public void progress() {

           }

         }));

     ExecutorService executor = Executors.newFixedThreadPool(numThreads);

     if (runAllChecks) {

       for (Vertex<I, V, E> vertex : graph) {

         V value = conf.createVertexValue();

         WritableUtils.copyInto(vertex.getValue(), value);

         vertex.setValue(value);

         vertex.setEdges((Iterable) WritableUtils.createCopy(

             (Writable) vertex.getEdges(), conf.getOutEdgesClass(), conf));

       }

     }

     final AtomicBoolean anyVertexAlive = new AtomicBoolean(true);

     for (int superstep = 0;; superstep++) {

       // serialize master to test continuable computation

       if (serializeMaster) {

         blockMasterLogic = (BlockMasterLogic) WritableUtils.createCopy(

             new KryoWritableWrapper<>(blockMasterLogic),

             KryoWritableWrapper.class,

             conf).get();

         blockMasterLogic.initializeAfterRead(internalApi);

       }

       if (!anyVertexAlive.get()) {

         break;

       }

       final BlockWorkerPieces workerPieces =

           blockMasterLogic.computeNext(superstep);

       if (workerPieces == null) {

         if (!conf.doOutputDuringComputation()) {

           List<Partition<I, V, E>> partitions = internalApi.getPartitions();

           for (Partition<I, V, E> partition : partitions) {

             for (Vertex<I, V, E> vertex : partition) {

               try {

                 vertexSaver.writeVertex(vertex);

               } catch (IOException | InterruptedException e) {

                 throw new RuntimeException(e);

               }

             }

           }

         }

         int left = executor.shutdownNow().size();

         Preconditions.checkState(0 == left, "Some work still left to be done?");

         break;

       } else {

         internalApi.afterMasterBeforeWorker(workerPieces);

         List<Partition<I, V, E>> partitions = internalApi.getPartitions();

         workerContextLogic.preSuperstep(

             internalWorkerApi,

             internalWorkerApi,

             KryoWritableWrapper.wrapAndCopy(workerPieces), superstep,

             internalApi.takeWorkerMessages());

         final CountDownLatch latch = new CountDownLatch(numPartitions);

         final AtomicReference<Throwable> exception = new AtomicReference<>();

         anyVertexAlive.set(false);

         for (final Partition<I, V, E> partition : partitions) {

           executor.execute(new Runnable() {

             @Override

             public void run() {

               try {

                 boolean anyCurVertexAlive = false;

                 BlockWorkerPieces localPieces =

                     KryoWritableWrapper.wrapAndCopy(workerPieces);

                 BlockWorkerLogic localLogic = new BlockWorkerLogic(localPieces);

                 localLogic.preSuperstep(internalWorkerApi, internalWorkerApi);

                 if (internalApi.ignoreExistingVertices()) {

                   Iterable<I> destinations =

                       internalApi.getPartitionDestinationVertices(

                           partition.getId());

                   if (!Iterables.isEmpty(destinations)) {

                     OnlyIdVertex<I> vertex = new OnlyIdVertex<>();

                     for (I vertexId : destinations) {

                       Iterable messages = internalApi.takeMessages(vertexId);

                       Preconditions.checkState(!Iterables.isEmpty(messages));

                       vertex.setId(vertexId);

                       localLogic.compute(vertex, messages);

                       anyCurVertexAlive = true;

                     }

                   }

                 } else {

                   for (Vertex<I, V, E> vertex : partition) {

                     Iterable messages =

                         internalApi.takeMessages(vertex.getId());

                     if (vertex.isHalted() && !Iterables.isEmpty(messages)) {

                       vertex.wakeUp();

                     }

                     // Equivalent of ComputeCallable.computePartition

                     if (!vertex.isHalted()) {

                       localLogic.compute(vertex, messages);

                       // Need to unwrap the mutated edges (possibly)

                       vertex.unwrapMutableEdges();

                       //Compact edges representation if possible

                       if (vertex instanceof Trimmable) {

                         ((Trimmable) vertex).trim();

                       }

                       // Write vertex to superstep output

                       // (no-op if it is not used)

                       if (doOutputDuringComputation) {

                         vertexSaver.writeVertex(vertex);

                       }

                       // Need to save the vertex changes (possibly)

                       partition.saveVertex(vertex);

                     }

                     if (!vertex.isHalted()) {

                       anyCurVertexAlive = true;

                     }

                   }

                 }

                 if (anyCurVertexAlive) {

                   anyVertexAlive.set(true);

                 }

                 localLogic.postSuperstep();

               // CHECKSTYLE: stop IllegalCatch

               // Need to propagate all exceptions within test

               } catch (Throwable t) {

               // CHECKSTYLE: resume IllegalCatch

                 t.printStackTrace();

                 exception.set(t);

               }

               latch.countDown();

             }

           });

         }

         try {

           latch.await();

         } catch (InterruptedException e) {

           throw new RuntimeException("Thread intentionally interrupted", e);

         }

         if (exception.get() != null) {

           throw new RuntimeException("Worker failed", exception.get());

         }

         workerContextLogic.postSuperstep();

         internalApi.afterWorkerBeforeMaster();

       }

     }

     workerContextLogic.postApplication();

     internalApi.postApplication();

   }

   private static

   <I extends WritableComparable, E extends Writable, V extends Writable>

   SimpleVertexWriter<I, V, E> noOpVertexSaver() {

     return new SimpleVertexWriter<I, V, E>() {

       @Override

       public void writeVertex(Vertex<I, V, E> vertex)

           throws IOException, InterruptedException {

         // No-op

       }

     };

   }

 }