/*

  * 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.ooc.data;

 import com.google.common.collect.Maps;

 import org.apache.giraph.bsp.BspService;

 import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;

 import org.apache.giraph.edge.OutEdges;

 import org.apache.giraph.graph.Vertex;

 import org.apache.giraph.ooc.OutOfCoreEngine;

 import org.apache.giraph.ooc.persistence.DataIndex;

 import org.apache.giraph.ooc.persistence.OutOfCoreDataAccessor;

 import org.apache.giraph.partition.Partition;

 import org.apache.giraph.partition.PartitionStore;

 import org.apache.giraph.utils.ExtendedDataOutput;

 import org.apache.giraph.utils.WritableUtils;

 import org.apache.giraph.worker.BspServiceWorker;

 import org.apache.hadoop.io.Writable;

 import org.apache.hadoop.io.WritableComparable;

 import org.apache.hadoop.mapreduce.Mapper;

 import org.apache.log4j.Logger;

 import java.io.DataInput;

 import java.io.DataOutput;

 import java.io.IOException;

 import java.util.Map;

 import static com.google.common.base.Preconditions.checkNotNull;

 /**

  * Implementation of a partition-store used for out-of-core mechanism.

  * Partition store is responsible for partition data, as well as data buffers in

  * INPUT_SUPERSTEP ("raw data buffer" -- defined in DiskBackedDataStore --

  * refers to vertex buffers in INPUT_SUPERSTEP).

  *

  * @param <I> Vertex id

  * @param <V> Vertex data

  * @param <E> Edge data

  */

 public class DiskBackedPartitionStore<I extends WritableComparable,

     V extends Writable, E extends Writable>

     extends DiskBackedDataStore<ExtendedDataOutput>

     implements PartitionStore<I, V, E> {

   /** Class logger. */

   private static final Logger LOG =

       Logger.getLogger(DiskBackedPartitionStore.class);

   /** Configuration */

   private final ImmutableClassesGiraphConfiguration<I, V, E> conf;

   /** Job context (for progress) */

   private final Mapper<?, ?, ?, ?>.Context context;

   /** In-memory partition store */

   private final PartitionStore<I, V, E> partitionStore;

   /**

    * Keeps number of vertices in partitions, right when they are last spilled

    * to the disk. This value may be inaccurate for in-memory partitions, but

    * is accurate for out-of-core partitions.

    */

   private final Map<Integer, Long> partitionVertexCount =

       Maps.newConcurrentMap();

   /**

    * Keeps number of edges in partitions, right when they are last spilled

    * to the disk. This value may be inaccurate for in-memory partitions, but

    * is accurate for out-of-core partitions.

    */

   private final Map<Integer, Long> partitionEdgeCount =

       Maps.newConcurrentMap();

   /**

    * Constructor.

    *

    * @param partitionStore In-memory partition store for which out-of-code

    *                       partition store would be a wrapper

    * @param conf Configuration

    * @param context Job context

    * @param oocEngine Out-of-core engine

    */

   public DiskBackedPartitionStore(

       PartitionStore<I, V, E> partitionStore,

       ImmutableClassesGiraphConfiguration<I, V, E> conf,

       Mapper<?, ?, ?, ?>.Context context,

       OutOfCoreEngine oocEngine) {

     super(conf, oocEngine);

     this.partitionStore = partitionStore;

     this.conf = conf;

     this.context = context;

   }

   @Override

   public boolean addPartition(Partition<I, V, E> partition) {

     boolean added = partitionStore.addPartition(partition);

     if (added) {

       oocEngine.getMetaPartitionManager()

           .addPartition(partition.getId());

     }

     return added;

   }

   @Override

   public Partition<I, V, E> removePartition(Integer partitionId) {

     // Set the partition as 'in process' so its data and messages do not get

     // spilled to disk until the remove is complete.

     oocEngine.getMetaPartitionManager().markPartitionAsInProcess(partitionId);

     oocEngine.retrievePartition(partitionId);

     Partition<I, V, E> partition = partitionStore.removePartition(partitionId);

     checkNotNull(partition, "removePartition: partition " + partitionId +

         " is not in memory for removal!");

     oocEngine.getMetaPartitionManager().removePartition(partitionId);

     return partition;

   }

   @Override

   public boolean hasPartition(Integer partitionId) {

     return oocEngine.getMetaPartitionManager().hasPartition(partitionId);

   }

   @Override

   public Iterable<Integer> getPartitionIds() {

     return oocEngine.getMetaPartitionManager().getPartitionIds();

   }

   @Override

   public int getNumPartitions() {

     return oocEngine.getMetaPartitionManager().getNumPartitions();

   }

   @Override

   public long getPartitionVertexCount(Integer partitionId) {

     if (partitionStore.hasPartition(partitionId)) {

       return partitionStore.getPartitionVertexCount(partitionId);

     } else {

       return partitionVertexCount.get(partitionId);

     }

   }

   @Override

   public long getPartitionEdgeCount(Integer partitionId) {

     if (partitionStore.hasPartition(partitionId)) {

       return partitionStore.getPartitionEdgeCount(partitionId);

     } else {

       return partitionEdgeCount.get(partitionId);

     }

   }

   @Override

   public boolean isEmpty() {

     return getNumPartitions() == 0;

   }

   @Override

   public void startIteration() {

     oocEngine.startIteration();

   }

   @Override

   public Partition<I, V, E> getNextPartition() {

     Integer partitionId = oocEngine.getNextPartition();

     if (partitionId == null) {

       return null;

     }

     Partition<I, V, E> partition = partitionStore.removePartition(partitionId);

     if (partition == null) {

       if (LOG.isInfoEnabled()) {

         LOG.info("getNextPartition: partition " + partitionId + " is not in " +

             "the partition store. Creating an empty partition for it.");

       }

       partition = conf.createPartition(partitionId, context);

     }

     partitionStore.addPartition(partition);

     return partition;

   }

   @Override

   public void putPartition(Partition<I, V, E> partition) {

     oocEngine.doneProcessingPartition(partition.getId());

   }

   @Override

   public void addPartitionVertices(Integer partitionId,

                                    ExtendedDataOutput extendedDataOutput) {

     addEntry(partitionId, extendedDataOutput);

   }

   @Override

   public void shutdown() {

     oocEngine.shutdown();

   }

   @Override

   public void initialize() {

     oocEngine.initialize();

   }

   /**

    * Read vertex data from an input and initialize the vertex.

    *

    * @param in     The input stream

    * @param vertex The vertex to initialize

    * @throws IOException

    */

   private void readVertexData(DataInput in, Vertex<I, V, E> vertex)

       throws IOException {

     I id = conf.createVertexId();

     id.readFields(in);

     V value = null;

     boolean hasNullValue = in.readBoolean();

     if (!hasNullValue) {

       value = conf.createVertexValue();

       value.readFields(in);

     }

     OutEdges<I, E> edges = conf.createAndInitializeOutEdges(0);

     vertex.initialize(id, value, edges);

     if (in.readBoolean()) {

       vertex.voteToHalt();

     } else {

       vertex.wakeUp();

     }

   }

   /**

    * Read vertex edges from an input and set them to the vertex.

    *

    * @param in        The input stream

    * @param partition The partition owning the vertex

    * @throws IOException

    */

   private void readOutEdges(DataInput in, Partition<I, V, E> partition)

       throws IOException {

     I id = conf.createVertexId();

     id.readFields(in);

     Vertex<I, V, E> v = partition.getVertex(id);

     if (v == null) {

       throw new IllegalStateException("Vertex with ID " + id +

         " not found in partition " + partition.getId() +

         " which has " + partition.getVertexCount() + " vertices and " +

         partition.getEdgeCount() + " edges.");

     }

     OutEdges<I, E> edges = (OutEdges<I, E>) v.getEdges();

     edges.readFields(in);

     partition.saveVertex(v);

   }

   @Override

   protected long loadInMemoryPartitionData(int partitionId, int ioThreadId,

                                            DataIndex index) throws IOException {

     long numBytes = 0;

     // Load vertices

     if (hasPartitionDataOnFile.remove(partitionId)) {

       Partition<I, V, E> partition = conf.createPartition(partitionId, context);

       OutOfCoreDataAccessor dataAccessor = oocEngine.getDataAccessor();

       index.addIndex(DataIndex.TypeIndexEntry.PARTITION_VERTICES);

       OutOfCoreDataAccessor.DataInputWrapper inputWrapper =

           dataAccessor.prepareInput(ioThreadId, index.copy());

       DataInput dataInput = inputWrapper.getDataInput();

       long numVertices = dataInput.readLong();

       for (long i = 0; i < numVertices; ++i) {

         Vertex<I, V, E> vertex = conf.createVertex();

         readVertexData(dataInput, vertex);

         partition.putVertex(vertex);

       }

       numBytes += inputWrapper.finalizeInput(true);

       // Load edges

       index.removeLastIndex()

           .addIndex(DataIndex.TypeIndexEntry.PARTITION_EDGES);

       inputWrapper = dataAccessor.prepareInput(ioThreadId, index.copy());

       dataInput = inputWrapper.getDataInput();

       for (int i = 0; i < numVertices; ++i) {

         readOutEdges(dataInput, partition);

       }

       // If the graph is static and it is not INPUT_SUPERSTEP, keep the file

       // around.

       boolean shouldDeleteEdges = false;

       if (!conf.isStaticGraph() ||

           oocEngine.getSuperstep() == BspService.INPUT_SUPERSTEP) {

         shouldDeleteEdges = true;

       }

       numBytes += inputWrapper.finalizeInput(shouldDeleteEdges);

       index.removeLastIndex();

       partitionStore.addPartition(partition);

     }

     return numBytes;

   }

   @Override

   protected ExtendedDataOutput readNextEntry(DataInput in) throws IOException {

     return WritableUtils.readExtendedDataOutput(in, conf);

   }

   @Override

   protected void addEntryToInMemoryPartitionData(int partitionId,

                                                  ExtendedDataOutput vertices) {

     if (!partitionStore.hasPartition(partitionId)) {

       oocEngine.getMetaPartitionManager().addPartition(partitionId);

     }

     partitionStore.addPartitionVertices(partitionId, vertices);

   }

   @Override

   public long loadPartitionData(int partitionId)

       throws IOException {

     return loadPartitionDataProxy(partitionId,

         new DataIndex().addIndex(DataIndex.TypeIndexEntry.PARTITION));

   }

   @Override

   public long offloadPartitionData(int partitionId)

       throws IOException {

     return offloadPartitionDataProxy(partitionId,

         new DataIndex().addIndex(DataIndex.TypeIndexEntry.PARTITION));

   }

   /**

    * Writes vertex data (Id, value and halted state) to stream.

    *

    * @param output The output stream

    * @param vertex The vertex to serialize

    * @throws IOException

    */

   private void writeVertexData(DataOutput output, Vertex<I, V, E> vertex)

       throws IOException {

     vertex.getId().write(output);

     V value = vertex.getValue();

     if (value != null) {

       output.writeBoolean(false);

       value.write(output);

     } else {

       output.writeBoolean(true);

     }

     output.writeBoolean(vertex.isHalted());

   }

   /**

    * Writes vertex edges (Id, edges) to stream.

    *

    * @param output The output stream

    * @param vertex The vertex to serialize

    * @throws IOException

    */

   private void writeOutEdges(DataOutput output, Vertex<I, V, E> vertex)

       throws IOException {

     vertex.getId().write(output);

     OutEdges<I, E> edges = (OutEdges<I, E>) vertex.getEdges();

     edges.write(output);

   }

   @Override

   protected long offloadInMemoryPartitionData(

       int partitionId, int ioThreadId, DataIndex index) throws IOException {

     long numBytes = 0;

     if (partitionStore.hasPartition(partitionId)) {

       OutOfCoreDataAccessor dataAccessor = oocEngine.getDataAccessor();

       partitionVertexCount.put(partitionId,

           partitionStore.getPartitionVertexCount(partitionId));

       partitionEdgeCount.put(partitionId,

           partitionStore.getPartitionEdgeCount(partitionId));

       Partition<I, V, E> partition =

           partitionStore.removePartition(partitionId);

       LOG.debug(

           "Offloading partition " + partition + " DataIndex[" + index + "]");

       index.addIndex(DataIndex.TypeIndexEntry.PARTITION_VERTICES);

       OutOfCoreDataAccessor.DataOutputWrapper outputWrapper =

           dataAccessor.prepareOutput(ioThreadId, index.copy(), false);

       DataOutput dataOutput = outputWrapper.getDataOutput();

       dataOutput.writeLong(partition.getVertexCount());

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

         writeVertexData(dataOutput, vertex);

       }

       numBytes += outputWrapper.finalizeOutput();

       index.removeLastIndex();

       // Avoid writing back edges if we have already written them once and

       // the graph is not changing.

       // If we are in the input superstep, we need to write the files

       // at least the first time, even though the graph is static.

       index.addIndex(DataIndex.TypeIndexEntry.PARTITION_EDGES);

       if (oocEngine.getSuperstep() == BspServiceWorker.INPUT_SUPERSTEP ||

           !conf.isStaticGraph() ||

           !dataAccessor.dataExist(ioThreadId, index)) {

         outputWrapper = dataAccessor.prepareOutput(ioThreadId, index.copy(),

             false);

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

           writeOutEdges(outputWrapper.getDataOutput(), vertex);

         }

         numBytes += outputWrapper.finalizeOutput();

       }

       index.removeLastIndex();

       hasPartitionDataOnFile.add(partitionId);

     }

     return numBytes;

   }

   @Override

   protected void writeEntry(ExtendedDataOutput vertices, DataOutput out)

       throws IOException {

     WritableUtils.writeExtendedDataOutput(vertices, out);

   }

   @Override

   public long offloadBuffers(int partitionId)

       throws IOException {

     return offloadBuffersProxy(partitionId,

         new DataIndex().addIndex(DataIndex.TypeIndexEntry.PARTITION));

   }

   @Override

   protected int entrySerializedSize(ExtendedDataOutput vertices) {

     return vertices.getPos();

   }

 }