TraceContext: Add generic method to start request tracing

At the moment there are 3 places where we start request tracing
(RestApiServlet, ReceiveCommits and SshCommand). In all these places we:

- generate a trace ID
- set a TRACE_ID tag with the generated trace ID
- enable force logging

Instead of repeating this code in several places offer one generic
method in TraceContext for this and use it from all places where we
start request tracing.

Change-Id: I73268c35123086e51967fd7b08f11a202df22d24
Signed-off-by: Edwin Kempin <ekempin@google.com>

java/com/google/gerrit/httpd/restapi/RestApiServlet.java

@@ -109,7 +109,6 @@ import com.google.gerrit.server.audit.ExtendedHttpAuditEvent;
 import com.google.gerrit.server.cache.PerThreadCache;
 import com.google.gerrit.server.config.GerritServerConfig;
 import com.google.gerrit.server.git.LockFailureException;
-import com.google.gerrit.server.logging.RequestId;
 import com.google.gerrit.server.logging.TraceContext;
 import com.google.gerrit.server.permissions.GlobalPermission;
 import com.google.gerrit.server.permissions.PermissionBackend;
@@ -1321,12 +1320,9 @@ public class RestApiServlet extends HttpServlet {
 
   private TraceContext enableTracing(HttpServletRequest req, HttpServletResponse res) {
     String v = req.getParameter(ParameterParser.TRACE_PARAMETER);
-    if (v != null && (v.isEmpty() || Boolean.parseBoolean(v))) {
-      RequestId traceId = new RequestId();
-      res.setHeader(X_GERRIT_TRACE, traceId.toString());
-      return TraceContext.open().forceLogging().addTag(RequestId.Type.TRACE_ID, traceId);
-    }
-    return TraceContext.DISABLED;
+    return TraceContext.newTrace(
+        v != null && (v.isEmpty() || Boolean.parseBoolean(v)),
+        (tagName, traceId) -> res.setHeader(X_GERRIT_TRACE, traceId.toString()));
   }
 
   private boolean isDelete(HttpServletRequest req) {

java/com/google/gerrit/server/git/receive/ReceiveCommits.java

@@ -577,14 +577,10 @@ class ReceiveCommits {
       Collection<ReceiveCommand> commands, MultiProgressMonitor progress) {
     parsePushOptions();
     try (TraceContext traceContext =
-        TraceContext.open()
-            .addTag(RequestId.Type.RECEIVE_ID, RequestId.forProject(project.getNameKey()))) {
-      if (tracePushOption.orElse(false)) {
-        RequestId traceId = new RequestId();
-        traceContext.forceLogging().addTag(RequestId.Type.TRACE_ID, traceId);
-        addMessage(RequestId.Type.TRACE_ID.name() + ": " + traceId);
-      }
-
+        TraceContext.newTrace(
+            tracePushOption.orElse(false),
+            (tagName, traceId) -> addMessage(tagName + ": " + traceId))) {
+      traceContext.addTag(RequestId.Type.RECEIVE_ID, RequestId.forProject(project.getNameKey()));
       try {
         if (!projectState.getProject().getState().permitsWrite()) {
           for (ReceiveCommand cmd : commands) {

java/com/google/gerrit/server/logging/TraceContext.java

@@ -20,12 +20,41 @@ import com.google.common.collect.HashBasedTable;
 import com.google.common.collect.Table;
 
 public class TraceContext implements AutoCloseable {
-  public static final TraceContext DISABLED = new TraceContext();
-
   public static TraceContext open() {
     return new TraceContext();
   }
 
+  /**
+   * Opens a new trace context for request tracing.
+   *
+   * <ul>
+   *   <li>sets a tag with a generated trace ID
+   *   <li>enables force logging
+   * </ul>
+   *
+   * <p>No-op if {@code trace} is {@code false}.
+   *
+   * @param trace whether tracing should be started
+   * @param traceIdConsumer consumer for the trace ID, should be used to return the generated trace
+   *     ID to the client, not invoked if {@code trace} is {@code false}
+   * @return the trace context
+   */
+  public static TraceContext newTrace(boolean trace, TraceIdConsumer traceIdConsumer) {
+    if (!trace) {
+      // Create an empty trace context.
+      return open();
+    }
+
+    RequestId traceId = new RequestId();
+    traceIdConsumer.accept(RequestId.Type.TRACE_ID.name(), traceId.toString());
+    return open().addTag(RequestId.Type.TRACE_ID, traceId).forceLogging();
+  }
+
+  @FunctionalInterface
+  public interface TraceIdConsumer {
+    void accept(String tagName, String traceId);
+  }
+
   // Table<TAG_NAME, TAG_VALUE, REMOVE_ON_CLOSE>
   private final Table<String, String, Boolean> tags = HashBasedTable.create();
 

java/com/google/gerrit/sshd/SshCommand.java

@@ -14,7 +14,6 @@
 
 package com.google.gerrit.sshd;
 
-import com.google.gerrit.server.logging.RequestId;
 import com.google.gerrit.server.logging.TraceContext;
 import java.io.IOException;
 import java.io.PrintWriter;
@@ -50,11 +49,7 @@ public abstract class SshCommand extends BaseCommand {
   protected abstract void run() throws UnloggedFailure, Failure, Exception;
 
   private TraceContext enableTracing() {
-    if (trace) {
-      RequestId traceId = new RequestId();
-      stderr.println(String.format("%s: %s", RequestId.Type.TRACE_ID, traceId));
-      return TraceContext.open().forceLogging().addTag(RequestId.Type.TRACE_ID, traceId);
-    }
-    return TraceContext.DISABLED;
+    return TraceContext.newTrace(
+        trace, (tagName, traceId) -> stderr.println(String.format("%s: %s", tagName, traceId)));
   }
 }

javatests/com/google/gerrit/server/logging/TraceContextTest.java

@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.gerrit.server.logging.TraceContext.TraceIdConsumer;
 import java.util.Map;
 import java.util.SortedMap;
 import java.util.SortedSet;
@@ -154,6 +155,29 @@ public class TraceContextTest {
     assertForceLogging(false);
   }
 
+  @Test
+  public void newTrace() {
+    TestTraceIdConsumer traceIdConsumer = new TestTraceIdConsumer();
+    try (TraceContext traceContext = TraceContext.newTrace(true, traceIdConsumer)) {
+      assertForceLogging(true);
+      assertThat(LoggingContext.getInstance().getTagsAsMap().keySet())
+          .containsExactly(RequestId.Type.TRACE_ID.name());
+    }
+    assertThat(traceIdConsumer.tagName).isEqualTo(RequestId.Type.TRACE_ID.name());
+    assertThat(traceIdConsumer.traceId).isNotNull();
+  }
+
+  @Test
+  public void newTraceDisabled() {
+    TestTraceIdConsumer traceIdConsumer = new TestTraceIdConsumer();
+    try (TraceContext traceContext = TraceContext.newTrace(false, traceIdConsumer)) {
+      assertForceLogging(false);
+      assertTags(ImmutableMap.of());
+    }
+    assertThat(traceIdConsumer.tagName).isNull();
+    assertThat(traceIdConsumer.traceId).isNull();
+  }
+
   private void assertTags(ImmutableMap<String, ImmutableSet<String>> expectedTagMap) {
     SortedMap<String, SortedSet<Object>> actualTagMap =
         LoggingContext.getInstance().getTags().asMap();
@@ -168,4 +192,15 @@ public class TraceContextTest {
     assertThat(LoggingContext.getInstance().shouldForceLogging(null, null, false))
         .isEqualTo(expected);
   }
+
+  private static class TestTraceIdConsumer implements TraceIdConsumer {
+    String tagName;
+    String traceId;
+
+    @Override
+    public void accept(String tagName, String traceId) {
+      this.tagName = tagName;
+      this.traceId = traceId;
+    }
+  }
 }