content: /*
   * 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.tez.runtime.library.common.shuffle.orderedgrouped;
  
  import java.io.IOException;
  import java.util.concurrent.atomic.AtomicInteger;
  
  import com.google.common.annotations.VisibleForTesting;
  import org.apache.tez.common.CallableWithNdc;
  import org.apache.tez.common.counters.TezCounter;
  import org.apache.tez.common.counters.TezCounters;
  import org.apache.tez.runtime.library.common.InputAttemptIdentifier;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.client.CelebornTezReader;
  import org.apache.celeborn.common.exception.CelebornIOException;
  
  public class CelebornTezShuffleDataFetcher extends CallableWithNdc<Void> {
    private static final Logger LOG = LoggerFactory.getLogger(CelebornTezShuffleDataFetcher.class);
  
    private enum ShuffleErrors {
      IO_ERROR,
      WRONG_LENGTH,
      BAD_ID,
      WRONG_MAP,
      CONNECTION,
      WRONG_REDUCE
    }
  
    private static final String SHUFFLE_ERR_GRP_NAME = "Shuffle Errors";
  
    private final TezCounter ioErrs;
    private final MergeManager merger;
  
    private long copyBlockCount = 0;
    private volatile boolean stopped = false;
  
    private final CelebornTezReader celebornTezReader;
    private long readTime = 0;
    private long decompressTime = 0;
    private long serializeTime = 0;
    private long waitTime = 0;
    private long copyTime = 0; // the sum of readTime + decompressTime + serializeTime + waitTime
    private final InputAttemptIdentifier inputAttemptIdentifier;
    private static int uniqueMapId = 0;
  
    private boolean hasPendingData = false;
    private long startWait;
    private int waitCount = 0;
    private byte[] shuffleData = null;
    private Integer partitionId;
    private final ExceptionReporter exceptionReporter;
  
    private final AtomicInteger issuedCnt = new AtomicInteger(0);
  
    public CelebornTezShuffleDataFetcher(
        InputAttemptIdentifier inputAttemptIdentifier,
        Integer partitionId,
        MergeManager merger,
        TezCounters tezCounters,
        CelebornTezReader celebornTezReader,
        ExceptionReporter exceptionReporter) {
      this.merger = merger;
      this.partitionId = partitionId;
      this.inputAttemptIdentifier = inputAttemptIdentifier;
      this.exceptionReporter = exceptionReporter;
      ioErrs = tezCounters.findCounter(SHUFFLE_ERR_GRP_NAME, ShuffleErrors.IO_ERROR.toString());
      this.celebornTezReader = celebornTezReader;
  
      LOG.info(
          "CelebornTezShuffleDataFetcher, partitionId:{}, inputAttemptIdentifier:{}.",
          this.partitionId,
          this.inputAttemptIdentifier);
    }
  
    @Override
    public Void callInternal() {
      try {
        celebornTezReader.init();
        fetchAllRssBlocks();
      } catch (InterruptedException ie) {
        // might not be respected when fetcher is in progress / server is busy.  TEZ-711
        // Set the status back
        LOG.warn(ie.getMessage(), ie);
        Thread.currentThread().interrupt();
        return null;
      } catch (Throwable t) {
        LOG.warn(t.getMessage(), t);
        exceptionReporter.reportException(t);
        // Shuffle knows how to deal with failures post shutdown via the onFailure hook
      }
      return null;
    }
  
    public void fetchAllRssBlocks() throws IOException, InterruptedException {
      while (!stopped) {
        try {
          // If merge is on, block
          merger.waitForInMemoryMerge();
          // Do shuffle
          copyFromRssServer();
        } catch (Exception e) {
          LOG.warn(e.getMessage(), e);
          throw e;
        }
      }
    }
  
    @VisibleForTesting
    public void copyFromRssServer() throws IOException {
      // fetch a block
      if (!hasPendingData) {
        final long startFetch = System.currentTimeMillis();
        shuffleData = celebornTezReader.getShuffleBlock();
        long fetchDuration = System.currentTimeMillis() - startFetch;
        readTime += fetchDuration;
      }
  
      if (shuffleData != null) {
        // start to merge
        final long startSerialization = System.currentTimeMillis();
        if (issueMapOutputMerge()) {
          long serializationDuration = System.currentTimeMillis() - startSerialization;
          serializeTime += serializationDuration;
          // if reserve successes, reset status for next fetch
          if (hasPendingData) {
            waitTime += System.currentTimeMillis() - startWait;
          }
          hasPendingData = false;
        } else {
          // if reserve fail, return and wait
          startWait = System.currentTimeMillis();
          return;
        }
  
        // update some status
        copyBlockCount++;
        copyTime = readTime + decompressTime + serializeTime + waitTime;
        updateStatus();
      } else {
        // finish reading data, close related reader and check data consistent
        celebornTezReader.close();
        LOG.info(
            "Reduce task "
                + inputAttemptIdentifier
                + " read block cnt: "
                + copyBlockCount
                + " cost "
                + readTime
                + " ms to fetch and "
                + decompressTime
                + " ms to decompress "
                + serializeTime
                + " ms to serialize and "
                + waitTime
                + " ms to wait resource"
                + ", copy time:"
                + copyTime);
        stopFetch();
      }
    }
  
    public Integer getPartitionId() {
      return partitionId;
    }
  
    private boolean issueMapOutputMerge() throws IOException {
      // Allocate a MapOutput (either in-memory or on-disk) to put uncompressed block
      // In Rss, a MapOutput is sent as multiple blocks, so the reducer needs to
      // treat each "block" as a faked "mapout".
      // To avoid name conflicts, we use getNextUniqueTaskAttemptID instead.
      // It will generate a unique TaskAttemptID(increased_seq++, 0).
      InputAttemptIdentifier uniqueInputAttemptIdentifier = getNextUniqueInputAttemptIdentifier();
      MapOutput mapOutput = null;
      try {
        issuedCnt.incrementAndGet();
        mapOutput = merger.reserve(uniqueInputAttemptIdentifier, shuffleData.length, 0, 1);
      } catch (IOException ioe) {
        // kill this reduce attempt
        ioErrs.increment(1);
        throw ioe;
      }
      // Check if we can shuffle *now* ...
      if (mapOutput == null || mapOutput.getType() == MapOutput.Type.WAIT) {
        LOG.info("RssMRFetcher - MergeManager returned status WAIT ...");
        // Not an error but wait to process data.
        // Use a retry flag to avoid re-fetch and re-uncompress.
        hasPendingData = true;
        waitCount++;
        return false;
      }
  
      // write data to mapOutput
      try {
        CelebornTezBypassWriter.write(mapOutput, shuffleData);
        // let the merger knows this block is ready for merging
        mapOutput.commit();
      } catch (Throwable t) {
        ioErrs.increment(1);
        mapOutput.abort();
        throw new CelebornIOException(
            "Reduce: "
                + inputAttemptIdentifier
                + " cannot write block to "
                + mapOutput.getClass().getSimpleName()
                + " due to: "
                + t.getClass().getName());
      }
      return true;
    }
  
    private InputAttemptIdentifier getNextUniqueInputAttemptIdentifier() {
      return new InputAttemptIdentifier(uniqueMapId++, 0);
    }
  
    private void updateStatus() {}
  
    @VisibleForTesting
    public int getRetryCount() {
      return waitCount;
    }
  
    private void stopFetch() {
      LOG.info("CelebornTezShuffleDataFetcher stop fetch");
      stopped = true;
    }
  
    public void shutDown() {
      stopFetch();
      LOG.info("CelebornTezShuffleDataFetcher shutdown");
    }
  }
  