[SPARK-53064][CORE] Rewrite MDC LogKey in Java

[pan3793]

What changes were proposed in this pull request?

This PR proposes to rewrite a few classes used by the Structured Logging API from Scala to Java.

Why are the changes needed?

Previously (before 3.5), modules under common were pure Java, and were easy to embed into other services, for example, the YARN External Shuffle Service will be run as a plugin of the YARN Resource Manager daemon process. With recent years' changes, some pure Java modules also require scala-library to be present at runtime, i.e., SPARK-52942 reports that YARN ESS causes YARN RM to fail to start due to missing scala-library in the classpath.

Instead of bundling scala-library into YARN ESS jar, #51650 (comment) suggests making it scala-free again.

This also makes Java's invocation of Structured Logging API much cleaner, now, it can be called without ugly $.MODULE$.

- MDC.of(LogKeys.HOST_PORT$.MODULE$, address);
+ MDC.of(LogKeys.HOST_PORT, address);

Does this PR introduce any user-facing change?

No, they are internal APIs, for those plugin developers who want to provide custom LogKeys, there still possible to make it compatible with both Spark 4.0 and the new API proposed by this PR, see https://github.com/pan3793/SPARK-53064

import org.apache.spark.internal.LogKey;

// CUSTOM_LOG_KEY is compatible with both Spark 4.0 and SPARK-53064
public class JavaCustomLogKeys {
  // Custom `LogKey` must be `implements LogKey`
  public static class CUSTOM_LOG_KEY implements LogKey {
    @Override
    public String name() {
      return "custom_log_key";
    }
  }

  // Singleton
  public static final CUSTOM_LOG_KEY CUSTOM_LOG_KEY = new CUSTOM_LOG_KEY();
}

How was this patch tested?

Pass GHA, and verified YARN ESS works without scala-library.

Was this patch authored or co-authored using generative AI tooling?

No.

[pan3793]

Since it's unable to write a class in Java that allows Scala to construct instances as same as Scala's case class, I create this method instead.

So, for logging in Scala

logInfo(log"Application ID is ${MDC(APP_ID, appId)}")

previously MDC(APP_ID, appId) calls case class's constructor, now it calls this method.

[pan3793]

cc @LuciferYang @gengliangwang @mridulm

[pan3793]

For reviewers, you can check only the first commit of this PR, the second one only modifies import to fix "unused import" compilation warnings

[pan3793]

this seems to be an exception previously

[gengliangwang]

@pan3793 Great work! LGTM overall. The changes are non-trivial, so I’ll take a closer look over the weekend.

[LuciferYang]

This is a great change, and overall it looks ok. I'll continue to review the details.

[gengliangwang]

So the customer log key from Spark 4.0 won't work anymore?
cc @panbingkun since you are the author of this one.

[pan3793]

I think so.

if we override the name(), then both 4.0 and this PR work. I updated to PR description to provide a demo project to help reviewers to verify that.

public class JavaCustomLogKeys {
  // Custom `LogKey` must be `implements LogKey`
  public static class CUSTOM_LOG_KEY implements LogKey {
+   @Override
+   public String name() {
+     return "custom_log_key";
+   }
  }

  // Singleton
  public static final CUSTOM_LOG_KEY CUSTOM_LOG_KEY = new CUSTOM_LOG_KEY();
}

---

BTW, the previous example have some issues on compilation

class JavaCustomLogKeys {
  // Custom `LogKey` must be `implements LogKey`
  public static class CUSTOM_LOG_KEY implements LogKey { }

  // Singleton
  public static final CUSTOM_LOG_KEY CUSTOM_LOG_KEY = new CUSTOM_LOG_KEY();
}

I can use javac to compile it, but scalac raises error

[INFO] --- scala:4.9.5:compile (scala-compile-first) @ SPARK-53064 ---
[INFO] Compiler bridge file: /Users/chengpan/.sbt/1.0/zinc/org.scala-sbt/org.scala-sbt-compiler-bridge_2.13-1.10.8-bin_2.13.16__61.0-1.10.8_20250303T164041.jar
[INFO] compiling 2 Scala sources and 1 Java source to /Users/chengpan/Projects/SPARK-53064/target/classes ...
[ERROR] /Users/chengpan/Projects/SPARK-53064/src/main/java/io/github/pan3793/JavaCustomLogKeys.java:10: CUSTOM_LOG_KEY is already defined as class CUSTOM_LOG_KEY
[ERROR] one error found

