/*
 * 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.service.deploy.master;

import static org.junit.Assert.*;

import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

import scala.Tuple2;
import scala.Tuple3;

import com.google.common.collect.ImmutableMap;
import org.junit.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.CsvSource;

import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.meta.DiskInfo;
import org.apache.celeborn.common.meta.WorkerInfo;
import org.apache.celeborn.common.protocol.PartitionLocation;
import org.apache.celeborn.common.protocol.StorageInfo;

public class SlotsAllocatorSuiteJ {
  private List<WorkerInfo> prepareWorkers(boolean hasDisks) {
    long assumedPartitionSize = 64 * 1024 * 1024;
    Random random = new Random();
    return basePrepareWorkers(
        3,
        hasDisks,
        ImmutableMap.of(
            "/mnt/disk1", random.nextInt() + 100 * 1024 * 1024 * 1024L,
            "/mnt/disk2", random.nextInt() + 95 * 1024 * 1024 * 1024L,
            "/mnt/disk3", random.nextInt() + 90 * 1024 * 1024 * 1024L),
        assumedPartitionSize,
        3,
        false,
        random);
  }

  @Test
  public void testAllocateSlotsForEmptyPartitionId() {
    final List<WorkerInfo> workers = prepareWorkers(true);
    final List<Integer> partitionIds = Collections.emptyList();
    final boolean shouldReplicate = true;

    check(workers, partitionIds, shouldReplicate, true);
  }

  @Test
  public void testAllocateSlotsForSinglePartitionId() {
    final List<WorkerInfo> workers = prepareWorkers(true);
    final List<Integer> partitionIds = Collections.singletonList(0);
    final boolean shouldReplicate = true;

    check(workers, partitionIds, shouldReplicate, true);
  }

  @Test
  public void testAllocateSlotsForSinglePartitionIdWithoutReplicate() {
    final List<WorkerInfo> workers = prepareWorkers(true);
    final List<Integer> partitionIds = Collections.singletonList(0);
    final boolean shouldReplicate = false;

    check(workers, partitionIds, shouldReplicate, true);
  }

  @Test
  public void testAllocateSlotsForTwoPartitionIdsWithoutReplicate() {
    final List<WorkerInfo> workers = prepareWorkers(true);
    final List<Integer> partitionIds = Arrays.asList(0, 1);
    final boolean shouldReplicate = false;

    check(workers, partitionIds, shouldReplicate, true);
  }

  @Test
  public void testAllocateSlotsForThreePartitionIdsWithoutReplicate() {
    final List<WorkerInfo> workers = prepareWorkers(true);
    final List<Integer> partitionIds = Arrays.asList(0, 1, 2);
    final boolean shouldReplicate = false;

    check(workers, partitionIds, shouldReplicate, true);
  }

  @Test
  public void testAllocateSlotsForThreeReduceIdsWithReplicate() {
    final List<WorkerInfo> workers = prepareWorkers(true);
    final List<Integer> partitionIds = Arrays.asList(0, 1, 2);
    final boolean shouldReplicate = true;

    check(workers, partitionIds, shouldReplicate, true);
  }

  @Test
  public void testAllocate3000ReduceIdsWithReplicate() {
    final List<WorkerInfo> workers = prepareWorkers(true);
    final List<Integer> partitionIds = new ArrayList<>();
    for (int i = 0; i < 3000; i++) {
      partitionIds.add(i);
    }
    final boolean shouldReplicate = true;

    check(workers, partitionIds, shouldReplicate, true);
  }

  @Test
  public void testAllocate3000ReduceIdsWithoutReplicate() {
    final List<WorkerInfo> workers = prepareWorkers(true);
    final List<Integer> partitionIds = new ArrayList<>();
    for (int i = 0; i < 3000; i++) {
      partitionIds.add(i);
    }
    final boolean shouldReplicate = true;

    check(workers, partitionIds, shouldReplicate, true);
  }

  @Test
  public void testAllocate3000ReduceIdsWithReplicateOnRoundRobin() {
    final List<WorkerInfo> workers = prepareWorkers(true);
    final List<Integer> partitionIds = new ArrayList<>();
    for (int i = 0; i < 3000; i++) {
      partitionIds.add(i);
    }
    final boolean shouldReplicate = true;

    check(workers, partitionIds, shouldReplicate, true, true, false, 0);
  }

  @Test
  public void testLoadAwarePrefersLowerActiveSlotsWhenConfigured() {
    final List<WorkerInfo> workers =
        basePrepareWorkers(
            2,
            true,
            ImmutableMap.of("/mnt/disk1", 100 * 1024 * 1024 * 1024L),
            64 * 1024 * 1024L,
            1,
            false,
            new Random(0));
    final DiskInfo overloadedDisk = workers.get(0).diskInfos().get("/mnt/disk1");
    final DiskInfo lightlyReservedDisk = workers.get(1).diskInfos().get("/mnt/disk1");
    overloadedDisk.activeSlots_$eq(1000);
    lightlyReservedDisk.activeSlots_$eq(0);

    final Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots =
        SlotsAllocator.offerSlotsLoadAware(
            workers,
            Collections.singletonList(0),
            false,
            false,
            2,
            1,
            0,
            0,
            1,
            StorageInfo.ALL_TYPES_AVAILABLE_MASK,
            false,
            0);

    assertTrue(slots.containsKey(workers.get(1)));
    assertFalse(slots.containsKey(workers.get(0)));
  }

  private void check(
      List<WorkerInfo> workers,
      List<Integer> partitionIds,
      boolean shouldReplicate,
      boolean expectSuccess) {
    check(workers, partitionIds, shouldReplicate, expectSuccess, false, false, 0);
  }

  private Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> check(
      List<WorkerInfo> workers,
      List<Integer> partitionIds,
      boolean shouldReplicate,
      boolean expectSuccess,
      boolean roundrobin,
      boolean interruptionAware,
      int interruptionAwareThreshold) {
    String shuffleKey = "appId-1";
    CelebornConf conf = new CelebornConf();
    conf.set(CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_DISKGROUP_NUM().key(), "2");
    conf.set(CelebornConf.MASTER_SLOT_ASSIGN_LOADAWARE_DISKGROUP_GRADIENT().key(), "1");
    Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots;
    if (roundrobin) {
      slots =
          SlotsAllocator.offerSlotsRoundRobin(
              workers,
              partitionIds,
              shouldReplicate,
              false,
              StorageInfo.ALL_TYPES_AVAILABLE_MASK,
              interruptionAware,
              interruptionAwareThreshold);
    } else {
      slots =
          SlotsAllocator.offerSlotsLoadAware(
              workers,
              partitionIds,
              shouldReplicate,
              false,
              conf.masterSlotAssignLoadAwareDiskGroupNum(),
              conf.masterSlotAssignLoadAwareDiskGroupGradient(),
              conf.masterSlotAssignLoadAwareFlushTimeWeight(),
              conf.masterSlotAssignLoadAwareFetchTimeWeight(),
              conf.masterSlotAssignLoadAwareActiveSlotsWeight(),
              StorageInfo.ALL_TYPES_AVAILABLE_MASK,
              interruptionAware,
              interruptionAwareThreshold);
    }
    if (expectSuccess) {
      if (shouldReplicate) {
        slots.forEach(
            (k, v) -> {
              Set<String> locationDuplicationSet = new HashSet<>();
              v._1.forEach(
                  i -> {
                    String uniqueId = i.getUniqueId();
                    assertFalse(locationDuplicationSet.contains(uniqueId));
                    locationDuplicationSet.add(uniqueId);
                    for (PartitionLocation location : v._1) {
                      assertTrue(
                          location.getHost().equals(k.host())
                              && location.getRpcPort() == k.rpcPort()
                              && location.getPushPort() == k.pushPort()
                              && location.getFetchPort() == k.fetchPort());
                    }
                  });
            });
      }
      Map<WorkerInfo, Map<String, Integer>> workerToAllocatedSlots =
          SlotsAllocator.slotsToDiskAllocations(slots);
      int unknownDiskSlots = 0;
      for (Map.Entry<WorkerInfo, Map<String, Integer>> entry : workerToAllocatedSlots.entrySet()) {
        WorkerInfo worker = entry.getKey();
        Map<String, Integer> allocationMap = entry.getValue();
        worker.allocateSlots(shuffleKey, allocationMap);
        if (allocationMap.containsKey("UNKNOWN_DISK")) {
          unknownDiskSlots += allocationMap.get("UNKNOWN_DISK");
        }
        if (roundrobin && !allocationMap.isEmpty()) {
          int maxSlots = Collections.max(allocationMap.values());
          int minSlots = Collections.min(allocationMap.values());
          assertTrue(
              "Worker "
                  + worker.host()
                  + " has unbalanced slot allocation. "
                  + "Max: "
                  + maxSlots
                  + ", Min: "
                  + minSlots,
              maxSlots - minSlots <= 1);
        }
      }
      int allocateToDiskSlots = 0;
      for (WorkerInfo worker : workers) {
        allocateToDiskSlots = (int) (allocateToDiskSlots + worker.usedSlots());
      }
      if (shouldReplicate) {
        assertTrue(partitionIds.size() * 2 >= unknownDiskSlots + allocateToDiskSlots);
      } else {
        assertTrue(partitionIds.size() >= unknownDiskSlots + allocateToDiskSlots);
      }
      assertEquals(0, unknownDiskSlots);
    } else {
      assertTrue(
          "Expect to fail to offer slots, but return " + slots.size() + " slots.", slots.isEmpty());
    }
    return slots;
  }

  private void checkSlotsOnDFS(
      List<WorkerInfo> workers,
      List<Integer> partitionIds,
      boolean shouldReplicate,
      boolean expectSuccess,
      boolean roundRobin,
      boolean enableS3) {
    checkSlotsOnDFS(
        workers, partitionIds, shouldReplicate, expectSuccess, roundRobin, enableS3, false);
  }

  private void checkSlotsOnDFS(
      List<WorkerInfo> workers,
      List<Integer> partitionIds,
      boolean shouldReplicate,
      boolean expectSuccess,
      boolean roundRobin,
      boolean enableS3,
      boolean enableOss) {
    CelebornConf conf = new CelebornConf();
    int availableStorageTypes;
    if (enableS3) {
      conf.set("celeborn.active.storage.levels", "S3");
      availableStorageTypes = StorageInfo.S3_MASK;
    } else if (enableOss) {
      conf.set("celeborn.active.storage.levels", "OSS");
      availableStorageTypes = StorageInfo.OSS_MASK;
    } else {
      conf.set("celeborn.active.storage.levels", "HDFS");
      availableStorageTypes = StorageInfo.HDFS_MASK;
    }
    Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>> slots;
    if (roundRobin) {
      slots =
          SlotsAllocator.offerSlotsRoundRobin(
              workers, partitionIds, shouldReplicate, false, availableStorageTypes, false, 0);
    } else {
      slots =
          SlotsAllocator.offerSlotsLoadAware(
              workers,
              partitionIds,
              shouldReplicate,
              false,
              3,
              0.1,
              0,
              1,
              0,
              StorageInfo.LOCAL_DISK_MASK | availableStorageTypes,
              false,
              0);
    }
    int allocatedPartitionCount = 0;
    for (Map.Entry<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>>
        workerToPartitions : slots.entrySet()) {
      List<PartitionLocation> primaryLocs = workerToPartitions.getValue()._1;
      List<PartitionLocation> replicaLocs = workerToPartitions.getValue()._2();
      allocatedPartitionCount += primaryLocs.size();
      allocatedPartitionCount += replicaLocs.size();
    }
    assertEquals(expectSuccess, !slots.isEmpty());
    assertEquals(
        shouldReplicate ? partitionIds.size() * 2 : partitionIds.size(), allocatedPartitionCount);
  }

  @Test
  public void testHDFSOnly() {
    final List<WorkerInfo> workers = prepareWorkers(false);
    final List<Integer> partitionIds = new ArrayList<>();
    for (int i = 0; i < 3000; i++) {
      partitionIds.add(i);
    }
    final boolean shouldReplicate = true;
    checkSlotsOnDFS(workers, partitionIds, shouldReplicate, true, true, false);
  }

  @Test
  public void testLocalDisksAndHDFSOnRoundRobin() {
    final List<WorkerInfo> workers = prepareWorkers(true);
    DiskInfo hdfs1 =
        new DiskInfo(
            "HDFS", Long.MAX_VALUE, 999999, 999999, Integer.MAX_VALUE, StorageInfo.Type.HDFS);
    DiskInfo hdfs2 =
        new DiskInfo(
            "HDFS", Long.MAX_VALUE, 999999, 999999, Integer.MAX_VALUE, StorageInfo.Type.HDFS);
    DiskInfo hdfs3 =
        new DiskInfo(
            "HDFS", Long.MAX_VALUE, 999999, 999999, Integer.MAX_VALUE, StorageInfo.Type.HDFS);
    hdfs1.maxSlots_$eq(Long.MAX_VALUE);
    hdfs2.maxSlots_$eq(Long.MAX_VALUE);
    hdfs3.maxSlots_$eq(Long.MAX_VALUE);
    workers.get(0).diskInfos().put("HDFS", hdfs1);
    workers.get(1).diskInfos().put("HDFS", hdfs2);
    workers.get(2).diskInfos().put("HDFS", hdfs3);
    final List<Integer> partitionIds = new ArrayList<>();
    for (int i = 0; i < 3000; i++) {
      partitionIds.add(i);
    }
    final boolean shouldReplicate = true;
    checkSlotsOnDFS(workers, partitionIds, shouldReplicate, true, true, false);
  }

  @Test
  public void testLocalDisksAndHDFSOnLoadAware() {
    final List<WorkerInfo> workers = prepareWorkers(true);
    DiskInfo hdfs1 =
        new DiskInfo(
            "HDFS", Long.MAX_VALUE, 999999, 999999, Integer.MAX_VALUE, StorageInfo.Type.HDFS);
    DiskInfo hdfs2 =
        new DiskInfo(
            "HDFS", Long.MAX_VALUE, 999999, 999999, Integer.MAX_VALUE, StorageInfo.Type.HDFS);
    //    DiskInfo hdfs3 = new DiskInfo("HDFS", Long.MAX_VALUE, 999999, 999999, Integer.MAX_VALUE,
    // StorageInfo.Type.HDFS);
    hdfs1.maxSlots_$eq(Long.MAX_VALUE);
    hdfs2.maxSlots_$eq(Long.MAX_VALUE);
    //    hdfs3.maxSlots_$eq(Long.MAX_VALUE);
    workers.get(0).diskInfos().put("HDFS", hdfs1);
    workers.get(1).diskInfos().put("HDFS", hdfs2);
    //    workers.get(2).diskInfos().put("HDFS", hdfs3);
    final List<Integer> partitionIds = new ArrayList<>();
    for (int i = 0; i < 3000; i++) {
      partitionIds.add(i);
    }
    final boolean shouldReplicate = true;
    checkSlotsOnDFS(workers, partitionIds, shouldReplicate, true, false, false);
  }

  @Test
  public void testLocalDisksAndHDFSOnLoadAwareWithInsufficientSlots() {
    final List<WorkerInfo> workers = prepareWorkers(true);
    DiskInfo hdfs1 =
        new DiskInfo(
            "HDFS", Long.MAX_VALUE, 999999, 999999, Integer.MAX_VALUE, StorageInfo.Type.HDFS);
    DiskInfo hdfs2 =
        new DiskInfo(
            "HDFS", Long.MAX_VALUE, 999999, 999999, Integer.MAX_VALUE, StorageInfo.Type.HDFS);
    //    DiskInfo hdfs3 = new DiskInfo("HDFS", Long.MAX_VALUE, 999999, 999999, Integer.MAX_VALUE,
    // StorageInfo.Type.HDFS);
    hdfs1.maxSlots_$eq(Long.MAX_VALUE);
    hdfs2.maxSlots_$eq(Long.MAX_VALUE);
    //    hdfs3.maxSlots_$eq(Long.MAX_VALUE);
    workers.get(0).diskInfos().put("HDFS", hdfs1);
    workers.get(1).diskInfos().put("HDFS", hdfs2);
    for (Map.Entry<String, DiskInfo> diskEntry : workers.get(2).diskInfos().entrySet()) {
      diskEntry.getValue().maxSlots_$eq(100);
    }
    //    workers.get(2).diskInfos().put("HDFS", hdfs3);
    final List<Integer> partitionIds = new ArrayList<>();
    for (int i = 0; i < 3000; i++) {
      partitionIds.add(i);
    }
    final boolean shouldReplicate = true;
    checkSlotsOnDFS(workers, partitionIds, shouldReplicate, true, false, false);
  }

  @Test
  public void testAllocateSlotsWithNoAvailableSlots() {
    final List<WorkerInfo> workers = prepareWorkers(true);
    // Simulates no available slots behavior with greatly changed estimatedPartitionSize for workers
    // with usable disks.
    workers.forEach(workerInfo -> workerInfo.updateDiskSlots(Long.MAX_VALUE));
    final List<Integer> partitionIds = Collections.singletonList(0);
    final boolean shouldReplicate = false;

    check(workers, partitionIds, shouldReplicate, true);
  }

  @Test
  public void testS3Only() {
    final List<WorkerInfo> workers = prepareWorkers(false);
    final List<Integer> partitionIds = new ArrayList<>();
    for (int i = 0; i < 3000; i++) {
      partitionIds.add(i);
    }
    final boolean shouldReplicate = true;
    checkSlotsOnDFS(workers, partitionIds, shouldReplicate, true, true, true);
  }

  @Test
  public void testLocalDisksAndS3() {
    final List<WorkerInfo> workers = prepareWorkers(true);
    DiskInfo s3DiskInfo1 =
        new DiskInfo("S3", Long.MAX_VALUE, 999999, 999999, Integer.MAX_VALUE, StorageInfo.Type.S3);
    DiskInfo s3DiskInfo2 =
        new DiskInfo("S3", Long.MAX_VALUE, 999999, 999999, Integer.MAX_VALUE, StorageInfo.Type.S3);
    s3DiskInfo1.maxSlots_$eq(Long.MAX_VALUE);
    s3DiskInfo2.maxSlots_$eq(Long.MAX_VALUE);
    workers.get(0).diskInfos().put("S3", s3DiskInfo1);
    workers.get(1).diskInfos().put("S3", s3DiskInfo2);
    final List<Integer> partitionIds = new ArrayList<>();
    for (int i = 0; i < 3000; i++) {
      partitionIds.add(i);
    }
    final boolean shouldReplicate = true;
    checkSlotsOnDFS(workers, partitionIds, shouldReplicate, true, true, true);
    checkSlotsOnDFS(workers, partitionIds, shouldReplicate, true, false, true);
  }

  @Test
  public void testOssOnly() {
    final List<WorkerInfo> workers = prepareWorkers(false);
    final List<Integer> partitionIds = new ArrayList<>();
    for (int i = 0; i < 3000; i++) {
      partitionIds.add(i);
    }
    final boolean shouldReplicate = true;
    checkSlotsOnDFS(workers, partitionIds, shouldReplicate, true, true, false, true);
  }

  @Test
  public void testLocalDisksAndOss() {
    final List<WorkerInfo> workers = prepareWorkers(true);
    DiskInfo ossDiskInfo1 =
        new DiskInfo(
            "OSS", Long.MAX_VALUE, 999999, 999999, Integer.MAX_VALUE, StorageInfo.Type.OSS);
    DiskInfo ossDiskInfo2 =
        new DiskInfo(
            "OSS", Long.MAX_VALUE, 999999, 999999, Integer.MAX_VALUE, StorageInfo.Type.OSS);
    ossDiskInfo1.maxSlots_$eq(Long.MAX_VALUE);
    ossDiskInfo2.maxSlots_$eq(Long.MAX_VALUE);
    workers.get(0).diskInfos().put("OSS", ossDiskInfo1);
    workers.get(1).diskInfos().put("OSS", ossDiskInfo2);
    final List<Integer> partitionIds = new ArrayList<>();
    for (int i = 0; i < 3000; i++) {
      partitionIds.add(i);
    }
    final boolean shouldReplicate = true;
    checkSlotsOnDFS(workers, partitionIds, shouldReplicate, true, true, false, true);
    checkSlotsOnDFS(workers, partitionIds, shouldReplicate, true, false, false, true);
  }

  @ParameterizedTest
  @CsvSource({"true, true", "true, false", "false, false", "false, true"})
  public void testInterruptionAwareSlotSelection(boolean shouldReplicate, boolean shouldRackAware) {
    long assumedPartitionSize = 64 * 1024 * 1024;
    double interruptionAwarePercentileThreshold = 50;
    Map<String, Long> diskPartitionToSize = new HashMap<>();
    diskPartitionToSize.put("/mnt/disk", 512 * 1024 * 1024L); // 0.5gb disk space
    // Cluster usable space is 50g, with 25g that will not be interrupted.
    List<WorkerInfo> workers =
        basePrepareWorkers(
            100, true, diskPartitionToSize, assumedPartitionSize, 20, true, new Random());
    Map<String, WorkerInfo> workersMap =
        workers.stream().collect(Collectors.toMap(WorkerInfo::host, worker -> worker));
    Tuple3<List<WorkerInfo>, List<WorkerInfo>, List<WorkerInfo>> prioritization =
        SlotsAllocator.prioritizeWorkersBasedOnInterruptionNotice(
            workers, shouldReplicate, shouldRackAware, interruptionAwarePercentileThreshold);
    List<WorkerInfo> workersWithoutInterruptions = prioritization._1();
    List<WorkerInfo> workersWithLateInterruptions = prioritization._2();
    List<WorkerInfo> workersWithEarlyInterruptions = prioritization._3();
    List<String> workersWithoutInterruptionsHosts = extractHosts(workersWithoutInterruptions);
    List<String> workersWithLateInterruptionsHosts = extractHosts(workersWithLateInterruptions);
    List<String> workersWithEarlyInterruptionsHosts = extractHosts(workersWithEarlyInterruptions);

    assertEquals(50, workersWithoutInterruptionsHosts.size());
    assertEquals(25, workersWithLateInterruptionsHosts.size());
    assertEquals(25, workersWithEarlyInterruptionsHosts.size());
    IntStream.range(0, 100)
        .forEach(
            i -> {
              String host = "host" + i;
              if (i % 2 == 0) {
                assertTrue(workersWithoutInterruptionsHosts.contains(host));
              } else if (i >= 51) {
                assertTrue(workersWithLateInterruptionsHosts.contains(host));
              } else {
                assertTrue(workersWithEarlyInterruptionsHosts.contains(host));
              }
            });

    // With replication enabled: 150 partitions * 128mb (64 primary, 64 replica) is roughly 19gb.
    // Both primaries and replicas should fit into workersWithoutInterruptions, since 19gb < 25gb
    // uninterrupted capacity.
    //
    // With replication disabled: 150 partitions * 64mb is roughly 9gb.
    // Similar to the above case, all primaries should fit into workersWithoutInterruptions since
    // 9gb < 25gb uninterrupted capacity.
    List<Integer> bestCasePartitionIds =
        IntStream.range(0, 150).boxed().collect(Collectors.toList());
    Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>>
        slotsFromBestCasePartitionIds =
            SlotsAllocator.offerSlotsRoundRobin(
                workers,
                bestCasePartitionIds,
                shouldReplicate,
                shouldRackAware,
                StorageInfo.ALL_TYPES_AVAILABLE_MASK,
                true,
                (int) interruptionAwarePercentileThreshold);
    slotsFromBestCasePartitionIds
        .values()
        .forEach(
            primaryReplicaSlots -> {
              List<PartitionLocation> primarySlots = primaryReplicaSlots._1;
              List<PartitionLocation> replicaSlots = primaryReplicaSlots._2;
              assertTrue(
                  primarySlots.stream()
                      .map(PartitionLocation::getHost)
                      .allMatch(workersWithoutInterruptionsHosts::contains));
              if (shouldReplicate) {
                assertTrue(
                    replicaSlots.stream()
                        .map(PartitionLocation::getHost)
                        .allMatch(workersWithoutInterruptionsHosts::contains));
                if (shouldRackAware) {
                  primarySlots.forEach(
                      slot -> {
                        WorkerInfo primary = workersMap.get(slot.getHost());
                        WorkerInfo replica = workersMap.get(slot.getPeer().getHost());
                        assertNotSame(primary.networkLocation(), replica.networkLocation());
                      });
                }
              }
            });

    List<WorkerInfo> primaryWorkerCandidates =
        combineWorkers(workersWithoutInterruptions, workersWithLateInterruptions);
    List<String> primaryWorkerCandidatesHosts = extractHosts(primaryWorkerCandidates);

    // With replication enabled: 300 partitions * 128mb (64 primary, 64 replica) is roughly 38gb.
    // In this case, primaries should be in workersWithoutInterruptions +
    // workersWithLateInterruptions, while
    // replicas can spill over into workersWithEarlyInterruptions.
    //
    // With replication disabled, we increase partitions to 600 to force this case:
    // 600 partitions * 64mb is roughly 38gb.
    // Similar to the above case, all primaries should be in workersWithoutInterruptions +
    // workersWithLateInterruptions.
    List<Integer> spillOverCasePartitionIds;
    if (shouldReplicate) {
      spillOverCasePartitionIds = IntStream.range(0, 300).boxed().collect(Collectors.toList());
    } else {
      spillOverCasePartitionIds = IntStream.range(0, 600).boxed().collect(Collectors.toList());
    }
    Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>>
        slotsFromSpillOverCasePartitionIds =
            SlotsAllocator.offerSlotsRoundRobin(
                workers,
                spillOverCasePartitionIds,
                shouldReplicate,
                shouldRackAware,
                StorageInfo.ALL_TYPES_AVAILABLE_MASK,
                true,
                (int) interruptionAwarePercentileThreshold);
    slotsFromSpillOverCasePartitionIds
        .values()
        .forEach(
            primaryReplicaSlots -> {
              List<PartitionLocation> primarySlots = primaryReplicaSlots._1;
              List<PartitionLocation> replicaSlots = primaryReplicaSlots._2;
              assertTrue(
                  primarySlots.stream()
                      .map(PartitionLocation::getHost)
                      .allMatch(primaryWorkerCandidatesHosts::contains));
              assertTrue(
                  primarySlots.stream()
                      .map(PartitionLocation::getHost)
                      .noneMatch(workersWithEarlyInterruptionsHosts::contains));
              if (shouldReplicate) {
                assertTrue(
                    replicaSlots.stream()
                        .map(PartitionLocation::getHost)
                        .allMatch(
                            host ->
                                primaryWorkerCandidatesHosts.contains(host)
                                    || workersWithEarlyInterruptionsHosts.contains(host)));
                if (shouldRackAware) {
                  primarySlots.forEach(
                      slot -> {
                        WorkerInfo primary = workersMap.get(slot.getHost());
                        WorkerInfo replica = workersMap.get(slot.getPeer().getHost());
                        assertNotSame(primary.networkLocation(), replica.networkLocation());
                      });
                }
              }
            });
    // With the slot restrictions in place for LoadAware, we expect to spill replicas into
    // workersWithEarlyInterruptionsHosts.
    // But primaries should be in workersWithoutInterruptions + workersWithLateInterruptions.
    Map<WorkerInfo, Tuple2<List<PartitionLocation>, List<PartitionLocation>>>
        loadAwareBestCasePartitionIdsSlots =
            check(
                workers,
                spillOverCasePartitionIds,
                shouldReplicate,
                true,
                false,
                true,
                (int) interruptionAwarePercentileThreshold);
    loadAwareBestCasePartitionIdsSlots
        .values()
        .forEach(
            primaryReplicaSlots -> {
              List<PartitionLocation> primarySlots = primaryReplicaSlots._1;
              List<PartitionLocation> replicaSlots = primaryReplicaSlots._2;
              assertTrue(
                  primarySlots.stream()
                      .map(PartitionLocation::getHost)
                      .allMatch(primaryWorkerCandidatesHosts::contains));
              assertTrue(
                  primarySlots.stream()
                      .map(PartitionLocation::getHost)
                      .noneMatch(workersWithEarlyInterruptionsHosts::contains));
              if (shouldReplicate) {
                assertTrue(
                    replicaSlots.stream()
                        .map(PartitionLocation::getHost)
                        .allMatch(
                            host ->
                                primaryWorkerCandidatesHosts.contains(host)
                                    || workersWithEarlyInterruptionsHosts.contains(host)));
                if (shouldRackAware) {
                  primarySlots.forEach(
                      slot -> {
                        WorkerInfo primary = workersMap.get(slot.getHost());
                        WorkerInfo replica = workersMap.get(slot.getPeer().getHost());
                        assertNotSame(primary.networkLocation(), replica.networkLocation());
                      });
                }
              }
            });
  }

  @Test
  public void testInterruptionAwareSlotSelectionWithNoInterruptions() {
    long assumedPartitionSize = 64 * 1024 * 1024;
    Map<String, Long> diskPartitionToSize = new HashMap<>();
    diskPartitionToSize.put("/mnt/disk", 512 * 1024 * 1024L); // 0.5gb disk space
    // Cluster usable space is 50g, with 25g that will not be interrupted.
    List<WorkerInfo> workers =
        basePrepareWorkers(
            100, true, diskPartitionToSize, assumedPartitionSize, 20, false, new Random());
    List<Integer> partitionIds = IntStream.range(0, 600).boxed().collect(Collectors.toList());
    check(workers, partitionIds, true, true, false, true, 50);
  }

  private List<String> extractHosts(List<WorkerInfo> workers) {
    return workers.stream().map(WorkerInfo::host).collect(Collectors.toList());
  }

  private List<WorkerInfo> combineWorkers(List<WorkerInfo>... workerLists) {
    return Arrays.stream(workerLists).flatMap(List::stream).collect(Collectors.toList());
  }

  static List<WorkerInfo> basePrepareWorkers(
      int numWorkers,
      boolean hasDisks,
      Map<String, Long> diskPartitionToSize,
      long assumedPartitionSize,
      int numNetworkLocations,
      boolean hasInterruptions,
      Random random) {
    return IntStream.range(0, numWorkers)
        .mapToObj(
            i -> {
              Map<String, DiskInfo> disks = new HashMap<>();
              if (hasDisks) {
                diskPartitionToSize.forEach(
                    (diskMountPoint, diskSize) -> {
                      DiskInfo diskInfo =
                          new DiskInfo(
                              diskMountPoint,
                              diskSize,
                              random.nextInt(1000),
                              random.nextInt(1000),
                              0);
                      diskInfo.availableSlots_$eq(
                          diskInfo.actualUsableSpace() / assumedPartitionSize);
                      disks.put(diskMountPoint, diskInfo);
                    });
              }
              WorkerInfo worker = new WorkerInfo("host" + i, i, i, i, i, i, disks, null);
              if (hasInterruptions) {
                if (i % 2 == 0) {
                  worker.nextInterruptionNotice_$eq(Long.MAX_VALUE);
                } else {
                  worker.nextInterruptionNotice_$eq(i);
                }
              }
              worker.networkLocation_$eq(String.valueOf(i % numNetworkLocations));
              return worker;
            })
        .collect(Collectors.toList());
  }
}
