giraph-core/src/main/java/org/apache/giraph/comm/SendMessageCache.java

/*
 * 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;
  }
}