/*

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

 import org.apache.giraph.combiner.MessageCombiner;

 import org.apache.giraph.comm.messages.primitives.IdByteArrayMessageStore;

 import org.apache.giraph.comm.messages.primitives.IdOneMessagePerVertexStore;

 import org.apache.giraph.comm.messages.primitives.IntFloatMessageStore;

 import org.apache.giraph.comm.messages.primitives.LongDoubleMessageStore;

 import org.apache.giraph.comm.messages.primitives.long_id.LongPointerListPerVertexStore;

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

 import org.apache.giraph.conf.GiraphConstants;

 import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;

 import org.apache.giraph.conf.MessageClasses;

 import org.apache.giraph.factories.MessageValueFactory;

 import org.apache.giraph.types.ops.PrimitiveIdTypeOps;

 import org.apache.giraph.types.ops.TypeOpsUtils;

 import org.apache.hadoop.io.DoubleWritable;

 import org.apache.hadoop.io.FloatWritable;

 import org.apache.hadoop.io.IntWritable;

 import org.apache.hadoop.io.LongWritable;

 import org.apache.hadoop.io.Writable;

 import org.apache.hadoop.io.WritableComparable;

 import org.apache.log4j.Logger;

 /**

  * Message store factory which produces message stores which hold all

  * messages in memory. Depending on whether or not combiner is currently used,

  * this factory creates {@link OneMessagePerVertexStore} or

  * {@link ByteArrayMessagesPerVertexStore}

  *

  * @param <I> Vertex id

  * @param <M> Message data

  */

 @SuppressWarnings("unchecked")

 public class InMemoryMessageStoreFactory<I extends WritableComparable,

     M extends Writable>

     implements MessageStoreFactory<I, M, MessageStore<I, M>> {

   /** Class logger */

   private static final Logger LOG =

       Logger.getLogger(InMemoryMessageStoreFactory.class);

   /** Partition info */

   protected PartitionSplitInfo<I> partitionInfo;

   /** Hadoop configuration */

   protected ImmutableClassesGiraphConfiguration<I, ?, ?> conf;

   /**

    * Default constructor allowing class invocation via Reflection.

    */

   public InMemoryMessageStoreFactory() {

   }

   /**

    * MessageStore to be used when combiner is enabled

    *

    * @param messageClass message class

    * @param messageValueFactory message value factory

    * @param messageCombiner message combiner

    * @return message store

    */

   protected MessageStore<I, M> newStoreWithCombiner(

       Class<M> messageClass,

       MessageValueFactory<M> messageValueFactory,

       MessageCombiner<? super I, M> messageCombiner) {

     MessageStore messageStore;

     Class<I> vertexIdClass = conf.getVertexIdClass();

     if (vertexIdClass.equals(IntWritable.class) &&

         messageClass.equals(FloatWritable.class)) {

       messageStore = new IntFloatMessageStore(

           (PartitionSplitInfo<IntWritable>) partitionInfo,

           (MessageCombiner<IntWritable, FloatWritable>) messageCombiner);

     } else if (vertexIdClass.equals(LongWritable.class) &&

         messageClass.equals(DoubleWritable.class)) {

       messageStore = new LongDoubleMessageStore(

           (PartitionSplitInfo<LongWritable>) partitionInfo,

           (MessageCombiner<LongWritable, DoubleWritable>) messageCombiner);

     } else {

       PrimitiveIdTypeOps<I> idTypeOps =

           TypeOpsUtils.getPrimitiveIdTypeOpsOrNull(vertexIdClass);

       if (idTypeOps != null) {

         messageStore = new IdOneMessagePerVertexStore<>(

           messageValueFactory, partitionInfo, messageCombiner, conf);

       } else {

         messageStore = new OneMessagePerVertexStore<I, M>(

           messageValueFactory, partitionInfo, messageCombiner, conf);

       }

     }

     return messageStore;

   }

   /**

    * MessageStore to be used when combiner is not enabled

    *

    * @param messageClass message class

    * @param messageValueFactory message value factory

    * @param encodeAndStore message encode and store type

    * @return message store

    */

   protected MessageStore<I, M> newStoreWithoutCombiner(

       Class<M> messageClass,

       MessageValueFactory<M> messageValueFactory,

       MessageEncodeAndStoreType encodeAndStore) {

     MessageStore messageStore = null;

     Class<I> vertexIdClass = conf.getVertexIdClass();

     // a special case for LongWritable with POINTER_LIST_PER_VERTEX

     if (vertexIdClass.equals(LongWritable.class) && encodeAndStore.equals(

         MessageEncodeAndStoreType.POINTER_LIST_PER_VERTEX)) {

       messageStore = new LongPointerListPerVertexStore(

         messageValueFactory, partitionInfo, conf);

     } else { // GENERAL

       if (encodeAndStore.equals(

           MessageEncodeAndStoreType.BYTEARRAY_PER_PARTITION) ||

           encodeAndStore.equals(

               MessageEncodeAndStoreType.EXTRACT_BYTEARRAY_PER_PARTITION)) {

         PrimitiveIdTypeOps<I> idTypeOps =

             TypeOpsUtils.getPrimitiveIdTypeOpsOrNull(vertexIdClass);

         if (idTypeOps != null) {

           messageStore = new IdByteArrayMessageStore<>(

               messageValueFactory, partitionInfo, conf);

         } else {

           messageStore = new ByteArrayMessagesPerVertexStore<>(

               messageValueFactory, partitionInfo, conf);

         }

       } else if (encodeAndStore.equals(

           MessageEncodeAndStoreType.POINTER_LIST_PER_VERTEX)) {

         messageStore = new PointerListPerVertexStore<>(

           messageValueFactory, partitionInfo, conf);

       }

     }

     return messageStore;

   }

   @Override

   public MessageStore<I, M> newStore(

       MessageClasses<I, M> messageClasses) {

     Class<M> messageClass = messageClasses.getMessageClass();

     MessageValueFactory<M> messageValueFactory =

         messageClasses.createMessageValueFactory(conf);

     MessageCombiner<? super I, M> messageCombiner =

         messageClasses.createMessageCombiner(conf);

     MessageStore messageStore;

     if (messageCombiner != null) {

       messageStore = newStoreWithCombiner(

           messageClass, messageValueFactory, messageCombiner);

     } else {

       messageStore = newStoreWithoutCombiner(

           messageClass, messageValueFactory,

           messageClasses.getMessageEncodeAndStoreType());

     }

     if (LOG.isInfoEnabled()) {

       LOG.info("newStore: Created " + messageStore.getClass() +

           " for vertex id " + conf.getVertexIdClass() +

           " and message value " + messageClass + " and" +

           (messageCombiner != null ? " message combiner " +

               messageCombiner.getClass() : " no combiner"));

     }

     int asyncMessageStoreThreads =

         GiraphConstants.ASYNC_MESSAGE_STORE_THREADS_COUNT.get(conf);

     if (asyncMessageStoreThreads > 0) {

       messageStore = new AsyncMessageStoreWrapper(

           messageStore,

           partitionInfo.getPartitionIds(),

           asyncMessageStoreThreads);

     }

     return messageStore;

   }

   @Override

   public void initialize(PartitionSplitInfo<I> partitionInfo,

       ImmutableClassesGiraphConfiguration<I, ?, ?> conf) {

     this.partitionInfo = partitionInfo;

     this.conf = conf;

   }

 }