giraph-block-app/src/main/java/org/apache/giraph/block_app/framework/internal/BlockWorkerPieces.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.block_app.framework.internal;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Objects;

import org.apache.giraph.block_app.framework.api.BlockApiHandle;
import org.apache.giraph.comm.messages.MessageEncodeAndStoreType;
import org.apache.giraph.conf.DefaultMessageClasses;
import org.apache.giraph.conf.GiraphConstants;
import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;
import org.apache.giraph.conf.MessageClasses;
import org.apache.giraph.counters.CustomCounter;
import org.apache.giraph.counters.CustomCounters;
import org.apache.giraph.factories.DefaultMessageValueFactory;
import org.apache.giraph.master.MasterCompute;
import org.apache.giraph.types.NoMessage;
import org.apache.giraph.utils.UnsafeReusableByteArrayInput;
import org.apache.giraph.utils.WritableUtils;
import org.apache.giraph.worker.WorkerGlobalCommUsage;
import org.apache.giraph.writable.kryo.KryoWritableWrapper;
import org.apache.hadoop.io.IntWritable;
import org.apache.hadoop.io.Writable;
import org.apache.log4j.Logger;

/**
 * Pair of pieces to be executed on workers in a superstep
 *
 * @param <S> Execution stage type
 */
@SuppressWarnings({ "rawtypes", "unchecked" })
public class BlockWorkerPieces<S> {
  private static final Logger LOG = Logger.getLogger(BlockWorkerPieces.class);

  /** Aggregator holding next worker computation */
  private static final
  String NEXT_WORKER_PIECES = "giraph.blocks.next_worker_pieces";

  /** Passed worker stats counter group */
  private static final String PASSED_WORKER_STATS_GROUP = "PassedWorker Stats";

  /** Total serialised size counter name */
  private static final String TOTAL_SERIALISED_SIZE_NAME =
          "total serialized size";

  /** Split parts counter name */
  private static final String SPLIT_PARTS_NAME = "split parts";

  private final PairedPieceAndStage<S> receiver;
  private final PairedPieceAndStage<S> sender;
  private final BlockApiHandle blockApiHandle;

  public BlockWorkerPieces(
      PairedPieceAndStage<S> receiver, PairedPieceAndStage<S> sender,
      BlockApiHandle blockApiHandle) {
    this.receiver = receiver;
    this.sender = sender;
    this.blockApiHandle = blockApiHandle;
  }

  public PairedPieceAndStage<S> getReceiver() {
    return receiver;
  }

  public PairedPieceAndStage<S> getSender() {
    return sender;
  }

  public BlockApiHandle getBlockApiHandle() {
    return blockApiHandle;
  }

  public MessageClasses getOutgoingMessageClasses(
      ImmutableClassesGiraphConfiguration conf) {
    MessageClasses messageClasses;
    if (sender == null || sender.getPiece() == null) {
      messageClasses = new DefaultMessageClasses(
          NoMessage.class,
          DefaultMessageValueFactory.class,
          null,
          MessageEncodeAndStoreType.BYTEARRAY_PER_PARTITION);
    } else {
      messageClasses = sender.getPiece().getMessageClasses(conf);
    }

    messageClasses.verifyConsistent(conf);
    return messageClasses;
  }

  @Override
  public String toString() {
    return "[receiver=" + receiver + ",sender=" + sender + "]";
  }

  public String toStringShort() {
    String receiverString =
        Objects.toString(receiver != null ? receiver.getPiece() : null);
    String senderString =
        Objects.toString(sender != null ? sender.getPiece() : null);
    if (receiverString.equals(senderString)) {
      return "[receiver&sender=" + receiverString + "]";
    } else {
      return "[receiver=" + receiverString + ",sender=" + senderString + "]";
    }
  }

  /**
   * Sets which WorkerComputation is going to be executed in the next superstep.
   */
  public static <S> void setNextWorkerPieces(
      MasterCompute master, BlockWorkerPieces<S> nextWorkerPieces) {
    Writable toBroadcast = new KryoWritableWrapper<>(nextWorkerPieces);
    byte[] data = WritableUtils.toByteArrayUnsafe(toBroadcast);

    // TODO: extract splitting logic into common utility
    int overhead = 4096;
    int singleSize = Math.max(
        overhead,
        GiraphConstants.MAX_MSG_REQUEST_SIZE.get(master.getConf()) - overhead);

    ArrayList<byte[]> splittedData = new ArrayList<>();
    if (data.length < singleSize) {
      splittedData.add(data);
    } else {
      for (int start = 0; start < data.length; start += singleSize) {
        splittedData.add(Arrays.copyOfRange(
            data, start, Math.min(data.length, start + singleSize)));
      }
    }

    LOG.info("Next worker piece - total serialized size: " + data.length +
        ", split into " + splittedData.size());
    CustomCounters.addCustomCounter(PASSED_WORKER_STATS_GROUP,
            TOTAL_SERIALISED_SIZE_NAME, CustomCounter.Aggregation.SUM);
    master.getContext().getCounter(PASSED_WORKER_STATS_GROUP,
            TOTAL_SERIALISED_SIZE_NAME)
        .increment(data.length);
    CustomCounters.addCustomCounter(PASSED_WORKER_STATS_GROUP,
            SPLIT_PARTS_NAME, CustomCounter.Aggregation.SUM);
    master.getContext().getCounter(PASSED_WORKER_STATS_GROUP, SPLIT_PARTS_NAME)
        .increment(splittedData.size());

    master.broadcast(NEXT_WORKER_PIECES, new IntWritable(splittedData.size()));

    for (int i = 0; i < splittedData.size(); i++) {
      master.broadcast(NEXT_WORKER_PIECES + "_part_" + i,
          KryoWritableWrapper.wrapIfNeeded(splittedData.get(i)));
    }

    master.setOutgoingMessageClasses(
        nextWorkerPieces.getOutgoingMessageClasses(master.getConf()));
  }

  public static <S> BlockWorkerPieces<S> getNextWorkerPieces(
      WorkerGlobalCommUsage worker) {
    int splits = worker.<IntWritable>getBroadcast(NEXT_WORKER_PIECES).get();

    int totalLength = 0;
    ArrayList<byte[]> splittedData = new ArrayList<>();
    for (int i = 0; i < splits; i++) {
      byte[] cur = KryoWritableWrapper.<byte[]>unwrapIfNeeded(
          worker.getBroadcast(NEXT_WORKER_PIECES + "_part_" + i));
      splittedData.add(cur);
      totalLength += cur.length;
    }

    byte[] merged;
    if (splits == 1) {
      merged = splittedData.get(0);
    } else {
      merged = new byte[totalLength];
      int index = 0;
      for (int i = 0; i < splits; i++) {
        System.arraycopy(
            splittedData.get(i), 0, merged, index, splittedData.get(i).length);
        index += splittedData.get(i).length;
      }
    }

    KryoWritableWrapper<BlockWorkerPieces<S>> wrapper =
        new KryoWritableWrapper<>();
    WritableUtils.fromByteArrayUnsafe(
        merged, wrapper, new UnsafeReusableByteArrayInput());
    return wrapper.get();
  }
}