Merge pull request square#601 from square/matthewdu/concurrent-java-g…

…eneration Add concurrent java generation.
Faire · Jun 17, 2016 · e151bfe · e151bfe
2 parents 3cb4da3 + f1ee6c0
commit e151bfe
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 24 deletions.
diff --git a/wire-compiler/src/main/java/com/squareup/wire/WireCompiler.java b/wire-compiler/src/main/java/com/squareup/wire/WireCompiler.java
@@ -35,6 +35,12 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Scanner;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 
 /**
  * Command line interface to the Wire Java generator.
@@ -88,6 +94,7 @@ public final class WireCompiler {
   public static final String NAMED_FILES_ONLY = "--named_files_only";
   public static final String ANDROID = "--android";
   public static final String COMPACT = "--compact";
+  public static final int MAX_WRITE_CONCURRENCY = 8;
 
   private static final String CODE_GENERATED_BY_WIRE =
       "Code generated by Wire protocol buffer compiler, do not edit.";
@@ -225,39 +232,73 @@ void compile() throws IOException {
         .withAndroid(emitAndroid)
         .withCompact(emitCompact);
 
+    ConcurrentLinkedQueue<Type> types = new ConcurrentLinkedQueue<>();
     for (ProtoFile protoFile : schema.protoFiles()) {
       // Check if we're skipping files not explicitly named.
       if (!sourceFileNames.isEmpty() && !sourceFileNames.contains(protoFile.location().path())) {
         if (namedFilesOnly || protoFile.location().path().equals(DESCRIPTOR_PROTO)) continue;
       }
+      types.addAll(protoFile.types());
+    }
 
-      for (Type type : protoFile.types()) {
-        TypeSpec typeSpec = javaGenerator.generateType(type);
-        ClassName javaTypeName = (ClassName) javaGenerator.typeName(type.type());
-        writeJavaFile(javaTypeName, typeSpec, type.location());
+    ExecutorService executor = Executors.newCachedThreadPool();
+    List<Future<Void>> futures = new ArrayList<>(MAX_WRITE_CONCURRENCY);
+    for (int i = 0; i < MAX_WRITE_CONCURRENCY; ++i) {
+      futures.add(i, executor.submit(new JavaFileWriter(javaGenerator, types)));
+    }
+    executor.shutdown();
+
+    try {
+      for (Future<Void> future : futures) {
+        future.get();
       }
+    } catch (ExecutionException e) {
+      throw new IOException(e.getMessage(), e);
+    } catch (InterruptedException e) {
+      throw new RuntimeException(e.getMessage(), e);
     }
   }
 
-  private void writeJavaFile(ClassName javaTypeName, TypeSpec typeSpec, Location location)
-      throws IOException {
-    JavaFile.Builder builder = JavaFile.builder(javaTypeName.packageName(), typeSpec)
-        .addFileComment("$L", CODE_GENERATED_BY_WIRE);
-    if (location != null) {
-      builder.addFileComment("\nSource file: $L", location.withoutBase());
+  class JavaFileWriter implements Callable<Void> {
+    private final JavaGenerator javaGenerator;
+    private final ConcurrentLinkedQueue<Type> queue;
+
+    public JavaFileWriter(JavaGenerator javaGenerator, ConcurrentLinkedQueue<Type> queue) {
+      this.javaGenerator = javaGenerator;
+      this.queue = queue;
     }
-    JavaFile javaFile = builder.build();
 
-    Path path = fs.getPath(javaOut);
-    log.artifact(path, javaFile);
+    @Override
+    public Void call() throws IOException {
+      while (true) {
+        Type type = queue.poll();
+        if (type == null) {
+          return null;
+        }
 
-    try {
-      if (!dryRun) {
-        javaFile.writeTo(path);
+        TypeSpec typeSpec = javaGenerator.generateType(type);
+        ClassName javaTypeName = (ClassName) javaGenerator.typeName(type.type());
+        Location location = type.location();
+
+        JavaFile.Builder builder = JavaFile.builder(javaTypeName.packageName(), typeSpec)
+            .addFileComment("$L", CODE_GENERATED_BY_WIRE);
+        if (location != null) {
+          builder.addFileComment("\nSource file: $L", location.withoutBase());
+        }
+        JavaFile javaFile = builder.build();
+
+        Path path = fs.getPath(javaOut);
+        log.artifact(path, javaFile);
+
+        try {
+          if (!dryRun) {
+            javaFile.writeTo(path);
+          }
+        } catch (IOException e) {
+          throw new IOException("Error emitting " + javaFile.packageName + "."
+              + javaFile.typeSpec.name + " to " + javaOut, e);
+        }
       }
-    } catch (IOException e) {
-      throw new IOException("Error emitting " + javaFile.packageName + "."
-          + javaFile.typeSpec.name + " to " + javaOut, e);
     }
   }
 }
diff --git a/wire-compiler/src/test/java/com/squareup/wire/StringWireLogger.java b/wire-compiler/src/test/java/com/squareup/wire/StringWireLogger.java
@@ -11,7 +11,7 @@ final class StringWireLogger implements WireLogger {
     this.quiet = quiet;
   }
 
-  @Override public void artifact(Path outputPath, JavaFile javaFile) {
+  @Override public synchronized void artifact(Path outputPath, JavaFile javaFile) {
     buffer.append(outputPath);
     buffer.append(" ");
     buffer.append(javaFile.packageName);
@@ -20,7 +20,7 @@ final class StringWireLogger implements WireLogger {
     buffer.append('\n');
   }
 
-  @Override public void info(String message) {
+  @Override public synchronized void info(String message) {
     if (!quiet) {
       buffer.append(message);
       buffer.append('\n');

diff --git a/wire-compiler/src/test/java/com/squareup/wire/WireCompilerTest.java b/wire-compiler/src/test/java/com/squareup/wire/WireCompilerTest.java
@@ -372,9 +372,9 @@ public class WireCompilerTest {
     };
     invokeCompiler(sources, "--includes=squareup.wire.protos.roots.TheService", "--dry_run", "--quiet");
 
-    assertThat(logger.getLog()).isEqualTo(""
-        + testDir.getAbsolutePath() + " com.squareup.wire.protos.roots.TheRequest\n"
-        + testDir.getAbsolutePath() + " com.squareup.wire.protos.roots.TheResponse\n");
+    assertThat(logger.getLog()).contains(
+        testDir.getAbsolutePath() + " com.squareup.wire.protos.roots.TheRequest\n",
+        testDir.getAbsolutePath() + " com.squareup.wire.protos.roots.TheResponse\n");
   }
 
   @Test public void noFiles() throws Exception {