/*

  * 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();

   }

 }