Skip to content

Commit

Permalink
multi-thread DexFileMerger
Browse files Browse the repository at this point in the history 
RELNOTES: speedup of incremental dexing tools

PiperOrigin-RevId: 164926895
  • Loading branch information
@kevin1e100 @hlopko
kevin1e100 authored and hlopko committed Aug 11, 2017
1 parent 53b9bab commit 6628d55
Show file tree
Hide file tree
Showing 4 changed files with 192 additions and 82 deletions.
42 changes: 34 additions & 8 deletions src/test/java/com/google/devtools/build/android/dexer/DexFileAggregatorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
package com.google.devtools.build.android.dexer;

import static com.google.common.truth.Truth.assertThat;
import static com.google.common.util.concurrent.MoreExecutors.newDirectExecutorService;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.times;
Expand Down Expand Up @@ -60,24 +61,37 @@ public void setUp() throws IOException {
@Test
public void testClose_emptyWritesNothing() throws Exception {
DexFileAggregator dexer =
new DexFileAggregator(new DxContext(), dest, MultidexStrategy.MINIMAL, DEX_LIMIT, WASTE);
new DexFileAggregator(
new DxContext(),
dest,
newDirectExecutorService(),
MultidexStrategy.MINIMAL,
DEX_LIMIT,
WASTE);
dexer.close();
verify(dest, times(0)).addFile(any(ZipEntry.class), any(Dex.class));
}

@Test
public void testAddAndClose_singleInputWritesThatInput() throws Exception {
DexFileAggregator dexer =
new DexFileAggregator(new DxContext(), dest, MultidexStrategy.MINIMAL, 0, WASTE);
new DexFileAggregator(
new DxContext(), dest, newDirectExecutorService(), MultidexStrategy.MINIMAL, 0, WASTE);
dexer.add(dex);
dexer.close();
verify(dest).addFile(any(ZipEntry.class), eq(dex));
}

@Test
public void testMultidex_underLimitWritesOneShard() throws Exception {
DexFileAggregator dexer = new DexFileAggregator(
new DxContext(), dest, MultidexStrategy.BEST_EFFORT, DEX_LIMIT, WASTE);
DexFileAggregator dexer =
new DexFileAggregator(
new DxContext(),
dest,
newDirectExecutorService(),
MultidexStrategy.BEST_EFFORT,
DEX_LIMIT,
WASTE);
Dex dex2 = DexFiles.toDex(convertClass(ByteStreams.class));
dexer.add(dex);
dexer.add(dex2);
Expand All @@ -89,8 +103,14 @@ public void testMultidex_underLimitWritesOneShard() throws Exception {

@Test
public void testMultidex_overLimitWritesSecondShard() throws Exception {
DexFileAggregator dexer = new DexFileAggregator(new DxContext(), dest,
MultidexStrategy.BEST_EFFORT, 2 /* dex has more than 2 methods and fields */, WASTE);
DexFileAggregator dexer =
new DexFileAggregator(
new DxContext(),
dest,
newDirectExecutorService(),
MultidexStrategy.BEST_EFFORT,
2 /* dex has more than 2 methods and fields */,
WASTE);
Dex dex2 = DexFiles.toDex(convertClass(ByteStreams.class));
dexer.add(dex); // classFile is already over limit but we take anything in empty shard
dexer.add(dex2); // this should start a new shard
Expand All @@ -103,8 +123,14 @@ public void testMultidex_overLimitWritesSecondShard() throws Exception {

@Test
public void testMonodex_alwaysWritesSingleShard() throws Exception {
DexFileAggregator dexer = new DexFileAggregator(new DxContext(), dest, MultidexStrategy.OFF,
2 /* dex has more than 2 methods and fields */, WASTE);
DexFileAggregator dexer =
new DexFileAggregator(
new DxContext(),
dest,
newDirectExecutorService(),
MultidexStrategy.OFF,
2 /* dex has more than 2 methods and fields */,
WASTE);
Dex dex2 = DexFiles.toDex(convertClass(ByteStreams.class));
dexer.add(dex);
dexer.add(dex2);
Expand Down
112 changes: 95 additions & 17 deletions src/tools/android/java/com/google/devtools/build/android/dexer/DexFileAggregator.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
// limitations under the License.
package com.google.devtools.build.android.dexer;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;

import com.android.dex.Dex;
Expand All @@ -26,13 +27,19 @@
import com.android.dx.merge.DexMerger;
import com.google.auto.value.AutoValue;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import java.io.Closeable;
import java.io.IOException;
import java.nio.BufferOverflowException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.zip.ZipEntry;

/**
Expand All @@ -59,23 +66,28 @@ class DexFileAggregator implements Closeable {
private final int wasteThresholdPerDex;
private final MultidexStrategy multidex;
private final DxContext context;
private DexFileArchive dest;
private final ListeningExecutorService executor;
private final DexFileArchive dest;

private int nextDexFileIndex = 0;
private ListenableFuture<Void> lastWriter = Futures.<Void>immediateFuture(null);

public DexFileAggregator(
DxContext context,
DexFileArchive dest,
ListeningExecutorService executor,
MultidexStrategy multidex,
int maxNumberOfIdxPerDex,
int wasteThresholdPerDex) {
this.context = context;
this.dest = dest;
this.executor = executor;
this.multidex = multidex;
this.maxNumberOfIdxPerDex = maxNumberOfIdxPerDex;
this.wasteThresholdPerDex = wasteThresholdPerDex;
}

public DexFileAggregator add(Dex dexFile) throws IOException {
public DexFileAggregator add(Dex dexFile) {
if (multidex.isMultidexAllowed()) {
// To determine whether currentShard is "full" we track unique field and method signatures,
// which predicts precisely the number of field and method indices.
Expand Down Expand Up @@ -114,13 +126,20 @@ public void close() throws IOException {
if (!currentShard.isEmpty()) {
rotateDexFile();
}
// Wait for last shard to be written before closing underlying archive
lastWriter.get();
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
} catch (ExecutionException e) {
Throwables.throwIfInstanceOf(e.getCause(), IOException.class);
Throwables.throwIfUnchecked(e.getCause());
throw new AssertionError("Unexpected execution exception", e);
} finally {
dest.close();
dest = null;
}
}

public void flush() throws IOException {
public void flush() {
checkState(multidex.isMultidexAllowed());
if (!currentShard.isEmpty()) {
rotateDexFile();
Expand All @@ -131,16 +150,24 @@ public int getDexFilesWritten() {
return nextDexFileIndex;
}

private void rotateDexFile() throws IOException {
private void rotateDexFile() {
writeMergedFile(currentShard.toArray(/* apparently faster than pre-sized array */ new Dex[0]));
currentShard.clear();
fieldsInCurrentShard.clear();
methodsInCurrentShard.clear();
}

private void writeMergedFile(Dex... dexes) throws IOException {
Dex merged = merge(dexes);
dest.addFile(nextArchiveEntry(), merged);
private void writeMergedFile(Dex... dexes) {
checkArgument(0 < dexes.length);
checkState(multidex.isMultidexAllowed() || nextDexFileIndex == 0);
String filename = getDexFileName(nextDexFileIndex++);
ListenableFuture<Dex> merged =
dexes.length == 1
? Futures.immediateFuture(dexes[0])
: executor.submit(new RunDexMerger(dexes));
lastWriter =
Futures.whenAllSucceed(lastWriter, merged)
.call(new WriteFile(filename, merged, dest), executor);
}

private Dex merge(Dex... dexes) throws IOException {
Expand All @@ -155,6 +182,9 @@ private Dex merge(Dex... dexes) throws IOException {
dexMerger.setCompactWasteThreshold(wasteThresholdPerDex);
return dexMerger.merge();
} catch (BufferOverflowException e) {
if (dexes.length <= 2) {
throw e;
}
// Bug in dx can cause this for ~1500 or more classes
Dex[] left = Arrays.copyOf(dexes, dexes.length / 2);
Dex[] right = Arrays.copyOfRange(dexes, left.length, dexes.length);
Expand All @@ -169,19 +199,16 @@ private Dex merge(Dex... dexes) throws IOException {
}
}

private ZipEntry nextArchiveEntry() {
checkState(multidex.isMultidexAllowed() || nextDexFileIndex == 0);
ZipEntry result = new ZipEntry(getDexFileName(nextDexFileIndex++));
result.setTime(0L); // Use simple stable timestamps for deterministic output
return result;
}

// More or less copied from from com.android.dx.command.dexer.Main
@VisibleForTesting
static String getDexFileName(int i) {
return i == 0 ? DexFormat.DEX_IN_JAR_NAME : DEX_PREFIX + (i + 1) + DEX_EXTENSION;
}

private static String typeName(Dex dex, int typeIndex) {
return dex.typeNames().get(typeIndex);
}

@AutoValue
abstract static class FieldDescriptor {
static FieldDescriptor fromDex(Dex dex, int fieldIndex) {
Expand Down Expand Up @@ -220,7 +247,58 @@ static MethodDescriptor fromDex(Dex dex, int methodIndex) {
abstract String returnType();
}

private static String typeName(Dex dex, int typeIndex) {
return dex.typeNames().get(typeIndex);
private class RunDexMerger implements Callable<Dex> {

private final Dex[] dexes;

public RunDexMerger(Dex... dexes) {
checkArgument(dexes.length >= 2, "Only got %s dex files to merge", dexes.length);
this.dexes = dexes;
}

@Override
public Dex call() throws IOException {
try {
return merge(dexes);
} catch (Throwable t) {
// Print out exceptions so they don't get swallowed completely
t.printStackTrace();
Throwables.throwIfInstanceOf(t, IOException.class);
Throwables.throwIfUnchecked(t);
throw new AssertionError(t); // shouldn't get here
}
}
}

private static class WriteFile implements Callable<Void> {

private final ListenableFuture<Dex> dex;
private final String filename;
private final DexFileArchive dest;

public WriteFile(String filename, ListenableFuture<Dex> dex, DexFileArchive dest) {
this.filename = filename;
this.dex = dex;
this.dest = dest;
}

@Override
public Void call() throws Exception {
try {
checkState(dex.isDone());
ZipEntry entry = new ZipEntry(filename);
entry.setTime(0L); // Use simple stable timestamps for deterministic output
dest.addFile(entry, dex.get());
return null;
} catch (Exception e) {
// Print out exceptions so they don't get swallowed completely
e.printStackTrace();
throw e;
} catch (Throwable t) {
t.printStackTrace();
Throwables.throwIfUnchecked(t);
throw new AssertionError(t); // shouldn't get here
}
}
}
}
25 changes: 13 additions & 12 deletions src/tools/android/java/com/google/devtools/build/android/dexer/DexFileArchive.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,49 +13,50 @@
// limitations under the License.
package com.google.devtools.build.android.dexer;

import static com.google.common.base.Preconditions.checkState;

import com.android.dex.Dex;
import com.google.common.io.ByteStreams;
import java.io.Closeable;
import java.io.IOException;
import java.io.InputStream;
import java.util.concurrent.atomic.AtomicReference;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;

/**
* Wrapper around a {@link ZipOutputStream} to simplify writing archives with {@code .dex} files.
* Adding files generally requires a {@link ZipEntry} in order to control timestamps.
*/
// TODO(kmb): Remove this class and inline into DexFileAggregator
class DexFileArchive implements Closeable {

private final ZipOutputStream out;

public DexFileArchive(ZipOutputStream out) {
this.out = out;
}

/**
* Copies the content of the given {@link InputStream} into an entry with the given details.
* Used to ensure writes from different threads are sequenced, which {@link DexFileAggregator}
* ensures by making the writer futures wait on each oter.
*/
public DexFileArchive copy(ZipEntry entry, InputStream in) throws IOException {
out.putNextEntry(entry);
ByteStreams.copy(in, out);
out.closeEntry();
return this;
private final AtomicReference<ZipEntry> inUse = new AtomicReference<>(null);

public DexFileArchive(ZipOutputStream out) {
this.out = out;
}

/**
* Adds a {@code .dex} file with the given details.
*/
public DexFileArchive addFile(ZipEntry entry, Dex dex) throws IOException {
checkState(inUse.compareAndSet(null, entry), "Already in use");
entry.setSize(dex.getLength());
out.putNextEntry(entry);
dex.writeTo(out);
out.closeEntry();
checkState(inUse.compareAndSet(entry, null), "Swooped in: ", inUse.get());
return this;
}

@Override
public void close() throws IOException {
checkState(inUse.get() == null, "Still in use: ", inUse.get());
out.close();
}
}
Loading

0 comments on commit 6628d55

Please sign in to comment.