[gengliangwang]

Thanks for the works!

[pan3793]

Rebased to resolve conflict

[LuciferYang]

Here, we employed a centralized Java enum to define all MDC keys. As far as I know, this approach results in the one-time initialization of all instances during class loading, thereby prolonging the class loading time and increasing memory usage.

Previously, because Scala's case objects are lazily loaded, unused keys did not consume memory or require initialization time.

Therefore, we may need to carry out some follow-up work later by splitting up this enum to mitigate the issues I've mentioned above.

If I'm mistaken, please correct me.

[pan3793]

I see your concern, but I am quite optimistic about this. The current approach avoids creating a class for each LogKey, creating hundreds of simple enum instances should not introduce much overhead for the JVM. I write a simple test to show that.

master branch

$ build/mvn -pl common/utils -am clean install -DskipTests
$ cat Test.java
import org.apache.spark.internal.LogKeys;

public class Main {
  public static void main(String[] args) {
    LogKeys.APP_ID$.MODULE$.name();
  }
}
$ time java -cp common/utils/target/spark-common-utils_2.13-4.1.0-SNAPSHOT.jar:scala-library-2.13.16.jar Test.java
java -cp  Test.java  0.73s user 0.04s system 215% cpu 0.360 total
$ time java -cp common/utils/target/spark-common-utils_2.13-4.1.0-SNAPSHOT.jar:scala-library-2.13.16.jar Test.java
java -cp  Test.java  0.74s user 0.04s system 217% cpu 0.362 total
$ time java -cp common/utils/target/spark-common-utils_2.13-4.1.0-SNAPSHOT.jar:scala-library-2.13.16.jar Test.java
java -cp  Test.java  0.74s user 0.04s system 217% cpu 0.359 total

this PR

$ build/mvn -pl common/utils -am clean install -DskipTests
$ cat Test.java
import org.apache.spark.internal.LogKeys;

public class Main {
  public static void main(String[] args) {
    LogKeys.APP_ID.name();
  }
}
$ time java -cp common/utils/target/spark-common-utils_2.13-4.1.0-SNAPSHOT.jar:scala-library-2.13.16.jar Test.java
java -cp  Test.java  0.60s user 0.05s system 165% cpu 0.396 total
$ time java -cp common/utils/target/spark-common-utils_2.13-4.1.0-SNAPSHOT.jar:scala-library-2.13.16.jar Test.java
java -cp  Test.java  0.61s user 0.04s system 201% cpu 0.321 total
$ time java -cp common/utils/target/spark-common-utils_2.13-4.1.0-SNAPSHOT.jar:scala-library-2.13.16.jar Test.java
java -cp  Test.java  0.60s user 0.04s system 199% cpu 0.316 total

it's even faster if I remove scala-library from the classpath

$ time java -cp common/utils/target/spark-common-utils_2.13-4.1.0-SNAPSHOT.jar Test.java
java -cp  Test.java  0.55s user 0.03s system 207% cpu 0.278 total
$ time java -cp common/utils/target/spark-common-utils_2.13-4.1.0-SNAPSHOT.jar Test.java
java -cp  Test.java  0.55s user 0.03s system 200% cpu 0.289 total
$ time java -cp common/utils/target/spark-common-utils_2.13-4.1.0-SNAPSHOT.jar Test.java
java -cp  Test.java  0.54s user 0.04s system 195% cpu 0.294 total

[pan3793]

and yes, we can split LogKeys into multiple enumerations if we observe a real perf bottleneck in the future :)

[LuciferYang]

Are the modifications in these lines strongly related to the intent of this pull request (PR)?

[pan3793]

it's necessary to make Java's Structured Logging API scala-free, otherwise, the Java SparkLogger still requires to access Scala Logging at runtime.

[LuciferYang]

Merged into master. Thank you @pan3793 and @gengliangwang