Make logging work correctly on Jetty (#2442)

util/src/main/java/google/registry/util/GcpJsonFormatter.java

@@ -14,7 +14,6 @@
 
 package google.registry.util;
 
-import com.google.common.collect.ImmutableMap;
 import com.google.common.flogger.LogSite;
 import com.google.common.flogger.backend.LogData;
 import com.google.common.flogger.backend.system.SimpleLogRecord;
@@ -52,6 +51,9 @@ public class GcpJsonFormatter extends Formatter {
    */
   private static final String SOURCE_LOCATION = "logging.googleapis.com/sourceLocation";
 
+  /** JSON field that stores the trace associated with the log entry, if any. */
+  private static final String TRACE = "logging.googleapis.com/trace";
+
   /** JSON field that contains the content, this will show up as the main entry in a log. */
   private static final String MESSAGE = "message";
 
@@ -63,6 +65,30 @@ public class GcpJsonFormatter extends Formatter {
 
   private static final Gson gson = new Gson();
 
+  private static final ThreadLocal<String> traceId = new ThreadLocal<>();
+
+  /**
+   * Set the Trace ID associated with any logging done by the current thread.
+   *
+   * @param id The traceID, in the form projects/[PROJECT_ID]/traces/[TRACE_ID]
+   */
+  public static void setCurrentTraceId(String id) {
+    if (id == null) {
+      traceId.remove();
+    } else {
+      traceId.set(id);
+    }
+  }
+
+  /**
+   * Get the Trace ID associated with any logging done by the current thread.
+   *
+   * @return id The traceID
+   */
+  public static String getCurrentTraceId() {
+    return traceId.get();
+  }
+
   @Override
   public String format(LogRecord record) {
     // Add an extra newline before the message for better displaying of multi-line logs. To see the
@@ -70,7 +96,6 @@ public class GcpJsonFormatter extends Formatter {
     // newline makes sure that the entire message starts on its own line, so that indentation within
     // the message is correct.
 
-    String message = '\n' + record.getMessage();
     String severity = severityFor(record.getLevel());
 
     // The rest is mostly lifted from java.util.logging.SimpleFormatter.
@@ -83,6 +108,7 @@ public class GcpJsonFormatter extends Formatter {
       }
       stacktrace = sw.toString();
     }
+    String message = '\n' + record.getMessage() + stacktrace;
 
     String function = "";
     if (record.getSourceClassName() != null) {
@@ -113,19 +139,18 @@ public class GcpJsonFormatter extends Formatter {
     if (!function.isEmpty()) {
       sourceLocation.put(FUNCTION, function);
     }
-    return gson.toJson(
-            ImmutableMap.of(
-                SEVERITY,
-                severity,
-                SOURCE_LOCATION,
-                sourceLocation,
-                // ImmutableMap.of(FILE, file, LINE, line, FUNCTION, function),
-                MESSAGE,
-                message + stacktrace))
-        // This trailing newline is required for the proxy because otherwise multiple logs might be
-        // sent to Stackdriver together (due to the async nature of the proxy), and not parsed
-        // correctly.
-        + '\n';
+
+    Map<String, Object> json = new LinkedHashMap<>();
+    json.put(SEVERITY, severity);
+    json.put(SOURCE_LOCATION, sourceLocation);
+    json.put(MESSAGE, message);
+    if (traceId.get() != null) {
+      json.put(TRACE, traceId.get());
+    }
+    // This trailing newline is required for the proxy because otherwise multiple logs might be
+    // sent to Stackdriver together (due to the async nature of the proxy), and not parsed
+    // correctly.
+    return gson.toJson(json) + '\n';
   }
 
   /**

util/src/test/java/google/registry/util/GcpJsonFormatterTest.java

@@ -69,6 +69,7 @@ class GcpJsonFormatterTest {
   @AfterEach
   void afterEach() {
     jdkLogger.removeHandler(handler);
+    GcpJsonFormatter.setCurrentTraceId(null);
   }
 
   @Test
@@ -76,7 +77,21 @@ class GcpJsonFormatterTest {
     logger.atInfo().log("Something I have to say");
     handler.close();
     String output = ostream.toString(StandardCharsets.US_ASCII);
-    assertThat(output).isEqualTo(makeJson("INFO", 76, "testSuccess", "Something I have to say"));
+    assertThat(output).isEqualTo(makeJson("INFO", 77, "testSuccess", "Something I have to say"));
+  }
+
+  @Test
+  void testSuccess_traceId() {
+    GcpJsonFormatter.setCurrentTraceId("trace_id");
+    logger.atInfo().log("Something I have to say");
+    handler.close();
+    String output = ostream.toString(StandardCharsets.US_ASCII);
+    String expected = makeJson("INFO", 86, "testSuccess_traceId", "Something I have to say");
+    // Remove the last two characters (}, \n) from the template and add the trace ID.
+    expected =
+        expected.substring(0, expected.length() - 2)
+            + ",\"logging.googleapis.com/trace\":\"trace_id\"}\n";
+    assertThat(output).isEqualTo(expected);
   }
 
   @Test
@@ -85,7 +100,7 @@ class GcpJsonFormatterTest {
     handler.close();
     String output = ostream.toString(StandardCharsets.US_ASCII);
     assertThat(output)
-        .isEqualTo(makeJson("ERROR", 84, "testSuccess_logLevel", "Something went terribly wrong"));
+        .isEqualTo(makeJson("ERROR", 99, "testSuccess_logLevel", "Something went terribly wrong"));
   }
 
   @Test
@@ -96,7 +111,7 @@ class GcpJsonFormatterTest {
     String prefix =
         makeJson(
             "ERROR",
-            93,
+            108,
             "testSuccess_withCause",
             "Something went terribly wrong\\njava.lang.RuntimeException: boom!");
     // Remove the last three characters (", }, \n) from the template as the actual output contains
@@ -113,7 +128,7 @@ class GcpJsonFormatterTest {
     String prefix =
         makeJson(
             "ERROR",
-            110,
+            125,
             "testSuccess_withStackTrace",
             "Something is worth checking\\ncom.google.common.flogger.LogSiteStackTrace: FULL");
     // Remove the last three characters (", }, \n) from the template as the actual output contains