/*
    * 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.netty;
  
  import java.io.IOException;
  import java.util.Iterator;
  import java.util.Map;
  import java.util.concurrent.ConcurrentMap;
  
  import org.apache.giraph.bsp.BspService;
  import org.apache.giraph.bsp.CentralizedServiceWorker;
  import org.apache.giraph.comm.SendEdgeCache;
  import org.apache.giraph.comm.SendMessageCache;
  import org.apache.giraph.comm.SendMutationsCache;
  import org.apache.giraph.comm.SendOneMessageToManyCache;
  import org.apache.giraph.comm.SendPartitionCache;
  import org.apache.giraph.comm.ServerData;
  import org.apache.giraph.comm.WorkerClient;
  import org.apache.giraph.comm.WorkerClientRequestProcessor;
  import org.apache.giraph.comm.messages.MessageStore;
  import org.apache.giraph.comm.requests.SendPartitionCurrentMessagesRequest;
  import org.apache.giraph.comm.requests.SendPartitionMutationsRequest;
  import org.apache.giraph.comm.requests.SendVertexRequest;
  import org.apache.giraph.comm.requests.SendWorkerEdgesRequest;
  import org.apache.giraph.comm.requests.SendWorkerVerticesRequest;
  import org.apache.giraph.comm.requests.WorkerRequest;
  import org.apache.giraph.comm.requests.WritableRequest;
  import org.apache.giraph.conf.GiraphConfiguration;
  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.graph.VertexMutations;
  import org.apache.giraph.metrics.GiraphMetrics;
  import org.apache.giraph.metrics.MetricNames;
  import org.apache.giraph.metrics.SuperstepMetricsRegistry;
  import org.apache.giraph.partition.Partition;
  import org.apache.giraph.partition.PartitionOwner;
  import org.apache.giraph.utils.ByteArrayVertexIdMessages;
  import org.apache.giraph.utils.ExtendedDataOutput;
  import org.apache.giraph.utils.PairList;
  import org.apache.giraph.utils.VertexIdEdges;
  import org.apache.giraph.worker.WorkerInfo;
  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.yammer.metrics.core.Counter;
  import com.yammer.metrics.core.Gauge;
  import com.yammer.metrics.util.PercentGauge;
  
  /**
   * Aggregate requests and sends them to the thread-safe NettyClient.  This
   * class is not thread-safe and expected to be used and then thrown away after
   * a phase of communication has completed.
   *
   * @param <I> Vertex id
   * @param <V> Vertex data
   * @param <E> Edge data
   */
  @SuppressWarnings("unchecked")
  public class NettyWorkerClientRequestProcessor<I extends WritableComparable,
      V extends Writable, E extends Writable> implements
      WorkerClientRequestProcessor<I, V, E> {
    /** Class logger */
    private static final Logger LOG =
        Logger.getLogger(NettyWorkerClientRequestProcessor.class);
    /** Cached partitions of vertices to send */
    private final SendPartitionCache<I, V, E> sendPartitionCache;
    /** Cached map of partitions to vertex indices to messages */
    private final SendMessageCache<I, Writable> sendMessageCache;
    /** Cache of edges to be sent. */
    private final SendEdgeCache<I, E> sendEdgeCache;
    /** Cached map of partitions to vertex indices to mutations */
    private final SendMutationsCache<I, V, E> sendMutationsCache =
        new SendMutationsCache<I, V, E>();
    /** NettyClient that could be shared among one or more instances */
    private final WorkerClient<I, V, E> workerClient;
    /** Maximum size of messages per remote worker to cache before sending */
    private final int maxMessagesSizePerWorker;
    /** Maximum size of vertices per remote worker to cache before sending. */
    private final int maxVerticesSizePerWorker;
    /** Maximum size of edges per remote worker to cache before sending. */
   private final int maxEdgesSizePerWorker;
   /** Maximum number of mutations per partition before sending */
   private final int maxMutationsPerPartition;
   /** Giraph configuration */
   private final ImmutableClassesGiraphConfiguration<I, V, E> configuration;
   /** Service worker */
   private final CentralizedServiceWorker<I, V, E> serviceWorker;
   /** Server data from the server (used for local requests) */
   private final ServerData<I, V, E> serverData;
 
   // Per-Superstep Metrics
   /** Number of requests that went on the wire */
   private final Counter localRequests;
   /** Number of requests that were handled locally */
   private final Counter remoteRequests;
   /** Cached message value factory */
   private final MessageValueFactory messageValueFactory;
 
   /**
    * Constructor.
    *
    * @param context Context
    * @param conf Configuration
    * @param serviceWorker Service worker
    * @param useOneMessageToManyIdsEncoding should use one message to many
    */
   public NettyWorkerClientRequestProcessor(
       Mapper<?, ?, ?, ?>.Context context,
       ImmutableClassesGiraphConfiguration<I, V, E> conf,
       CentralizedServiceWorker<I, V, E> serviceWorker,
       boolean useOneMessageToManyIdsEncoding) {
     this.workerClient = serviceWorker.getWorkerClient();
     this.configuration = conf;
 
 
     sendPartitionCache =
         new SendPartitionCache<I, V, E>(conf, serviceWorker);
     sendEdgeCache = new SendEdgeCache<I, E>(conf, serviceWorker);
     maxMessagesSizePerWorker =
         GiraphConfiguration.MAX_MSG_REQUEST_SIZE.get(conf);
     maxVerticesSizePerWorker =
         GiraphConfiguration.MAX_VERTEX_REQUEST_SIZE.get(conf);
     if (useOneMessageToManyIdsEncoding) {
       sendMessageCache =
         new SendOneMessageToManyCache<I, Writable>(conf, serviceWorker,
           this, maxMessagesSizePerWorker);
     } else {
       sendMessageCache =
         new SendMessageCache<I, Writable>(conf, serviceWorker,
           this, maxMessagesSizePerWorker);
     }
     maxEdgesSizePerWorker =
         GiraphConfiguration.MAX_EDGE_REQUEST_SIZE.get(conf);
     maxMutationsPerPartition =
         GiraphConfiguration.MAX_MUTATIONS_PER_REQUEST.get(conf);
     this.serviceWorker = serviceWorker;
     this.serverData = serviceWorker.getServerData();
 
     // Per-Superstep Metrics.
     // Since this object is not long lived we just initialize the metrics here.
     SuperstepMetricsRegistry smr = GiraphMetrics.get().perSuperstep();
     localRequests = smr.getCounter(MetricNames.LOCAL_REQUESTS);
     remoteRequests = smr.getCounter(MetricNames.REMOTE_REQUESTS);
     setupGauges(smr, localRequests, remoteRequests);
     messageValueFactory = configuration.createOutgoingMessageValueFactory();
   }
 
   @Override
   public void sendMessageRequest(I destVertexId, Writable message) {
     this.sendMessageCache.sendMessageRequest(destVertexId, message);
   }
 
   @Override
   public void sendMessageToAllRequest(
     Vertex<I, V, E> vertex, Writable message) {
     this.sendMessageCache.sendMessageToAllRequest(vertex, message);
   }
 
   @Override
   public void sendMessageToAllRequest(
     Iterator<I> vertexIdIterator, Writable message) {
     this.sendMessageCache.sendMessageToAllRequest(vertexIdIterator, message);
   }
 
   @Override
   public void sendPartitionRequest(WorkerInfo workerInfo,
                                    Partition<I, V, E> partition) {
     if (LOG.isTraceEnabled()) {
       LOG.trace("sendVertexRequest: Sending to " + workerInfo +
           ", with partition " + partition);
     }
 
     WritableRequest vertexRequest = new SendVertexRequest<I, V, E>(partition);
     doRequest(workerInfo, vertexRequest);
 
     // Messages are stored separately
     if (serviceWorker.getSuperstep() != BspService.INPUT_SUPERSTEP) {
       sendPartitionMessages(workerInfo, partition);
       ConcurrentMap<I, VertexMutations<I, V, E>> vertexMutationMap =
           serverData.getPartitionMutations().remove(partition.getId());
       WritableRequest partitionMutationsRequest =
           new SendPartitionMutationsRequest<I, V, E>(partition.getId(),
               vertexMutationMap);
       doRequest(workerInfo, partitionMutationsRequest);
     }
   }
 
   /**
    * Send all messages for a partition to another worker.
    *
    * @param workerInfo Worker to send the partition messages to
    * @param partition Partition whose messages to send
    */
   private void sendPartitionMessages(WorkerInfo workerInfo,
                                      Partition<I, V, E> partition) {
     final int partitionId = partition.getId();
     MessageStore<I, Writable> messageStore =
         serverData.getCurrentMessageStore();
     ByteArrayVertexIdMessages<I, Writable> vertexIdMessages =
         new ByteArrayVertexIdMessages<I, Writable>(
             messageValueFactory);
     vertexIdMessages.setConf(configuration);
     vertexIdMessages.initialize();
     for (I vertexId :
         messageStore.getPartitionDestinationVertices(partitionId)) {
       // Messages cannot be re-used from this iterable, but add()
       // serializes the message, making this safe
       Iterable<Writable> messages = messageStore.getVertexMessages(vertexId);
       for (Writable message : messages) {
         vertexIdMessages.add(vertexId, message);
       }
       if (vertexIdMessages.getSize() > maxMessagesSizePerWorker) {
         WritableRequest messagesRequest =
             new SendPartitionCurrentMessagesRequest<I, V, E, Writable>(
             partitionId, vertexIdMessages);
         doRequest(workerInfo, messagesRequest);
         vertexIdMessages =
             new ByteArrayVertexIdMessages<I, Writable>(
                 messageValueFactory);
         vertexIdMessages.setConf(configuration);
         vertexIdMessages.initialize();
       }
     }
     if (!vertexIdMessages.isEmpty()) {
       WritableRequest messagesRequest = new
           SendPartitionCurrentMessagesRequest<I, V, E, Writable>(
           partitionId, vertexIdMessages);
       doRequest(workerInfo, messagesRequest);
     }
     messageStore.clearPartition(partitionId);
   }
 
   @Override
   public boolean sendVertexRequest(PartitionOwner partitionOwner,
       Vertex<I, V, E> vertex) {
     // Add the vertex to the cache
     int workerMessageSize = sendPartitionCache.addVertex(
         partitionOwner, vertex);
 
     // Send a request if the cache of outgoing message to
     // the remote worker 'workerInfo' is full enough to be flushed
     if (workerMessageSize >= maxVerticesSizePerWorker) {
       PairList<Integer, ExtendedDataOutput>
           workerPartitionVertices =
           sendPartitionCache.removeWorkerData(partitionOwner.getWorkerInfo());
       WritableRequest writableRequest =
           new SendWorkerVerticesRequest<I, V, E>(
               configuration, workerPartitionVertices);
       doRequest(partitionOwner.getWorkerInfo(), writableRequest);
       return true;
     }
 
     return false;
   }
 
   @Override
   public void addEdgeRequest(I vertexIndex, Edge<I, E> edge) throws
       IOException {
     PartitionOwner partitionOwner =
         serviceWorker.getVertexPartitionOwner(vertexIndex);
     int partitionId = partitionOwner.getPartitionId();
     if (LOG.isTraceEnabled()) {
       LOG.trace("addEdgeRequest: Sending edge " + edge + " for index " +
           vertexIndex + " with partition " + partitionId);
     }
 
     // Add the message to the cache
     int partitionMutationCount =
         sendMutationsCache.addEdgeMutation(partitionId, vertexIndex, edge);
 
     sendMutationsRequestIfFull(
         partitionId, partitionOwner, partitionMutationCount);
   }
 
   @Override
   public boolean sendEdgeRequest(I sourceVertexId, Edge<I, E> edge)
     throws IOException {
     PartitionOwner owner =
         serviceWorker.getVertexPartitionOwner(sourceVertexId);
     WorkerInfo workerInfo = owner.getWorkerInfo();
     final int partitionId = owner.getPartitionId();
     if (LOG.isTraceEnabled()) {
       LOG.trace("sendEdgeRequest: Send bytes (" + edge.toString() +
           ") to " + sourceVertexId + " on worker " + workerInfo);
     }
 
     // Add the message to the cache
     int workerEdgesSize = sendEdgeCache.addEdge(
         workerInfo, partitionId, sourceVertexId, edge);
 
     // Send a request if the cache of outgoing edges to the remote worker is
     // full
     if (workerEdgesSize >= maxEdgesSizePerWorker) {
       PairList<Integer, VertexIdEdges<I, E>> workerEdges =
           sendEdgeCache.removeWorkerEdges(workerInfo);
       WritableRequest writableRequest =
           new SendWorkerEdgesRequest<I, E>(workerEdges);
       doRequest(workerInfo, writableRequest);
       return true;
     }
 
     return false;
   }
 
   /**
    * Send a mutations request if the maximum number of mutations per partition
    * was met.
    *
    * @param partitionId Partition id
    * @param partitionOwner Owner of the partition
    * @param partitionMutationCount Number of mutations for this partition
    */
   private void sendMutationsRequestIfFull(
       int partitionId, PartitionOwner partitionOwner,
       int partitionMutationCount) {
     // Send a request if enough mutations are there for a partition
     if (partitionMutationCount >= maxMutationsPerPartition) {
       Map<I, VertexMutations<I, V, E>> partitionMutations =
           sendMutationsCache.removePartitionMutations(partitionId);
       WritableRequest writableRequest =
           new SendPartitionMutationsRequest<I, V, E>(
               partitionId, partitionMutations);
       doRequest(partitionOwner.getWorkerInfo(), writableRequest);
     }
   }
 
   @Override
   public void removeEdgesRequest(I vertexIndex,
                                  I destinationVertexIndex) throws IOException {
     PartitionOwner partitionOwner =
         serviceWorker.getVertexPartitionOwner(vertexIndex);
     int partitionId = partitionOwner.getPartitionId();
     if (LOG.isTraceEnabled()) {
       LOG.trace("removeEdgesRequest: Removing edge " +
           destinationVertexIndex +
           " for index " + vertexIndex + " with partition " + partitionId);
     }
 
     // Add the message to the cache
     int partitionMutationCount =
         sendMutationsCache.removeEdgeMutation(
             partitionId, vertexIndex, destinationVertexIndex);
 
     sendMutationsRequestIfFull(
         partitionId, partitionOwner, partitionMutationCount);
   }
 
   @Override
   public void addVertexRequest(Vertex<I, V, E> vertex) throws IOException {
     PartitionOwner partitionOwner =
         serviceWorker.getVertexPartitionOwner(vertex.getId());
     int partitionId = partitionOwner.getPartitionId();
     if (LOG.isTraceEnabled()) {
       LOG.trace("addVertexRequest: Sending vertex " + vertex +
           " to partition " + partitionId);
     }
 
     // Add the message to the cache
     int partitionMutationCount =
         sendMutationsCache.addVertexMutation(partitionId, vertex);
 
     sendMutationsRequestIfFull(
         partitionId, partitionOwner, partitionMutationCount);
   }
 
   @Override
   public void removeVertexRequest(I vertexIndex) throws IOException {
     PartitionOwner partitionOwner =
         serviceWorker.getVertexPartitionOwner(vertexIndex);
     int partitionId = partitionOwner.getPartitionId();
     if (LOG.isTraceEnabled()) {
       LOG.trace("removeVertexRequest: Removing vertex index " +
           vertexIndex + " from partition " + partitionId);
     }
 
     // Add the message to the cache
     int partitionMutationCount =
         sendMutationsCache.removeVertexMutation(partitionId, vertexIndex);
 
     sendMutationsRequestIfFull(
         partitionId, partitionOwner, partitionMutationCount);
   }
 
   @Override
   public void flush() throws IOException {
     // Execute the remaining sends messages (if any)
     // including individual and compact messages.
     sendMessageCache.flush();
 
     // Execute the remaining sends vertices (if any)
     PairList<WorkerInfo, PairList<Integer, ExtendedDataOutput>>
         remainingVertexCache = sendPartitionCache.removeAllData();
     PairList<WorkerInfo,
         PairList<Integer, ExtendedDataOutput>>.Iterator
         vertexIterator = remainingVertexCache.getIterator();
     while (vertexIterator.hasNext()) {
       vertexIterator.next();
       WritableRequest writableRequest =
           new SendWorkerVerticesRequest(
               configuration, vertexIterator.getCurrentSecond());
       doRequest(vertexIterator.getCurrentFirst(), writableRequest);
     }
 
     // Execute the remaining sends edges (if any)
     PairList<WorkerInfo, PairList<Integer,
         VertexIdEdges<I, E>>>
         remainingEdgeCache = sendEdgeCache.removeAllEdges();
     PairList<WorkerInfo,
         PairList<Integer, VertexIdEdges<I, E>>>.Iterator
         edgeIterator = remainingEdgeCache.getIterator();
     while (edgeIterator.hasNext()) {
       edgeIterator.next();
       WritableRequest writableRequest =
           new SendWorkerEdgesRequest<I, E>(
               edgeIterator.getCurrentSecond());
       doRequest(edgeIterator.getCurrentFirst(), writableRequest);
     }
 
     // Execute the remaining sends mutations (if any)
     Map<Integer, Map<I, VertexMutations<I, V, E>>> remainingMutationsCache =
         sendMutationsCache.removeAllPartitionMutations();
     for (Map.Entry<Integer, Map<I, VertexMutations<I, V, E>>> entry :
         remainingMutationsCache.entrySet()) {
       WritableRequest writableRequest =
           new SendPartitionMutationsRequest<I, V, E>(
               entry.getKey(), entry.getValue());
       PartitionOwner partitionOwner =
           serviceWorker.getVertexPartitionOwner(
               entry.getValue().keySet().iterator().next());
       doRequest(partitionOwner.getWorkerInfo(), writableRequest);
     }
   }
 
   @Override
   public long resetMessageCount() {
     return this.sendMessageCache.resetMessageCount();
   }
 
   @Override
   public long resetMessageBytesCount() {
     return this.sendMessageCache.resetMessageBytesCount();
   }
 
   /**
    * When doing the request, short circuit if it is local
    *
    * @param workerInfo Worker info
    * @param writableRequest Request to either submit or run locally
    */
   public void doRequest(WorkerInfo workerInfo,
                          WritableRequest writableRequest) {
     // If this is local, execute locally
     if (serviceWorker.getWorkerInfo().getTaskId() ==
         workerInfo.getTaskId()) {
       ((WorkerRequest) writableRequest).doRequest(serverData);
       localRequests.inc();
     } else {
       workerClient.sendWritableRequest(
           workerInfo.getTaskId(), writableRequest);
       remoteRequests.inc();
     }
   }
 
   /**
    * Sets up gauges for superstep metrics.
    * This has to be static so that internal objects created here don't
    * hold references to this$0. Otherwise we memory leaking
    * NettyWorkerClientRequestProcessor objects.
    *
    * @param smr metric registry for current superstep
    * @param localRequests counter for local requests
    * @param remoteRequests counter for remote requests
    */
   private static void setupGauges(SuperstepMetricsRegistry smr,
                                   final Counter localRequests,
                                   final Counter remoteRequests) {
     final Gauge<Long> totalRequests = smr.getGauge(MetricNames.TOTAL_REQUESTS,
         new Gauge<Long>() {
           @Override
           public Long value() {
             return localRequests.count() + remoteRequests.count();
           }
         }
     );
     smr.getGauge(MetricNames.PERCENT_LOCAL_REQUESTS, new PercentGauge() {
       @Override protected double getNumerator() {
         return localRequests.count();
       }
 
       @Override protected double getDenominator() {
         return totalRequests.value();
       }
     });
   }
 }