[Java] Shade some widely used dependencies in bazel_jar_jar rule. (#21237)

These dependencies are widely used:
- com.google.common
- com.google.protobuf
- com.google.thirdparty

So that we need to shade them to avoid being conflict with jars introduced by user.

In this PR, we introduce a `bazel_jar_jar` rule for doing these and also shade them in maven pom files.

bazel/ray_deps_build_all.bzl

@@ -6,6 +6,7 @@ load("@com_github_jupp0r_prometheus_cpp//bazel:repositories.bzl", "prometheus_cp
 load("@com_github_grpc_grpc//third_party/py:python_configure.bzl", "python_configure")
 load("@com_github_grpc_grpc//bazel:grpc_deps.bzl", "grpc_deps")
 load("@rules_proto_grpc//:repositories.bzl", "rules_proto_grpc_toolchains")
+load("@com_github_johnynek_bazel_jar_jar//:jar_jar.bzl", "jar_jar_repositories")
 
 
 def ray_deps_build_all():
@@ -17,3 +18,4 @@ def ray_deps_build_all():
   python_configure(name = "local_config_python")
   grpc_deps()
   rules_proto_grpc_toolchains()
+  jar_jar_repositories()

bazel/ray_deps_setup.bzl

@@ -193,6 +193,12 @@ def ray_deps_setup():
         sha256 = "d21e155ac9a455831f81608bb06620e4a1d75012a630faf11f4c25ad10cfc9bb",
     )
 
+    auto_http_archive(
+        name = "com_github_johnynek_bazel_jar_jar",
+        url = "https://github.com/johnynek/bazel_jar_jar/archive/171f268569384c57c19474b04aebe574d85fde0d.tar.gz",
+        sha256 = "97c5f862482a05f385bd8f9d28a9bbf684b0cf3fae93112ee96f3fb04d34b193",
+    )
+
     auto_http_archive(
         name = "io_opencensus_cpp",
         url = "https://github.com/census-instrumentation/opencensus-cpp/archive/b14a5c0dcc2da8a7fc438fab637845c73438b703.zip",

java/BUILD.bazel

@@ -2,6 +2,10 @@ load("//bazel:ray.bzl", "define_java_module")
 load("//bazel:ray.bzl", "native_java_binary")
 load("//bazel:ray.bzl", "native_java_library")
 load("@rules_proto_grpc//java:defs.bzl", "java_proto_compile")
+load(
+    "@com_github_johnynek_bazel_jar_jar//:jar_jar.bzl",
+    "jar_jar",
+)
 
 exports_files([
     "testng.xml",
@@ -311,10 +315,23 @@ java_binary(
     ],
 )
 
+jar_jar(
+    name = "ray_dist_shaded",
+    input_jar = "//java:ray_dist_deploy.jar",
+    rules = "//java:shade_rule",
+)
+
+# Shade dependencies in tests fat jar.
+jar_jar(
+    name = "all_tests_shaded",
+    input_jar = "//java:all_tests_deploy.jar",
+    rules = "//java:shade_rule",
+)
+
 genrule(
     name = "ray_java_pkg",
     srcs = [
-        "//java:ray_dist_deploy.jar",
+        "//java:ray_dist_shaded.jar",
         "//java:gen_maven_deps",
         "//streaming/java:gen_maven_deps",
     ],
@@ -322,7 +339,7 @@ genrule(
     cmd = """
         WORK_DIR="$$(pwd)"
         rm -rf "$$WORK_DIR/python/ray/jars" && mkdir -p "$$WORK_DIR/python/ray/jars"
-        cp -f $(location //java:ray_dist_deploy.jar) "$$WORK_DIR/python/ray/jars/ray_dist.jar"
+        cp -f $(location //java:ray_dist_shaded.jar) "$$WORK_DIR/python/ray/jars/ray_dist.jar"
         date > $@
     """,
     local = 1,

java/runtime/pom_template.xml

@@ -80,6 +80,53 @@
           <outputDirectory>${output.directory}</outputDirectory>
         </configuration>
       </plugin>
+
+      <plugin>
+        <groupId>org.apache.maven.plugins</groupId>
+          <artifactId>maven-shade-plugin</artifactId>
+          <version>3.1.1</version>
+          <executions>
+            <execution>
+                    <phase>package</phase>
+                    <goals>
+                        <goal>shade</goal>
+                    </goals>
+                    <configuration>
+                        <minimizeJar>false</minimizeJar>
+                        <artifactSet>
+                            <includes>
+                                <include>com.google.guava</include>
+                                <include>com.google.protobuf</include>
+                            </includes>
+                        </artifactSet>
+                        <relocations>
+                            <relocation>
+                                <pattern>com.google.common</pattern>
+                                <shadedPattern>io.ray.shaded.com.google.common</shadedPattern>
+                            </relocation>
+                            <relocation>
+                                <pattern>com.google.protobuf</pattern>
+                                <shadedPattern>io.ray.shaded.com.google.protobuf</shadedPattern>
+                            </relocation>
+                            <relocation>
+                                <pattern>com.google.thirdparty</pattern>
+                                <shadedPattern>io.ray.shaded.com.google.thirdparty</shadedPattern>
+                            </relocation>
+                        </relocations>
+                        <filters>
+                            <filter>
+                                <artifact>*:*</artifact>
+                                <excludes>
+                                    <exclude>META-INF/*.SF</exclude>
+                                    <exclude>META-INF/*.DSA</exclude>
+                                    <exclude>META-INF/*.RSA</exclude>
+                                </excludes>
+                            </filter>
+                        </filters>
+                    </configuration>
+                </execution>
+            </executions>
+        </plugin>
     </plugins>
   </build>
 

java/shade_rule

@@ -0,0 +1,3 @@
+rule com.google.common.** io.ray.shaded.com.google.common.@1
+rule com.google.protobuf.** io.ray.shaded.com.google.protobuf.@1
+rule com.google.thirdparty.** io.ray.shaded.com.google.thirdparty.@1

java/test.sh

@@ -64,7 +64,7 @@ echo "Build java maven deps."
 bazel build //java:gen_maven_deps
 
 echo "Build test jar."
-bazel build //java:all_tests_deploy.jar
+bazel build //java:all_tests_shaded.jar
 
 java/generate_jni_header_files.sh
 
@@ -91,7 +91,7 @@ while true; do
   # TODO(hchen): Ideally, we should use the following bazel command to run Java tests. However, if there're skipped tests,
   # TestNG will exit with code 2. And bazel treats it as test failure.
   # bazel test //java:all_tests --config=ci || cluster_exit_code=$?
-  run_testng java -cp "$ROOT_DIR"/../bazel-bin/java/all_tests_deploy.jar org.testng.TestNG -d /tmp/ray_java_test_output "$ROOT_DIR"/testng.xml
+  run_testng java -cp "$ROOT_DIR"/../bazel-bin/java/all_tests_shaded.jar org.testng.TestNG -d /tmp/ray_java_test_output "$ROOT_DIR"/testng.xml
 
   echo Finished cluster mode test round $round
   date
@@ -103,7 +103,7 @@ done
 
 echo "Running tests under single-process mode."
 # bazel test //java:all_tests --jvmopt="-Dray.run-mode=SINGLE_PROCESS" --config=ci || single_exit_code=$?
-run_testng java -Dray.run-mode="SINGLE_PROCESS" -cp "$ROOT_DIR"/../bazel-bin/java/all_tests_deploy.jar org.testng.TestNG -d /tmp/ray_java_test_output "$ROOT_DIR"/testng.xml
+run_testng java -Dray.run-mode="SINGLE_PROCESS" -cp "$ROOT_DIR"/../bazel-bin/java/all_tests_shaded.jar org.testng.TestNG -d /tmp/ray_java_test_output "$ROOT_DIR"/testng.xml
 
 echo "Running connecting existing cluster tests."
 case "${OSTYPE}" in
@@ -112,8 +112,8 @@ case "${OSTYPE}" in
   *) echo "Can't get ip address for ${OSTYPE}"; exit 1;;
 esac
 RAY_BACKEND_LOG_LEVEL=debug ray start --head --port=6379 --redis-password=123456 --node-ip-address="$ip"
-RAY_BACKEND_LOG_LEVEL=debug java -cp bazel-bin/java/all_tests_deploy.jar -Dray.address="$ip:6379"\
- -Dray.redis.password='123456' -Dray.job.code-search-path="$PWD/bazel-bin/java/all_tests_deploy.jar" io.ray.test.MultiDriverTest
+RAY_BACKEND_LOG_LEVEL=debug java -cp bazel-bin/java/all_tests_shaded.jar -Dray.address="$ip:6379"\
+ -Dray.redis.password='123456' -Dray.job.code-search-path="$PWD/bazel-bin/java/all_tests_shaded.jar" io.ray.test.MultiDriverTest
 ray stop
 
 echo "Running documentation demo code."
@@ -122,7 +122,7 @@ for file in "$docdemo_path"*.java; do
   file=${file#"$docdemo_path"}
   class=${file%".java"}
   echo "Running $class"
-  java -cp bazel-bin/java/all_tests_deploy.jar -Dray.job.num-java-workers-per-process=1\
+  java -cp bazel-bin/java/all_tests_shaded.jar -Dray.job.num-java-workers-per-process=1\
    -Dray.raylet.startup-token=0 "io.ray.docdemo.$class"
 done
 popd
@@ -136,7 +136,7 @@ popd
 
 pushd "$ROOT_DIR"
 echo "Running performance test."
-run_timeout 60 java -cp "$ROOT_DIR"/../bazel-bin/java/all_tests_deploy.jar io.ray.performancetest.test.ActorPerformanceTestCase1
+run_timeout 60 java -cp "$ROOT_DIR"/../bazel-bin/java/all_tests_shaded.jar io.ray.performancetest.test.ActorPerformanceTestCase1
 # The performance process may be killed by run_timeout, so clear ray here.
 ray stop
 popd

src/ray/core_worker/lib/java/jni_init.cc

@@ -146,6 +146,20 @@ JavaVM *jvm;
 
 inline jclass LoadClass(JNIEnv *env, const char *class_name) {
   jclass tempLocalClassRef = env->FindClass(class_name);
+   if (tempLocalClassRef == nullptr) {
+    const std::string shaded_class_prefix = "io/ray/shaded/";
+    const auto class_name_str = std::string(class_name);
+    const auto this_prefix = class_name_str.substr(0, shaded_class_prefix.size());
+    if (this_prefix == shaded_class_prefix) {
+      // This is a shaded class, and try to load the original class.
+      env->ExceptionClear();
+      auto no_shaded_class_name =
+          class_name_str.substr(shaded_class_prefix.size(), class_name_str.size());
+      tempLocalClassRef = env->FindClass(no_shaded_class_name.c_str());
+    }
+  }
+  RAY_CHECK(tempLocalClassRef) << "Can't load Java class " << class_name;
+
   jclass ret = (jclass)env->NewGlobalRef(tempLocalClassRef);
   RAY_CHECK(ret) << "Can't load Java class " << class_name;
   env->DeleteLocalRef(tempLocalClassRef);
@@ -236,7 +250,8 @@ jint JNI_OnLoad(JavaVM *vm, void *reserved) {
   java_base_id_get_bytes = env->GetMethodID(java_base_id_class, "getBytes", "()[B");
 
   java_abstract_message_lite_class =
-      LoadClass(env, "com/google/protobuf/AbstractMessage");
+      LoadClass(env, "io/ray/shaded/com/google/protobuf/AbstractMessage");
+
   java_abstract_message_lite_to_byte_array =
       env->GetMethodID(java_abstract_message_lite_class, "toByteArray", "()[B");