bazel: Adopt proto_lang_toolchain for java_grpc_library() (#12994)

Author: AgraVator

compiler/BUILD.bazel

@@ -1,7 +1,8 @@
+load("@com_google_protobuf//bazel/toolchains:proto_lang_toolchain.bzl", "proto_lang_toolchain")
 load("@rules_cc//cc:defs.bzl", "cc_binary")
 load("@rules_java//java:defs.bzl", "java_library")
 load("@rules_jvm_external//:defs.bzl", "artifact")
-load("//:java_grpc_library.bzl", "java_rpc_toolchain")
+load("//:java_grpc_library.bzl", "java_rpc_java_config")
 
 # This should not generally be referenced. Users should use java_grpc_library
 cc_binary(
@@ -42,17 +43,32 @@ java_library(
     ],
 )
 
-java_rpc_toolchain(
+proto_lang_toolchain(
     name = "java_grpc_library_toolchain",
+    command_line = "--rpc-plugin_out=$(OUT)",
+    output_files = "single",
     plugin = "//compiler:grpc_java_plugin",
-    runtime = [":java_grpc_library_deps__do_not_reference"],
+    plugin_format_flag = "--plugin=protoc-gen-rpc-plugin=%s",
+    runtime = ":java_grpc_library_deps__do_not_reference",
     visibility = ["//visibility:public"],
 )
 
-java_rpc_toolchain(
+proto_lang_toolchain(
     name = "java_lite_grpc_library_toolchain",
+    command_line = "--rpc-plugin_out=lite:$(OUT)",
+    output_files = "single",
     plugin = "//compiler:grpc_java_plugin",
-    plugin_arg = "lite",
-    runtime = [":java_lite_grpc_library_deps__do_not_reference"],
+    plugin_format_flag = "--plugin=protoc-gen-rpc-plugin=%s",
+    runtime = ":java_lite_grpc_library_deps__do_not_reference",
+    visibility = ["//visibility:public"],
+)
+
+java_rpc_java_config(
+    name = "java_grpc_library_java_config",
+    visibility = ["//visibility:public"],
+)
+
+java_rpc_java_config(
+    name = "java_lite_grpc_library_java_config",
     visibility = ["//visibility:public"],
 )

java_grpc_library.bzl

@@ -1,123 +1,75 @@
 """Build rule for java_grpc_library."""
 
+load("@com_google_protobuf//bazel/common:proto_common.bzl", "proto_common")
 load("@com_google_protobuf//bazel/common:proto_info.bzl", "ProtoInfo")
+load("@com_google_protobuf//bazel/common:proto_lang_toolchain_info.bzl", "ProtoLangToolchainInfo")
 load("@rules_java//java:defs.bzl", "JavaInfo", "JavaPluginInfo", "java_common")
 
-_JavaRpcToolchainInfo = provider(
-    fields = [
-        "java_toolchain",
-        "java_plugins",
-        "plugin",
-        "plugin_arg",
-        "protoc",
-        "runtime",
-    ],
+# Building a java_grpc_library happens in two phases: protoc generates Java
+# sources, then those sources are compiled. `proto_lang_toolchain` configures
+# the first phase, but it is language-agnostic and so cannot express anything
+# about the second. This provider carries that Java-specific configuration.
+_JavaRpcJavaConfigInfo = provider(
+    doc = "Java compilation settings for the sources protoc generates.",
+    fields = {
+        "java_plugins": "(list[Target]) JavaPluginInfo targets to run as " +
+                        "annotation processors over the generated sources.",
+    },
 )
 
-def _java_rpc_toolchain_impl(ctx):
-    return [
-        _JavaRpcToolchainInfo(
-            java_toolchain = ctx.attr._java_toolchain,
-            java_plugins = ctx.attr.java_plugins,
-            plugin = ctx.attr.plugin,
-            plugin_arg = ctx.attr.plugin_arg,
-            protoc = ctx.attr._protoc,
-            runtime = ctx.attr.runtime,
-        ),
-        platform_common.ToolchainInfo(),  # Magic for b/78647825
-    ]
+def _java_rpc_java_config_impl(ctx):
+    return [_JavaRpcJavaConfigInfo(java_plugins = ctx.attr.java_plugins)]
 
-java_rpc_toolchain = rule(
+java_rpc_java_config = rule(
+    doc = """Java compilation settings to accompany a `proto_lang_toolchain`.
+
+A `java_grpc_library` runs protoc and then compiles the generated sources.
+`proto_lang_toolchain` describes the protoc invocation; this rule describes the
+Java compilation that follows it.""",
     attrs = {
-        # This attribute has a "magic" name recognized by the native DexArchiveAspect (b/78647825).
-        "runtime": attr.label_list(
-            cfg = "target",
-            providers = [JavaInfo],
-        ),
-        "plugin": attr.label(
-            cfg = "exec",
-            executable = True,
-        ),
-        "plugin_arg": attr.string(),
-        "_protoc": attr.label(
-            cfg = "exec",
-            default = Label("@com_google_protobuf//:protoc"),
-            executable = True,
-        ),
         "java_plugins": attr.label_list(
             default = [],
             providers = [JavaPluginInfo],
-        ),
-        "_java_toolchain": attr.label(
-            default = Label("@bazel_tools//tools/jdk:current_java_toolchain"),
+            doc = "Annotation processors to run over the generated sources.",
         ),
     },
-    provides = [
-        _JavaRpcToolchainInfo,
-        platform_common.ToolchainInfo,
-    ],
-    implementation = _java_rpc_toolchain_impl,
+    provides = [_JavaRpcJavaConfigInfo],
+    implementation = _java_rpc_java_config_impl,
 )
 
-# "repository" here is for Bazel builds that span multiple WORKSPACES.
-def _path_ignoring_repository(f):
-    # Bazel creates a _virtual_imports directory in case the .proto source files
-    # need to be accessed at a path that's different from their source path:
-    # https://github.com/bazelbuild/bazel/blob/0.27.1/src/main/java/com/google/devtools/build/lib/rules/proto/ProtoCommon.java#L289
-    #
-    # In that case, the import path of the .proto file is the path relative to
-    # the virtual imports directory of the rule in question.
-    virtual_imports = "/_virtual_imports/"
-    if virtual_imports in f.path:
-        return f.path.split(virtual_imports)[1].split("/", 1)[1]
-    elif len(f.owner.workspace_root) == 0:
-        # |f| is in the main repository
-        return f.short_path
-    else:
-        # If |f| is a generated file, it will have "bazel-out/*/genfiles" prefix
-        # before "external/workspace", so we need to add the starting index of "external/workspace"
-        return f.path[f.path.find(f.owner.workspace_root) + len(f.owner.workspace_root) + 1:]
-
 def _java_rpc_library_impl(ctx):
     if len(ctx.attr.srcs) != 1:
         fail("Exactly one src value supported", "srcs")
     if ctx.attr.srcs[0].label.package != ctx.label.package:
         print(("in srcs attribute of {0}: Proto source with label {1} should be in " +
                "same package as consuming rule").format(ctx.label, ctx.attr.srcs[0].label))
 
-    toolchain = ctx.attr._toolchain[_JavaRpcToolchainInfo]
-    srcs = ctx.attr.srcs[0][ProtoInfo].direct_sources
-    descriptor_set_in = ctx.attr.srcs[0][ProtoInfo].transitive_descriptor_sets
-
+    toolchain = ctx.attr._toolchain[ProtoLangToolchainInfo]
     srcjar = ctx.actions.declare_file("%s-proto-gensrc.jar" % ctx.label.name)
 
-    args = ctx.actions.args()
-    args.add(toolchain.plugin[DefaultInfo].files_to_run.executable, format = "--plugin=protoc-gen-rpc-plugin=%s")
-    args.add("--rpc-plugin_out={0}:{1}".format(toolchain.plugin_arg, srcjar.path))
-    args.add_joined("--descriptor_set_in", descriptor_set_in, join_with = ctx.configuration.host_path_separator)
-    args.add_all(srcs, map_each = _path_ignoring_repository)
-
-    ctx.actions.run(
-        inputs = depset(srcs, transitive = [descriptor_set_in, toolchain.plugin[DefaultInfo].files]),
-        outputs = [srcjar],
-        executable = toolchain.protoc[DefaultInfo].files_to_run,
-        arguments = [args],
-        use_default_shell_env = True,
-        toolchain = None,
+    proto_common.compile(
+        actions = ctx.actions,
+        proto_info = ctx.attr.srcs[0][ProtoInfo],
+        proto_lang_toolchain_info = toolchain,
+        generated_files = [srcjar],
     )
 
     deps_java_info = java_common.merge([dep[JavaInfo] for dep in ctx.attr.deps])
+    runtime_deps = [toolchain.runtime[JavaInfo]] if toolchain.runtime else []
 
     java_info = java_common.compile(
         ctx,
-        java_toolchain = toolchain.java_toolchain[java_common.JavaToolchainInfo],
+        java_toolchain = ctx.toolchains["@bazel_tools//tools/jdk:toolchain_type"].java,
         source_jars = [srcjar],
         output = ctx.outputs.jar,
         output_source_jar = ctx.outputs.srcjar,
-        plugins = [plugin[JavaPluginInfo] for plugin in toolchain.java_plugins],
+        plugins = [
+            plugin[JavaPluginInfo]
+            for plugin in ctx.attr._java_config[_JavaRpcJavaConfigInfo].java_plugins
+        ],
         deps = [
             java_common.make_non_strict(deps_java_info),
-        ] + [dep[JavaInfo] for dep in toolchain.runtime],
+        ] + runtime_deps,
     )
 
     return [java_info]
@@ -136,6 +88,11 @@ _java_grpc_library = rule(
         ),
         "_toolchain": attr.label(
             default = Label("//compiler:java_grpc_library_toolchain"),
+            providers = [ProtoLangToolchainInfo],
+        ),
+        "_java_config": attr.label(
+            default = Label("//compiler:java_grpc_library_java_config"),
+            providers = [_JavaRpcJavaConfigInfo],
         ),
     },
     toolchains = ["@bazel_tools//tools/jdk:toolchain_type"],
@@ -163,6 +120,11 @@ INTERNAL_java_grpc_library_for_xds = rule(
         ),
         "_toolchain": attr.label(
             default = Label("//xds:java_grpc_library_toolchain"),
+            providers = [ProtoLangToolchainInfo],
+        ),
+        "_java_config": attr.label(
+            default = Label("//xds:java_grpc_library_java_config"),
+            providers = [_JavaRpcJavaConfigInfo],
         ),
     },
     toolchains = ["@bazel_tools//tools/jdk:toolchain_type"],
@@ -190,6 +152,11 @@ _java_lite_grpc_library = rule(
         # This attribute has a "magic" name recognized by the native DexArchiveAspect (b/78647825).
         "_toolchain": attr.label(
             default = Label("//compiler:java_lite_grpc_library_toolchain"),
+            providers = [ProtoLangToolchainInfo],
+        ),
+        "_java_config": attr.label(
+            default = Label("//compiler:java_lite_grpc_library_java_config"),
+            providers = [_JavaRpcJavaConfigInfo],
         ),
     },
     toolchains = ["@bazel_tools//tools/jdk:toolchain_type"],

xds/BUILD.bazel

@@ -1,9 +1,10 @@
 load("@bazel_jar_jar//:jar_jar.bzl", "jar_jar")
 load("@com_google_protobuf//bazel:java_proto_library.bzl", "java_proto_library")
 load("@com_google_protobuf//bazel:proto_library.bzl", "proto_library")
+load("@com_google_protobuf//bazel/toolchains:proto_lang_toolchain.bzl", "proto_lang_toolchain")
 load("@rules_java//java:defs.bzl", "java_binary", "java_library", "java_test")
 load("@rules_jvm_external//:defs.bzl", "artifact")
-load("//:java_grpc_library.bzl", "INTERNAL_java_grpc_library_for_xds", "java_grpc_library", "java_rpc_toolchain")
+load("//:java_grpc_library.bzl", "INTERNAL_java_grpc_library_for_xds", "java_grpc_library", "java_rpc_java_config")
 
 # Mirrors the dependencies included in the artifact on Maven Central for usage
 # with maven_install's override_targets. Should only be used as a dep for
@@ -319,10 +320,17 @@ java_grpc_library(
     deps = [":xds_orca_java_proto"],
 )
 
-java_rpc_toolchain(
+proto_lang_toolchain(
     name = "java_grpc_library_toolchain",
+    command_line = "--rpc-plugin_out=$(OUT)",
+    output_files = "single",
     plugin = "//compiler:grpc_java_plugin",
-    runtime = [":java_grpc_library_deps"],
+    plugin_format_flag = "--plugin=protoc-gen-rpc-plugin=%s",
+    runtime = ":java_grpc_library_deps",
+)
+
+java_rpc_java_config(
+    name = "java_grpc_library_java_config",
 )
 
 java_library(