/*

  * 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 com.google.common.collect.Sets;

 import org.apache.commons.lang3.tuple.MutablePair;

 import org.apache.commons.lang3.tuple.Pair;

 import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;

 import org.apache.giraph.conf.IntConfOption;

 import org.apache.giraph.ooc.OutOfCoreEngine;

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

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

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

 import org.apache.log4j.Logger;

 import java.io.DataInput;

 import java.io.DataOutput;

 import java.io.IOException;

 import java.util.ArrayList;

 import java.util.HashSet;

 import java.util.List;

 import java.util.Map;

 import java.util.Set;

 import java.util.concurrent.ConcurrentMap;

 import java.util.concurrent.locks.ReadWriteLock;

 import java.util.concurrent.locks.ReentrantReadWriteLock;

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

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

 import static org.apache.giraph.conf.GiraphConstants.ONE_MB;

 /**

  * This class provides basic operations for data structures that have to

  * participate in out-of-core mechanism. Essential subclasses of this class are:

  *  - DiskBackedPartitionStore (for partition data)

  *  - DiskBackedMessageStore (for messages)

  *  - DiskBackedEdgeStore (for edges read in INPUT_SUPERSTEP)

  * Basically, any data structure that may cause OOM to happen can be implemented

  * as a subclass of this class.

  *

  * There are two different terms used in the rest of this class:

  *  - "data store" refers to in-memory representation of data. Usually this is

  *    stored per-partition in in-memory implementations of data structures. For

  *    instance, "data store" of a DiskBackedPartitionStore would collection of

  *    all partitions kept in the in-memory partition store within the

  *    DiskBackedPartitionStore.

  *  - "raw data buffer" refers to raw data which were supposed to be

  *    de-serialized and added to the data store, but they remain 'as is' in the

  *    memory because their corresponding partition is offloaded to disk and is

  *    not available in the data store.

  *

  * @param <T> raw data format of the data store subclassing this class

  */

 public abstract class DiskBackedDataStore<T> {

   /**

    * Minimum size of a buffer (in bytes) to flush to disk. This is used to

    * decide whether vertex/edge buffers are large enough to flush to disk.

    */

   public static final IntConfOption MINIMUM_BUFFER_SIZE_TO_FLUSH =

       new IntConfOption("giraph.flushBufferSize", 8 * ONE_MB,

           "Minimum size of a buffer (in bytes) to flush to disk.");

   /** Class logger. */

   private static final Logger LOG = Logger.getLogger(

       DiskBackedDataStore.class);

   /** Out-of-core engine */

   protected final OutOfCoreEngine oocEngine;

   /**

    * Set containing ids of all partitions where the partition data is in some

    * file on disk.

    * Note that the out-of-core mechanism may decide to put the data for a

    * partition on disk, while the partition data is empty. For instance, at the

    * beginning of a superstep, out-of-core mechanism may decide to put incoming

    * messages of a partition on disk, while the partition has not received any

    * messages. In such scenarios, the "out-of-core mechanism" thinks that the

    * partition data is on disk, while disk-backed data stores may want to

    * optimize for IO/metadata accesses and decide not to create/write anything

    * on files on disk.

    * In summary, there is a subtle difference between this field and

    * `hasPartitionOnDisk` field. Basically, this field is used for optimizing

    * IO (mainly metadata) accesses by disk-backed stores, while

    * `hasPartitionDataOnDisk` is the view that out-of-core mechanism has

    * regarding partition storage statuses. Since out-of-core mechanism does not

    * know about the actual data for a partition, these two fields have to be

    * separate.

    */

   protected final Set<Integer> hasPartitionDataOnFile =

       Sets.newConcurrentHashSet();

   /** Cached value for MINIMUM_BUFFER_SIZE_TO_FLUSH */

   private final int minBufferSizeToOffload;

   /** Set containing ids of all out-of-core partitions */

   private final Set<Integer> hasPartitionDataOnDisk =

       Sets.newConcurrentHashSet();

   /**

    * Map of partition ids to list of raw data buffers. The map will have entries

    * only for partitions that their in-memory data structures are currently

    * offloaded to disk. We keep the aggregate size of buffers for each partition

    * as part of the values in the map to estimate how much memory we can free up

    * if we offload data buffers of a particular partition to disk.

    */

   private final ConcurrentMap<Integer, Pair<Integer, List<T>>> dataBuffers =

       Maps.newConcurrentMap();

   /**

    * Map of partition ids to number of raw data buffers offloaded to disk for

    * each partition. The map will have entries only for partitions that their

    * in-memory data structures are currently out of core. It is necessary to

    * know the number of data buffers on disk for a particular partition when we

    * are loading all these buffers back in memory.

    */

   private final ConcurrentMap<Integer, Integer> numDataBuffersOnDisk =

       Maps.newConcurrentMap();

   /**

    * Lock to avoid overlapping of read and write on data associated with each

    * partition.

    * */

   private final ConcurrentMap<Integer, ReadWriteLock> locks =

       Maps.newConcurrentMap();

   /**

    * Constructor.

    *

    * @param conf Configuration

    * @param oocEngine Out-of-core engine

    */

   DiskBackedDataStore(ImmutableClassesGiraphConfiguration conf,

                       OutOfCoreEngine oocEngine) {

     this.minBufferSizeToOffload = MINIMUM_BUFFER_SIZE_TO_FLUSH.get(conf);

     this.oocEngine = oocEngine;

   }

   /**

    * Retrieves a lock for a given partition. If the lock for the given partition

    * does not exist, creates a new lock.

    *

    * @param partitionId id of the partition the lock is needed for

    * @return lock for a given partition

    */

   private ReadWriteLock getPartitionLock(int partitionId) {

     ReadWriteLock readWriteLock = locks.get(partitionId);

     if (readWriteLock == null) {

       readWriteLock = new ReentrantReadWriteLock();

       ReadWriteLock temp = locks.putIfAbsent(partitionId, readWriteLock);

       if (temp != null) {

         readWriteLock = temp;

       }

     }

     return readWriteLock;

   }

   /**

    * Adds a data entry for a given partition to the current data store. If data

    * of a given partition in data store is already offloaded to disk, adds the

    * data entry to appropriate raw data buffer list.

    *

    * @param partitionId id of the partition to add the data entry to

    * @param entry data entry to add

    */

   protected void addEntry(int partitionId, T entry) {

     // Addition of data entries to a data store is much more common than

     // out-of-core operations. Besides, in-memory data store implementations

     // existing in the code base already account for parallel addition to data

     // stores. Therefore, using read lock would optimize for parallel addition

     // to data stores, specially for cases where the addition should happen for

     // partitions that are entirely in memory.

     ReadWriteLock rwLock = getPartitionLock(partitionId);

     rwLock.readLock().lock();

     if (hasPartitionDataOnDisk.contains(partitionId)) {

       List<T> entryList = new ArrayList<>();

       entryList.add(entry);

       int entrySize = entrySerializedSize(entry);

       MutablePair<Integer, List<T>> newPair =

           new MutablePair<>(entrySize, entryList);

       Pair<Integer, List<T>> oldPair =

           dataBuffers.putIfAbsent(partitionId, newPair);

       if (oldPair != null) {

         synchronized (oldPair) {

           newPair = (MutablePair<Integer, List<T>>) oldPair;

           newPair.setLeft(oldPair.getLeft() + entrySize);

           newPair.getRight().add(entry);

         }

       }

     } else {

       addEntryToInMemoryPartitionData(partitionId, entry);

     }

     rwLock.readLock().unlock();

   }

   /**

    * Loads and assembles all data for a given partition, and put it into the

    * data store. Returns the number of bytes transferred from disk to memory in

    * the loading process.

    *

    * @param partitionId id of the partition to load and assemble all data for

    * @return number of bytes loaded from disk to memory

    * @throws IOException

    */

   public abstract long loadPartitionData(int partitionId) throws IOException;

   /**

    * The proxy method that does the actual operation for `loadPartitionData`,

    * but uses the data index given by the caller.

    *

    * @param partitionId id of the partition to load and assemble all data for

    * @param index data index chain for the data to load

    * @return number of bytes loaded from disk to memory

    * @throws IOException

    */

   protected long loadPartitionDataProxy(int partitionId, DataIndex index)

       throws IOException {

     long numBytes = 0;

     ReadWriteLock rwLock = getPartitionLock(partitionId);

     rwLock.writeLock().lock();

     if (hasPartitionDataOnDisk.contains(partitionId)) {

       int ioThreadId =

           oocEngine.getMetaPartitionManager().getOwnerThreadId(partitionId);

       numBytes += loadInMemoryPartitionData(partitionId, ioThreadId,

           index.addIndex(NumericIndexEntry.createPartitionEntry(partitionId)));

       hasPartitionDataOnDisk.remove(partitionId);

       // Loading raw data buffers from disk if there is any and applying those

       // to already loaded in-memory data.

       Integer numBuffers = numDataBuffersOnDisk.remove(partitionId);

       if (numBuffers != null) {

         checkState(numBuffers > 0);

         index.addIndex(DataIndex.TypeIndexEntry.BUFFER);

         OutOfCoreDataAccessor.DataInputWrapper inputWrapper =

             oocEngine.getDataAccessor().prepareInput(ioThreadId, index.copy());

         DataInput dataInput = inputWrapper.getDataInput();

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

           T entry = readNextEntry(dataInput);

           addEntryToInMemoryPartitionData(partitionId, entry);

         }

         numBytes += inputWrapper.finalizeInput(true);

         index.removeLastIndex();

       }

       index.removeLastIndex();

       // Applying in-memory raw data buffers to in-memory partition data.

       Pair<Integer, List<T>> pair = dataBuffers.remove(partitionId);

       if (pair != null) {

         for (T entry : pair.getValue()) {

           addEntryToInMemoryPartitionData(partitionId, entry);

         }

       }

     }

     rwLock.writeLock().unlock();

     return numBytes;

   }

   /**

    * Offloads partition data of a given partition in the data store to disk, and

    * returns the number of bytes offloaded from memory to disk.

    *

    * @param partitionId id of the partition to offload its data

    * @return number of bytes offloaded from memory to disk

    * @throws IOException

    */

   public abstract long offloadPartitionData(int partitionId) throws IOException;

   /**

    * The proxy method that does the actual operation for `offloadPartitionData`,

    * but uses the data index given by the caller.

    *

    * @param partitionId id of the partition to offload its data

    * @param index data index chain for the data to offload

    * @return number of bytes offloaded from memory to disk

    * @throws IOException

    */

   @edu.umd.cs.findbugs.annotations.SuppressWarnings(

       "UL_UNRELEASED_LOCK_EXCEPTION_PATH")

   protected long offloadPartitionDataProxy(

       int partitionId, DataIndex index) throws IOException {

     ReadWriteLock rwLock = getPartitionLock(partitionId);

     rwLock.writeLock().lock();

     hasPartitionDataOnDisk.add(partitionId);

     rwLock.writeLock().unlock();

     int ioThreadId =

         oocEngine.getMetaPartitionManager().getOwnerThreadId(partitionId);

     long numBytes = offloadInMemoryPartitionData(partitionId, ioThreadId,

         index.addIndex(NumericIndexEntry.createPartitionEntry(partitionId)));

     index.removeLastIndex();

     return numBytes;

   }

   /**

    * Offloads raw data buffers of a given partition to disk, and returns the

    * number of bytes offloaded from memory to disk.

    *

    * @param partitionId id of the partition to offload its raw data buffers

    * @return number of bytes offloaded from memory to disk

    * @throws IOException

    */

   public abstract long offloadBuffers(int partitionId) throws IOException;

   /**

    * The proxy method that does the actual operation for `offloadBuffers`,

    * but uses the data index given by the caller.

    *

    * @param partitionId id of the partition to offload its raw data buffers

    * @param index data index chain for the data to offload its buffers

    * @return number of bytes offloaded from memory to disk

    * @throws IOException

    */

   protected long offloadBuffersProxy(int partitionId, DataIndex index)

       throws IOException {

     Pair<Integer, List<T>> pair = dataBuffers.get(partitionId);

     if (pair == null || pair.getLeft() < minBufferSizeToOffload) {

       return 0;

     }

     ReadWriteLock rwLock = getPartitionLock(partitionId);

     rwLock.writeLock().lock();

     pair = dataBuffers.remove(partitionId);

     rwLock.writeLock().unlock();

     checkNotNull(pair);

     checkState(!pair.getRight().isEmpty());

     int ioThreadId =

         oocEngine.getMetaPartitionManager().getOwnerThreadId(partitionId);

     index.addIndex(NumericIndexEntry.createPartitionEntry(partitionId))

         .addIndex(DataIndex.TypeIndexEntry.BUFFER);

     OutOfCoreDataAccessor.DataOutputWrapper outputWrapper =

         oocEngine.getDataAccessor().prepareOutput(ioThreadId, index.copy(),

             true);

     for (T entry : pair.getRight()) {

       writeEntry(entry, outputWrapper.getDataOutput());

     }

     long numBytes = outputWrapper.finalizeOutput();

     index.removeLastIndex().removeLastIndex();

     int numBuffers = pair.getRight().size();

     Integer oldNumBuffersOnDisk =

         numDataBuffersOnDisk.putIfAbsent(partitionId, numBuffers);

     if (oldNumBuffersOnDisk != null) {

       numDataBuffersOnDisk.replace(partitionId,

           oldNumBuffersOnDisk + numBuffers);

     }

     return numBytes;

   }

   /**

    * Looks through all partitions that their data is not in the data store (is

    * offloaded to disk), and sees if any of them has enough raw data buffer in

    * memory. If so, puts that partition in a list to return.

    *

    * @param ioThreadId Id of the IO thread who would offload the buffers

    * @return Set of partition ids of all partition raw buffers where the

    *         aggregate size of buffers are large enough and it is worth flushing

    *         those buffers to disk

    */

   public Set<Integer> getCandidateBuffersToOffload(int ioThreadId) {

     Set<Integer> result = new HashSet<>();

     for (Map.Entry<Integer, Pair<Integer, List<T>>> entry :

         dataBuffers.entrySet()) {

       int partitionId = entry.getKey();

       long aggregateBufferSize = entry.getValue().getLeft();

       if (aggregateBufferSize > minBufferSizeToOffload &&

           oocEngine.getMetaPartitionManager().getOwnerThreadId(partitionId) ==

               ioThreadId) {

         result.add(partitionId);

       }

     }

     return result;

   }

   /**

    * Writes a single raw entry to a given output stream.

    *

    * @param entry entry to write to output

    * @param out output stream to write the entry to

    * @throws IOException

    */

   protected abstract void writeEntry(T entry, DataOutput out)

       throws IOException;

   /**

    * Reads the next available raw entry from a given input stream.

    *

    * @param in input stream to read the entry from

    * @return entry read from an input stream

    * @throws IOException

    */

   protected abstract T readNextEntry(DataInput in) throws IOException;

   /**

    * Loads data of a partition into data store. Returns number of bytes loaded.

    *

    * @param partitionId id of the partition to load its data

    * @param ioThreadId id of the IO thread performing the load

    * @param index data index chain for the data to load

    * @return number of bytes loaded from disk to memory

    * @throws IOException

    */

   protected abstract long loadInMemoryPartitionData(

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

   /**

    * Offloads data of a partition in data store to disk. Returns the number of

    * bytes offloaded to disk

    *

    * @param partitionId id of the partition to offload to disk

    * @param ioThreadId id of the IO thread performing the offload

    * @param index data index chain for the data to offload

    * @return number of bytes offloaded from memory to disk

    * @throws IOException

    */

   protected abstract long offloadInMemoryPartitionData(

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

   /**

    * Gets the size of a given entry in bytes.

    *

    * @param entry input entry to find its size

    * @return size of given input entry in bytes

    */

   protected abstract int entrySerializedSize(T entry);

   /**

    * Adds a single entry for a given partition to the in-memory data store.

    *

    * @param partitionId id of the partition to add the data to

    * @param entry input entry to add to the data store

    */

   protected abstract void addEntryToInMemoryPartitionData(int partitionId,

                                                           T entry);

 }