/*

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

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

 import java.util.ArrayList;

 import java.util.Collections;

 import java.util.List;

 import java.util.Map;

 import java.util.concurrent.ConcurrentMap;

 import org.apache.giraph.bsp.CentralizedServiceWorker;

 import org.apache.giraph.comm.aggregators.AllAggregatorServerData;

 import org.apache.giraph.comm.aggregators.OwnerAggregatorServerData;

 import org.apache.giraph.comm.messages.MessageStore;

 import org.apache.giraph.comm.messages.MessageStoreFactory;

 import org.apache.giraph.comm.messages.queue.AsyncMessageStoreWrapper;

 import org.apache.giraph.conf.GiraphConstants;

 import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;

 import org.apache.giraph.edge.EdgeStore;

 import org.apache.giraph.edge.EdgeStoreFactory;

 import org.apache.giraph.graph.Vertex;

 import org.apache.giraph.graph.VertexMutations;

 import org.apache.giraph.graph.VertexResolver;

 import org.apache.giraph.ooc.OutOfCoreEngine;

 import org.apache.giraph.ooc.data.DiskBackedEdgeStore;

 import org.apache.giraph.ooc.data.DiskBackedMessageStore;

 import org.apache.giraph.ooc.data.DiskBackedPartitionStore;

 import org.apache.giraph.partition.Partition;

 import org.apache.giraph.partition.PartitionStore;

 import org.apache.giraph.partition.SimplePartitionStore;

 import org.apache.giraph.utils.ReflectionUtils;

 import org.apache.hadoop.io.Writable;

 import org.apache.hadoop.io.WritableComparable;

 import org.apache.hadoop.mapreduce.Mapper;

 import org.apache.log4j.Logger;

 import com.google.common.collect.Iterables;

 import com.google.common.collect.Maps;

 /**

  * Anything that the server stores

  *

  * @param <I> Vertex id

  * @param <V> Vertex data

  * @param <E> Edge data

  */

 @SuppressWarnings("rawtypes")

 public class ServerData<I extends WritableComparable,

     V extends Writable, E extends Writable> {

   /** Class logger */

   private static final Logger LOG = Logger.getLogger(ServerData.class);

   /** Configuration */

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

   /** Partition store for this worker. */

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

   /** Edge store for this worker. */

   private final EdgeStore<I, V, E> edgeStore;

   /** Message store factory */

   private final MessageStoreFactory<I, Writable, MessageStore<I, Writable>>

       messageStoreFactory;

   /**

    * Message store for incoming messages (messages which will be consumed

    * in the next super step)

    */

   private volatile MessageStore<I, Writable> incomingMessageStore;

   /**

    * Message store for current messages (messages which we received in

    * previous super step and which will be consumed in current super step)

    */

   private volatile MessageStore<I, Writable> currentMessageStore;

   /**

    * Map of partition ids to vertex mutations from other workers. These are

    * mutations that should be applied before execution of *current* super step.

    * (accesses to keys should be thread-safe as multiple threads may resolve

    * mutations of different partitions at the same time)

    */

   private ConcurrentMap<Integer,

       ConcurrentMap<I, VertexMutations<I, V, E>>>

       oldPartitionMutations = Maps.newConcurrentMap();

   /**

    * Map of partition ids to vertex mutations from other workers. These are

    * mutations that are coming from other workers as the execution goes one in a

    * super step. These mutations should be applied in the *next* super step.

    * (this should be thread-safe)

    */

   private ConcurrentMap<Integer,

       ConcurrentMap<I, VertexMutations<I, V, E>>>

       partitionMutations = Maps.newConcurrentMap();

   /**

    * Holds aggregators which current worker owns from current superstep

    */

   private final OwnerAggregatorServerData ownerAggregatorData;

   /**

    * Holds old aggregators from previous superstep

    */

   private final AllAggregatorServerData allAggregatorData;

   /** Service worker */

   private final CentralizedServiceWorker<I, V, E> serviceWorker;

   /** Store for current messages from other workers to this worker */

   private volatile List<Writable> currentWorkerToWorkerMessages =

       Collections.synchronizedList(new ArrayList<Writable>());

   /** Store for message from other workers to this worker for next superstep */

   private volatile List<Writable> incomingWorkerToWorkerMessages =

       Collections.synchronizedList(new ArrayList<Writable>());

   /** Job context (for progress) */

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

   /** Out-of-core engine */

   private final OutOfCoreEngine oocEngine;

   /**

    * Constructor.

    *

    * @param service Service worker

    * @param workerServer Worker server

    * @param conf Configuration

    * @param context Mapper context

    */

   public ServerData(

       CentralizedServiceWorker<I, V, E> service,

       WorkerServer workerServer,

       ImmutableClassesGiraphConfiguration<I, V, E> conf,

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

     this.serviceWorker = service;

     this.conf = conf;

     this.messageStoreFactory = createMessageStoreFactory();

     EdgeStoreFactory<I, V, E> edgeStoreFactory = conf.createEdgeStoreFactory();

     edgeStoreFactory.initialize(service, conf, context);

     EdgeStore<I, V, E> inMemoryEdgeStore = edgeStoreFactory.newStore();

     PartitionStore<I, V, E> inMemoryPartitionStore =

         new SimplePartitionStore<I, V, E>(conf, context);

     if (GiraphConstants.USE_OUT_OF_CORE_GRAPH.get(conf)) {

       oocEngine = new OutOfCoreEngine(conf, service, workerServer);

       partitionStore =

           new DiskBackedPartitionStore<I, V, E>(inMemoryPartitionStore,

               conf, context, oocEngine);

       edgeStore =

           new DiskBackedEdgeStore<I, V, E>(inMemoryEdgeStore, conf, oocEngine);

     } else {

       partitionStore = inMemoryPartitionStore;

       edgeStore = inMemoryEdgeStore;

       oocEngine = null;

     }

     ownerAggregatorData = new OwnerAggregatorServerData(context);

     allAggregatorData = new AllAggregatorServerData(context, conf);

     this.context = context;

   }

   /**

    * Decide which message store should be used for current application,

    * and create the factory for that store

    *

    * @return Message store factory

    */

   private MessageStoreFactory<I, Writable, MessageStore<I, Writable>>

   createMessageStoreFactory() {

     Class<? extends MessageStoreFactory> messageStoreFactoryClass =

         MESSAGE_STORE_FACTORY_CLASS.get(conf);

     MessageStoreFactory messageStoreFactoryInstance =

         ReflectionUtils.newInstance(messageStoreFactoryClass);

     messageStoreFactoryInstance.initialize(serviceWorker, conf);

     return messageStoreFactoryInstance;

   }

   /**

    * Return the out-of-core engine for this worker.

    *

    * @return The out-of-core engine

    */

   public OutOfCoreEngine getOocEngine() {

     return oocEngine;

   }

   /**

    * Return the edge store for this worker.

    *

    * @return The edge store

    */

   public EdgeStore<I, V, E> getEdgeStore() {

     return edgeStore;

   }

   /**

    * Return the partition store for this worker.

    *

    * @return The partition store

    */

   public PartitionStore<I, V, E> getPartitionStore() {

     return partitionStore;

   }

   /**

    * Get message store for incoming messages (messages which will be consumed

    * in the next super step)

    *

    * @param <M> Message data

    * @return Incoming message store

    */

   public <M extends Writable> MessageStore<I, M> getIncomingMessageStore() {

     return (MessageStore<I, M>) incomingMessageStore;

   }

   /**

    * Get message store for current messages (messages which we received in

    * previous super step and which will be consumed in current super step)

    *

    * @param <M> Message data

    * @return Current message store

    */

   public <M extends Writable> MessageStore<I, M> getCurrentMessageStore() {

     return (MessageStore<I, M>) currentMessageStore;

   }

   /**

    * Re-initialize message stores.

    * Discards old values if any.

    */

   public void resetMessageStores() {

     if (currentMessageStore != null) {

       currentMessageStore.clearAll();

       currentMessageStore = null;

     }

     if (incomingMessageStore != null) {

       incomingMessageStore.clearAll();

       incomingMessageStore = null;

     }

     prepareSuperstep();

   }

   /** Prepare for next superstep */

   public void prepareSuperstep() {

     if (currentMessageStore != null) {

       currentMessageStore.clearAll();

     }

     MessageStore<I, Writable> nextCurrentMessageStore;

     MessageStore<I, Writable> nextIncomingMessageStore;

     MessageStore<I, Writable> messageStore;

     // First create the necessary in-memory message stores. If out-of-core

     // mechanism is enabled, we wrap the in-memory message stores within

     // disk-backed messages stores.

     if (incomingMessageStore != null) {

       nextCurrentMessageStore = incomingMessageStore;

     } else {

       messageStore = messageStoreFactory.newStore(

           conf.getIncomingMessageClasses());

       if (oocEngine == null) {

         nextCurrentMessageStore = messageStore;

       } else {

         nextCurrentMessageStore = new DiskBackedMessageStore<>(

             conf, oocEngine, messageStore,

             conf.getIncomingMessageClasses().useMessageCombiner(),

             serviceWorker.getSuperstep());

       }

     }

     messageStore = messageStoreFactory.newStore(

         conf.getOutgoingMessageClasses());

     if (oocEngine == null) {

       nextIncomingMessageStore = messageStore;

     } else {

       nextIncomingMessageStore = new DiskBackedMessageStore<>(

           conf, oocEngine, messageStore,

           conf.getOutgoingMessageClasses().useMessageCombiner(),

           serviceWorker.getSuperstep() + 1);

     }

     // If out-of-core engine is enabled, we avoid overlapping of out-of-core

     // decisions with change of superstep. This avoidance is done to simplify

     // the design and reduce excessive use of synchronization primitives.

     if (oocEngine != null) {

       oocEngine.getSuperstepLock().writeLock().lock();

     }

     currentMessageStore = nextCurrentMessageStore;

     incomingMessageStore = nextIncomingMessageStore;

     if (oocEngine != null) {

       oocEngine.reset();

       oocEngine.getSuperstepLock().writeLock().unlock();

     }

     currentMessageStore.finalizeStore();

     currentWorkerToWorkerMessages = incomingWorkerToWorkerMessages;

     incomingWorkerToWorkerMessages =

         Collections.synchronizedList(new ArrayList<Writable>());

   }

   /**

    * Get the vertex mutations (synchronize on the values)

    *

    * @return Vertex mutations

    */

   public ConcurrentMap<Integer, ConcurrentMap<I, VertexMutations<I, V, E>>>

   getPartitionMutations() {

     return partitionMutations;

   }

   /**

    * Get holder for aggregators which current worker owns

    *

    * @return Holder for aggregators which current worker owns

    */

   public OwnerAggregatorServerData getOwnerAggregatorData() {

     return ownerAggregatorData;

   }

   /**

    * Get holder for aggregators from previous superstep

    *

    * @return Holder for aggregators from previous superstep

    */

   public AllAggregatorServerData getAllAggregatorData() {

     return allAggregatorData;

   }

   /**

    * Get the reference of the service worker.

    *

    * @return CentralizedServiceWorker

    */

   public CentralizedServiceWorker<I, V, E> getServiceWorker() {

     return this.serviceWorker;

   }

   /**

    * Get and clear worker to worker messages for this superstep. Can be

    * called only once per superstep.

    *

    * @return List of messages for this worker

    */

   public List<Writable> getAndClearCurrentWorkerToWorkerMessages() {

     List<Writable> ret = currentWorkerToWorkerMessages;

     currentWorkerToWorkerMessages = null;

     return ret;

   }

   /**

    * Add incoming message to this worker for next superstep. Thread-safe.

    *

    * @param message Message received

    */

   public void addIncomingWorkerToWorkerMessage(Writable message) {

     incomingWorkerToWorkerMessages.add(message);

   }

   /**

    * Get worker to worker messages received in previous superstep.

    * @return list of current worker to worker messages.

    */

   public List<Writable> getCurrentWorkerToWorkerMessages() {

     return currentWorkerToWorkerMessages;

   }

   /**

    * Prepare resolving mutation.

    */

   public void prepareResolveMutations() {

     oldPartitionMutations = partitionMutations;

     partitionMutations = Maps.newConcurrentMap();

   }

   /**

    * Resolve mutations specific for a partition. This method is called once

    * per partition, before the computation for that partition starts.

    * @param partition The partition to resolve mutations for

    */

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

     Integer partitionId = partition.getId();

     VertexResolver<I, V, E> vertexResolver = conf.createVertexResolver();

     ConcurrentMap<I, VertexMutations<I, V, E>> prevPartitionMutations =

         oldPartitionMutations.get(partitionId);

     boolean ignoreExistingVertices =

         conf.getIncomingMessageClasses().ignoreExistingVertices();

     // Resolve mutations that are explicitly sent for this partition

     if (prevPartitionMutations != null) {

       for (Map.Entry<I, VertexMutations<I, V, E>> entry : prevPartitionMutations

           .entrySet()) {

         I vertexId = entry.getKey();

         Vertex<I, V, E> originalVertex = partition.getVertex(vertexId);

         VertexMutations<I, V, E> vertexMutations = entry.getValue();

         Vertex<I, V, E> vertex = vertexResolver.resolve(vertexId,

             originalVertex, vertexMutations,

             !ignoreExistingVertices &&

             getCurrentMessageStore().hasMessagesForVertex(entry.getKey()));

         if (LOG.isDebugEnabled()) {

           LOG.debug("resolvePartitionMutations: Resolved vertex index " +

               vertexId + " in partition index " + partitionId +

               " with original vertex " + originalVertex +

               ", returned vertex " + vertex + " on superstep " +

               serviceWorker.getSuperstep() + " with mutations " +

               vertexMutations);

         }

         if (vertex != null) {

           partition.putVertex(vertex);

         } else if (originalVertex != null) {

           partition.removeVertex(vertexId);

           if (!ignoreExistingVertices) {

             getCurrentMessageStore().clearVertexMessages(vertexId);

           }

         }

         context.progress();

       }

     }

     if (!ignoreExistingVertices) {

       // Keep track of vertices which are not here in the partition, but have

       // received messages

       Iterable<I> destinations = getCurrentMessageStore().

           getPartitionDestinationVertices(partitionId);

       if (!Iterables.isEmpty(destinations)) {

         for (I vertexId : destinations) {

           if (partition.getVertex(vertexId) == null) {

             Vertex<I, V, E> vertex =

                 vertexResolver.resolve(vertexId, null, null, true);

             if (LOG.isDebugEnabled()) {

               LOG.debug(

                   "resolvePartitionMutations: A non-existing vertex has " +

                   "message(s). Added vertex index " + vertexId +

                   " in partition index " + partitionId +

                   ", vertex = " + vertex + ", on superstep " +

                   serviceWorker.getSuperstep());

             }

             if (vertex != null) {

               partition.putVertex(vertex);

             }

             context.progress();

           }

         }

       }

     }

   }

   /**

    * In case of async message store we have to wait for all messages

    * to be processed before going into next superstep.

    */

   public void waitForComplete() {

     if (incomingMessageStore instanceof AsyncMessageStoreWrapper) {

       ((AsyncMessageStoreWrapper) incomingMessageStore).waitToComplete();

     }

   }

 }