giraph-core/src/main/java/org/apache/giraph/graph/GiraphTransferRegulator.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.graph;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.io.WritableComparable;
import org.apache.giraph.partition.PartitionOwner;
import com.google.common.collect.Maps;
import java.util.Map;


/** Utility class to manage data transfers from
 * a local worker reading InputSplits.
 * Currently, this class measures # of vertices and edges
 * per outgoing Collection of graph data (destined for a
 * particular Partition and remote worker node, preselected
 * by the master.)
 *
 * TODO: implement defaults and configurable options for
 * measuring the size of input <V> or <E> data
 * per read vertex, and setting limits on totals per outgoing
 * graph data Collection etc. (See GIRAPH-260)
 */
public class GiraphTransferRegulator {
  /** Maximum vertices to read from an InputSplit locally that are
   * to be routed to a remote worker, before sending them. */
  public static final String MAX_VERTICES_PER_TRANSFER =
    "giraph.maxVerticesPerTransfer";
  /** Default maximum number of vertices per
   * temp partition before sending. */
  public static final int MAX_VERTICES_PER_TRANSFER_DEFAULT = 10000;
  /**
   * Maximum edges to read from an InputSplit locally that are
   * to be routed to a remote worker, before sending them.
   */
  public static final String MAX_EDGES_PER_TRANSFER =
    "giraph.maxEdgesPerTransfer";
  /** Default maximum number of vertices per
   * temp partition before sending. */
  public static final int MAX_EDGES_PER_TRANSFER_DEFAULT = 80000;

  /** Internal state to measure when
   * the next data transfer of a Collection
   * of vertices read by the local worker that
   * owns this regulator is ready to be sent
   * to the remote worker node that the master
   * has assigned the vertices to */
  private Map<Integer, Integer> edgeAccumulator;

  /** Internal state to measure when
   * the next data transfer of a Collection
   * of vertices read by the local worker that
   * owns this regulator is ready to be sent
   * to the remote worker node that the master
   * has assigned the vertices to */
  private Map<Integer, Integer> vertexAccumulator;

  /** Number of vertices per data transfer */
  private final int maxVerticesPerTransfer;

  /** Number of edges per data transfer */
  private final int maxEdgesPerTransfer;

  /** Vertex count total for this InputSplit */
  private long totalVertexCount;

  /** Edge count total for this InputSplit */
  private long totalEdgeCount;

  /** Default constructor
   * @param conf the Configuration for this job
   */
  public GiraphTransferRegulator(Configuration conf) {
    vertexAccumulator = Maps.<Integer, Integer>newHashMap();
    edgeAccumulator = Maps.<Integer, Integer>newHashMap();
    maxVerticesPerTransfer = conf.getInt(
        MAX_VERTICES_PER_TRANSFER,
        MAX_VERTICES_PER_TRANSFER_DEFAULT);
    maxEdgesPerTransfer = conf.getInt(
        MAX_EDGES_PER_TRANSFER,
        MAX_EDGES_PER_TRANSFER_DEFAULT);
    totalEdgeCount = 0;
    totalVertexCount = 0;
  }

  /** Is this outbound data Collection full,
   * and ready to transfer?
   * @param owner the partition owner for the outbound data
   * @return 'true' if the temp partition data is ready to transfer
   */
  public boolean transferThisPartition(PartitionOwner owner) {
    final int partitionId = owner.getPartitionId();
    if (getEdgesForPartition(partitionId) >=
      maxEdgesPerTransfer ||
      getVerticesForPartition(partitionId) >=
      maxVerticesPerTransfer) {
      vertexAccumulator.put(partitionId, 0);
      edgeAccumulator.put(partitionId, 0);
      return true;
    }
    return false;
  }

  /** get current vertex count for a given Collection of
   * data soon to be transfered to its permanent home.
   * @param partId the partition id to check the count on.
   * @return the count of vertices.
   */
  private int getVerticesForPartition(final int partId) {
    return vertexAccumulator.get(partId) == null ?
      0 : vertexAccumulator.get(partId);
  }

  /** get current edge count for a given Collection of
   * data soon to be transfered to its permanent home.
   * @param partId the partition id to check the count on.
   * @return the count of edges.
   */
  private int getEdgesForPartition(final int partId) {
    return edgeAccumulator.get(partId) == null ?
      0 : edgeAccumulator.get(partId);
  }

  /** Clear storage to reset for reading new InputSplit */
  public void clearCounters() {
    totalEdgeCount = 0;
    totalVertexCount = 0;
    vertexAccumulator.clear();
    edgeAccumulator.clear();
  }

  /** Increment V &amp; E counts for new vertex read, store values
   * for that outgoing _temporary_ Partition, which shares the
   * Partition ID for the actual remote Partition the collection
   * will eventually be processed in.
   * @param partitionOwner the owner of the Partition this data
   *  will eventually belong to.
   * @param vertex the vertex to extract counts from.
   * @param <I> the vertex id type.
   * @param <V> the vertex value type.
   * @param <E> the edge value type.
   */
  public <I extends WritableComparable, V extends Writable,
  E extends Writable> void
  incrementCounters(PartitionOwner partitionOwner,
    Vertex<I, V, E> vertex) {
    final int id = partitionOwner.getPartitionId();
    // vertex counts
    vertexAccumulator
      .put(id, getVerticesForPartition(id) + 1);
    totalVertexCount++;
    // edge counts
    totalEdgeCount += vertex.getNumEdges();
    edgeAccumulator.put(id, getEdgesForPartition(id) +
      vertex.getNumEdges());
  }

  /** Getter for MAX edge count to initiate a transfer
    * @return max edge count per transfer */
  public long getMaxEdgesPerTransfer() {
    return maxEdgesPerTransfer;
  }

  /** Getter for MAX vertex count to initiate a transfer
   * @return max edge count per transfer */
  public long getMaxVerticesPerTransfer() {
    return maxVerticesPerTransfer;
  }

  /** Getter for total edge count for the current InputSplit
    * @return the # of total edges counted in this InputSplit */
  public long getTotalEdges() {
    return totalEdgeCount;
  }

  /** Getter for total vetex count for the current InputSplit
   * @return the total # of vertices counted in this InputSplit */
  public long getTotalVertices() {
    return totalVertexCount;
  }
}