aboutsummaryrefslogtreecommitdiffhomepage
path: root/src/tools/android/java/com/google/devtools/build/android/ziputils/Splitter.java
blob: d476581109a6725071dba4995838892a5aac2753 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
// 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;
  }
}