/*
    * 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.ADDITIONAL_MSG_REQUEST_SIZE;
  import static org.apache.giraph.conf.GiraphConstants.MAX_MSG_REQUEST_SIZE;
  
  import java.util.Iterator;
  
  import org.apache.giraph.bsp.CentralizedServiceWorker;
  import org.apache.giraph.comm.netty.NettyWorkerClientRequestProcessor;
  import org.apache.giraph.comm.requests.SendWorkerMessagesRequest;
  import org.apache.giraph.comm.requests.WritableRequest;
  import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;
  import org.apache.giraph.edge.Edge;
  import org.apache.giraph.factories.MessageValueFactory;
  import org.apache.giraph.graph.Vertex;
  import org.apache.giraph.partition.PartitionOwner;
  import org.apache.giraph.utils.ByteArrayVertexIdMessages;
  import org.apache.giraph.utils.PairList;
  import org.apache.giraph.utils.VertexIdMessages;
  import org.apache.giraph.worker.WorkerInfo;
  import org.apache.hadoop.io.Writable;
  import org.apache.hadoop.io.WritableComparable;
  import org.apache.log4j.Logger;
  
  /**
   * Aggregates the messages to be sent to workers so they can be sent
   * in bulk.  Not thread-safe.
   *
   * @param <I> Vertex id
   * @param <M> Message data
   */
  @SuppressWarnings("unchecked")
  public class SendMessageCache<I extends WritableComparable, M extends Writable>
      extends SendVertexIdDataCache<I, M, VertexIdMessages<I, M>> {
    /** Class logger */
    private static final Logger LOG =
        Logger.getLogger(SendMessageCache.class);
    /** Messages sent during the last superstep */
    protected long totalMsgsSentInSuperstep = 0;
    /** Message bytes sent during the last superstep */
    protected long totalMsgBytesSentInSuperstep = 0;
    /** Max message size sent to a worker */
    protected final int maxMessagesSizePerWorker;
    /** NettyWorkerClientRequestProcessor for message sending */
    protected final NettyWorkerClientRequestProcessor<I, ?, ?> clientProcessor;
    /** Cached message value factory */
    protected MessageValueFactory<M> messageValueFactory;
    /**
     * Constructor
     *
     * @param conf Giraph configuration
     * @param serviceWorker Service worker
     * @param processor NettyWorkerClientRequestProcessor
     * @param maxMsgSize Max message size sent to a worker
     */
    public SendMessageCache(ImmutableClassesGiraphConfiguration conf,
        CentralizedServiceWorker<?, ?, ?> serviceWorker,
        NettyWorkerClientRequestProcessor<I, ?, ?> processor,
        int maxMsgSize) {
      super(conf, serviceWorker, MAX_MSG_REQUEST_SIZE.get(conf),
          ADDITIONAL_MSG_REQUEST_SIZE.get(conf));
      maxMessagesSizePerWorker = maxMsgSize;
      clientProcessor = processor;
      messageValueFactory =
              conf.createOutgoingMessageValueFactory();
    }
  
    @Override
    public VertexIdMessages<I, M> createVertexIdData() {
      return new ByteArrayVertexIdMessages<I, M>(messageValueFactory);
    }
  
    /**
     * Add a message to the cache.
     *
     * @param workerInfo the remote worker destination
     * @param partitionId the remote Partition this message belongs to
     * @param destVertexId vertex id that is ultimate destination
     * @param message Message to send to remote worker
     * @return Size of messages for the worker.
     */
   public int addMessage(WorkerInfo workerInfo,
                         int partitionId, I destVertexId, M message) {
     return addData(workerInfo, partitionId, destVertexId, message);
   }
 
   /**
    * Add a message to the cache with serialized ids.
    *
    * @param workerInfo The remote worker destination
    * @param partitionId The remote Partition this message belongs to
    * @param serializedId Serialized vertex id that is ultimate destination
    * @param idSerializerPos The end position of serialized id in the byte array
    * @param message Message to send to remote worker
    * @return Size of messages for the worker.
    */
   protected int addMessage(WorkerInfo workerInfo, int partitionId,
       byte[] serializedId, int idSerializerPos, M message) {
     return addData(
       workerInfo, partitionId, serializedId,
       idSerializerPos, message);
   }
 
   /**
    * Gets the messages for a worker and removes it from the cache.
    *
    * @param workerInfo the address of the worker who owns the data
    *                   partitions that are receiving the messages
    * @return List of pairs (partitionId, ByteArrayVertexIdMessages),
    *         where all partition ids belong to workerInfo
    */
   protected PairList<Integer, VertexIdMessages<I, M>>
   removeWorkerMessages(WorkerInfo workerInfo) {
     return removeWorkerData(workerInfo);
   }
 
   /**
    * Gets all the messages and removes them from the cache.
    *
    * @return All vertex messages for all partitions
    */
   private PairList<WorkerInfo, PairList<
       Integer, VertexIdMessages<I, M>>> removeAllMessages() {
     return removeAllData();
   }
 
   /**
    * Send a message to a target vertex id.
    *
    * @param destVertexId Target vertex id
    * @param message The message sent to the target
    */
   public void sendMessageRequest(I destVertexId, M message) {
     PartitionOwner owner =
       getServiceWorker().getVertexPartitionOwner(destVertexId);
     WorkerInfo workerInfo = owner.getWorkerInfo();
     final int partitionId = owner.getPartitionId();
     if (LOG.isTraceEnabled()) {
       LOG.trace("sendMessageRequest: Send bytes (" + message.toString() +
         ") to " + destVertexId + " on worker " + workerInfo);
     }
     ++totalMsgsSentInSuperstep;
     // Add the message to the cache
     int workerMessageSize = addMessage(
       workerInfo, partitionId, destVertexId, message);
     // Send a request if the cache of outgoing message to
     // the remote worker 'workerInfo' is full enough to be flushed
     if (workerMessageSize >= maxMessagesSizePerWorker) {
       PairList<Integer, VertexIdMessages<I, M>>
         workerMessages = removeWorkerMessages(workerInfo);
       WritableRequest writableRequest =
         new SendWorkerMessagesRequest<I, M>(workerMessages);
       totalMsgBytesSentInSuperstep += writableRequest.getSerializedSize();
       clientProcessor.doRequest(workerInfo, writableRequest);
       // Notify sending
       getServiceWorker().getGraphTaskManager().notifySentMessages();
     }
   }
 
   /**
    * An iterator wrapper on edges to return
    * target vertex ids.
    */
   public static class TargetVertexIdIterator<I extends WritableComparable>
       implements Iterator<I> {
     /** An edge iterator */
     private final Iterator<Edge<I, Writable>> edgesIterator;
 
     /**
      * Constructor.
      *
      * @param vertex The source vertex of the out edges
      */
     public TargetVertexIdIterator(Vertex<I, ?, ?> vertex) {
       edgesIterator =
         ((Vertex<I, Writable, Writable>) vertex).getEdges().iterator();
     }
 
     @Override
     public boolean hasNext() {
       return edgesIterator.hasNext();
     }
 
     @Override
     public I next() {
       return edgesIterator.next().getTargetVertexId();
     }
 
     @Override
     public void remove() {
       throw new UnsupportedOperationException();
     }
   }
 
   /**
    * Send message to all its neighbors
    *
    * @param vertex The source vertex
    * @param message The message sent to a worker
    */
   public void sendMessageToAllRequest(Vertex<I, ?, ?> vertex, M message) {
     TargetVertexIdIterator targetVertexIterator =
       new TargetVertexIdIterator(vertex);
     sendMessageToAllRequest(targetVertexIterator, message);
   }
 
   /**
    * Send message to the target ids in the iterator
    *
    * @param vertexIdIterator The iterator of target vertex ids
    * @param message The message sent to a worker
    */
   public void sendMessageToAllRequest(Iterator<I> vertexIdIterator, M message) {
     while (vertexIdIterator.hasNext()) {
       sendMessageRequest(vertexIdIterator.next(), message);
     }
   }
 
   /**
    * Flush the rest of the messages to the workers.
    */
   public void flush() {
     PairList<WorkerInfo, PairList<Integer,
         VertexIdMessages<I, M>>>
     remainingMessageCache = removeAllMessages();
     PairList<WorkerInfo, PairList<
         Integer, VertexIdMessages<I, M>>>.Iterator
     iterator = remainingMessageCache.getIterator();
     while (iterator.hasNext()) {
       iterator.next();
       WritableRequest writableRequest =
         new SendWorkerMessagesRequest<I, M>(
           iterator.getCurrentSecond());
       totalMsgBytesSentInSuperstep += writableRequest.getSerializedSize();
       clientProcessor.doRequest(
         iterator.getCurrentFirst(), writableRequest);
     }
   }
 
   /**
    * Reset the message count per superstep.
    *
    * @return The message count sent in last superstep
    */
   public long resetMessageCount() {
     long messagesSentInSuperstep = totalMsgsSentInSuperstep;
     totalMsgsSentInSuperstep = 0;
     return messagesSentInSuperstep;
   }
 
   /**
    * Reset the message bytes count per superstep.
    *
    * @return The message count sent in last superstep
    */
   public long resetMessageBytesCount() {
     long messageBytesSentInSuperstep = totalMsgBytesSentInSuperstep;
     totalMsgBytesSentInSuperstep = 0;
     return messageBytesSentInSuperstep;
   }
 }