Splitter.java

// Copyright 2015 The Bazel Authors. All rights reserved.
//
// 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 com.google.devtools.build.android.ziputils;

import com.google.common.base.Preconditions;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

class Splitter {

  private static final String ARCHIVE_FILE_SEPARATOR = "/";

  private final int numberOfShards;
  private final Map<String, Integer> assigned;
  private int size = 0;
  private int shard = 0;
  private String currPath = null;
  private int remaining;
  private int idealSize;
  private int almostFull;

  /**
   * Creates a splitter for splitting an expected number of entries into
   * a given number of shards. The current shard is shard 0.
   */
  public Splitter(int numberOfShards, int expectedSize) {
    this.numberOfShards = numberOfShards;
    this.remaining = expectedSize;
    this.assigned = new HashMap<>();
    idealSize = remaining / (numberOfShards - shard);
    // Before you change this, please do the math.
    // It's not always perfect, but designed to keep shards reasonably balanced in most cases.
    int limit = Math.min(Math.min(10, (idealSize + 3) / 4), (int) Math.log(numberOfShards));
    almostFull = idealSize - limit;
  }

  /**
   * Forces mapping of the given entries to be that of the current shard.
   * The estimated number of remaining entries to process is adjusted,
   * by subtracting the number of as-of-yet unassigned entries from the
   * filter.
   */
  public void assign(Collection<String> filter) {
    if (filter != null) {
      for (String s : filter) {
        if (!assigned.containsKey(s)) {
          remaining--;
        }
        assigned.put(s, shard);
      }
      size = filter.size();
    }
  }

  /**
   * Forces increment of the current shard. May be called externally.
   * Typically right after calling {@link #assign(java.util.Collection)}.
   */
  public void nextShard() {
    if (shard < numberOfShards - 1) {
      shard++;
      size = 0;
      addEntries(0);
    }
  }

  /**
   * Adjusts the number of estimated entries to be process by the given count.
   */
  public void addEntries(int count) {
    this.remaining += count;
    idealSize = numberOfShards > shard ? remaining / (numberOfShards - shard) : remaining;
    // Before you change this, please do the math.
    // It's not always perfect, but designed to keep shards reasonably balanced in most cases.
    int limit = Math.min(Math.min(10, (idealSize + 3) / 4), (int) Math.log(numberOfShards));
    almostFull = idealSize -  limit;
  }

  /**
   * Assigns the given entry to an output file.
   */
  public int assign(String path) {
    Preconditions.checkState(shard < numberOfShards, "Too many shards!");
    Integer assignment = assigned.get(path);
    if (assignment != null) {
      return assignment;
    }
    remaining--;

    // last shard, no choice
    if (shard == numberOfShards - 1) {
      size++;
      assigned.put(currPath, shard);
      return shard;
    }

    // Forced split to try to avoid empty shards
    if (remaining < numberOfShards - shard - 1) {
      if (size > 0) {
        nextShard();
      }
      size++;
      assigned.put(currPath, shard);
      return shard;
    }

    String prevPath = currPath;
    currPath = path;
    // If current shard is at least "almost full", check for package boundary?
    if (prevPath != null && size >= almostFull) {
      int i = currPath.lastIndexOf(ARCHIVE_FILE_SEPARATOR);
      String dir = i > 0 ? currPath.substring(0, i) : ".";
      i = prevPath.lastIndexOf(ARCHIVE_FILE_SEPARATOR);
      String prevDir = i > 0 ? prevPath.substring(0, i) : ".";
      if (!dir.equals(prevDir)) {
        nextShard();
      }
    }
    assigned.put(currPath, shard);
    size++;
    return shard;
  }
}