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.plugin.flink;
  
  import java.util.*;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.ExecutorService;
  
  import org.apache.flink.api.common.JobID;
  import org.apache.flink.runtime.io.network.partition.ResultPartitionID;
  import org.apache.flink.runtime.shuffle.JobShuffleContext;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import org.apache.celeborn.client.LifecycleManager;
  import org.apache.celeborn.client.listener.WorkerStatusListener;
  import org.apache.celeborn.client.listener.WorkersStatus;
  import org.apache.celeborn.common.meta.ShufflePartitionLocationInfo;
  import org.apache.celeborn.common.meta.WorkerInfo;
  import org.apache.celeborn.common.util.JavaUtils;
  
  public class ShuffleResourceTracker implements WorkerStatusListener {
    private static final Logger LOG = LoggerFactory.getLogger(ShuffleResourceTracker.class);
    private final ExecutorService executorService;
    private final LifecycleManager lifecycleManager;
    // JobID -> ShuffleResourceListener
    private final Map<JobID, JobShuffleResourceListener> shuffleResourceListeners =
        JavaUtils.newConcurrentHashMap();
    private static final int MAX_RETRY_TIMES = 3;
  
    public ShuffleResourceTracker(
        ExecutorService executorService, LifecycleManager lifecycleManager) {
      this.executorService = executorService;
      this.lifecycleManager = lifecycleManager;
      lifecycleManager.registerWorkerStatusListener(this);
    }
  
    public void registerJob(JobShuffleContext jobShuffleContext) {
      shuffleResourceListeners.put(
          jobShuffleContext.getJobId(),
          new JobShuffleResourceListener(jobShuffleContext, executorService));
    }
  
    public void addPartitionResource(
        JobID jobId, int shuffleId, int partitionId, ResultPartitionID partitionID) {
      JobShuffleResourceListener shuffleResourceListener = shuffleResourceListeners.get(jobId);
      shuffleResourceListener.addPartitionResource(shuffleId, partitionId, partitionID);
    }
  
    public void removePartitionResource(JobID jobID, int shuffleId, int partitionId) {
      JobShuffleResourceListener shuffleResourceListener = shuffleResourceListeners.get(jobID);
      if (shuffleResourceListener != null) {
        shuffleResourceListener.removePartitionResource(shuffleId, partitionId);
      }
    }
  
    public JobShuffleResourceListener getJobResourceListener(JobID jobID) {
      return shuffleResourceListeners.get(jobID);
    }
  
    public void unRegisterJob(JobID jobID) {
      shuffleResourceListeners.remove(jobID);
    }
  
    @Override
    public void notifyChangedWorkersStatus(WorkersStatus workersStatus) {
      try {
        List<WorkerInfo> unknownWorkers = workersStatus.unknownWorkers;
        if (unknownWorkers != null && !unknownWorkers.isEmpty()) {
          // untrack by job
          for (Map.Entry<JobID, JobShuffleResourceListener> entry :
              shuffleResourceListeners.entrySet()) {
            Set<ResultPartitionID> partitionIds = new HashSet<>();
            JobShuffleResourceListener shuffleResourceListener = entry.getValue();
            for (Map.Entry<Integer, Map<Integer, ResultPartitionID>> mapEntry :
                shuffleResourceListener.getResultPartitionMap().entrySet()) {
              int shuffleId = mapEntry.getKey();
              if (!mapEntry.getValue().isEmpty()) {
                for (WorkerInfo unknownWorker : unknownWorkers) {
                  Map<String, ShufflePartitionLocationInfo> shuffleAllocateInfo =
                      lifecycleManager.workerSnapshots(shuffleId);
                  // shuffleResourceListener may release when the shuffle is ended
                  if (shuffleAllocateInfo != null) {
                    ShufflePartitionLocationInfo shufflePartitionLocationInfo =
                        shuffleAllocateInfo.get(unknownWorker.toUniqueId());
                    if (shufflePartitionLocationInfo != null) {
                      // TODO if we support partition replica for map partition we need refactor this
                      //  Currently we only untrack primary partitions for map partition
                      shufflePartitionLocationInfo
                          .removeAndGetAllPrimaryPartitionIds()
                          .forEach(
                              id -> {
                                ResultPartitionID resultPartitionId =
                                    shuffleResourceListener.removePartitionResource(shuffleId, id);
                                if (resultPartitionId != null) {
                                  partitionIds.add(resultPartitionId);
                                }
                              });
                    }
                  }
                }
              }
            }
  
            shuffleResourceListener.notifyStopTrackingPartitions(partitionIds, MAX_RETRY_TIMES);
          }
        }
      } catch (Throwable e) {
        // listener never throw exception
        LOG.error("Failed to handle unknown workers, message: {}.", e.getMessage(), e);
      }
    }
  
    public static class JobShuffleResourceListener {
  
      private final JobShuffleContext context;
      private final ExecutorService executorService;
      // celeborn shuffleId -> partitionId -> Flink ResultPartitionID
      private Map<Integer, Map<Integer, ResultPartitionID>> resultPartitionMap =
          JavaUtils.newConcurrentHashMap();
  
      public JobShuffleResourceListener(
          JobShuffleContext jobShuffleContext, ExecutorService executorService) {
        this.context = jobShuffleContext;
        this.executorService = executorService;
      }
  
      public void addPartitionResource(
          int shuffleId, int partitionId, ResultPartitionID partitionID) {
        Map<Integer, ResultPartitionID> shufflePartitionMap =
            resultPartitionMap.computeIfAbsent(shuffleId, (s) -> JavaUtils.newConcurrentHashMap());
        shufflePartitionMap.put(partitionId, partitionID);
      }
  
      private void notifyStopTrackingPartitions(
          Set<ResultPartitionID> partitionIDS, int remainingRetries) {
        if (partitionIDS == null || partitionIDS.isEmpty()) {
          return;
        }
  
        LOG.info(
            "jobId: {}, stop tracking partitions {}.",
            context.getJobId(),
            Arrays.toString(partitionIDS.toArray()));
  
        int count = remainingRetries - 1;
        try {
          CompletableFuture<?> future = context.stopTrackingAndReleasePartitions(partitionIDS);
          future.whenCompleteAsync(
              (ignored, throwable) -> {
                if (throwable == null) {
                  return;
                }
  
                if (count == 0) {
                  LOG.error(
                      "jobId: {}, Failed to stop tracking partitions {}.",
                      context.getJobId(),
                      Arrays.toString(partitionIDS.toArray()));
                  return;
                }
                notifyStopTrackingPartitions(partitionIDS, count);
              },
              executorService);
        } catch (Throwable throwable) {
          if (count == 0) {
            LOG.error(
                "jobId: {}, Failed to stop tracking partitions {}.",
                context.getJobId(),
                Arrays.toString(partitionIDS.toArray()),
                throwable);
            return;
          }
          notifyStopTrackingPartitions(partitionIDS, count);
        }
      }
  
      public Map<Integer, Map<Integer, ResultPartitionID>> getResultPartitionMap() {
        return resultPartitionMap;
      }
  
      public ResultPartitionID removePartitionResource(int shuffleId, int partitionId) {
        Map<Integer, ResultPartitionID> partitionIDMap = resultPartitionMap.get(shuffleId);
        if (partitionIDMap != null) {
          return partitionIDMap.remove(partitionId);
        }
  
        return null;
      }
    }
  }
  