giraph-core/src/main/java/org/apache/giraph/comm/ServerData.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.MESSAGE_STORE_FACTORY_CLASS;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentMap;

import org.apache.giraph.bsp.CentralizedServiceWorker;
import org.apache.giraph.comm.aggregators.AllAggregatorServerData;
import org.apache.giraph.comm.aggregators.OwnerAggregatorServerData;
import org.apache.giraph.comm.messages.MessageStore;
import org.apache.giraph.comm.messages.MessageStoreFactory;
import org.apache.giraph.comm.messages.queue.AsyncMessageStoreWrapper;
import org.apache.giraph.conf.GiraphConstants;
import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;
import org.apache.giraph.edge.EdgeStore;
import org.apache.giraph.edge.EdgeStoreFactory;
import org.apache.giraph.graph.Vertex;
import org.apache.giraph.graph.VertexMutations;
import org.apache.giraph.graph.VertexResolver;
import org.apache.giraph.ooc.OutOfCoreEngine;
import org.apache.giraph.ooc.data.DiskBackedEdgeStore;
import org.apache.giraph.ooc.data.DiskBackedMessageStore;
import org.apache.giraph.ooc.data.DiskBackedPartitionStore;
import org.apache.giraph.partition.Partition;
import org.apache.giraph.partition.PartitionStore;
import org.apache.giraph.partition.SimplePartitionStore;
import org.apache.giraph.utils.ReflectionUtils;
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.google.common.collect.Iterables;
import com.google.common.collect.Maps;

/**
 * Anything that the server stores
 *
 * @param <I> Vertex id
 * @param <V> Vertex data
 * @param <E> Edge data
 */
