Make ServiceTalk protoc plugin extensible

servicetalk-grpc-protoc/build.gradle

@@ -78,6 +78,8 @@ protobuf {
   }
   generateProtoTasks {
     all().each { task ->
+      task.generateDescriptorSet = true
+      task.descriptorSetOptions.includeImports = true
       task.inputs
           .file(pluginJar)
           .withNormalizer(ClasspathNormalizer)
@@ -113,6 +115,7 @@ task testJavadoc(type: Javadoc) {
   dependsOn tasks.compileTestJava
 
   source = protobuf.generatedFilesBaseDir
+  exclude "**/*.desc"
   classpath = sourceSets.test.compileClasspath
   destinationDir = file("$buildDir/tmp/testjavadoc")
   options.addBooleanOption("Xdoclint:syntax", true)
@@ -122,3 +125,7 @@ task testJavadoc(type: Javadoc) {
 }
 
 test.finalizedBy(testJavadoc)
+
+test {
+  systemProperty 'generatedFilesBaseDir', protobuf.generatedFilesBaseDir
+}

servicetalk-grpc-protoc/src/main/java/io/servicetalk/grpc/protoc/FileDescriptor.java

@@ -66,6 +66,7 @@ final class FileDescriptor implements GenerationContext {
     private final String typeNameSuffix;
     private final List<TypeSpec.Builder> serviceClassBuilders;
     private final Set<String> reservedJavaTypeName = new HashSet<>();
+    private final Map<TypeSpec.Builder, ServiceDescriptorProto> protoForServiceBuilder;
 
     /**
      * A single protoc file for which we will be generating classes
@@ -98,12 +99,18 @@ final class FileDescriptor implements GenerationContext {
         reservedJavaTypeName.add(outerClassName);
 
         serviceClassBuilders = new ArrayList<>(protoFile.getServiceCount());
+        protoForServiceBuilder = new HashMap<>();
     }
 
     String protoFileName() {
         return protoFile.getName();
     }
 
+    @Nullable
+    String getProtoPackageName() {
+        return protoPackageName;
+    }
+
     List<ServiceDescriptorProto> protoServices() {
         return protoFile.getServiceList();
     }
@@ -120,10 +127,10 @@ DescriptorProtos.SourceCodeInfo sourceCodeInfo() {
         return protoFile.getSourceCodeInfo();
     }
 
-    private void addMessageTypes(final List<DescriptorProto> messageTypes,
-                                 final @Nullable String parentProtoScope,
-                                 final String parentJavaScope,
-                                 final Map<String, ClassName> messageTypesMap) {
+    private static void addMessageTypes(final List<DescriptorProto> messageTypes,
+                                        @Nullable final String parentProtoScope,
+                                        final String parentJavaScope,
+                                        final Map<String, ClassName> messageTypesMap) {
         messageTypes.forEach(t -> {
             final String protoTypeName = parentProtoScope != null ?
                     (parentProtoScope + '.' + t.getName()) : '.' + t.getName();
@@ -170,6 +177,7 @@ public TypeSpec.Builder newServiceClassBuilder(final ServiceDescriptorProto serv
         }
 
         serviceClassBuilders.add(builder);
+        protoForServiceBuilder.put(builder, serviceProto);
         return builder;
     }
 
@@ -196,7 +204,11 @@ void writeTo(final CodeGeneratorResponse.Builder responseBuilder) {
 
             insertSingleFileContent("// " + GENERATED_BY_COMMENT, fileName, responseBuilder);
             for (final TypeSpec.Builder builder : serviceClassBuilders) {
-                insertSingleFileContent(builder.addModifiers(STATIC).build().toString(), fileName, responseBuilder);
+                String content = addInsertionPoint(
+                    builder.addModifiers(STATIC).build().toString(),
+                    protoForServiceBuilder.get(builder).getName()
+                );
+                insertSingleFileContent(content, fileName, responseBuilder);
             }
             return;
         }
@@ -205,6 +217,7 @@ void writeTo(final CodeGeneratorResponse.Builder responseBuilder) {
         final String packageName = javaPackageName();
         for (final TypeSpec.Builder builder : serviceClassBuilders) {
             final TypeSpec serviceType = builder.build();
+            ServiceDescriptorProto serviceDescriptorProto = protoForServiceBuilder.get(builder);
             final File.Builder fileBuilder = File.newBuilder();
             fileBuilder.setName(calculateFileName(packageName, serviceType.name));
 
@@ -213,11 +226,21 @@ void writeTo(final CodeGeneratorResponse.Builder responseBuilder) {
                     .addFileComment(GENERATED_BY_COMMENT)
                     .build();
 
-            fileBuilder.setContent(javaFile.toString());
+            fileBuilder.setContent(addInsertionPoint(javaFile.toString(), serviceDescriptorProto.getName()));
             responseBuilder.addFile(fileBuilder.build());
         }
     }
 
+    private String addInsertionPoint(String content, String name) {
+        String fqn = protoPackageName != null ? protoPackageName + '.' + name : name;
+        content = content.replaceAll("class __" + fqn + " \\{\n *}", insertionPoint(fqn));
+        return content;
+    }
+
+    static String insertionPoint(final String fqn) {
+        return "// @@protoc_insertion_point(service_scope:" + fqn + ')';
+    }
+
     private static void insertSingleFileContent(final String content, String fileName,
                                                 final CodeGeneratorResponse.Builder responseBuilder) {
         final File.Builder fileBuilder = File.newBuilder();

servicetalk-grpc-protoc/src/main/java/io/servicetalk/grpc/protoc/Generator.java

@@ -242,11 +242,12 @@ private State(final ServiceDescriptorProto serviceProto, String name, int servic
     /**
      * Generate Service class for the provided proto service descriptor.
      *
+     *
      * @param serviceProto The service descriptor.
      * @param serviceIndex The index of the service within the current file (0 based).
      * @return The service class builder
      */
-    TypeSpec.Builder generate(final ServiceDescriptorProto serviceProto, final int serviceIndex) {
+    TypeSpec.Builder generate(FileDescriptor f, final ServiceDescriptorProto serviceProto, final int serviceIndex) {
         final String name = context.deconflictJavaTypeName(
                 sanitizeIdentifier(serviceProto.getName(), false) + Service);
         final State state = new State(serviceProto, name, serviceIndex);
@@ -265,10 +266,18 @@ TypeSpec.Builder generate(final ServiceDescriptorProto serviceProto, final int s
         addClientMetadata(state, serviceClassBuilder);
         addClientInterfaces(state, serviceClassBuilder);
         addClientFactory(state, serviceClassBuilder);
+        // this empty class is a placeholder and get replaced with insertion point comment
+        serviceClassBuilder.addType(TypeSpec.classBuilder("__" + serviceFQN(f, serviceProto)).build());
 
         return serviceClassBuilder;
     }
 
+    private String serviceFQN(FileDescriptor f, ServiceDescriptorProto serviceDescriptorProto) {
+        return f.getProtoPackageName() != null ?
+            f.getProtoPackageName() + "." + serviceDescriptorProto.getName() :
+            serviceDescriptorProto.getName();
+    }
+
     private TypeSpec.Builder addSerializationProviderInit(final State state,
                                                           final TypeSpec.Builder serviceClassBuilder) {
         final CodeBlock.Builder staticInitBlockBuilder = CodeBlock.builder()

servicetalk-grpc-protoc/src/main/java/io/servicetalk/grpc/protoc/Main.java

@@ -138,7 +138,7 @@ private static CodeGeneratorResponse generate(final CodeGeneratorRequest request
                 List<ServiceDescriptorProto> serviceDescriptorProtoList = f.protoServices();
                 for (int i = 0; i < serviceDescriptorProtoList.size(); ++i) {
                     ServiceDescriptorProto serviceDescriptor = serviceDescriptorProtoList.get(i);
-                    generator.generate(serviceDescriptor, i);
+                    generator.generate(f, serviceDescriptor, i);
                 }
                 f.writeTo(responseBuilder);
             }

servicetalk-grpc-protoc/src/test/java/io/servicetalk/grpc/protoc/InsertionPointTest.java

@@ -0,0 +1,107 @@
+/*
+ * Copyright © 2020 Apple Inc. and the ServiceTalk project authors
+ *
+ * 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 io.servicetalk.grpc.protoc;
+
+import com.google.protobuf.DescriptorProtos;
+import com.google.protobuf.compiler.PluginProtos.CodeGeneratorRequest;
+import com.google.protobuf.compiler.PluginProtos.CodeGeneratorResponse;
+import org.junit.jupiter.api.Test;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.File;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.PrintStream;
+import java.nio.file.Files;
+import java.util.List;
+
+import static io.servicetalk.grpc.protoc.FileDescriptor.insertionPoint;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+class InsertionPointTest {
+
+    @Test
+    void insertionPointExistsInMultiFiles() throws IOException {
+        List<CodeGeneratorResponse.File> files = generate("test_multi.proto");
+        assertEquals(2, files.size());
+
+        assertTrue(files.get(0).getContent().contains(insertionPoint("test.multi.Tester")));
+        assertTrue(files.get(1).getContent().contains(insertionPoint("test.multi.Tester2")));
+    }
+
+    @Test
+    void insertionPointExistsInSingleFile() throws IOException {
+        List<CodeGeneratorResponse.File> files = generate("test_single.proto");
+        assertEquals(3, files.size());
+        assertTrue(files.get(1).getContent().contains(insertionPoint("test.single.Greeter")));
+        assertTrue(files.get(2).getContent().contains(insertionPoint("test.single.Fareweller")));
+    }
+
+    private static List<CodeGeneratorResponse.File> generate(String file) throws IOException {
+        final CodeGeneratorRequest.Builder reqBuilder = CodeGeneratorRequest.newBuilder();
+        reqBuilder.addAllProtoFile(getFileDescriptorSet().getFileList());
+        reqBuilder.addFileToGenerate(file);
+        CodeGeneratorResponse codeGeneratorResponse = executePlugin(reqBuilder.build());
+        return codeGeneratorResponse.getFileList();
+    }
+
+    /**
+     * Executes protoc servicetalk plugin
+     * @param request generator request
+     * @return generator response
+     * @throws IOException throws exception if it fails to get hold of descriptor set
+     */
+    private static CodeGeneratorResponse executePlugin(CodeGeneratorRequest request) throws IOException {
+        // make stdin with request
+        ByteArrayOutputStream tempCollector = new ByteArrayOutputStream();
+        request.writeTo(tempCollector);
+        ByteArrayInputStream stdinWithRequest = new ByteArrayInputStream(tempCollector.toByteArray());
+
+        // hold stdout into this after plugin execution
+        ByteArrayOutputStream stdoutWithResponse = new ByteArrayOutputStream();
+
+        InputStream stdin = System.in;
+        PrintStream stdout = System.out;
+        try {
+            // prepare stdin and stdout for the plugin execution
+            System.setIn(stdinWithRequest);
+            System.setOut(new PrintStream(stdoutWithResponse));
+            // execute plugin
+            Main.main();
+            return CodeGeneratorResponse.parseFrom(stdoutWithResponse.toByteArray());
+        } finally {
+            // restore
+            System.setIn(stdin);
+            System.setOut(stdout);
+        }
+    }
+
+    /**
+     * This descriptor file is expected to be generated at build time by the protoc plugin
+     * @return descriptor fileset
+     * @throws IOException throws exception if it fails to locate descriptor set on file system
+     */
+    static DescriptorProtos.FileDescriptorSet getFileDescriptorSet() throws IOException {
+        String baseDir = System.getProperty("generatedFilesBaseDir",
+                "servicetalk-grpc-protoc/build/generated/sources/proto");
+        File descriptorSet = new File(baseDir + "/test/descriptor_set.desc");
+        assertTrue(descriptorSet.exists());
+        byte[] data = Files.readAllBytes(descriptorSet.toPath());
+        return DescriptorProtos.FileDescriptorSet.parseFrom(data);
+    }
+}