content: /*
   * Copyright 2021 Alibaba Group Holding Limited.
   *
   * Licensed 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 com.alibaba.graphscope.parallel.mm.impl;
  
  import static org.apache.giraph.conf.GiraphConstants.MAX_OUT_MSG_CACHE_SIZE;
  import static org.apache.giraph.utils.ByteUtils.SIZE_OF_LONG;
  
  import com.alibaba.fastffi.llvm4jni.runtime.JavaRuntime;
  import com.alibaba.graphscope.communication.FFICommunicator;
  import com.alibaba.graphscope.ds.PropertyNbrUnit;
  import com.alibaba.graphscope.fragment.BaseArrowProjectedFragment;
  import com.alibaba.graphscope.fragment.IFragment;
  import com.alibaba.graphscope.fragment.adaptor.AbstractArrowProjectedAdaptor;
  import com.alibaba.graphscope.graph.GiraphVertexIdManager;
  import com.alibaba.graphscope.graph.impl.VertexImpl;
  import com.alibaba.graphscope.parallel.DefaultMessageManager;
  import com.alibaba.graphscope.serialization.FFIByteVectorOutputStream;
  import com.alibaba.graphscope.stdcxx.FFIByteVector;
  import com.alibaba.graphscope.stdcxx.FFIByteVectorFactory;
  import com.alibaba.graphscope.utils.LongIdParser;
  
  import org.apache.giraph.conf.ImmutableClassesGiraphConfiguration;
  import org.apache.giraph.graph.Vertex;
  import org.apache.hadoop.io.Writable;
  import org.apache.hadoop.io.WritableComparable;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.IOException;
  
  public class GiraphMpiMessageManager<
                  OID_T extends WritableComparable,
                  VDATA_T extends Writable,
                  EDATA_T extends Writable,
                  IN_MSG_T extends Writable,
                  OUT_MSG_T extends Writable,
                  GS_VID_T,
                  GS_OID_T>
          extends AbstractMessageManager<
                  OID_T, VDATA_T, EDATA_T, IN_MSG_T, OUT_MSG_T, GS_VID_T, GS_OID_T> {
  
      public static long THRESHOLD;
      private static Logger logger = LoggerFactory.getLogger(GiraphMpiMessageManager.class);
      private FFIByteVectorOutputStream[] cacheOut;
  
      private LongIdParser idParser;
      private PropertyNbrUnit<Long> nbrUnit;
      private long offsetBeginPtrFirstAddr;
      private long offsetEndPtrFirstAddr;
      private long nbrUnitEleSize;
      private long nbrUnitInitAddress;
      private BaseArrowProjectedFragment<GS_OID_T, GS_VID_T, ?, ?> projectedFragment;
      private GiraphVertexIdManager<GS_VID_T, OID_T> idManager;
  
      public GiraphMpiMessageManager(
              IFragment fragment,
              DefaultMessageManager defaultMessageManager,
              ImmutableClassesGiraphConfiguration configuration,
              FFICommunicator communicator,
              GiraphVertexIdManager<GS_VID_T, OID_T> idManager) {
          super(fragment, defaultMessageManager, configuration, communicator);
          this.idManager = idManager;
          THRESHOLD = MAX_OUT_MSG_CACHE_SIZE.get(configuration);
          this.projectedFragment =
                  ((AbstractArrowProjectedAdaptor<GS_OID_T, GS_VID_T, ?, ?>) fragment)
                          .getBaseArrayProjectedFragment();
  
          this.cacheOut = new FFIByteVectorOutputStream[fragment.fnum()];
          for (int i = 0; i < fragment.fnum(); ++i) {
              this.cacheOut[i] = new FFIByteVectorOutputStream();
              this.cacheOut[i].resize(THRESHOLD);
          }
          idParser = new LongIdParser(fragment.fnum(), 1);
          nbrUnit = (PropertyNbrUnit<Long>) this.projectedFragment.getOutEdgesPtr();
          offsetEndPtrFirstAddr = this.projectedFragment.getOEOffsetsEndPtr();
          offsetBeginPtrFirstAddr = this.projectedFragment.getOEOffsetsBeginPtr();
          nbrUnitEleSize = nbrUnit.elementSize();
          nbrUnitInitAddress = nbrUnit.getAddress();
      }
  
      /**
       * Called by our frame work, to deserialize the messages from c++ to java. Must be called before
       * getMessages
       */
      @Override
      public void receiveMessages() {
          // put message to currentIncoming message store
          FFIByteVector tmpVector = (FFIByteVector) FFIByteVectorFactory.INSTANCE.create();
          long bytesOfReceivedMsg = 0;
          while (grapeMessager.getPureMessage(tmpVector)) {
              // The retrieved tmp vector has been resized, so the cached objAddress is not available.
              // trigger the refresh
              tmpVector.touch();
              // OutArchive will do the resize;
              if (logger.isDebugEnabled()) {
                  logger.debug("Frag [{}] digest message of size {}", fragId, tmpVector.size());
              }
              ///////////////////////////////////////////
              currentIncomingMessageStore.digest(tmpVector);
              ///////////////////////////////////////////
              bytesOfReceivedMsg += tmpVector.size();
          }
          tmpVector.delete();
          logger.info(
                  "Frag [{}] totally Received [{}] bytes from others starting deserialization",
                  fragId,
                  bytesOfReceivedMsg);
      }
  
      /**
       * Send one message to dstOid.
       *
       * @param dstOid  vertex to receive this message.
       * @param message message.
       */
      @Override
      public void sendMessage(OID_T dstOid, OUT_MSG_T message) {
          GS_VID_T lid = idManager.oid2Lid(dstOid);
          grapeVertex.setValue(lid);
          sendMessage(grapeVertex, message);
      }
  
      private void sendMessage(com.alibaba.graphscope.ds.Vertex<GS_VID_T> vertex, OUT_MSG_T msg) {
          int dstfragId = fragment.getFragId(vertex);
          if (cacheOut[dstfragId].bytesWritten() >= THRESHOLD && dstfragId != fragId) {
              cacheOut[dstfragId].writeLong(
                      0, cacheOut[dstfragId].bytesWritten() - 8); // minus size_of_long
              cacheOut[dstfragId].finishSetting();
              // the vertex will be swapped. so this vector is empty;
              grapeMessager.sendToFragment(dstfragId, cacheOut[dstfragId].getVector());
              //            cacheOut[dstfragId] = new FFIByteVectorOutputStream();
              //            cacheOut[dstfragId].resize(THRESHOLD);
              cacheOut[dstfragId].reset();
              cacheOut[dstfragId].writeLong(0, 0);
          }
          try {
              cacheOut[dstfragId].writeLong((Long) fragment.vertex2Gid(vertex));
              msg.write(cacheOut[dstfragId]);
          } catch (IOException e) {
              e.printStackTrace();
          }
      }
  
      /**
       * Send message to neighbor vertices.
       *
       * @param vertex
       */
      @Override
      public void sendMessageToAllEdges(Vertex<OID_T, VDATA_T, EDATA_T> vertex, OUT_MSG_T message) {
          VertexImpl<GS_VID_T, OID_T, VDATA_T, EDATA_T> vertexImpl =
                  (VertexImpl<GS_VID_T, OID_T, VDATA_T, EDATA_T>) vertex;
  
          long lid = vertexImpl.getLocalId();
          long offset = idParser.getOffset(lid);
          long oeBeginOffset = JavaRuntime.getLong(offsetBeginPtrFirstAddr + offset * 8);
          long oeEndOffset = JavaRuntime.getLong(offsetEndPtrFirstAddr + offset * 8);
          long curAddress = nbrUnitInitAddress + nbrUnitEleSize * oeBeginOffset;
          long endAddress = nbrUnitInitAddress + nbrUnitEleSize * oeEndOffset;
  
          while (curAddress < endAddress) {
              nbrUnit.setAddress(curAddress);
              grapeVertex.setValue((GS_VID_T) nbrUnit.vid());
              sendMessage(grapeVertex, message);
              curAddress += nbrUnitEleSize;
          }
  
          // send msg through outgoing adjlist
          //        AdjList<GS_VID_T, ?> adaptorAdjList = fragment.getOutgoingAdjList(grapeVertex);
          //
          //        for (NbrBase<GS_VID_T, ?> nbr : adaptorAdjList.nbrBases()) {
          //            com.alibaba.graphscope.ds.Vertex<GS_VID_T> curVertex = nbr.neighbor();
          //            sendMessage(curVertex, message);
          //        }
  
      }
  
      /**
       * Make sure all messages has been sent. Clean outputstream buffer
       */
      @Override
      public void finishMessageSending() {
          for (int i = 0; i < fragNum; ++i) {
              long bytesWritten = cacheOut[i].bytesWritten();
              cacheOut[i].finishSetting();
              cacheOut[i].writeLong(0, bytesWritten - SIZE_OF_LONG);
  
              if (bytesWritten == SIZE_OF_LONG) {
                  logger.debug(
                          "[Finish msg] sending skip msg from {} -> {}, since msg size: {}",
                          fragId,
                          i,
                          bytesWritten);
                  continue;
              }
              if (i == fragId) {
                  nextIncomingMessageStore.digest(cacheOut[i].getVector());
                  logger.info(
                          "In final step, Frag [{}] digest msg to self of size: {}",
                          fragId,
                          bytesWritten);
              } else {
                  grapeMessager.sendToFragment(i, cacheOut[i].getVector());
                  logger.info(
                          "In final step, Frag [{}] send msg to [{}] of size: {}",
                          fragId,
                          i,
                          bytesWritten);
              }
          }
          //        if (maxSuperStep > 0) {
          //            grapeMessager.ForceContinue();
          //            maxSuperStep -= 1;
          //        }
  
          //        logger.debug("[Unused res] {}", unused);
          //        logger.debug("adaptor hasNext {}, grape hasNext{}", adaptorHasNext, grapeHasNext);
          //        logger.debug("adaptor next {}, grape next {}", adaptorNext, grapeNext);
          //        logger.debug("adaptor neighbor {}, grape neighbor {}", adaptorNeighbor,
          // grapeNeighbor);
      }
  
      @Override
      public void preSuperstep() {
          for (int i = 0; i < fragNum; ++i) {
              //            cacheOut[i].resize(THRESHOLD);
              cacheOut[i].reset();
              //            if (i != fragId) {
              // only write size info for mpi messages, local message don't need size.
              try {
                  cacheOut[i].writeLong(0);
              } catch (IOException e) {
                  e.printStackTrace();
              }
              //            }
          }
      }
  
      @Override
      public void postSuperstep() {
          currentIncomingMessageStore.swap(nextIncomingMessageStore);
          nextIncomingMessageStore.clearAll();
      }
  
      @Override
      public void postApplication() {
          for (int i = 0; i < fragNum; ++i) {
              cacheOut[i].close();
          }
      }
  }
  