giraph-core/src/main/java/org/apache/giraph/graph/AddressesAndPartitionsWritable.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.giraph.partition.PartitionOwner;
import org.apache.giraph.master.MasterInfo;
import org.apache.giraph.utils.ReflectionUtils;
import org.apache.giraph.utils.WritableUtils;
import org.apache.giraph.worker.WorkerInfo;
import org.apache.hadoop.io.Writable;

import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;

import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;
import java.util.Collection;
import java.util.List;
import java.util.Map;

/**
 * Helper class to write descriptions of master, workers and partition owners
 */
public class AddressesAndPartitionsWritable implements Writable {
  /** Master information */
  private MasterInfo masterInfo;
  /** List of all workers */
  private List<WorkerInfo> workerInfos;
  /** Collection of partitions */
  private Collection<PartitionOwner> partitionOwners;

  /**
   * Constructor when we want to serialize object
   *
   * @param masterInfo Master information
   * @param workerInfos List of all workers
   * @param partitionOwners Collection of partitions
   */
  public AddressesAndPartitionsWritable(MasterInfo masterInfo,
      List<WorkerInfo> workerInfos,
      Collection<PartitionOwner> partitionOwners) {
    this.masterInfo = masterInfo;
    this.workerInfos = workerInfos;
    this.partitionOwners = partitionOwners;
  }

  /** Constructor for reflection */
  public AddressesAndPartitionsWritable() {
  }

  /**
   * Get master information
   *
   * @return Master information
   */
  public MasterInfo getMasterInfo() {
    return masterInfo;
  }

  /**
   * Get all workers
   *
   * @return List of all workers
   */
  public List<WorkerInfo> getWorkerInfos() {
    return workerInfos;
  }

  /**
   * Get partition owners
   *
   * @return Collection of partition owners
   */
  public Collection<PartitionOwner> getPartitionOwners() {
    return partitionOwners;
  }

  @Override
  public void write(DataOutput output) throws IOException {
    masterInfo.write(output);

    output.writeInt(workerInfos.size());
    for (WorkerInfo workerInfo : workerInfos) {
      workerInfo.write(output);
    }

    Map<Integer, WorkerInfo> workerInfoMap = getAsWorkerInfoMap(workerInfos);
    // Also write out the previous worker information that are used
    // in the partition owners
    List<WorkerInfo> previousWorkerInfos = Lists.newArrayList();
    for (PartitionOwner partitionOwner : partitionOwners) {
      if (partitionOwner.getPreviousWorkerInfo() != null) {
        if (!workerInfoMap.containsKey(
            partitionOwner.getPreviousWorkerInfo().getTaskId())) {
          previousWorkerInfos.add(partitionOwner.getPreviousWorkerInfo());
        }
      }
    }
    output.writeInt(previousWorkerInfos.size());
    for (WorkerInfo workerInfo : previousWorkerInfos) {
      workerInfo.write(output);
    }

    output.writeInt(partitionOwners.size());
    if (partitionOwners.size() > 0) {
      WritableUtils.writeClass(
          partitionOwners.iterator().next().getClass(), output);
    }
    for (PartitionOwner partitionOwner : partitionOwners) {
      partitionOwner.writeWithWorkerIds(output);
    }
  }

  @Override
  public void readFields(DataInput input) throws IOException {
    masterInfo = new MasterInfo();
    masterInfo.readFields(input);

    int workerInfosSize = input.readInt();
    workerInfos = Lists.newArrayListWithCapacity(workerInfosSize);
    for (int i = 0; i < workerInfosSize; i++) {
      WorkerInfo workerInfo = new WorkerInfo();
      workerInfo.readFields(input);
      workerInfos.add(workerInfo);
    }

    Map<Integer, WorkerInfo> workerInfoMap = getAsWorkerInfoMap(workerInfos);
    int additionalWorkerInfos = input.readInt();
    for (int i = 0; i < additionalWorkerInfos; i++) {
      WorkerInfo workerInfo = new WorkerInfo();
      workerInfo.readFields(input);
      workerInfoMap.put(workerInfo.getTaskId(), workerInfo);
    }

    int partitionOwnersSize = input.readInt();
    Class<PartitionOwner> partitionOwnerClass = null;
    if (partitionOwnersSize > 0) {
      partitionOwnerClass = WritableUtils.readClass(input);
    }
    partitionOwners = Lists.newArrayListWithCapacity(partitionOwnersSize);
    for (int i = 0; i < partitionOwnersSize; i++) {
      PartitionOwner partitionOwner =
          ReflectionUtils.newInstance(partitionOwnerClass);
      partitionOwner.readFieldsWithWorkerIds(input, workerInfoMap);
      partitionOwners.add(partitionOwner);
    }
  }

  /**
   * Convert Iterable of WorkerInfos to the map from task id to WorkerInfo.
   *
   * @param workerInfos Iterable of WorkerInfos
   * @return The map from task id to WorkerInfo
   */
  private static Map<Integer, WorkerInfo> getAsWorkerInfoMap(
      Iterable<WorkerInfo> workerInfos) {
    Map<Integer, WorkerInfo> workerInfoMap =
        Maps.newHashMapWithExpectedSize(Iterables.size(workerInfos));
    for (WorkerInfo workerInfo : workerInfos) {
      workerInfoMap.put(workerInfo.getTaskId(), workerInfo);
    }
    return workerInfoMap;
  }
}