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.hadoop.mapreduce.task.reduce;
  
  import java.io.IOException;
  
  import org.apache.hadoop.mapred.*;
  import org.apache.hadoop.mapreduce.task.reduce.Shuffle.ShuffleError;
  import org.apache.hadoop.util.Progress;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.client.ShuffleClient;
  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.identity.UserIdentifier;
  import org.apache.celeborn.reflect.DynConstructors;
  import org.apache.celeborn.reflect.DynMethods;
  import org.apache.celeborn.util.HadoopUtils;
  
  public class CelebornShuffleConsumer<K, V>
      implements ShuffleConsumerPlugin<K, V>, ExceptionReporter {
    private static final Logger logger = LoggerFactory.getLogger(CelebornShuffleConsumer.class);
    private JobConf mrJobConf;
    private MergeManager<K, V> merger;
    private Throwable throwable = null;
    private Progress copyPhase;
    private TaskStatus taskStatus;
    private org.apache.hadoop.mapreduce.TaskAttemptID reduceId;
    private TaskUmbilicalProtocol umbilical;
    private Reporter reporter;
    private ShuffleClientMetrics metrics;
    private Task reduceTask;
    private ShuffleClient shuffleClient;
  
    @Override
    public void init(Context<K, V> context) {
  
      reduceId = context.getReduceId();
      mrJobConf = context.getJobConf();
      JobConf celebornJobConf = new JobConf(HadoopUtils.MR_CELEBORN_CONF);
  
      umbilical = context.getUmbilical();
      reporter = context.getReporter();
      try {
        this.metrics = createMetrics(reduceId, mrJobConf);
      } catch (Exception e) {
        logger.error("Fatal error occurred, failed to get shuffle client metrics.", e);
        reportException(e);
      }
      copyPhase = context.getCopyPhase();
      taskStatus = context.getStatus();
      reduceTask = context.getReduceTask();
  
      String appId = celebornJobConf.get(HadoopUtils.MR_CELEBORN_APPLICATION_ID);
      String lmHost = celebornJobConf.get(HadoopUtils.MR_CELEBORN_LM_HOST);
      int lmPort = Integer.parseInt(celebornJobConf.get(HadoopUtils.MR_CELEBORN_LM_PORT));
      logger.info("Reducer initialized with celeborn {} {} {}", appId, lmHost, lmPort);
      CelebornConf celebornConf = HadoopUtils.fromYarnConf(mrJobConf);
      shuffleClient =
          ShuffleClient.get(
              appId,
              lmHost,
              lmPort,
              celebornConf,
              new UserIdentifier(
                  celebornConf.userSpecificTenant(), celebornConf.userSpecificUserName()));
      this.merger =
          new MergeManagerImpl<>(
              reduceId,
              mrJobConf,
              context.getLocalFS(),
              context.getLocalDirAllocator(),
              reporter,
              context.getCodec(),
              context.getCombinerClass(),
              context.getCombineCollector(),
              context.getSpilledRecordsCounter(),
              context.getReduceCombineInputCounter(),
              context.getMergedMapOutputsCounter(),
              this,
              context.getMergePhase(),
              context.getMapOutputFile());
    }
  
    private ShuffleClientMetrics createMetrics(
        org.apache.hadoop.mapreduce.TaskAttemptID taskAttemptID, JobConf jobConf)
        throws NoSuchMethodException {
      // for hadoop 3.1+ see MAPREDUCE-6861
      try {
        return DynMethods.builder("create")
            .impl(
                ShuffleClientMetrics.class,
                org.apache.hadoop.mapreduce.TaskAttemptID.class,
                JobConf.class)
            .buildStaticChecked()
            .invoke(taskAttemptID, jobConf);
      } catch (Exception e) {
        // ignore this exception because the createMetrics might use hadoop2
      }
  
      // for hadoop 3.1 see MAPREDUCE-6526
      try {
        return DynMethods.builder("create")
            .impl(ShuffleClientMetrics.class)
            .buildStaticChecked()
            .invoke(taskAttemptID, jobConf);
      } catch (Exception e) {
      }
  
      // for hadoop 2
      return DynConstructors.builder(ShuffleClientMetrics.class)
          .hiddenImpl(new Class[] {org.apache.hadoop.mapreduce.TaskAttemptID.class, JobConf.class})
          .buildChecked()
          .invoke(null, taskAttemptID, jobConf);
    }
  
    @Override
    public RawKeyValueIterator run() throws IOException {
      logger.info(
          "In reduce:{}, Celeborn mr client start to read shuffle data."
              + " Create inputstream with params: shuffleId 0 reduceId:{} attemptId:{}",
          reduceId,
          reduceId.getTaskID().getId(),
          reduceId.getId());
  
      MetricsCallback metricsCallback =
          new MetricsCallback() {
            @Override
            public void incBytesRead(long bytesRead) {}
  
            @Override
            public void incReadTime(long time) {}
          };
  
      CelebornInputStream shuffleInputStream =
          shuffleClient.readPartition(
              0,
              reduceId.getTaskID().getId(),
              reduceId.getId(),
              0,
              0,
              Integer.MAX_VALUE,
              metricsCallback);
      CelebornShuffleFetcher<K, V> shuffleReader =
          new CelebornShuffleFetcher(
              reduceId, taskStatus, merger, copyPhase, reporter, metrics, shuffleInputStream);
      shuffleReader.fetchAndMerge();
  
      copyPhase.complete();
      taskStatus.setPhase(TaskStatus.Phase.SORT);
      reduceTask.statusUpdate(umbilical);
  
      RawKeyValueIterator kvIter;
      try {
        kvIter = merger.close();
      } catch (Throwable e) {
        throw new ShuffleError("Error while doing final merge ", e);
      }
  
      logger.info("In reduce: {} Celeborn mr client read shuffle data complete", reduceId);
  
      return kvIter;
    }
  
    @Override
    public void close() {}
  
    @Override
    public void reportException(Throwable throwable) {
      if (this.throwable == null) {
        this.throwable = throwable;
      }
    }
  }
  