@SuppressWarnings("rawtypes")
public class ServerData<I extends WritableComparable,
    V extends Writable, E extends Writable> {
  /** Class logger */
  private static final Logger LOG = Logger.getLogger(ServerData.class);
  /** Configuration */
  private final ImmutableClassesGiraphConfiguration<I, V, E> conf;
  /** Partition store for this worker. */
  private volatile PartitionStore<I, V, E> partitionStore;
  /** Edge store for this worker. */
  private final EdgeStore<I, V, E> edgeStore;
  /** Message store factory */
  private final MessageStoreFactory<I, Writable, MessageStore<I, Writable>>
      messageStoreFactory;
  /**
   * Message store for incoming messages (messages which will be consumed
   * in the next super step)
   */
  private volatile MessageStore<I, Writable> incomingMessageStore;
  /**
   * Message store for current messages (messages which we received in
   * previous super step and which will be consumed in current super step)
   */
  private volatile MessageStore<I, Writable> currentMessageStore;
  /**
   * Map of partition ids to vertex mutations from other workers. These are
   * mutations that should be applied before execution of *current* super step.
   * (accesses to keys should be thread-safe as multiple threads may resolve
   * mutations of different partitions at the same time)
   */
  private ConcurrentMap<Integer,
      ConcurrentMap<I, VertexMutations<I, V, E>>>
      oldPartitionMutations = Maps.newConcurrentMap();
  /**
   * Map of partition ids to vertex mutations from other workers. These are
   * mutations that are coming from other workers as the execution goes one in a
   * super step. These mutations should be applied in the *next* super step.
   * (this should be thread-safe)
   */
  private ConcurrentMap<Integer,
      ConcurrentMap<I, VertexMutations<I, V, E>>>
      partitionMutations = Maps.newConcurrentMap();
  /**
   * Holds aggregators which current worker owns from current superstep
   */
  private final OwnerAggregatorServerData ownerAggregatorData;
  /**
   * Holds old aggregators from previous superstep
   */
  private final AllAggregatorServerData allAggregatorData;
  /** Service worker */
  private final CentralizedServiceWorker<I, V, E> serviceWorker;

  /** Store for current messages from other workers to this worker */
  private volatile List<Writable> currentWorkerToWorkerMessages =
      Collections.synchronizedList(new ArrayList<Writable>());
  /** Store for message from other workers to this worker for next superstep */
  private volatile List<Writable> incomingWorkerToWorkerMessages =
      Collections.synchronizedList(new ArrayList<Writable>());

  /** Job context (for progress) */
  private final Mapper<?, ?, ?, ?>.Context context;
  /** Out-of-core engine */
  private final OutOfCoreEngine oocEngine;

  /**
   * Constructor.
   *
   * @param service Service worker
   * @param workerServer Worker server
   * @param conf Configuration
   * @param context Mapper context
   */
  public ServerData(
      CentralizedServiceWorker<I, V, E> service,
      WorkerServer workerServer,
      ImmutableClassesGiraphConfiguration<I, V, E> conf,
      Mapper<?, ?, ?, ?>.Context context) {
    this.serviceWorker = service;
    this.conf = conf;
    this.messageStoreFactory = createMessageStoreFactory();
    EdgeStoreFactory<I, V, E> edgeStoreFactory = conf.createEdgeStoreFactory();
    edgeStoreFactory.initialize(service, conf, context);
    EdgeStore<I, V, E> inMemoryEdgeStore = edgeStoreFactory.newStore();
    PartitionStore<I, V, E> inMemoryPartitionStore =
        new SimplePartitionStore<I, V, E>(conf, context);
    if (GiraphConstants.USE_OUT_OF_CORE_GRAPH.get(conf)) {
      oocEngine = new OutOfCoreEngine(conf, service, workerServer);
      partitionStore =
          new DiskBackedPartitionStore<I, V, E>(inMemoryPartitionStore,
              conf, context, oocEngine);
      edgeStore =
          new DiskBackedEdgeStore<I, V, E>(inMemoryEdgeStore, conf, oocEngine);
    } else {
      partitionStore = inMemoryPartitionStore;
      edgeStore = inMemoryEdgeStore;
      oocEngine = null;
    }
    ownerAggregatorData = new OwnerAggregatorServerData(context);
    allAggregatorData = new AllAggregatorServerData(context, conf);
    this.context = context;
  }

  /**
   * Decide which message store should be used for current application,
   * and create the factory for that store
   *
   * @return Message store factory
   */
  private MessageStoreFactory<I, Writable, MessageStore<I, Writable>>
  createMessageStoreFactory() {
    Class<? extends MessageStoreFactory> messageStoreFactoryClass =
        MESSAGE_STORE_FACTORY_CLASS.get(conf);

    MessageStoreFactory messageStoreFactoryInstance =
        ReflectionUtils.newInstance(messageStoreFactoryClass);
    messageStoreFactoryInstance.initialize(serviceWorker, conf);

    return messageStoreFactoryInstance;
  }

  /**
   * Return the out-of-core engine for this worker.
   *
   * @return The out-of-core engine
   */
  public OutOfCoreEngine getOocEngine() {
    return oocEngine;
  }

  /**
   * Return the edge store for this worker.
   *
   * @return The edge store
   */
  public EdgeStore<I, V, E> getEdgeStore() {
    return edgeStore;
  }

  /**
   * Return the partition store for this worker.
   *
   * @return The partition store
   */
  public PartitionStore<I, V, E> getPartitionStore() {
    return partitionStore;
  }

  /**
   * Get message store for incoming messages (messages which will be consumed
   * in the next super step)
   *
   * @param <M> Message data
   * @return Incoming message store
   */
  public <M extends Writable> MessageStore<I, M> getIncomingMessageStore() {
    return (MessageStore<I, M>) incomingMessageStore;
  }

  /**
   * Get message store for current messages (messages which we received in
   * previous super step and which will be consumed in current super step)
   *
   * @param <M> Message data
   * @return Current message store
   */
  public <M extends Writable> MessageStore<I, M> getCurrentMessageStore() {
    return (MessageStore<I, M>) currentMessageStore;
  }

  /**
   * Re-initialize message stores.
   * Discards old values if any.
   */
  public void resetMessageStores() {
    if (currentMessageStore != null) {
      currentMessageStore.clearAll();
      currentMessageStore = null;
    }
    if (incomingMessageStore != null) {
      incomingMessageStore.clearAll();
      incomingMessageStore = null;
    }
    prepareSuperstep();
  }

  /** Prepare for next superstep */
  public void prepareSuperstep() {
    if (currentMessageStore != null) {
      currentMessageStore.clearAll();
    }

    MessageStore<I, Writable> nextCurrentMessageStore;
    MessageStore<I, Writable> nextIncomingMessageStore;
    MessageStore<I, Writable> messageStore;

    // First create the necessary in-memory message stores. If out-of-core
    // mechanism is enabled, we wrap the in-memory message stores within
    // disk-backed messages stores.
    if (incomingMessageStore != null) {
      nextCurrentMessageStore = incomingMessageStore;
    } else {
      messageStore = messageStoreFactory.newStore(
          conf.getIncomingMessageClasses());
      if (oocEngine == null) {
        nextCurrentMessageStore = messageStore;
      } else {
        nextCurrentMessageStore = new DiskBackedMessageStore<>(
            conf, oocEngine, messageStore,
            conf.getIncomingMessageClasses().useMessageCombiner(),
            serviceWorker.getSuperstep());
      }
    }

    messageStore = messageStoreFactory.newStore(
        conf.getOutgoingMessageClasses());
    if (oocEngine == null) {
      nextIncomingMessageStore = messageStore;
    } else {
      nextIncomingMessageStore = new DiskBackedMessageStore<>(
          conf, oocEngine, messageStore,
          conf.getOutgoingMessageClasses().useMessageCombiner(),
          serviceWorker.getSuperstep() + 1);
    }

    // If out-of-core engine is enabled, we avoid overlapping of out-of-core
    // decisions with change of superstep. This avoidance is done to simplify
    // the design and reduce excessive use of synchronization primitives.
    if (oocEngine != null) {
      oocEngine.getSuperstepLock().writeLock().lock();
    }
    currentMessageStore = nextCurrentMessageStore;
    incomingMessageStore = nextIncomingMessageStore;
    if (oocEngine != null) {
      oocEngine.reset();
      oocEngine.getSuperstepLock().writeLock().unlock();
    }
    currentMessageStore.finalizeStore();

    currentWorkerToWorkerMessages = incomingWorkerToWorkerMessages;
    incomingWorkerToWorkerMessages =
        Collections.synchronizedList(new ArrayList<Writable>());
  }

  /**
   * Get the vertex mutations (synchronize on the values)
   *
   * @return Vertex mutations
   */
  public ConcurrentMap<Integer, ConcurrentMap<I, VertexMutations<I, V, E>>>
  getPartitionMutations() {
    return partitionMutations;
  }

  /**
   * Get holder for aggregators which current worker owns
   *
   * @return Holder for aggregators which current worker owns
   */
  public OwnerAggregatorServerData getOwnerAggregatorData() {
    return ownerAggregatorData;
  }

  /**
   * Get holder for aggregators from previous superstep
   *
   * @return Holder for aggregators from previous superstep
   */
  public AllAggregatorServerData getAllAggregatorData() {
    return allAggregatorData;
  }

  /**
   * Get the reference of the service worker.
   *
   * @return CentralizedServiceWorker
   */
  public CentralizedServiceWorker<I, V, E> getServiceWorker() {
    return this.serviceWorker;
  }

  /**
   * Get and clear worker to worker messages for this superstep. Can be
   * called only once per superstep.
   *
   * @return List of messages for this worker
   */
  public List<Writable> getAndClearCurrentWorkerToWorkerMessages() {
    List<Writable> ret = currentWorkerToWorkerMessages;
    currentWorkerToWorkerMessages = null;
    return ret;
  }

  /**
   * Add incoming message to this worker for next superstep. Thread-safe.
   *
   * @param message Message received
   */
  public void addIncomingWorkerToWorkerMessage(Writable message) {
    incomingWorkerToWorkerMessages.add(message);
  }


  /**
   * Get worker to worker messages received in previous superstep.
   * @return list of current worker to worker messages.
   */
  public List<Writable> getCurrentWorkerToWorkerMessages() {
    return currentWorkerToWorkerMessages;
  }

  /**
   * Prepare resolving mutation.
   */
  public void prepareResolveMutations() {
    oldPartitionMutations = partitionMutations;
    partitionMutations = Maps.newConcurrentMap();
  }

  /**
   * Resolve mutations specific for a partition. This method is called once
   * per partition, before the computation for that partition starts.
   * @param partition The partition to resolve mutations for
   */
  public void resolvePartitionMutation(Partition<I, V, E> partition) {
    Integer partitionId = partition.getId();
    VertexResolver<I, V, E> vertexResolver = conf.createVertexResolver();
    ConcurrentMap<I, VertexMutations<I, V, E>> prevPartitionMutations =
        oldPartitionMutations.get(partitionId);

    boolean ignoreExistingVertices =
        conf.getIncomingMessageClasses().ignoreExistingVertices();

    // Resolve mutations that are explicitly sent for this partition
    if (prevPartitionMutations != null) {
      for (Map.Entry<I, VertexMutations<I, V, E>> entry : prevPartitionMutations
          .entrySet()) {
        I vertexId = entry.getKey();
        Vertex<I, V, E> originalVertex = partition.getVertex(vertexId);
        VertexMutations<I, V, E> vertexMutations = entry.getValue();
        Vertex<I, V, E> vertex = vertexResolver.resolve(vertexId,
            originalVertex, vertexMutations,
            !ignoreExistingVertices &&
            getCurrentMessageStore().hasMessagesForVertex(entry.getKey()));

        if (LOG.isDebugEnabled()) {
          LOG.debug("resolvePartitionMutations: Resolved vertex index " +
              vertexId + " in partition index " + partitionId +
              " with original vertex " + originalVertex +
              ", returned vertex " + vertex + " on superstep " +
              serviceWorker.getSuperstep() + " with mutations " +
              vertexMutations);
        }

        if (vertex != null) {
          partition.putVertex(vertex);
        } else if (originalVertex != null) {
          partition.removeVertex(vertexId);
          if (!ignoreExistingVertices) {
            getCurrentMessageStore().clearVertexMessages(vertexId);
          }
        }
        context.progress();
      }
    }

    if (!ignoreExistingVertices) {
      // Keep track of vertices which are not here in the partition, but have
      // received messages
      Iterable<I> destinations = getCurrentMessageStore().
          getPartitionDestinationVertices(partitionId);
      if (!Iterables.isEmpty(destinations)) {
        for (I vertexId : destinations) {
          if (partition.getVertex(vertexId) == null) {
            Vertex<I, V, E> vertex =
                vertexResolver.resolve(vertexId, null, null, true);

            if (LOG.isDebugEnabled()) {
              LOG.debug(
                  "resolvePartitionMutations: A non-existing vertex has " +
                  "message(s). Added vertex index " + vertexId +
                  " in partition index " + partitionId +
                  ", vertex = " + vertex + ", on superstep " +
                  serviceWorker.getSuperstep());
            }

            if (vertex != null) {
              partition.putVertex(vertex);
            }
            context.progress();
          }
        }
      }
    }
  }

  /**
   * In case of async message store we have to wait for all messages
   * to be processed before going into next superstep.
   */
  public void waitForComplete() {
    if (incomingMessageStore instanceof AsyncMessageStoreWrapper) {
      ((AsyncMessageStoreWrapper) incomingMessageStore).waitToComplete();
    }
  }
}