/** * Copyright 2019 Airbnb. Licensed under Apache-2.0. See LICENSE in the project root for license * information. */ package com.airbnb.dynein.scheduler.partition; import com.airbnb.dynein.scheduler.Constants; import io.fabric8.kubernetes.api.model.Pod; import io.fabric8.kubernetes.api.model.PodList; import io.fabric8.kubernetes.api.model.apps.ReplicaSet; import io.fabric8.kubernetes.client.AppsAPIGroupClient; import io.fabric8.kubernetes.client.DefaultKubernetesClient; import io.fabric8.kubernetes.client.NamespacedKubernetesClient; import io.fabric8.kubernetes.client.Watch; import io.fabric8.kubernetes.client.Watcher; import io.fabric8.kubernetes.client.dsl.FilterWatchListDeletable; import java.text.Collator; import java.util.ArrayList; import java.util.List; import java.util.stream.Collectors; import javax.inject.Inject; import javax.inject.Named; import lombok.AllArgsConstructor; @AllArgsConstructor(onConstructor = @__(@Inject)) public class K8SConsecutiveAllocationPolicy implements PartitionPolicy { private static final String K8S_CONDITION_READY = "Ready"; private static final String K8S_CONDITION_TRUE = "True"; private final DefaultKubernetesClient kubernetesClient; private final AppsAPIGroupClient appsClient; @Named(Constants.MAX_PARTITIONS) private final Integer numPartitions; public List<Integer> getPartitions() { String podName = System.getenv("K8S_POD_NAME"); String namespace = System.getenv("K8S_NAMESPACE"); NamespacedKubernetesClient namespacedClient = kubernetesClient.inNamespace(namespace); ReplicaSet replicaSet; Pod pod = namespacedClient.pods().withName(podName).get(); String replicaSetName = pod.getMetadata().getOwnerReferences().get(0).getName(); replicaSet = appsClient.replicaSets().inNamespace(namespace).withName(replicaSetName).get(); FilterWatchListDeletable<Pod, PodList, Boolean, Watch, Watcher<Pod>> deploymentPods = namespacedClient.pods().withLabelSelector(replicaSet.getSpec().getSelector()); List<String> activePods = deploymentPods .list() .getItems() .stream() .filter( isActive -> isActive .getStatus() .getConditions() .stream() .anyMatch( condition -> condition.getType().equals(K8S_CONDITION_READY) && condition.getStatus().equals(K8S_CONDITION_TRUE))) .map(it -> it.getMetadata().getName()) .sorted(Collator.getInstance()) .collect(Collectors.toList()); int podIndex = activePods.indexOf(podName); int numPods = activePods.size(); if (numPods == 0 || podIndex == -1) { return new ArrayList<>(); } List<Integer> partitions = new ArrayList<>(); int split = numPartitions / numPods; int start = podIndex * split; int end = (podIndex == numPods - 1) ? numPartitions - 1 : ((podIndex + 1) * split) - 1; for (int i = start; i <= end; i++) { partitions.add(i); } return partitions; } }