/*

  * 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 java.io.DataInput;

 import java.io.DataOutput;

 import java.io.IOException;

 import java.util.Collections;

 import java.util.concurrent.ConcurrentMap;

 import org.apache.giraph.bsp.CentralizedServiceWorker;

 import org.apache.giraph.combiner.MessageCombiner;

 import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;

 import org.apache.giraph.conf.MessageClasses;

 import org.apache.giraph.factories.MessageValueFactory;

 import org.apache.giraph.utils.VertexIdMessageIterator;

 import org.apache.giraph.utils.VertexIdMessages;

 import org.apache.giraph.utils.WritableUtils;

 import org.apache.hadoop.io.Writable;

 import org.apache.hadoop.io.WritableComparable;

 /**

  * Implementation of {@link SimpleMessageStore} where we have a single

  * message per vertex.

  * Used when {@link org.apache.giraph.combiner.MessageCombiner} is provided.

  *

  * @param <I> Vertex id

  * @param <M> Message data

  */

 public class OneMessagePerVertexStore<I extends WritableComparable,

     M extends Writable> extends SimpleMessageStore<I, M, M> {

   /** MessageCombiner for messages */

   private final MessageCombiner<? super I, M> messageCombiner;

   /**

    * @param messageValueFactory Message class held in the store

    * @param partitionInfo Partition split info

    * @param messageCombiner MessageCombiner for messages

    * @param config Hadoop configuration

    */

   public OneMessagePerVertexStore(

     MessageValueFactory<M> messageValueFactory,

     PartitionSplitInfo<I> partitionInfo,

     MessageCombiner<? super I, M> messageCombiner,

     ImmutableClassesGiraphConfiguration<I, ?, ?> config

   ) {

     super(messageValueFactory, partitionInfo, config);

     this.messageCombiner =

         messageCombiner;

   }

   @Override

   public boolean isPointerListEncoding() {

     return false;

   }

   @Override

   public void addPartitionMessages(

       int partitionId,

       VertexIdMessages<I, M> messages) {

     ConcurrentMap<I, M> partitionMap =

         getOrCreatePartitionMap(partitionId);

     VertexIdMessageIterator<I, M> vertexIdMessageIterator =

       messages.getVertexIdMessageIterator();

     // This loop is a little complicated as it is optimized to only create

     // the minimal amount of vertex id and message objects as possible.

     while (vertexIdMessageIterator.hasNext()) {

       vertexIdMessageIterator.next();

       I vertexId = vertexIdMessageIterator.getCurrentVertexId();

       M currentMessage =

           partitionMap.get(vertexIdMessageIterator.getCurrentVertexId());

       if (currentMessage == null) {

         M newMessage = messageCombiner.createInitialMessage();

         currentMessage = partitionMap.putIfAbsent(

             vertexIdMessageIterator.releaseCurrentVertexId(), newMessage);

         if (currentMessage == null) {

           currentMessage = newMessage;

         }

       }

       synchronized (currentMessage) {

         messageCombiner.combine(vertexId, currentMessage,

             vertexIdMessageIterator.getCurrentMessage());

       }

     }

   }

   @Override

   public void addMessage(I vertexId, M message) throws IOException {

     ConcurrentMap<I, M> partitionMap =

       getOrCreatePartitionMap(getPartitionId(vertexId));

     M currentMessage = partitionMap.get(vertexId);

     if (currentMessage == null) {

       M newMessage = messageCombiner.createInitialMessage();

       // need a copy as vertexId might be reusable

       I copyId = WritableUtils.createCopy(vertexId);

       currentMessage = partitionMap.putIfAbsent(copyId, newMessage);

       if (currentMessage == null) {

         currentMessage = newMessage;

       }

     }

     synchronized (currentMessage) {

       messageCombiner.combine(vertexId, currentMessage, message);

     }

   }

   @Override

   protected Iterable<M> getMessagesAsIterable(M message) {

     return Collections.singleton(message);

   }

   @Override

   protected int getNumberOfMessagesIn(ConcurrentMap<I, M> partitionMap) {

     return partitionMap.size();

   }

   @Override

   protected void writeMessages(M messages, DataOutput out) throws IOException {

     messages.write(out);

   }

   @Override

   protected M readFieldsForMessages(DataInput in) throws IOException {

     M message = messageValueFactory.newInstance();

     message.readFields(in);

     return message;

   }

   /**

    * Create new factory for this message store

    *

    * @param service Worker service

    * @param config  Hadoop configuration

    * @param <I>     Vertex id

    * @param <M>     Message data

    * @return Factory

    */

   public static <I extends WritableComparable, M extends Writable>

   MessageStoreFactory<I, M, MessageStore<I, M>> newFactory(

       CentralizedServiceWorker<I, ?, ?> service,

       ImmutableClassesGiraphConfiguration<I, ?, ?> config) {

     return new Factory<I, M>(service, config);

   }

   /**

    * Factory for {@link OneMessagePerVertexStore}

    *

    * @param <I> Vertex id

    * @param <M> Message data

    */

   private static class Factory<I extends WritableComparable,

       M extends Writable>

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

     /** Service worker */

     private PartitionSplitInfo<I> partitionInfo;

     /** Hadoop configuration */

     private ImmutableClassesGiraphConfiguration<I, ?, ?> config;

     /**

      * @param partitionInfo Partition split info

      * @param config  Hadoop configuration

      */

     public Factory(

       PartitionSplitInfo<I> partitionInfo,

       ImmutableClassesGiraphConfiguration<I, ?, ?> config

     ) {

       this.partitionInfo = partitionInfo;

       this.config = config;

     }

     @Override

     public MessageStore<I, M> newStore(

         MessageClasses<I, M> messageClasses) {

       return new OneMessagePerVertexStore<I, M>(

           messageClasses.createMessageValueFactory(config), partitionInfo,

           messageClasses.createMessageCombiner(config), config);

     }

     @Override

     public void initialize(PartitionSplitInfo<I> partitionInfo,

         ImmutableClassesGiraphConfiguration<I, ?, ?> conf) {

       this.partitionInfo = partitionInfo;

       this.config = conf;

     }

   }

 }