org.apache.aurora.scheduler.updater.strategy.VariableBatchStrategy.java Source code

Introduction

Here is the source code for org.apache.aurora.scheduler.updater.strategy.VariableBatchStrategy.java

Source

/**
 * Licensed 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.aurora.scheduler.updater.strategy;

import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import com.google.common.collect.Ordering;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * An update strategy that will only add more work when the current active group is empty.
 * Size of the groups are picked from the supplied list.
 * Last element is picked multiple times if necessary.
 *
 * @param <T> Instance type.
 */
public class VariableBatchStrategy<T extends Comparable<T>> implements UpdateStrategy<T> {
    private final Ordering<T> ordering;
    protected final ImmutableList<Integer> groupSizes;
    private final boolean rollingForward;
    private Optional<Integer> totalModInstanceCount;

    private static final Logger LOG = LoggerFactory.getLogger(VariableBatchStrategy.class);

    /**
     * Creates a variable active-limited strategy that applies an upper bound to all results.
     *
     * @param maxActiveGroups  List of Maximum group sizes. Each group size represents a step.
     * {@link #getNextGroup(Set, Set)}.
     */
    public VariableBatchStrategy(Ordering<T> ordering, List<Integer> maxActiveGroups, boolean rollingForward) {

        this.ordering = Objects.requireNonNull(ordering);
        this.rollingForward = rollingForward;

        maxActiveGroups.forEach(x -> Preconditions.checkArgument(x > 0));

        this.groupSizes = ImmutableList.copyOf(maxActiveGroups);
        this.totalModInstanceCount = Optional.empty();
    }

    // Determine how far we're into the update based upon how many instances are waiting
    // to be modified.
    private int determineCurGroupSize(int remaining) {
        // Calculate which groupIndex we are in by finding out how many instances we have left to update
        int modified = totalModInstanceCount.get() - remaining;
        int finalGroupSize = Iterables.getLast(groupSizes);

        LOG.debug(
                "Variable Batch Update progress: {} instances have been modified, "
                        + "{} instances remain unmodified, and {} overall instances to be modified.",
                modified, remaining, totalModInstanceCount.get());

        if (rollingForward) {
            int sum = 0;
            for (Integer groupSize : groupSizes) {
                sum += groupSize;

                if (sum > modified) {
                    return groupSize;
                }
            }
            // Return last step when number of instances > sum of all groups
            return finalGroupSize;
        } else {
            // To perform the update in reverse, we use the number of remaining tasks left to update
            // instead of using the number of already modified instances. In a rollback, the remaining
            // count represents the number of instances that were already modified while rolling forward
            // and need to be reverted.
            int curGroupSize = remaining;

            for (Integer groupSize : groupSizes) {
                // This handles an in between step. i.e.: updated instances = 4, update groups = [2,3]
                // which results in update groups 2 and 2 rolling forward at the time of failure.
                if (curGroupSize <= groupSize) {
                    return curGroupSize;
                }

                curGroupSize -= groupSize;
            }

            // Handle the case where number of instances update were
            // greater than the sum of all update groups
            // Calculate the size of the last update group size performed while rolling forward.
            curGroupSize = curGroupSize % finalGroupSize;
            if (curGroupSize == 0) {
                return finalGroupSize;
            } else {
                return curGroupSize;
            }
        }
    }

    @Override
    public final Set<T> getNextGroup(Set<T> idle, Set<T> active) {
        // Get the size for the idle set on the first run only. This is representative of the number
        // of overall instance modifications this update will trigger.
        if (!totalModInstanceCount.isPresent()) {
            totalModInstanceCount = Optional.of(idle.size());
        }

        // Limit group size to the current size of the group minus the number of instances currently
        // being modified.
        return ordering.sortedCopy(doGetNextGroup(idle, active)).stream()
                .limit(Math.max(0, determineCurGroupSize(idle.size()) - active.size())).collect(Collectors.toSet());
    }

    /**
     * Return a list of instances to be updated.
     * Returns an empty list if the current active group has not completed.
     *
     * @param idle Idle instances, candidate for being updated.
     * @param active Instances currently being updated.
     * @return all idle instances to start updating.
     */
    Set<T> doGetNextGroup(Set<T> idle, Set<T> active) {
        return active.isEmpty() ? idle : ImmutableSet.of();
    }
}