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.celeborn.client;
  
  import java.io.IOException;
  import java.nio.ByteBuffer;
  import java.util.*;
  import java.util.concurrent.Callable;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.TimeUnit;
  
  import scala.Tuple2;
  import scala.Tuple3;
  import scala.reflect.ClassTag$;
  
  import com.google.common.annotations.VisibleForTesting;
  import com.google.common.collect.Lists;
  import com.google.protobuf.InvalidProtocolBufferException;
  import io.netty.buffer.CompositeByteBuf;
  import io.netty.buffer.Unpooled;
  import org.apache.commons.lang3.StringUtils;
  import org.apache.commons.lang3.tuple.Pair;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.client.compress.Compressor;
  import org.apache.celeborn.client.read.CelebornInputStream;
  import org.apache.celeborn.client.read.MetricsCallback;
  import org.apache.celeborn.common.CelebornConf;
  import org.apache.celeborn.common.exception.CelebornBroadcastException;
  import org.apache.celeborn.common.exception.CelebornIOException;
  import org.apache.celeborn.common.exception.CelebornRuntimeException;
  import org.apache.celeborn.common.identity.UserIdentifier;
  import org.apache.celeborn.common.metrics.source.Role;
  import org.apache.celeborn.common.network.TransportContext;
  import org.apache.celeborn.common.network.buffer.NettyManagedBuffer;
  import org.apache.celeborn.common.network.client.*;
  import org.apache.celeborn.common.network.protocol.*;
  import org.apache.celeborn.common.network.protocol.SerdeVersion;
  import org.apache.celeborn.common.network.sasl.SaslClientBootstrap;
  import org.apache.celeborn.common.network.sasl.SaslCredentials;
  import org.apache.celeborn.common.network.server.BaseMessageHandler;
  import org.apache.celeborn.common.network.util.TransportConf;
  import org.apache.celeborn.common.protocol.*;
  import org.apache.celeborn.common.protocol.message.ControlMessages.*;
  import org.apache.celeborn.common.protocol.message.StatusCode;
  import org.apache.celeborn.common.rpc.RpcAddress;
  import org.apache.celeborn.common.rpc.RpcEndpointRef;
  import org.apache.celeborn.common.rpc.RpcEnv;
  import org.apache.celeborn.common.unsafe.Platform;
  import org.apache.celeborn.common.util.*;
  import org.apache.celeborn.common.write.DataBatches;
  import org.apache.celeborn.common.write.LocationPushFailedBatches;
  import org.apache.celeborn.common.write.PushState;
  
  public class ShuffleClientImpl extends ShuffleClient {
    private static final Logger logger = LoggerFactory.getLogger(ShuffleClientImpl.class);
  
    protected static final byte PRIMARY_MODE = PartitionLocation.Mode.PRIMARY.mode();
  
    private static final Random RND = new Random();
  
    protected final CelebornConf conf;
  
    private final UserIdentifier userIdentifier;
  
    private final int registerShuffleMaxRetries;
    private final long registerShuffleRetryWaitMs;
    private final int rpcMaxRetries;
    private final long rpcRetryWait;
    private final int maxReviveTimes;
    private final boolean testRetryRevive;
    private final int pushBufferMaxSize;
    protected final long pushDataTimeout;
  
    private final RpcEnv rpcEnv;
  
    protected RpcEndpointRef lifecycleManagerRef;
  
    private TransportContext transportContext;
    protected TransportClientFactory dataClientFactory;
  
    protected final int BATCH_HEADER_SIZE = 4 * 4;
  
    protected byte[] extension;
  
    // key: appShuffleIdentifier, value: shuffleId
    protected Map<String, Tuple2<Integer, Boolean>> shuffleIdCache = JavaUtils.newConcurrentHashMap();
  
    // key: shuffleId, value: (partitionId, PartitionLocation)
    final Map<Integer, ConcurrentHashMap<Integer, PartitionLocation>> reducePartitionMap =
        JavaUtils.newConcurrentHashMap();
  
    // key: shuffleId, value: Set(mapId)
    protected final ConcurrentHashMap<Integer, Set<Integer>> mapperEndMap =
        JavaUtils.newConcurrentHashMap();
  
    // shuffleIds which have finished all map tasks
    protected final Set<Integer> stageEndShuffleSet = ConcurrentHashMap.newKeySet();
  
    // key: shuffleId-mapId-attemptId
    protected final Map<String, PushState> pushStates = JavaUtils.newConcurrentHashMap();
  
    private final boolean pushExcludeWorkerOnFailureEnabled;
    private final boolean shuffleCompressionEnabled;
    private final boolean shuffleIntegrityCheckEnabled;
  
    private final Set<String> pushExcludedWorkers = ConcurrentHashMap.newKeySet();
    private final ConcurrentHashMap<String, Long> fetchExcludedWorkers =
        JavaUtils.newConcurrentHashMap();
    private boolean pushReplicateEnabled;
    private boolean fetchExcludeWorkerOnFailureEnabled;
  
    private final ExecutorService pushDataRetryPool;
  
    private final Map<Integer, Set<Integer>> splitting = JavaUtils.newConcurrentHashMap();
  
    protected final String appUniqueId;
    private final boolean authEnabled;
    private final TransportConf dataTransportConf;
  
    @SuppressWarnings("ThreadLocalUsage")
    private final ThreadLocal<Compressor> compressorThreadLocal =
        new ThreadLocal<Compressor>() {
          @Override
          protected Compressor initialValue() {
            return Compressor.getCompressor(conf);
          }
        };
  
    private final ReviveManager reviveManager;
  
    private final boolean dataPushFailureTrackingEnabled;
  
    public static class ReduceFileGroups {
      public Map<Integer, Set<PartitionLocation>> partitionGroups;
      public Map<String, LocationPushFailedBatches> pushFailedBatches;
      public int[] mapAttempts;
      public Set<Integer> partitionIds;
  
      ReduceFileGroups(
          Map<Integer, Set<PartitionLocation>> partitionGroups,
          int[] mapAttempts,
          Set<Integer> partitionIds,
          Map<String, LocationPushFailedBatches> pushFailedBatches) {
        this.partitionGroups = partitionGroups;
        this.mapAttempts = mapAttempts;
        this.partitionIds = partitionIds;
        this.pushFailedBatches = pushFailedBatches;
      }
  
      public ReduceFileGroups() {
        this.partitionGroups = null;
        this.mapAttempts = null;
        this.partitionIds = null;
        this.pushFailedBatches = null;
      }
  
      public void update(ReduceFileGroups fileGroups) {
        partitionGroups = fileGroups.partitionGroups;
        mapAttempts = fileGroups.mapAttempts;
        partitionIds = fileGroups.partitionIds;
        pushFailedBatches = fileGroups.pushFailedBatches;
      }
    }
  
    // key: shuffleId
    protected final Map<Integer, Tuple3<ReduceFileGroups, String, Exception>> reduceFileGroupsMap =
        JavaUtils.newConcurrentHashMap();
  
    private final TransportMessagesHelper messagesHelper = new TransportMessagesHelper();
  
    public ShuffleClientImpl(String appUniqueId, CelebornConf conf, UserIdentifier userIdentifier) {
      super();
      this.appUniqueId = appUniqueId;
      this.conf = conf;
      this.userIdentifier = userIdentifier;
      registerShuffleMaxRetries = conf.clientRegisterShuffleMaxRetry();
      registerShuffleRetryWaitMs = conf.clientRegisterShuffleRetryWaitMs();
      rpcMaxRetries = conf.clientRpcMaxRetries();
      rpcRetryWait = conf.clientRpcRetryWait();
      maxReviveTimes = conf.clientPushMaxReviveTimes();
      testRetryRevive = conf.testRetryRevive();
      pushBufferMaxSize = conf.clientPushBufferMaxSize();
      pushExcludeWorkerOnFailureEnabled = conf.clientPushExcludeWorkerOnFailureEnabled();
      shuffleCompressionEnabled = !conf.shuffleCompressionCodec().equals(CompressionCodec.NONE);
      pushReplicateEnabled = conf.clientPushReplicateEnabled();
      fetchExcludeWorkerOnFailureEnabled = conf.clientFetchExcludeWorkerOnFailureEnabled();
      shuffleIntegrityCheckEnabled = conf.clientShuffleIntegrityCheckEnabled();
      if (conf.clientPushReplicateEnabled()) {
        pushDataTimeout = conf.pushDataTimeoutMs() * 2;
      } else {
        pushDataTimeout = conf.pushDataTimeoutMs();
      }
      authEnabled = conf.authEnabledOnClient();
      dataPushFailureTrackingEnabled = conf.clientAdaptiveOptimizeSkewedPartitionReadEnabled();
  
      // init rpc env
      rpcEnv =
          RpcEnv.create(
              RpcNameConstants.SHUFFLE_CLIENT_SYS,
              TransportModuleConstants.RPC_APP_CLIENT_MODULE,
              Utils.localHostName(conf),
              0,
              conf,
              Role.CLIENT(),
              scala.None$.empty());
  
      String module = TransportModuleConstants.DATA_MODULE;
      dataTransportConf = Utils.fromCelebornConf(conf, module, conf.networkIoThreads(module));
      initDataClientFactoryIfNeeded();
      int pushDataRetryThreads = conf.clientPushRetryThreads();
      pushDataRetryPool =
          ThreadUtils.newDaemonCachedThreadPool("celeborn-retry-sender", pushDataRetryThreads, 60);
  
      reviveManager = new ReviveManager(this, conf);
  
      logger.info("Created ShuffleClientImpl, appUniqueId: {}", appUniqueId);
    }
  
    protected List<TransportClientBootstrap> createBootstraps() {
      if (authEnabled && null != lifecycleManagerRef) {
        PbApplicationMetaRequest pbApplicationMetaRequest =
            PbApplicationMetaRequest.newBuilder().setAppId(appUniqueId).build();
        PbApplicationMeta pbApplicationMeta =
            lifecycleManagerRef.askSync(
                pbApplicationMetaRequest,
                conf.clientRpcRegisterShuffleAskTimeout(),
                ClassTag$.MODULE$.apply(PbApplicationMeta.class));
        List<TransportClientBootstrap> bootstraps = Lists.newArrayList();
        bootstraps.add(
            new SaslClientBootstrap(
                dataTransportConf,
                appUniqueId,
                new SaslCredentials(appUniqueId, pbApplicationMeta.getSecret())));
        return Collections.unmodifiableList(bootstraps);
      } else {
        return Collections.emptyList();
      }
    }
  
    private void initDataClientFactoryIfNeeded() {
      if (dataClientFactory != null) {
        return;
      }
      this.transportContext =
          new TransportContext(
              dataTransportConf, new BaseMessageHandler(), conf.clientCloseIdleConnections());
      if (!authEnabled) {
        logger.info("Initializing data client factory for {}.", appUniqueId);
        dataClientFactory = transportContext.createClientFactory();
      } else if (lifecycleManagerRef != null) {
        logger.info("Initializing data client factory for secured {}.", appUniqueId);
        List<TransportClientBootstrap> bootstraps = createBootstraps();
        dataClientFactory = transportContext.createClientFactory(bootstraps);
      }
    }
  
    private boolean isPushTargetWorkerExcluded(
        PartitionLocation location, RpcResponseCallback wrappedCallback) {
      // If pushExcludeWorkerOnFailureEnabled = false, pushExcludedWorkers should be empty.
      if (pushExcludedWorkers.contains(location.hostAndPushPort())) {
        wrappedCallback.onFailure(
            new CelebornIOException(StatusCode.PUSH_DATA_PRIMARY_WORKER_EXCLUDED));
        return true;
      } else if (location.hasPeer()
          && pushExcludedWorkers.contains(location.getPeer().hostAndPushPort())) {
        wrappedCallback.onFailure(
            new CelebornIOException(StatusCode.PUSH_DATA_REPLICA_WORKER_EXCLUDED));
        return true;
      } else {
        return false;
      }
    }
  
    private void submitRetryPushData(
        int shuffleId,
        byte[] body,
        int batchId,
        PushDataRpcResponseCallback pushDataRpcResponseCallback,
        PushState pushState,
        ReviveRequest request,
        int remainReviveTimes,
        long dueTime) {
      int mapId = request.mapId;
      int attemptId = request.attemptId;
      PartitionLocation loc = request.loc;
      StatusCode cause = request.cause;
      int partitionId = loc.getId();
      long reviveWaitTime = dueTime - System.currentTimeMillis();
      final long delta = 50;
      long accumulatedTime = 0;
      while (request.reviveStatus == StatusCode.REVIVE_INITIALIZED.getValue()
          && accumulatedTime <= reviveWaitTime) {
        try {
          Thread.sleep(delta);
          accumulatedTime += delta;
        } catch (InterruptedException e) {
          logger.error("Interrupted while waiting for Revive result!");
          Thread.currentThread().interrupt();
        }
      }
      if (mapperEnded(shuffleId, mapId)) {
        logger.debug(
            "Revive for push data success, but the mapper already ended for shuffle {} map {} attempt {} partition {} batch {} location {}.",
            shuffleId,
            mapId,
            attemptId,
            partitionId,
            batchId,
            loc);
        pushState.removeBatch(batchId, loc.hostAndPushPort());
      } else if (request.reviveStatus != StatusCode.SUCCESS.getValue()) {
        pushDataRpcResponseCallback.onFailure(
            new CelebornIOException(
                cause
                    + " then revive but "
                    + StatusCode.REVIVE_FAILED
                    + ", revive status "
                    + request.reviveStatus
                    + "("
                    + StatusCode.fromValue(request.reviveStatus)
                    + ")"
                    + ", old location: "
                    + request.loc));
      } else {
        PartitionLocation newLoc = reducePartitionMap.get(shuffleId).get(partitionId);
        logger.info(
            "Revive for push data success, new location for shuffle {} map {} attempt {} partition {} batch {} is location {}.",
            shuffleId,
            mapId,
            attemptId,
            partitionId,
            batchId,
            newLoc);
        pushDataRpcResponseCallback.updateLatestPartition(newLoc);
        try {
          if (!isPushTargetWorkerExcluded(newLoc, pushDataRpcResponseCallback)) {
            if (!testRetryRevive || remainReviveTimes < 1) {
              assert dataClientFactory != null;
              TransportClient client =
                  dataClientFactory.createClient(newLoc.getHost(), newLoc.getPushPort(), partitionId);
              NettyManagedBuffer newBuffer = new NettyManagedBuffer(Unpooled.wrappedBuffer(body));
              String shuffleKey = Utils.makeShuffleKey(appUniqueId, shuffleId);
              PushData newPushData =
                  new PushData(PRIMARY_MODE, shuffleKey, newLoc.getUniqueId(), newBuffer);
              client.pushData(newPushData, pushDataTimeout, pushDataRpcResponseCallback);
            } else {
              throw new RuntimeException(
                  "Mock push data submit retry failed. remainReviveTimes = "
                      + remainReviveTimes
                      + ".");
            }
          }
        } catch (Exception e) {
          logger.error(
              "Exception raised while pushing data for shuffle {} map {} attempt {} partition {} batch {} location {}.",
              shuffleId,
              mapId,
              attemptId,
              partitionId,
              batchId,
              newLoc,
              e);
          if (e instanceof InterruptedException) {
            pushDataRpcResponseCallback.onFailure(e);
          } else {
            pushDataRpcResponseCallback.onFailure(
                new CelebornIOException(StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_PRIMARY, e));
          }
        }
      }
    }
  
    public ReviveRequest[] addAndGetReviveRequests(
        int shuffleId,
        int mapId,
        int attemptId,
        ArrayList<DataBatches.DataBatch> batches,
        StatusCode cause) {
      ReviveRequest[] reviveRequests = new ReviveRequest[batches.size()];
      for (int i = 0; i < batches.size(); i++) {
        DataBatches.DataBatch batch = batches.get(i);
        PartitionLocation loc = batch.loc;
        ReviveRequest reviveRequest =
            new ReviveRequest(shuffleId, mapId, attemptId, loc.getId(), loc.getEpoch(), loc, cause);
        reviveManager.addRequest(reviveRequest);
        reviveRequests[i] = reviveRequest;
      }
      return reviveRequests;
    }
  
    private void submitRetryPushMergedData(
        PushState pushState,
        int shuffleId,
        int mapId,
        int attemptId,
        ArrayList<DataBatches.DataBatch> batches,
        StatusCode cause,
        Integer oldGroupedBatchId,
        ReviveRequest[] reviveRequests,
        int remainReviveTimes,
        long reviveResponseDueTime) {
      HashMap<Pair<String, String>, DataBatches> newDataBatchesMap = new HashMap<>();
      ArrayList<DataBatches.DataBatch> reviveFailedBatchesMap = new ArrayList<>();
  
      long reviveWaitTime = reviveResponseDueTime - System.currentTimeMillis();
      final long delta = 50;
      long accumulatedTime = 0;
      int index = 0;
      while (index < reviveRequests.length && accumulatedTime <= reviveWaitTime) {
        ReviveRequest request = reviveRequests[index];
        DataBatches.DataBatch batch = batches.get(index);
        if (request.reviveStatus != StatusCode.REVIVE_INITIALIZED.getValue()) {
          if (mapperEnded(shuffleId, mapId)) {
            logger.debug(
                "Revive for push merged data success, but the mapper already ended for shuffle {} map {} attempt {} partition {} batch {}.",
                shuffleId,
                mapId,
                attemptId,
                request.partitionId,
                oldGroupedBatchId);
          } else if (request.reviveStatus == StatusCode.SUCCESS.getValue()) {
            PartitionLocation newLoc = reducePartitionMap.get(shuffleId).get(request.partitionId);
            DataBatches newDataBatches =
                newDataBatchesMap.computeIfAbsent(genAddressPair(newLoc), (s) -> new DataBatches());
            newDataBatches.addDataBatch(newLoc, batch.batchId, batch.body);
          } else {
            if (remainReviveTimes > 0) {
              reviveFailedBatchesMap.add(batch);
            } else {
              String errorMsg =
                  String.format(
                      "Revive failed while pushing merged for shuffle %d map %d attempt %d partition %d batch %d location %s.",
                      shuffleId, mapId, attemptId, request.partitionId, oldGroupedBatchId, batch.loc);
              pushState.exception.compareAndSet(
                  null,
                  new CelebornIOException(
                      errorMsg,
                      new CelebornIOException(
                          cause
                              + " then revive but "
                              + request.reviveStatus
                              + "("
                              + StatusCode.fromValue(request.reviveStatus)
                              + ")")));
              return;
            }
          }
          index++;
        } else {
          try {
            Thread.sleep(delta);
          } catch (InterruptedException e) {
            logger.error("Interrupted while waiting for Revive result!");
            Thread.currentThread().interrupt();
          }
          accumulatedTime += delta;
        }
      }
  
      for (int i = index; i < reviveRequests.length; i++) {
        ReviveRequest request = reviveRequests[index];
        DataBatches.DataBatch batch = batches.get(i);
        if (remainReviveTimes > 0) {
          reviveFailedBatchesMap.add(batch);
        } else {
          String errorMsg =
              String.format(
                  "Revive failed while pushing merged for shuffle %d map %d attempt %d partition %d batch %d location %s.",
                  shuffleId, mapId, attemptId, request.partitionId, oldGroupedBatchId, batch.loc);
          pushState.exception.compareAndSet(
              null,
              new CelebornIOException(
                  errorMsg,
                  new CelebornIOException(
                      cause
                          + " then revive but "
                          + request.reviveStatus
                          + "("
                          + StatusCode.fromValue(request.reviveStatus)
                          + ")")));
          return;
        }
      }
  
      for (Map.Entry<Pair<String, String>, DataBatches> entry : newDataBatchesMap.entrySet()) {
        Pair<String, String> addressPair = entry.getKey();
        DataBatches newDataBatches = entry.getValue();
        doPushMergedData(
            addressPair,
            shuffleId,
            mapId,
            attemptId,
            newDataBatches.requireBatches(),
            pushState,
            remainReviveTimes);
      }
      if (reviveFailedBatchesMap.isEmpty()) {
        pushState.removeBatch(oldGroupedBatchId, batches.get(0).loc.hostAndPushPort());
      } else {
        ReviveRequest[] requests =
            addAndGetReviveRequests(shuffleId, mapId, attemptId, reviveFailedBatchesMap, cause);
        pushDataRetryPool.submit(
            () ->
                submitRetryPushMergedData(
                    pushState,
                    shuffleId,
                    mapId,
                    attemptId,
                    reviveFailedBatchesMap,
                    cause,
                    oldGroupedBatchId,
                    requests,
                    remainReviveTimes - 1,
                    System.currentTimeMillis()
                        + conf.clientRpcRequestPartitionLocationAskTimeout().duration().toMillis()));
      }
    }
  
    private Pair<String, String> genAddressPair(PartitionLocation loc) {
      if (loc.hasPeer()) {
        return Pair.of(loc.hostAndPushPort(), loc.getPeer().hostAndPushPort());
      } else {
        return Pair.of(loc.hostAndPushPort(), null);
      }
    }
  
    private ConcurrentHashMap<Integer, PartitionLocation> registerShuffle(
        int shuffleId, int numMappers, int numPartitions) throws CelebornIOException {
      return registerShuffleInternal(
          shuffleId,
          numMappers,
          numPartitions,
          () ->
              lifecycleManagerRef.askSync(
                  RegisterShuffle$.MODULE$.apply(shuffleId, numMappers, numPartitions),
                  conf.clientRpcRegisterShuffleAskTimeout(),
                  rpcMaxRetries,
                  rpcRetryWait,
                  ClassTag$.MODULE$.apply(PbRegisterShuffleResponse.class)));
    }
  
    @Override
    public PartitionLocation registerMapPartitionTask(
        int shuffleId, int numMappers, int mapId, int attemptId, int partitionId) throws IOException {
      return registerMapPartitionTask(shuffleId, numMappers, mapId, attemptId, partitionId, false);
    }
  
    public PartitionLocation registerMapPartitionTask(
        int shuffleId,
        int numMappers,
        int mapId,
        int attemptId,
        int partitionId,
        boolean isSegmentGranularityVisible)
        throws IOException {
      logger.info(
          "Register MapPartition task for shuffle {} map {} attempt {} partition {} with {} mapper.",
          shuffleId,
          mapId,
          attemptId,
          partitionId,
          numMappers);
      ConcurrentHashMap<Integer, PartitionLocation> partitionLocationMap =
          registerShuffleInternal(
              shuffleId,
              numMappers,
              numMappers,
              () ->
                  lifecycleManagerRef.askSync(
                      RegisterMapPartitionTask$.MODULE$.apply(
                          shuffleId,
                          numMappers,
                          mapId,
                          attemptId,
                          partitionId,
                          isSegmentGranularityVisible),
                      conf.clientRpcRegisterShuffleAskTimeout(),
                      ClassTag$.MODULE$.apply(PbRegisterShuffleResponse.class)));
  
      return partitionLocationMap.get(partitionId);
    }
  
    @Override
    public ConcurrentHashMap<Integer, PartitionLocation> getPartitionLocation(
        int shuffleId, int numMappers, int numPartitions) throws CelebornIOException {
      try {
        return reducePartitionMap.computeIfAbsent(
            shuffleId,
            (id) -> {
              try {
                return registerShuffle(shuffleId, numMappers, numPartitions);
              } catch (CelebornIOException e) {
                throw new RuntimeException(e);
              }
            });
      } catch (RuntimeException e) {
        if (e.getCause() instanceof CelebornIOException) {
          throw (CelebornIOException) e.getCause();
        } else {
          throw e;
        }
      }
    }
  
    @Override
    public PushState getPushState(String mapKey) {
      return pushStates.computeIfAbsent(mapKey, (s) -> new PushState(conf));
    }
  
    @Override
    public Tuple2<Integer, Boolean> getShuffleId(
        int appShuffleId, String appShuffleIdentifier, boolean isWriter, boolean isBarrierStage) {
      return shuffleIdCache.computeIfAbsent(
          appShuffleIdentifier,
          (id) -> {
            PbGetShuffleId pbGetShuffleId =
                PbGetShuffleId.newBuilder()
                    .setAppShuffleId(appShuffleId)
                    .setAppShuffleIdentifier(appShuffleIdentifier)
                    .setIsShuffleWriter(isWriter)
                    .setIsBarrierStage(isBarrierStage)
                    .build();
            PbGetShuffleIdResponse pbGetShuffleIdResponse =
                lifecycleManagerRef.askSync(
                    pbGetShuffleId,
                    conf.clientRpcRegisterShuffleAskTimeout(),
                    ClassTag$.MODULE$.apply(PbGetShuffleIdResponse.class));
            return Tuple2.apply(
                pbGetShuffleIdResponse.getShuffleId(), pbGetShuffleIdResponse.getSuccess());
          });
    }
  
    @Override
    public void readReducerPartitionEnd(
        int shuffleId,
        int partitionId,
        int startMapIndex,
        int endMapIndex,
        int crc32,
        long bytesWritten)
        throws IOException {
      PbReadReducerPartitionEnd pbReadReducerPartitionEnd =
          PbReadReducerPartitionEnd.newBuilder()
              .setShuffleId(shuffleId)
              .setPartitionId(partitionId)
              .setStartMaxIndex(startMapIndex)
              .setEndMapIndex(endMapIndex)
              .setCrc32(crc32)
              .setBytesWritten(bytesWritten)
              .build();
  
      PbReadReducerPartitionEndResponse pbReducerPartitionEndResponse =
          lifecycleManagerRef.askSync(
              pbReadReducerPartitionEnd,
              conf.clientRpcRegisterShuffleAskTimeout(),
              ClassTag$.MODULE$.apply(PbReadReducerPartitionEndResponse.class));
      if (pbReducerPartitionEndResponse.getStatus() != StatusCode.SUCCESS.getValue()) {
        throw new CelebornIOException(pbReducerPartitionEndResponse.getErrorMsg());
      }
    }
  
    @Override
    public boolean reportShuffleFetchFailure(int appShuffleId, int shuffleId, long taskId) {
      PbReportShuffleFetchFailure pbReportShuffleFetchFailure =
          PbReportShuffleFetchFailure.newBuilder()
              .setAppShuffleId(appShuffleId)
              .setShuffleId(shuffleId)
              .setTaskId(taskId)
              .build();
      PbReportShuffleFetchFailureResponse pbReportShuffleFetchFailureResponse =
          lifecycleManagerRef.askSync(
              pbReportShuffleFetchFailure,
              conf.clientRpcRegisterShuffleAskTimeout(),
              ClassTag$.MODULE$.apply(PbReportShuffleFetchFailureResponse.class));
      return pbReportShuffleFetchFailureResponse.getSuccess();
    }
  
    @Override
    public boolean reportBarrierTaskFailure(int appShuffleId, String appShuffleIdentifier) {
      PbReportBarrierStageAttemptFailure pbReportBarrierStageAttemptFailure =
          PbReportBarrierStageAttemptFailure.newBuilder()
              .setAppShuffleId(appShuffleId)
              .setAppShuffleIdentifier(appShuffleIdentifier)
              .build();
      PbReportBarrierStageAttemptFailureResponse pbReportBarrierStageAttemptFailureResponse =
          lifecycleManagerRef.askSync(
              pbReportBarrierStageAttemptFailure,
              conf.clientRpcRegisterShuffleAskTimeout(),
              ClassTag$.MODULE$.apply(PbReportBarrierStageAttemptFailureResponse.class));
      return pbReportBarrierStageAttemptFailureResponse.getSuccess();
    }
  
    private ConcurrentHashMap<Integer, PartitionLocation> registerShuffleInternal(
        int shuffleId,
        int numMappers,
        int numPartitions,
        Callable<PbRegisterShuffleResponse> callable)
        throws CelebornIOException {
      int numRetries = registerShuffleMaxRetries;
      StatusCode lastFailedStatusCode = null;
      while (numRetries > 0) {
        try {
          PbRegisterShuffleResponse response = callable.call();
          StatusCode respStatus = StatusCode.fromValue(response.getStatus());
          if (StatusCode.SUCCESS.equals(respStatus)) {
            ConcurrentHashMap<Integer, PartitionLocation> result = JavaUtils.newConcurrentHashMap();
            Tuple2<List<PartitionLocation>, List<PartitionLocation>> locations =
                PbSerDeUtils.fromPbPackedPartitionLocationsPair(
                    response.getPackedPartitionLocationsPair());
            for (PartitionLocation location : locations._1) {
              pushExcludedWorkers.remove(location.hostAndPushPort());
              if (location.hasPeer()) {
                pushExcludedWorkers.remove(location.getPeer().hostAndPushPort());
              }
              result.put(location.getId(), location);
            }
            return result;
          } else if (StatusCode.SLOT_NOT_AVAILABLE.equals(respStatus)) {
            lastFailedStatusCode = respStatus;
            logger.error(
                "LifecycleManager request slots return {}, retry again, remain retry times {}.",
                StatusCode.SLOT_NOT_AVAILABLE,
                numRetries - 1);
          } else if (StatusCode.RESERVE_SLOTS_FAILED.equals(respStatus)) {
            lastFailedStatusCode = respStatus;
            logger.error(
                "LifecycleManager request slots return {}, retry again, remain retry times {}.",
                StatusCode.RESERVE_SLOTS_FAILED,
                numRetries - 1);
          } else {
            lastFailedStatusCode = respStatus;
            logger.error(
                "LifecycleManager request slots return {}, retry again, remain retry times {}.",
                StatusCode.REQUEST_FAILED,
                numRetries - 1);
          }
        } catch (Exception e) {
          if (e instanceof InterruptedException) {
            Thread.currentThread().interrupt();
          }
          logger.error(
              "Exception raised while registering shuffle {} with {} mapper and {} partitions.",
              shuffleId,
              numMappers,
              numPartitions,
              e);
          throw new CelebornIOException("Register shuffle failed for shuffle " + shuffleId + ".", e);
        }
  
        try {
          TimeUnit.MILLISECONDS.sleep(registerShuffleRetryWaitMs);
        } catch (InterruptedException e) {
          break;
        }
        numRetries--;
      }
      throw new CelebornIOException(
          "Register shuffle failed for shuffle " + shuffleId + ", reason: " + lastFailedStatusCode);
    }
  
    protected void limitMaxInFlight(String mapKey, PushState pushState, String hostAndPushPort)
        throws IOException {
      boolean reachLimit = pushState.limitMaxInFlight(hostAndPushPort);
  
      if (reachLimit) {
        throw new CelebornIOException(
            String.format(
                "Waiting timeout for task %s while limiting max in-flight requests to %s",
                mapKey, hostAndPushPort),
            pushState.exception.get());
      }
    }
  
    protected void limitZeroInFlight(String mapKey, PushState pushState) throws IOException {
      boolean reachLimit = pushState.limitZeroInFlight();
  
      if (reachLimit) {
        throw new CelebornIOException(
            String.format(
                "Waiting timeout for task %s while limiting zero in-flight requests", mapKey),
            pushState.exception.get());
      }
    }
  
    /**
     * Check if a newer PartitionLocation(with larger epoch) exists in local cache.
     *
     * @param shuffleMap The mapping between shuffle id and partition location.
     * @param partitionId The id of partition.
     * @param epoch The epoch of revive.
     * @param wait Whether to wait for some time for a newer partition location.
     * @return whether newer partition location exists in local cache.
     */
    boolean newerPartitionLocationExists(
        Map<Integer, PartitionLocation> shuffleMap, int partitionId, int epoch, boolean wait) {
      PartitionLocation currentLocation = shuffleMap.get(partitionId);
      if (currentLocation != null && currentLocation.getEpoch() > epoch) {
        return true;
      } else if (wait) {
        long sleepTimeMs = RND.nextInt(50);
        if (sleepTimeMs > 30) {
          try {
            TimeUnit.MILLISECONDS.sleep(sleepTimeMs);
          } catch (InterruptedException e) {
            logger.error("Waiting revived location was interrupted.", e);
            Thread.currentThread().interrupt();
          }
        }
  
        currentLocation = shuffleMap.get(partitionId);
        return currentLocation != null && currentLocation.getEpoch() > epoch;
      } else {
        return false;
      }
    }
  
    void excludeWorkerByCause(StatusCode cause, PartitionLocation oldLocation) {
      if (pushExcludeWorkerOnFailureEnabled && oldLocation != null) {
        switch (cause) {
          case PUSH_DATA_CREATE_CONNECTION_FAIL_PRIMARY:
          case PUSH_DATA_CONNECTION_EXCEPTION_PRIMARY:
          case PUSH_DATA_TIMEOUT_PRIMARY:
            pushExcludedWorkers.add(oldLocation.hostAndPushPort());
            break;
          case PUSH_DATA_CREATE_CONNECTION_FAIL_REPLICA:
          case PUSH_DATA_CONNECTION_EXCEPTION_REPLICA:
          case PUSH_DATA_TIMEOUT_REPLICA:
            pushExcludedWorkers.add(oldLocation.getPeer().hostAndPushPort());
            break;
          default: // fall out
        }
      }
    }
  
    private boolean revive(
        int shuffleId,
        int mapId,
        int attemptId,
        int partitionId,
        int epoch,
        PartitionLocation oldLocation,
        StatusCode cause) {
      excludeWorkerByCause(cause, oldLocation);
  
      Set<Integer> mapIds = new HashSet<>();
      mapIds.add(mapId);
      List<ReviveRequest> requests = new ArrayList<>();
      ReviveRequest req =
          new ReviveRequest(shuffleId, mapId, attemptId, partitionId, epoch, oldLocation, cause);
      requests.add(req);
      Map<Integer, Integer> results = reviveBatch(shuffleId, mapIds, requests);
  
      if (mapperEnded(shuffleId, mapId)) {
        logger.debug(
            "Revive success, but the mapper ended for shuffle {} map {} attempt {} partition {}, just return true(Assume revive successfully).",
            shuffleId,
            mapId,
            attemptId,
            partitionId);
        return true;
      } else {
        return results != null
            && results.containsKey(partitionId)
            && results.get(partitionId) == StatusCode.SUCCESS.getValue();
      }
    }
  
    /** @return partitionId -> StatusCode#getValue */
    Map<Integer, Integer> reviveBatch(
        int shuffleId, Set<Integer> mapIds, Collection<ReviveRequest> requests) {
      // partitionId -> StatusCode#getValue
      Map<Integer, Integer> results = new HashMap<>();
  
      // Local cached map of (partitionId -> PartitionLocation)
      ConcurrentHashMap<Integer, PartitionLocation> partitionLocationMap =
          reducePartitionMap.get(shuffleId);
  
      Map<Integer, PartitionLocation> oldLocMap = new HashMap<>();
      Iterator<ReviveRequest> iter = requests.iterator();
      while (iter.hasNext()) {
        ReviveRequest req = iter.next();
        oldLocMap.put(req.partitionId, req.loc);
      }
      try {
        PbChangeLocationResponse response =
            lifecycleManagerRef.askSync(
                Revive$.MODULE$.apply(shuffleId, mapIds, requests),
                conf.clientRpcRequestPartitionLocationAskTimeout(),
                ClassTag$.MODULE$.apply(PbChangeLocationResponse.class));
  
        for (int i = 0; i < response.getEndedMapIdCount(); i++) {
          int mapId = response.getEndedMapId(i);
          mapperEndMap.computeIfAbsent(shuffleId, (id) -> ConcurrentHashMap.newKeySet()).add(mapId);
        }
  
        for (int i = 0; i < response.getPartitionInfoCount(); i++) {
          PbChangeLocationPartitionInfo partitionInfo = response.getPartitionInfo(i);
          int partitionId = partitionInfo.getPartitionId();
          int statusCode = partitionInfo.getStatus();
          if (partitionInfo.getOldAvailable()) {
            PartitionLocation oldLoc = oldLocMap.get(partitionId);
            // Currently, revive only check if main location available, here won't remove peer loc.
            pushExcludedWorkers.remove(oldLoc.hostAndPushPort());
          }
  
          if (StatusCode.SUCCESS.getValue() == statusCode) {
            PartitionLocation loc =
                PbSerDeUtils.fromPbPartitionLocation(partitionInfo.getPartition());
            partitionLocationMap.put(partitionId, loc);
            pushExcludedWorkers.remove(loc.hostAndPushPort());
            if (loc.hasPeer()) {
              pushExcludedWorkers.remove(loc.getPeer().hostAndPushPort());
            }
          } else if (StatusCode.STAGE_ENDED.getValue() == statusCode) {
            stageEndShuffleSet.add(shuffleId);
            return results;
          } else if (StatusCode.SHUFFLE_UNREGISTERED.getValue() == statusCode) {
            logger.error("SHUFFLE_NOT_REGISTERED!");
            return null;
          }
          results.put(partitionId, statusCode);
        }
  
        return results;
      } catch (Exception e) {
        StringBuilder partitionIds = new StringBuilder();
        StringBuilder epochs = new StringBuilder();
        requests.forEach(
            (req) -> {
              partitionIds.append(req.partitionId).append(",");
              epochs.append(req.epoch).append(",");
            });
        logger.error(
            "Exception raised while reviving for shuffle {} partitionIds {} epochs {}.",
            shuffleId,
            partitionIds,
            epochs,
            e);
        if (e instanceof InterruptedException) {
          Thread.currentThread().interrupt();
        }
        return null;
      }
    }
  
    private interface PushDataRpcResponseCallback extends RpcResponseCallback {
      default void updateLatestPartition(PartitionLocation latest) {}
    }
  
    public int pushOrMergeData(
        int shuffleId,
        int mapId,
        int attemptId,
        int partitionId,
        byte[] data,
        int offset,
        int length,
        int numMappers,
        int numPartitions,
        boolean doPush,
        boolean skipCompress)
        throws IOException {
      // mapKey
      final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
      // return if shuffle stage already ended
      if (mapperEnded(shuffleId, mapId)) {
        logger.debug(
            "Push or merge data ignored because mapper already ended for shuffle {} map {} attempt {} partition {}.",
            shuffleId,
            mapId,
            attemptId,
            partitionId);
        PushState pushState = pushStates.get(mapKey);
        if (pushState != null) {
          pushState.cleanup();
        }
        return 0;
      }
      // register shuffle if not registered
      final ConcurrentHashMap<Integer, PartitionLocation> map =
          getPartitionLocation(shuffleId, numMappers, numPartitions);
  
      // get location
      // If rerun or speculation task running after LifecycleManager call stageEnd,
      // register shuffle will return an empty location map, client need revive for a new location.
      if (!map.containsKey(partitionId)) {
        if (!revive(
            shuffleId,
            mapId,
            attemptId,
            partitionId,
            -1,
            null,
            StatusCode.PUSH_DATA_FAIL_NON_CRITICAL_CAUSE_PRIMARY)) {
          throw new CelebornIOException(
              String.format("Revive for shuffle %s partition %d failed.", shuffleId, partitionId));
        }
      }
  
      if (mapperEnded(shuffleId, mapId)) {
        logger.debug(
            "Push or merge data ignored because mapper already ended for shuffle {} map {} attempt {} partition {}.",
            shuffleId,
            mapId,
            attemptId,
            partitionId);
        PushState pushState = pushStates.get(mapKey);
        if (pushState != null) {
          pushState.cleanup();
        }
        return 0;
      }
  
      final PartitionLocation loc = map.get(partitionId);
      if (loc == null) {
        throw new CelebornIOException(
            String.format(
                "Partition location for shuffle %s partition %d is NULL!", shuffleId, partitionId));
      }
  
      PushState pushState = getPushState(mapKey);
  
      // increment batchId
      final int nextBatchId = pushState.nextBatchId();
  
      // Track commit metadata if shuffle compression and integrity check are enabled and this request
      // is not for pushing metadata itself.
      if (shuffleIntegrityCheckEnabled) {
        pushState.addDataWithOffsetAndLength(partitionId, data, offset, length);
      }
  
      if (shuffleCompressionEnabled && !skipCompress) {
        // compress data
        final Compressor compressor = compressorThreadLocal.get();
        compressor.compress(data, offset, length);
  
        data = compressor.getCompressedBuffer();
        offset = 0;
        length = compressor.getCompressedTotalSize();
      }
  
      final byte[] body = new byte[BATCH_HEADER_SIZE + length];
      Platform.putInt(body, Platform.BYTE_ARRAY_OFFSET, mapId);
      Platform.putInt(body, Platform.BYTE_ARRAY_OFFSET + 4, attemptId);
      Platform.putInt(body, Platform.BYTE_ARRAY_OFFSET + 8, nextBatchId);
      Platform.putInt(body, Platform.BYTE_ARRAY_OFFSET + 12, length);
      System.arraycopy(data, offset, body, BATCH_HEADER_SIZE, length);
  
      if (doPush) {
        // check limit
        limitMaxInFlight(mapKey, pushState, loc.hostAndPushPort());
  
        // add inFlight requests
        pushState.addBatch(nextBatchId, body.length, loc.hostAndPushPort());
  
        // build PushData request
        NettyManagedBuffer buffer = new NettyManagedBuffer(Unpooled.wrappedBuffer(body));
        final String shuffleKey = Utils.makeShuffleKey(appUniqueId, shuffleId);
        PushData pushData = new PushData(PRIMARY_MODE, shuffleKey, loc.getUniqueId(), buffer);
  
        // build callback
        RpcResponseCallback callback =
            new RpcResponseCallback() {
              @Override
              public void onSuccess(ByteBuffer response) {
                if (response.remaining() > 0 && response.get() == StatusCode.MAP_ENDED.getValue()) {
                  mapperEndMap
                      .computeIfAbsent(shuffleId, (id) -> ConcurrentHashMap.newKeySet())
                      .add(mapId);
                }
                logger.debug(
                    "Push data to {} success for shuffle {} map {} attempt {} partition {} batch {}.",
                    loc.hostAndPushPort(),
                    shuffleId,
                    mapId,
                    attemptId,
                    partitionId,
                    nextBatchId);
              }
  
              @Override
              public void onFailure(Throwable e) {
                String errorMsg =
                    String.format(
                        "Push data to %s failed for shuffle %d map %d attempt %d partition %d batch %d.",
                        loc, shuffleId, mapId, attemptId, partitionId, nextBatchId);
                pushState.exception.compareAndSet(null, new CelebornIOException(errorMsg, e));
              }
            };
  
        RpcResponseCallback wrappedCallback =
            new PushDataRpcResponseCallback() {
              int remainReviveTimes = maxReviveTimes;
              PartitionLocation latest = loc;
  
              @Override
              public void updateLatestPartition(PartitionLocation newloc) {
                pushState.addBatch(nextBatchId, body.length, newloc.hostAndPushPort());
                pushState.removeBatch(nextBatchId, this.latest.hostAndPushPort());
                this.latest = newloc;
              }
  
              @Override
              public void onSuccess(ByteBuffer response) {
                if (response.remaining() > 0) {
                  byte reason = response.get();
                  if (reason == StatusCode.SOFT_SPLIT.getValue()) {
                    logger.debug(
                        "Push data to {} soft split required for shuffle {} map {} attempt {} partition {} batch {}.",
                        latest.hostAndPushPort(),
                        shuffleId,
                        mapId,
                        attemptId,
                        partitionId,
                        nextBatchId);
                    if (!newerPartitionLocationExists(
                        reducePartitionMap.get(shuffleId), partitionId, latest.getEpoch(), false)) {
                      ReviveRequest reviveRequest =
                          new ReviveRequest(
                              shuffleId,
                              mapId,
                              attemptId,
                              partitionId,
                              latest.getEpoch(),
                              latest,
                              StatusCode.SOFT_SPLIT);
                      reviveManager.addRequest(reviveRequest);
                    }
                    pushState.onSuccess(latest.hostAndPushPort());
                    pushState.removeBatch(nextBatchId, latest.hostAndPushPort());
                    callback.onSuccess(response);
                  } else if (reason == StatusCode.HARD_SPLIT.getValue()) {
                    logger.debug(
                        "Push data to {} hard split required for shuffle {} map {} attempt {} partition {} batch {}.",
                        latest.hostAndPushPort(),
                        shuffleId,
                        mapId,
                        attemptId,
                        partitionId,
                        nextBatchId);
                    if (dataPushFailureTrackingEnabled && pushReplicateEnabled) {
                      pushState.recordFailedBatch(
                          latest.getUniqueId(), mapId, attemptId, nextBatchId);
                    }
                    ReviveRequest reviveRequest =
                        new ReviveRequest(
                            shuffleId,
                            mapId,
                            attemptId,
                            partitionId,
                            latest.getEpoch(),
                            latest,
                            StatusCode.HARD_SPLIT);
                    reviveManager.addRequest(reviveRequest);
                    long dueTime =
                        System.currentTimeMillis()
                            + conf.clientRpcRequestPartitionLocationAskTimeout()
                                .duration()
                                .toMillis();
                    pushDataRetryPool.submit(
                        () ->
                            submitRetryPushData(
                                shuffleId,
                                body,
                                nextBatchId,
                                this,
                                pushState,
                                reviveRequest,
                                remainReviveTimes,
                                dueTime));
                  } else if (reason == StatusCode.PUSH_DATA_SUCCESS_PRIMARY_CONGESTED.getValue()) {
                    logger.debug(
                        "Push data to {} primary congestion required for shuffle {} map {} attempt {} partition {} batch {}.",
                        latest.hostAndPushPort(),
                        shuffleId,
                        mapId,
                        attemptId,
                        partitionId,
                        nextBatchId);
                    pushState.onCongestControl(latest.hostAndPushPort());
                    pushState.removeBatch(nextBatchId, latest.hostAndPushPort());
                    callback.onSuccess(response);
                  } else if (reason == StatusCode.PUSH_DATA_SUCCESS_REPLICA_CONGESTED.getValue()) {
                    logger.debug(
                        "Push data to {} replica congestion required for shuffle {} map {} attempt {} partition {} batch {}.",
                        latest.hostAndPushPort(),
                        shuffleId,
                        mapId,
                        attemptId,
                        partitionId,
                        nextBatchId);
                    pushState.onCongestControl(latest.hostAndPushPort());
                    pushState.removeBatch(nextBatchId, latest.hostAndPushPort());
                    callback.onSuccess(response);
                  } else {
                    // StageEnd.
                    response.rewind();
                    pushState.onSuccess(latest.hostAndPushPort());
                    pushState.removeBatch(nextBatchId, latest.hostAndPushPort());
                    callback.onSuccess(response);
                  }
                } else {
                  pushState.onSuccess(latest.hostAndPushPort());
                  pushState.removeBatch(nextBatchId, latest.hostAndPushPort());
                  callback.onSuccess(response);
                }
              }
  
              @Override
              public void onFailure(Throwable e) {
                if (dataPushFailureTrackingEnabled) {
                  pushState.recordFailedBatch(latest.getUniqueId(), mapId, attemptId, nextBatchId);
                }
                if (pushState.exception.get() != null) {
                  return;
                }
                if (e instanceof InterruptedException) {
                  Thread.currentThread().interrupt();
                  callback.onFailure(e);
                  return;
                }
                StatusCode cause = getPushDataFailCause(e.getMessage());
                if (remainReviveTimes <= 0) {
                  if (e instanceof CelebornIOException) {
                    callback.onFailure(e);
                  } else {
                    callback.onFailure(new CelebornIOException(cause, e));
                  }
                  return;
                }
  
                logger.error(
                    "Push data to {} failed for shuffle {} map {} attempt {} partition {} batch {}, remain revive times {}.",
                    latest.hostAndPushPort(),
                    shuffleId,
                    mapId,
                    attemptId,
                    partitionId,
                    nextBatchId,
                    remainReviveTimes,
                    e);
                // async retry push data
                if (!mapperEnded(shuffleId, mapId)) {
                  remainReviveTimes = remainReviveTimes - 1;
                  ReviveRequest reviveRequest =
                      new ReviveRequest(
                          shuffleId, mapId, attemptId, partitionId, latest.getEpoch(), latest, cause);
                  reviveManager.addRequest(reviveRequest);
                  long dueTime =
                      System.currentTimeMillis()
                          + conf.clientRpcRequestPartitionLocationAskTimeout().duration().toMillis();
                  pushDataRetryPool.submit(
                      () ->
                          submitRetryPushData(
                              shuffleId,
                              body,
                              nextBatchId,
                              this,
                              pushState,
                              reviveRequest,
                              remainReviveTimes,
                              dueTime));
                } else {
                  pushState.removeBatch(nextBatchId, latest.hostAndPushPort());
                  logger.info(
                      "Push data to {} failed but mapper already ended for shuffle {} map {} attempt {} partition {} batch {}, remain revive times {}.",
                      latest.hostAndPushPort(),
                      shuffleId,
                      mapId,
                      attemptId,
                      partitionId,
                      nextBatchId,
                      remainReviveTimes);
                }
              }
            };
  
        // do push data
        try {
          if (!isPushTargetWorkerExcluded(loc, wrappedCallback)) {
            if (!testRetryRevive) {
              assert dataClientFactory != null;
              TransportClient client =
                  dataClientFactory.createClient(loc.getHost(), loc.getPushPort(), partitionId);
              client.pushData(pushData, pushDataTimeout, wrappedCallback);
            } else {
              wrappedCallback.onFailure(
                  new CelebornIOException(
                      StatusCode.PUSH_DATA_FAIL_NON_CRITICAL_CAUSE_PRIMARY,
                      new RuntimeException("Mock push data first time failed.")));
            }
          }
        } catch (Exception e) {
          logger.error(
              "Exception raised while pushing data for shuffle {} map {} attempt {} partition {} batch {} location {}.",
              shuffleId,
              mapId,
              attemptId,
              partitionId,
              nextBatchId,
              loc,
              e);
          if (e instanceof InterruptedException) {
            wrappedCallback.onFailure(e);
          } else {
            wrappedCallback.onFailure(
                new CelebornIOException(StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_PRIMARY, e));
          }
        }
      } else {
        // add batch data
        logger.debug("Merge batch {}.", nextBatchId);
        Pair<String, String> addressPair = genAddressPair(loc);
        boolean shouldPush = pushState.addBatchData(addressPair, loc, nextBatchId, body);
        if (shouldPush) {
          limitMaxInFlight(mapKey, pushState, loc.hostAndPushPort());
          DataBatches dataBatches = pushState.takeDataBatches(addressPair);
          if (dataBatches != null) {
            doPushMergedData(
                addressPair,
                shuffleId,
                mapId,
                attemptId,
                dataBatches.requireBatches(),
                pushState,
                maxReviveTimes);
          }
        }
      }
  
      return body.length;
    }
  
    @Override
    public int pushData(
        int shuffleId,
        int mapId,
        int attemptId,
        int partitionId,
        byte[] data,
        int offset,
        int length,
        int numMappers,
        int numPartitions)
        throws IOException {
      return pushOrMergeData(
          shuffleId,
          mapId,
          attemptId,
          partitionId,
          data,
          offset,
          length,
          numMappers,
          numPartitions,
          true,
          false);
    }
  
    @Override
    public void prepareForMergeData(int shuffleId, int mapId, int attemptId) throws IOException {
      final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
      PushState pushState = pushStates.get(mapKey);
      if (pushState != null) {
        limitZeroInFlight(mapKey, pushState);
      }
    }
  
    @Override
    public int mergeData(
        int shuffleId,
        int mapId,
        int attemptId,
        int partitionId,
        byte[] data,
        int offset,
        int length,
        int numMappers,
        int numPartitions)
        throws IOException {
      return pushOrMergeData(
          shuffleId,
          mapId,
          attemptId,
          partitionId,
          data,
          offset,
          length,
          numMappers,
          numPartitions,
          false,
          false);
    }
  
    @Override
    public void pushMergedData(int shuffleId, int mapId, int attemptId) throws IOException {
      final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
      PushState pushState = pushStates.get(mapKey);
      if (pushState == null) {
        return;
      }
      ArrayList<Map.Entry<Pair<String, String>, DataBatches>> batchesArr =
          new ArrayList<>(pushState.batchesMap.entrySet());
      while (!batchesArr.isEmpty()) {
        Map.Entry<Pair<String, String>, DataBatches> entry =
            batchesArr.get(RND.nextInt(batchesArr.size()));
        limitMaxInFlight(mapKey, pushState, entry.getKey().getLeft());
        ArrayList<DataBatches.DataBatch> batches = entry.getValue().requireBatches(pushBufferMaxSize);
        if (entry.getValue().getTotalSize() == 0) {
          batchesArr.remove(entry);
        }
        doPushMergedData(
            entry.getKey(), shuffleId, mapId, attemptId, batches, pushState, maxReviveTimes);
      }
    }
  
    private void doPushMergedData(
        Pair<String, String> addressPair,
        int shuffleId,
        int mapId,
        int attemptId,
        ArrayList<DataBatches.DataBatch> batches,
        PushState pushState,
        int remainReviveTimes) {
      String hostPort = addressPair.getLeft();
      String[] hostPortArr = Utils.parseColonSeparatedHostPorts(hostPort, 1);
      final String host = hostPortArr[0];
      final int port = Integer.parseInt(hostPortArr[1]);
  
      int groupedBatchId = pushState.nextBatchId();
      int groupedBatchBytesSize = batches.stream().mapToInt(batch -> batch.body.length).sum();
      pushState.addBatch(groupedBatchId, groupedBatchBytesSize, hostPort);
  
      final int numBatches = batches.size();
      final Integer[] partitionIds = new Integer[numBatches];
      final String[] partitionUniqueIds = new String[numBatches];
      final int[] offsets = new int[numBatches];
      final int[] batchIds = new int[numBatches];
      int currentSize = 0;
      CompositeByteBuf byteBuf = Unpooled.compositeBuffer();
      for (int i = 0; i < numBatches; i++) {
        DataBatches.DataBatch batch = batches.get(i);
        partitionIds[i] = batch.loc.getId();
        partitionUniqueIds[i] = batch.loc.getUniqueId();
        offsets[i] = currentSize;
        batchIds[i] = batch.batchId;
        currentSize += batch.body.length;
        byteBuf.addComponent(true, Unpooled.wrappedBuffer(batch.body));
      }
      NettyManagedBuffer buffer = new NettyManagedBuffer(byteBuf);
      String shuffleKey = Utils.makeShuffleKey(appUniqueId, shuffleId);
      PushMergedData mergedData =
          new PushMergedData(PRIMARY_MODE, shuffleKey, partitionUniqueIds, offsets, buffer);
  
      RpcResponseCallback callback =
          new RpcResponseCallback() {
            @Override
            public void onSuccess(ByteBuffer response) {
              logger.debug(
                  "Push merged data to {} success for shuffle {} map {} attempt {} partition {} groupedBatch {} batch {}.",
                  addressPair,
                  shuffleId,
                  mapId,
                  attemptId,
                  Arrays.toString(partitionIds),
                  groupedBatchId,
                  Arrays.toString(batchIds));
              pushState.removeBatch(groupedBatchId, hostPort);
              if (response.remaining() > 0 && response.get() == StatusCode.MAP_ENDED.getValue()) {
                mapperEndMap
                    .computeIfAbsent(shuffleId, (id) -> ConcurrentHashMap.newKeySet())
                    .add(mapId);
              }
            }
  
            @Override
            public void onFailure(Throwable e) {
              String errorMsg =
                  String.format(
                      "Push merged data to %s failed for shuffle %d map %d attempt %d partition %s groupedBatch %d batch %s, remain revive times %d.",
                      addressPair,
                      shuffleId,
                      mapId,
                      attemptId,
                      Arrays.toString(partitionIds),
                      groupedBatchId,
                      Arrays.toString(batchIds),
                      remainReviveTimes);
              pushState.exception.compareAndSet(null, new CelebornIOException(errorMsg, e));
              if (logger.isDebugEnabled()) {
                for (int i = 0; i < numBatches; i++) {
                  logger.debug(
                      "Push merged data to {} failed for shuffle {} map {} attempt {} partition {} groupedBatch {} batch {}, remain revive times {}.",
                      addressPair,
                      shuffleId,
                      mapId,
                      attemptId,
                      partitionIds[i],
                      groupedBatchId,
                      batchIds[i],
                      remainReviveTimes);
                }
              }
            }
          };
  
      RpcResponseCallback wrappedCallback =
          new RpcResponseCallback() {
            @Override
            public void onSuccess(ByteBuffer response) {
              byte reason = response.get();
              if (reason == StatusCode.HARD_SPLIT.getValue()
                  || reason == StatusCode.SOFT_SPLIT.getValue()) {
                ArrayList<DataBatches.DataBatch> batchesNeedResubmit;
                if (response.remaining() > 0) {
                  batchesNeedResubmit = new ArrayList<>();
                  PbPushMergedDataSplitPartitionInfo partitionInfo;
                  try {
                    partitionInfo = TransportMessage.fromByteBuffer(response).getParsedPayload();
                  } catch (CelebornIOException | InvalidProtocolBufferException e) {
                    callback.onFailure(
                        new CelebornIOException("parse pushMergedData response failed", e));
                    return;
                  }
                  List<Integer> splitPartitionIndexes = partitionInfo.getSplitPartitionIndexesList();
                  List<Integer> statusCodeList = partitionInfo.getStatusCodesList();
                  StringBuilder dataBatchReviveInfos = new StringBuilder();
                  for (int i = 0; i < splitPartitionIndexes.size(); i++) {
                    int partitionIndex = splitPartitionIndexes.get(i);
                    int batchId = batches.get(partitionIndex).batchId;
                    dataBatchReviveInfos.append(
                        String.format(
                            "(batchId=%d, partitionId=%d, cause=%s)",
                            batchId,
                            partitionIds[partitionIndex],
                            StatusCode.fromValue(statusCodeList.get(i).byteValue())));
                    if (statusCodeList.get(i) == StatusCode.SOFT_SPLIT.getValue()) {
                      PartitionLocation loc = batches.get(partitionIndex).loc;
                      if (!newerPartitionLocationExists(
                          reducePartitionMap.get(shuffleId), loc.getId(), loc.getEpoch(), false)) {
                        ReviveRequest reviveRequest =
                            new ReviveRequest(
                                shuffleId,
                                mapId,
                                attemptId,
                                loc.getId(),
                                loc.getEpoch(),
                                loc,
                                StatusCode.SOFT_SPLIT);
                        reviveManager.addRequest(reviveRequest);
                      }
                    } else {
                      batchesNeedResubmit.add(batches.get(partitionIndex));
                    }
                  }
                  logger.info(
                      "Push merged data to {} partial success required for shuffle {} map {} attempt {} groupedBatch {}. split batches {}.",
                      addressPair,
                      shuffleId,
                      mapId,
                      attemptId,
                      groupedBatchId,
                      dataBatchReviveInfos);
                } else {
                  // Workers that do not incorporate changes from [CELEBORN-1721]
                  // will respond with a status of HARD_SPLIT,
                  // but will not include a PbPushMergedDataSplitPartitionInfo.
                  // For backward compatibility, all batches must be resubmitted.
                  batchesNeedResubmit = batches;
                  logger.info(
                      "Push merged data to {} hard split required for shuffle {} map {} attempt {} partition {} groupedBatch {} batch {}.",
                      addressPair,
                      shuffleId,
                      mapId,
                      attemptId,
                      Arrays.toString(partitionIds),
                      groupedBatchId,
                      Arrays.toString(batchIds));
                }
                if (batchesNeedResubmit.isEmpty()) {
                  pushState.onSuccess(hostPort);
                  callback.onSuccess(ByteBuffer.wrap(new byte[] {StatusCode.SOFT_SPLIT.getValue()}));
                } else {
                  if (dataPushFailureTrackingEnabled && pushReplicateEnabled) {
                    for (DataBatches.DataBatch resubmitBatch : batchesNeedResubmit) {
                      pushState.recordFailedBatch(
                          resubmitBatch.loc.getUniqueId(), mapId, attemptId, resubmitBatch.batchId);
                    }
                  }
                  ReviveRequest[] requests =
                      addAndGetReviveRequests(
                          shuffleId, mapId, attemptId, batchesNeedResubmit, StatusCode.HARD_SPLIT);
                  pushDataRetryPool.submit(
                      () ->
                          submitRetryPushMergedData(
                              pushState,
                              shuffleId,
                              mapId,
                              attemptId,
                              batchesNeedResubmit,
                              StatusCode.HARD_SPLIT,
                              groupedBatchId,
                              requests,
                              remainReviveTimes,
                              System.currentTimeMillis()
                                  + conf.clientRpcRequestPartitionLocationAskTimeout()
                                      .duration()
                                      .toMillis()));
                }
              } else if (reason == StatusCode.PUSH_DATA_SUCCESS_PRIMARY_CONGESTED.getValue()) {
                logger.debug(
                    "Push merged data to {} primary congestion required for shuffle {} map {} attempt {} partition {} groupedBatch {} batch {}.",
                    addressPair,
                    shuffleId,
                    mapId,
                    attemptId,
                    Arrays.toString(partitionIds),
                    groupedBatchId,
                    Arrays.toString(batchIds));
                pushState.onCongestControl(hostPort);
                callback.onSuccess(response);
              } else if (reason == StatusCode.PUSH_DATA_SUCCESS_REPLICA_CONGESTED.getValue()) {
                logger.debug(
                    "Push merged data to {} replica congestion required for shuffle {} map {} attempt {} partition {} groupedBatch {} batch {}.",
                    addressPair,
                    shuffleId,
                    mapId,
                    attemptId,
                    Arrays.toString(partitionIds),
                    groupedBatchId,
                    Arrays.toString(batchIds));
                pushState.onCongestControl(hostPort);
                callback.onSuccess(response);
              } else if (reason == StatusCode.MAP_ENDED.getValue()) {
                pushState.onSuccess(hostPort);
                callback.onSuccess(ByteBuffer.wrap(new byte[] {StatusCode.MAP_ENDED.getValue()}));
              } else { // success
                pushState.onSuccess(hostPort);
                callback.onSuccess(ByteBuffer.wrap(new byte[] {StatusCode.SUCCESS.getValue()}));
              }
            }
  
            @Override
            public void onFailure(Throwable e) {
              if (dataPushFailureTrackingEnabled) {
                for (int i = 0; i < numBatches; i++) {
                  pushState.recordFailedBatch(partitionUniqueIds[i], mapId, attemptId, batchIds[i]);
                }
              }
              if (pushState.exception.get() != null) {
                return;
              }
              if (e instanceof InterruptedException) {
                Thread.currentThread().interrupt();
                callback.onFailure(e);
                return;
              }
              StatusCode cause = getPushDataFailCause(e.getMessage());
              if (remainReviveTimes <= 0) {
                if (e instanceof CelebornIOException) {
                  callback.onFailure(e);
                } else {
                  callback.onFailure(new CelebornIOException(cause, e));
                }
                return;
              }
              logger.error(
                  "Push merged data to {} failed for shuffle {} map {} attempt {} partition {} groupedBatch {} batch {}, remain revive times {}.",
                  addressPair,
                  shuffleId,
                  mapId,
                  attemptId,
                  Arrays.toString(partitionIds),
                  groupedBatchId,
                  Arrays.toString(batchIds),
                  remainReviveTimes,
                  e);
              if (!mapperEnded(shuffleId, mapId)) {
                ReviveRequest[] requests =
                    addAndGetReviveRequests(shuffleId, mapId, attemptId, batches, cause);
                pushDataRetryPool.submit(
                    () ->
                        submitRetryPushMergedData(
                            pushState,
                            shuffleId,
                            mapId,
                            attemptId,
                            batches,
                            cause,
                            groupedBatchId,
                            requests,
                            remainReviveTimes - 1,
                            System.currentTimeMillis()
                                + conf.clientRpcRequestPartitionLocationAskTimeout()
                                    .duration()
                                    .toMillis()));
              } else {
                pushState.removeBatch(groupedBatchId, hostPort);
                logger.info(
                    "Push merged data to {} failed but mapper already ended for shuffle {} map {} attempt {} partition {} groupedBatch {} batch {}, remain revive times {}.",
                    hostPort,
                    shuffleId,
                    mapId,
                    attemptId,
                    Arrays.toString(partitionIds),
                    groupedBatchId,
                    Arrays.toString(batchIds),
                    remainReviveTimes);
              }
            }
          };
  
      // do push merged data
      try {
        if (!isPushTargetWorkerExcluded(batches.get(0).loc, wrappedCallback)) {
          if (!testRetryRevive || remainReviveTimes < 1) {
            assert dataClientFactory != null;
            TransportClient client = dataClientFactory.createClient(host, port);
            client.pushMergedData(mergedData, pushDataTimeout, wrappedCallback);
          } else {
            wrappedCallback.onFailure(
                new CelebornIOException(
                    StatusCode.PUSH_DATA_FAIL_NON_CRITICAL_CAUSE_PRIMARY,
                    new RuntimeException("Mock push merge data failed.")));
          }
        }
      } catch (Exception e) {
        logger.error(
            "Exception raised while pushing merged data for shuffle {} map {} attempt {} partition {} groupedBatch {} batch {} location {}.",
            shuffleId,
            mapId,
            attemptId,
            Arrays.toString(partitionIds),
            groupedBatchId,
            Arrays.toString(batchIds),
            addressPair,
            e);
        if (e instanceof InterruptedException) {
          wrappedCallback.onFailure(e);
        } else {
          wrappedCallback.onFailure(
              new CelebornIOException(StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_PRIMARY, e));
        }
      }
    }
  
    @Override
    public void mapperEnd(int shuffleId, int mapId, int attemptId, int numMappers, int numPartitions)
        throws IOException {
      mapEndInternal(shuffleId, mapId, attemptId, numMappers, numPartitions, -1);
    }
  
    @Override
    public void mapPartitionMapperEnd(
        int shuffleId, int mapId, int attemptId, int numMappers, int numPartitions, int partitionId)
        throws IOException {
      mapEndInternal(shuffleId, mapId, attemptId, numMappers, numPartitions, partitionId);
    }
  
    private void mapEndInternal(
        int shuffleId,
        int mapId,
        int attemptId,
        int numMappers,
        int numPartitions,
        Integer partitionId)
        throws IOException {
      final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
      PushState pushState = getPushState(mapKey);
  
      try {
        limitZeroInFlight(mapKey, pushState);
  
        // send CRC32 and num bytes per partition if e2e checks are enabled
        int[] crc32PerPartition =
            pushState.getCRC32PerPartition(shuffleIntegrityCheckEnabled, numPartitions);
        long[] bytesPerPartition =
            pushState.getBytesWrittenPerPartition(shuffleIntegrityCheckEnabled, numPartitions);
  
        MapperEndResponse response =
            lifecycleManagerRef.askSync(
                new MapperEnd(
                    shuffleId,
                    mapId,
                    attemptId,
                    numMappers,
                    partitionId,
                    pushState.getFailedBatches(),
                    numPartitions,
                    crc32PerPartition,
                    bytesPerPartition),
                rpcMaxRetries,
                rpcRetryWait,
                ClassTag$.MODULE$.apply(MapperEndResponse.class));
        if (response.status() != StatusCode.SUCCESS) {
          throw new CelebornIOException("MapperEnd failed! StatusCode: " + response.status());
        }
      } finally {
        pushStates.remove(mapKey);
      }
    }
  
    @Override
    public void cleanup(int shuffleId, int mapId, int attemptId) {
      final String mapKey = Utils.makeMapKey(shuffleId, mapId, attemptId);
      PushState pushState = pushStates.remove(mapKey);
      if (pushState != null) {
        pushState.exception.compareAndSet(null, new CelebornIOException("Cleaned Up"));
        pushState.cleanup();
      }
    }
  
    @Override
    public boolean cleanupShuffle(int shuffleId) {
      // clear status
      reducePartitionMap.remove(shuffleId);
      reduceFileGroupsMap.remove(shuffleId);
      mapperEndMap.remove(shuffleId);
      stageEndShuffleSet.remove(shuffleId);
      splitting.remove(shuffleId);
  
      logger.info("Unregistered shuffle {}.", shuffleId);
      return true;
    }
  
    protected Tuple3<ReduceFileGroups, String, Exception> loadFileGroupInternal(
        int shuffleId, boolean isSegmentGranularityVisible) {
      long getReducerFileGroupStartTime = System.nanoTime();
      String exceptionMsg = null;
      Exception exception = null;
      if (lifecycleManagerRef == null) {
        exceptionMsg = "Driver endpoint is null!";
        logger.warn(exceptionMsg);
        return Tuple3.apply(null, exceptionMsg, exception);
      }
      try {
        GetReducerFileGroup getReducerFileGroup =
            new GetReducerFileGroup(shuffleId, isSegmentGranularityVisible, SerdeVersion.V1);
  
        GetReducerFileGroupResponse response =
            lifecycleManagerRef.askSync(
                getReducerFileGroup,
                conf.clientRpcGetReducerFileGroupAskTimeout(),
                rpcMaxRetries,
                rpcRetryWait,
                ClassTag$.MODULE$.apply(GetReducerFileGroupResponse.class));
        switch (response.status()) {
          case SUCCESS:
            if (response.broadcast() != null && response.broadcast().length > 0) {
              response =
                  ShuffleClient.deserializeReducerFileGroupResponse(shuffleId, response.broadcast());
              if (response == null) {
                throw new CelebornBroadcastException(
                    "Failed to get GetReducerFileGroupResponse broadcast for shuffle: " + shuffleId);
              }
            }
            logger.info(
                "Shuffle {} request reducer file group success using {} ms, result partition size {}.",
                shuffleId,
                TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - getReducerFileGroupStartTime),
                response.fileGroup().size());
            return Tuple3.apply(
                new ReduceFileGroups(
                    response.fileGroup(),
                    response.attempts(),
                    response.partitionIds(),
                    response.pushFailedBatches()),
                null,
                null);
          case SHUFFLE_UNREGISTERED:
            logger.warn(
                "Request {} return {} for {}.", getReducerFileGroup, response.status(), shuffleId);
            // return empty result
            return Tuple3.apply(
                new ReduceFileGroups(
                    response.fileGroup(),
                    response.attempts(),
                    response.partitionIds(),
                    response.pushFailedBatches()),
                null,
                null);
          case STAGE_END_TIMEOUT:
          case SHUFFLE_DATA_LOST:
            exceptionMsg =
                String.format(
                    "Request %s return %s for %s.",
                    getReducerFileGroup, response.status(), shuffleId);
            logger.warn(exceptionMsg);
            break;
          default: // fall out
        }
      } catch (Exception e) {
        if (e instanceof InterruptedException) {
          Thread.currentThread().interrupt();
        }
        logger.error("Exception raised while call GetReducerFileGroup for {}.", shuffleId, e);
        exceptionMsg = e.getMessage();
        exception = e;
      }
      return Tuple3.apply(null, exceptionMsg, exception);
    }
  
    @Override
    public ReduceFileGroups updateFileGroup(int shuffleId, int partitionId)
        throws CelebornIOException {
      return updateFileGroup(shuffleId, partitionId, false);
    }
  
    @Override
    public boolean isShuffleStageEnd(int shuffleId) throws Exception {
      if (null != lifecycleManagerRef) {
        PbGetStageEnd request = PbGetStageEnd.newBuilder().setShuffleId(shuffleId).build();
        PbGetStageEndResponse response =
            lifecycleManagerRef.askSync(
                request,
                rpcMaxRetries,
                rpcRetryWait,
                ClassTag$.MODULE$.apply(PbGetStageEndResponse.class));
        return response.getStageEnd();
      } else {
        throw new RuntimeException("Driver endpoint is null!");
      }
    }
  
    public ReduceFileGroups updateFileGroup(
        int shuffleId, int partitionId, boolean isSegmentGranularityVisible)
        throws CelebornIOException {
      Tuple3<ReduceFileGroups, String, Exception> fileGroupTuple =
          reduceFileGroupsMap.compute(
              shuffleId,
              (id, existsTuple) -> {
                if (existsTuple == null || existsTuple._1() == null) {
                  return loadFileGroupInternal(shuffleId, isSegmentGranularityVisible);
                } else {
                  return existsTuple;
                }
              });
      if (fileGroupTuple._1() == null) {
        throw new CelebornIOException(
            loadFileGroupException(shuffleId, partitionId, (fileGroupTuple._2())),
            fileGroupTuple._3());
      } else {
        return fileGroupTuple._1();
      }
    }
  
    protected String loadFileGroupException(int shuffleId, int partitionId, String exceptionMsg) {
      return String.format(
          "Failed to load file group of shuffle %d partition %d! %s",
          shuffleId,
          partitionId,
          StringUtils.isEmpty(exceptionMsg) ? StringUtils.EMPTY : exceptionMsg);
    }
  
    @Override
    public CelebornInputStream readPartition(
        int shuffleId,
        int appShuffleId,
        int partitionId,
        int attemptNumber,
        long taskId,
        int startMapIndex,
        int endMapIndex,
        ExceptionMaker exceptionMaker,
        ArrayList<PartitionLocation> locations,
        ArrayList<PbStreamHandler> streamHandlers,
        Map<String, LocationPushFailedBatches> failedBatchSetMap,
        Map<String, Pair<Integer, Integer>> chunksRange,
        int[] mapAttempts,
        MetricsCallback metricsCallback,
        boolean needDecompress)
        throws IOException {
      if (shuffleId == Utils$.MODULE$.UNKNOWN_APP_SHUFFLE_ID()) {
        logger.warn("Shuffle data is empty for shuffle {}: UNKNOWN_APP_SHUFFLE_ID.", shuffleId);
        return CelebornInputStream.empty();
      }
  
      // When `mapAttempts` is not null, it's guaranteed that the code path comes from
      // CelebornShuffleReader, which means `updateFileGroup` is already called and
      // batch open stream has been tried
      if (mapAttempts == null) {
        ReduceFileGroups fileGroups = updateFileGroup(shuffleId, partitionId, false);
        mapAttempts = fileGroups.mapAttempts;
        if (fileGroups.partitionGroups.containsKey(partitionId)) {
          locations = new ArrayList(fileGroups.partitionGroups.get(partitionId));
        }
      }
  
      if (locations == null || locations.size() == 0) {
        logger.warn("Shuffle data is empty for shuffle {} partition {}.", shuffleId, partitionId);
        return CelebornInputStream.empty();
      } else {
        String shuffleKey = Utils.makeShuffleKey(appUniqueId, shuffleId);
        assert dataClientFactory != null;
        return CelebornInputStream.create(
            conf,
            dataClientFactory,
            shuffleKey,
            locations,
            streamHandlers,
            mapAttempts,
            failedBatchSetMap,
            chunksRange,
            attemptNumber,
            taskId,
            startMapIndex,
            endMapIndex,
            fetchExcludedWorkers,
            this,
            appShuffleId,
            shuffleId,
            partitionId,
            exceptionMaker,
            metricsCallback,
            needDecompress);
      }
    }
  
    @VisibleForTesting
    public Map<Integer, Tuple3<ReduceFileGroups, String, Exception>> getReduceFileGroupsMap() {
      return reduceFileGroupsMap;
    }
  
    @Override
    public void shutdown() {
      if (null != reviveManager) {
        reviveManager.close();
      }
      if (null != rpcEnv) {
        rpcEnv.shutdown();
      }
      if (null != dataClientFactory) {
        dataClientFactory.close();
      }
      if (null != transportContext) {
        transportContext.close();
      }
      if (null != pushDataRetryPool) {
        pushDataRetryPool.shutdown();
      }
      if (null != lifecycleManagerRef) {
        lifecycleManagerRef = null;
      }
  
      shuffleIdCache.clear();
      pushExcludedWorkers.clear();
      fetchExcludedWorkers.clear();
      messagesHelper.close();
      logger.warn("Shuffle client has been shutdown!");
    }
  
    @Override
    public void setupLifecycleManagerRef(String host, int port) {
      logger.info("setupLifecycleManagerRef: host = {}, port = {}", host, port);
      try {
        lifecycleManagerRef =
            rpcEnv.setupEndpointRef(
                new RpcAddress(host, port),
                RpcNameConstants.LIFECYCLE_MANAGER_EP,
                rpcMaxRetries,
                rpcRetryWait);
      } catch (Exception e) {
        throw new CelebornRuntimeException("setupLifecycleManagerRef failed!", e);
      }
  
      initDataClientFactoryIfNeeded();
    }
  
    @Override
    public void setupLifecycleManagerRef(RpcEndpointRef endpointRef) {
      lifecycleManagerRef = endpointRef;
      initDataClientFactoryIfNeeded();
    }
  
    @Override
    public void setExtension(byte[] extension) {
      this.extension = extension;
    }
  
    boolean mapperEnded(int shuffleId, int mapId) {
      return (mapperEndMap.containsKey(shuffleId) && mapperEndMap.get(shuffleId).contains(mapId))
          || isStageEnded(shuffleId);
    }
  
    protected boolean isStageEnded(int shuffleId) {
      return stageEndShuffleSet.contains(shuffleId);
    }
  
    private StatusCode getPushDataFailCause(String message) {
      logger.debug("Push data failed cause message: {}", message);
      StatusCode cause;
      if (message == null) {
        logger.error("Push data throw unexpected exception");
        cause = StatusCode.PUSH_DATA_FAIL_NON_CRITICAL_CAUSE_PRIMARY;
      } else if (message.startsWith(StatusCode.PUSH_DATA_FAIL_NON_CRITICAL_CAUSE_REPLICA.name())) {
        cause = StatusCode.PUSH_DATA_FAIL_NON_CRITICAL_CAUSE_REPLICA;
      } else if (message.startsWith(StatusCode.PUSH_DATA_WRITE_FAIL_REPLICA.name())) {
        cause = StatusCode.PUSH_DATA_WRITE_FAIL_REPLICA;
      } else if (message.startsWith(StatusCode.PUSH_DATA_WRITE_FAIL_PRIMARY.name())) {
        cause = StatusCode.PUSH_DATA_WRITE_FAIL_PRIMARY;
      } else if (message.startsWith(StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_PRIMARY.name())) {
        cause = StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_PRIMARY;
      } else if (message.startsWith(StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_REPLICA.name())) {
        cause = StatusCode.PUSH_DATA_CREATE_CONNECTION_FAIL_REPLICA;
      } else if (message.startsWith(StatusCode.PUSH_DATA_CONNECTION_EXCEPTION_PRIMARY.name())) {
        cause = StatusCode.PUSH_DATA_CONNECTION_EXCEPTION_PRIMARY;
      } else if (message.startsWith(StatusCode.PUSH_DATA_CONNECTION_EXCEPTION_REPLICA.name())) {
        cause = StatusCode.PUSH_DATA_CONNECTION_EXCEPTION_REPLICA;
      } else if (message.startsWith(StatusCode.PUSH_DATA_TIMEOUT_PRIMARY.name())) {
        cause = StatusCode.PUSH_DATA_TIMEOUT_PRIMARY;
      } else if (message.startsWith(StatusCode.PUSH_DATA_TIMEOUT_REPLICA.name())) {
        cause = StatusCode.PUSH_DATA_TIMEOUT_REPLICA;
      } else if (message.startsWith(StatusCode.REPLICATE_DATA_FAILED.name())) {
        cause = StatusCode.REPLICATE_DATA_FAILED;
      } else if (message.startsWith(StatusCode.PUSH_DATA_PRIMARY_WORKER_EXCLUDED.name())) {
        cause = StatusCode.PUSH_DATA_PRIMARY_WORKER_EXCLUDED;
      } else if (message.startsWith(StatusCode.PUSH_DATA_REPLICA_WORKER_EXCLUDED.name())) {
        cause = StatusCode.PUSH_DATA_REPLICA_WORKER_EXCLUDED;
      } else if (message.startsWith(StatusCode.PUSH_DATA_FAIL_PARTITION_NOT_FOUND.name())) {
        cause = StatusCode.PUSH_DATA_FAIL_PARTITION_NOT_FOUND;
      } else if (ExceptionUtils.connectFail(message)) {
        // Throw when push to primary worker connection causeException.
        cause = StatusCode.PUSH_DATA_CONNECTION_EXCEPTION_PRIMARY;
      } else {
        cause = StatusCode.PUSH_DATA_FAIL_NON_CRITICAL_CAUSE_PRIMARY;
      }
      return cause;
    }
  
    @VisibleForTesting
    @Override
    public TransportClientFactory getDataClientFactory() {
      return dataClientFactory;
    }
  
    @Override
    public void excludeFailedFetchLocation(String hostAndFetchPort, Exception e) {
      if (pushReplicateEnabled
          && fetchExcludeWorkerOnFailureEnabled
          && Utils.isCriticalCauseForFetch(e)) {
        fetchExcludedWorkers.put(hostAndFetchPort, System.currentTimeMillis());
      }
    }
  }
  