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
   *
   * <p>http://www.apache.org/licenses/LICENSE-2.0
   *
   * <p>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.impl;
  
  import static org.apache.celeborn.tez.plugin.util.CelebornTezUtils.getParentPrivateField;
  
  import java.io.IOException;
  import java.util.*;
  import java.util.concurrent.BlockingQueue;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.LinkedBlockingQueue;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicBoolean;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.concurrent.locks.Condition;
  import java.util.concurrent.locks.ReentrantLock;
  
  import com.google.common.collect.Sets;
  import com.google.common.util.concurrent.FutureCallback;
  import com.google.common.util.concurrent.Futures;
  import com.google.common.util.concurrent.ListenableFuture;
  import com.google.common.util.concurrent.ListeningExecutorService;
  import org.apache.hadoop.conf.Configuration;
  import org.apache.hadoop.io.compress.CompressionCodec;
  import org.apache.hadoop.yarn.api.records.ApplicationAttemptId;
  import org.apache.tez.common.CallableWithNdc;
  import org.apache.tez.common.GuavaShim;
  import org.apache.tez.common.Preconditions;
  import org.apache.tez.common.counters.TezCounter;
  import org.apache.tez.dag.records.TezTaskAttemptID;
  import org.apache.tez.runtime.api.InputContext;
  import org.apache.tez.runtime.api.TaskFailureType;
  import org.apache.tez.runtime.library.common.CompositeInputAttemptIdentifier;
  import org.apache.tez.runtime.library.common.InputAttemptIdentifier;
  import org.apache.tez.runtime.library.common.shuffle.*;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.client.CelebornTezReader;
  import org.apache.celeborn.client.ShuffleClient;
  import org.apache.celeborn.tez.plugin.util.CelebornTezUtils;
  
  public class CelebornShuffleManager extends ShuffleManager {
  
    private static final Logger LOG = LoggerFactory.getLogger(CelebornShuffleManager.class);
  
    private final InputContext inputContext;
    private final int numInputs;
    private final ListeningExecutorService schedulerExecutor;
    private final int shuffleId;
    private final ApplicationAttemptId applicationAttemptId;
    private final ShuffleClient shuffleClient;
    private final FetchedInputAllocator inputManager;
  
    private final Set<CelebornFetcher> celebornRunningFetchers;
    private final BlockingQueue<FetchedInput> completedInputs;
    private final BlockingQueue<InputHost> pendingHosts;
  
    private final AtomicInteger numCompletedInputs = new AtomicInteger(0);
    // Required to be held when manipulating pendingHosts
    private final ReentrantLock lock = new ReentrantLock();
    private final Condition wakeLoop = lock.newCondition();
  
    private final int numFetchers;
    private final int maxTimeToWaitForReportMillis;
    private final String sourceDestNameTrimmed;
    private final TezCounter shufflePhaseTime;
  
    private final AtomicBoolean isShutdown;
    private final long startTime;
  
    private volatile Throwable shuffleError;
    private final Set<Integer> successRssPartitionSet = new HashSet<>();
    private final Set<Integer> runningRssPartitionMap = new HashSet<>();
    private final Set<Integer> allRssPartition = Sets.newConcurrentHashSet();
    private final BlockingQueue<Integer> pendingPartition = new LinkedBlockingQueue<>();
    Map<Integer, List<InputAttemptIdentifier>> partitionToInput = new HashMap<>();
    private final AtomicInteger numNoDataInput = new AtomicInteger(0);
    private final AtomicInteger numWithDataInput = new AtomicInteger(0);
    private final boolean broadcastOrOneToOne;
  
    public CelebornShuffleManager(
        InputContext inputContext,
        Configuration conf,
        int numInputs,
        ShuffleClient shuffleClient,
        int bufferSize,
        boolean ifileReadAheadEnabled,
        int ifileReadAheadLength,
        CompressionCodec codec,
        FetchedInputAllocator inputAllocator,
        int shuffleId,
        ApplicationAttemptId applicationAttemptId,
        boolean broadcastOrOneToOne)
        throws IOException {
      super(
          inputContext,
          conf,
          numInputs,
          bufferSize,
          ifileReadAheadEnabled,
          ifileReadAheadLength,
          codec,
          inputAllocator);
      this.inputContext = inputContext;
      this.shuffleId = shuffleId;
      this.applicationAttemptId = applicationAttemptId;
      this.shuffleClient = shuffleClient;
      this.broadcastOrOneToOne = broadcastOrOneToOne;
      this.inputManager = inputAllocator;
      this.celebornRunningFetchers =
          Collections.newSetFromMap(new ConcurrentHashMap<CelebornFetcher, Boolean>());
      this.numInputs = numInputs;
      this.numFetchers = (int) getParentPrivateField(this, "numFetchers");
      this.schedulerExecutor =
          (ListeningExecutorService) getParentPrivateField(this, "schedulerExecutor");
      this.maxTimeToWaitForReportMillis =
          (int) getParentPrivateField(this, "maxTimeToWaitForReportMillis");
      this.isShutdown = (AtomicBoolean) getParentPrivateField(this, "isShutdown");
      this.sourceDestNameTrimmed = (String) getParentPrivateField(this, "sourceDestNameTrimmed");
      this.completedInputs =
          (BlockingQueue<FetchedInput>) getParentPrivateField(this, "completedInputs");
      this.pendingHosts = (BlockingQueue<InputHost>) getParentPrivateField(this, "pendingHosts");
      this.shufflePhaseTime = (TezCounter) getParentPrivateField(this, "shufflePhaseTime");
      this.startTime = (long) getParentPrivateField(this, "startTime");
    }
  
    @Override
    public void run() throws IOException {
      Preconditions.checkState(inputManager != null, "InputManager must be configured");
      // todo maxTimeToWaitForReportMillis
      ListenableFuture<Void> runShuffleFuture =
          schedulerExecutor.submit(new CelebornRunShuffleCallable());
      Futures.addCallback(
          runShuffleFuture, new SchedulerFutureCallback(), GuavaShim.directExecutor());
      // Shutdown this executor once this task, and the callback complete.
      schedulerExecutor.shutdown();
    }
  
    private boolean isAllInputFetched() {
      LOG.info(
          "Check isAllInputFetched, numNoDataInput:{}, numWithDataInput:{},numInputs:{},  "
              + "successRssPartitionSet:{},  allRssPartition:{}.",
          numNoDataInput,
          numWithDataInput,
          numInputs,
          successRssPartitionSet,
          allRssPartition);
      return (numNoDataInput.get() + numWithDataInput.get() >= numInputs)
          && (successRssPartitionSet.size() >= allRssPartition.size());
    }
  
    private boolean isAllInputAdded() {
      LOG.info(
          "Check isAllInputAdded, numNoDataInput:{}, numWithDataInput:{},numInputs:{},  "
              + "successRssPartitionSet:{}, allRssPartition:{}.",
          numNoDataInput,
          numWithDataInput,
          numInputs,
          successRssPartitionSet,
          allRssPartition);
      return numNoDataInput.get() + numWithDataInput.get() >= numInputs;
    }
  
    private class CelebornRunShuffleCallable extends CallableWithNdc<Void> {
  
      public CelebornRunShuffleCallable() {}
  
      @Override
      protected Void callInternal() throws Exception {
        while (!isShutdown.get() && !isAllInputFetched()) {
          lock.lock();
          try {
            while (((celebornRunningFetchers.size() >= numFetchers || pendingPartition.isEmpty())
                    && !isAllInputFetched())
                || !isAllInputAdded()) {
              inputContext.notifyProgress();
              boolean ret = wakeLoop.await(1000, TimeUnit.MILLISECONDS);
              if (isShutdown.get()) {
                break;
              }
            }
          } finally {
            lock.unlock();
          }
  
          if (shuffleError != null) {
            // InputContext has already been informed of a fatal error. Relying on
            // tez to kill the task.
            break;
          }
  
          LOG.debug("{}: NumCompletedInputs: {}", sourceDestNameTrimmed, numCompletedInputs);
          if (!isAllInputFetched() && !isShutdown.get()) {
            lock.lock();
            try {
              int maxFetchersToRun = numFetchers - celebornRunningFetchers.size();
              int count = 0;
              while (pendingPartition.peek() != null && !isShutdown.get()) {
                Integer partition = null;
                try {
                  partition = pendingPartition.take();
                } catch (InterruptedException e) {
                  if (isShutdown.get()) {
                    LOG.info(
                        sourceDestNameTrimmed
                            + ": "
                            + "Interrupted and hasBeenShutdown, Breaking out of ShuffleScheduler Loop");
                    Thread.currentThread().interrupt();
                    break;
                  } else {
                    throw e;
                  }
                }
                if (LOG.isDebugEnabled()) {
                  LOG.debug(
                      sourceDestNameTrimmed + ": " + "Processing pending partition: " + partition);
                }
                if (!isShutdown.get()
                    && (!successRssPartitionSet.contains(partition)
                        && !runningRssPartitionMap.contains(partition))) {
  
                  runningRssPartitionMap.add(partition);
  
                  CelebornFetcher fetcher = constructFetcherForCeleborn(partition);
                  celebornRunningFetchers.add(fetcher);
                  if (isShutdown.get()) {
                    LOG.info(
                        sourceDestNameTrimmed
                            + ": "
                            + "hasBeenShutdown,"
                            + "Breaking out of ShuffleScheduler Loop");
                    break;
                  }
                  ListenableFuture<FetchResult> future = fetcherExecutor.submit(fetcher);
                  Futures.addCallback(
                      future, new FetchFutureCallback(fetcher), GuavaShim.directExecutor());
                  if (++count >= maxFetchersToRun) {
                    break;
                  }
                } else {
                  if (LOG.isDebugEnabled()) {
                    LOG.debug(
                        sourceDestNameTrimmed
                            + ": "
                            + "Skipping partition: "
                            + partition
                            + " since is shutdown");
                  }
                }
              }
            } finally {
              lock.unlock();
            }
          }
        }
        shufflePhaseTime.setValue(System.currentTimeMillis() - startTime);
        LOG.info(
            sourceDestNameTrimmed
                + ": "
                + "Shutting down FetchScheduler, Was Interrupted: "
                + Thread.currentThread().isInterrupted());
        if (!fetcherExecutor.isShutdown()) {
          fetcherExecutor.shutdownNow();
        }
        return null;
      }
    }
  
    CelebornFetcher constructFetcherForCeleborn(int partition) {
      CelebornTezReader reader =
          new CelebornTezReader(
              shuffleClient, shuffleId, partition, applicationAttemptId.getAttemptId());
      CelebornFetcher celebornFetcher = new CelebornFetcher(this, inputManager, reader);
      return celebornFetcher;
    }
  
    /////////////////// Methods for InputEventHandler
  
    public void addKnownInput(
        String hostName,
        int port,
        CompositeInputAttemptIdentifier srcAttemptIdentifier,
        int srcPhysicalIndex) {
      HostPort identifier = new HostPort(hostName, port);
      super.addKnownInput(hostName, port, srcAttemptIdentifier, srcPhysicalIndex);
      lock.lock();
      try {
        // pendingHosts is useless in celeborn
        pendingHosts.remove(identifier);
        if (broadcastOrOneToOne) {
          String pathComponent = srcAttemptIdentifier.getPathComponent();
          TezTaskAttemptID taskAttemptId =
              TezTaskAttemptID.fromString(
                  CelebornTezUtils.uniqueIdentifierToAttemptId(pathComponent));
          int partitionId = taskAttemptId.getTaskID().getId();
          LOG.info("Read partition, shuffleId {} partitionId {}", shuffleId, partitionId);
          if (!allRssPartition.contains(partitionId)) {
            pendingPartition.add(partitionId);
          }
          allRssPartition.add(partitionId);
          partitionToInput.putIfAbsent(partitionId, new ArrayList<>());
          partitionToInput.get(partitionId).add(srcAttemptIdentifier);
        } else {
          for (int i = 0; i < srcAttemptIdentifier.getInputIdentifierCount(); i++) {
            int p = srcPhysicalIndex + i;
            LOG.info(
                "PartitionToInput, original:{}, add:{},  now:{}",
                srcAttemptIdentifier,
                srcAttemptIdentifier.expand(i),
                partitionToInput.get(p));
            if (!allRssPartition.contains(srcPhysicalIndex + i)) {
              pendingPartition.add(p);
            }
            allRssPartition.add(p);
            partitionToInput.putIfAbsent(p, new ArrayList<>());
            partitionToInput.get(p).add(srcAttemptIdentifier);
            LOG.info("Add partition:{}, after add, now partition:{}", p, allRssPartition);
          }
        }
  
        numWithDataInput.incrementAndGet();
        LOG.info("numWithDataInput:{}.", numWithDataInput.get());
      } finally {
        lock.unlock();
      }
    }
  
    @Override
    public void addCompletedInputWithNoData(InputAttemptIdentifier srcAttemptIdentifier) {
      super.addCompletedInputWithNoData(srcAttemptIdentifier);
      numNoDataInput.incrementAndGet();
      LOG.info(
          "AddCompletedInputWithNoData, numNoDataInput:{}, numWithDataInput:{},numInputs:{},  "
              + "successRssPartitionSet:{}, allRssPartition:{}.",
          numNoDataInput,
          numWithDataInput,
          numInputs,
          successRssPartitionSet,
          allRssPartition);
    }
  
    public void shutdown() throws InterruptedException {
  
      if (!isShutdown.get()) {
        super.shutdown();
        lock.lock();
        try {
          wakeLoop.signal(); // signal the fetch-scheduler
          for (CelebornFetcher fetcher : celebornRunningFetchers) {
            try {
              fetcher.shutdown(); // This could be parallelized.
            } catch (Exception e) {
              LOG.warn(
                  "Error while stopping fetcher during shutdown. Ignoring and continuing. Message={}",
                  e.getMessage());
            }
          }
        } finally {
          lock.unlock();
        }
      }
    }
  
    public boolean isAllPartitionFetched() {
      lock.lock();
      try {
        if (!allRssPartition.containsAll(successRssPartitionSet)) {
          LOG.error(
              "Failed to check partition, all partition:{}, success partiton:{}",
              allRssPartition,
              successRssPartitionSet);
        }
        return isAllInputFetched();
      } finally {
        lock.unlock();
      }
    }
  
    /**
     * @return the next available input, or null if there are no available inputs. This method will
     *     block if there are currently no available inputs, but more may become available.
     */
    public FetchedInput getNextInput() throws InterruptedException {
      // Check for no additional inputs
      FetchedInput fetchedInput = null;
      if (completedInputs.peek() == null) {
        while (true) {
          fetchedInput = completedInputs.poll(2000, TimeUnit.MICROSECONDS);
          if (fetchedInput != null) {
            break;
          } else if (isAllPartitionFetched()) {
            fetchedInput = completedInputs.poll(100, TimeUnit.MICROSECONDS);
            LOG.info("GetNextInput, enter isAllPartitionFetched");
            break;
          }
          LOG.info("GetNextInput, out loop");
          Thread.sleep(1000);
        }
      } else {
        fetchedInput = completedInputs.take();
      }
  
      if (fetchedInput instanceof NullFetchedInput) {
        LOG.info("getNextInput, NullFetchedInput is null:{}", fetchedInput);
        fetchedInput = null;
      }
      LOG.info("getNextInput, fetchedInput:{}", fetchedInput);
      return fetchedInput;
    }
  
    private class SchedulerFutureCallback implements FutureCallback<Void> {
  
      @Override
      public void onSuccess(Void result) {
        LOG.info(sourceDestNameTrimmed + ": " + "Scheduler thread completed");
      }
  
      @Override
      public void onFailure(Throwable t) {
        if (isShutdown.get()) {
          LOG.debug("{}: Already shutdown. Ignoring error.", sourceDestNameTrimmed, t);
        } else {
          LOG.error(sourceDestNameTrimmed + ": " + "Scheduler failed with error: ", t);
          inputContext.reportFailure(TaskFailureType.NON_FATAL, t, "Shuffle Scheduler Failed");
        }
      }
    }
  
    private class FetchFutureCallback implements FutureCallback<FetchResult> {
  
      private final CelebornFetcher fetcher;
  
      public FetchFutureCallback(CelebornFetcher fetcher) {
        this.fetcher = fetcher;
      }
  
      private void doBookKeepingForFetcherComplete() {
        lock.lock();
        try {
          celebornRunningFetchers.remove(fetcher);
          wakeLoop.signal();
        } finally {
          lock.unlock();
        }
      }
  
      @Override
      public void onSuccess(FetchResult result) {
        fetcher.shutdown();
        if (isShutdown.get()) {
          LOG.debug("{}: Already shutdown. Ignoring event from fetcher", sourceDestNameTrimmed);
        } else {
          lock.lock();
          try {
            successRssPartitionSet.add(fetcher.getPartitionId());
            runningRssPartitionMap.remove(fetcher.getPartitionId());
            LOG.info(
                "FetchFutureCallback allRssPartition:{}, successRssPartitionSet:{}, runningRssPartitionMap:{}.",
                allRssPartition,
                successRssPartitionSet,
                runningRssPartitionMap);
            doBookKeepingForFetcherComplete();
          } finally {
            lock.unlock();
          }
        }
      }
  
      @Override
      public void onFailure(Throwable t) {
        // Unsuccessful - the fetcher may not have shutdown correctly. Try shutting it down.
        fetcher.shutdown();
        if (isShutdown.get()) {
          LOG.debug("{}: Already shutdown. Ignoring error from fetcher.", sourceDestNameTrimmed, t);
        } else {
          LOG.error(sourceDestNameTrimmed + ": " + "Fetcher failed with error: ", t);
          shuffleError = t;
          inputContext.reportFailure(TaskFailureType.NON_FATAL, t, "Fetch failed");
          doBookKeepingForFetcherComplete();
        }
      }
    }
  }
  