Support request tracing for REST calls by setting a header in the request

At the moment request tracing for REST calls is enabled by setting the
'trace' or 'trace=<trace-id>' request parameter. This is good for manual
use since adding the request parameter to the URL is very easy and can
be done quickly.

For providing copy-pastable examples for how to do tracing this is not
ideal, since it depends on the concrete URL how the request parameter
would need to be set, e.g. it can be that '?trace' or '&trace' needs to
be appended depending on whether the URL already contains request
parameters or not.

With this change we now support enabling request tracing for REST calls
also by setting a 'X-Gerrit-Trace' header in the request. For manual use
this is less easy but it makes providing copy-pastable examples for how
to do tracing easier as one can now do:

  curl -D /tmp/gerrit -H X-Gerrit-Trace URL
  grep X-Gerrit-Trace /tmp/gerrit

Change-Id: I793ca9fff83ef23f5720390931599a9a85e868c7
Signed-off-by: Edwin Kempin <ekempin@google.com>
diff --git a/Documentation/rest-api.txt b/Documentation/rest-api.txt
index 97cca38..8f6a47b 100644
--- a/Documentation/rest-api.txt
+++ b/Documentation/rest-api.txt
@@ -210,6 +210,15 @@
   GET /changes/myProject~master~I8473b95934b5732ac55d26311a706c9c2bde9940/suggest_reviewers?trace&q=J
 ----
 
+Alternatively request tracing can also be enabled by setting the
+`X-Gerrit-Trace` header:
+
+.Example Request
+----
+  GET /changes/myProject~master~I8473b95934b5732ac55d26311a706c9c2bde9940/suggest_reviewers?q=J
+  X-Gerrit-Trace: issue/123
+----
+
 Enabling tracing results in additional logs with debug information that
 are written to the `error_log`. All logs that correspond to the traced
 request are associated with the trace ID. The trace ID is returned with
diff --git a/Documentation/user-request-tracing.txt b/Documentation/user-request-tracing.txt
index ac09d3c..8430e97 100644
--- a/Documentation/user-request-tracing.txt
+++ b/Documentation/user-request-tracing.txt
@@ -11,10 +11,10 @@
 request type:
 
 * REST API: For REST calls tracing can be enabled by setting the
-  `trace` request parameter, the trace ID is returned as
-  `X-Gerrit-Trace` header. More information about this can be found in
-  the link:rest-api.html#tracing[Request Tracing] section of the
-  link:rest-api.html[REST API documentation].
+  `trace` request parameter or the `X-Gerrit-Trace` header, the trace
+  ID is returned as `X-Gerrit-Trace` header. More information about
+  this can be found in the link:rest-api.html#tracing[Request Tracing]
+  section of the link:rest-api.html[REST API documentation].
 * SSH API: For SSH calls tracing can be enabled by setting the
   `--trace` option. More information about this can be found in
   the link:cmd-index.html#trace[Trace] section of the
diff --git a/java/com/google/gerrit/acceptance/HttpResponse.java b/java/com/google/gerrit/acceptance/HttpResponse.java
index 6c03793..3e98d71 100644
--- a/java/com/google/gerrit/acceptance/HttpResponse.java
+++ b/java/com/google/gerrit/acceptance/HttpResponse.java
@@ -14,13 +14,16 @@
 
 package com.google.gerrit.acceptance;
 
+import static com.google.common.collect.ImmutableList.toImmutableList;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
 import java.io.IOException;
 import java.io.InputStreamReader;
 import java.io.Reader;
 import java.nio.ByteBuffer;
+import java.util.Arrays;
 import org.apache.http.Header;
 import org.eclipse.jgit.util.IO;
 import org.eclipse.jgit.util.RawParseUtils;
@@ -61,6 +64,13 @@
     return hdr != null ? hdr.getValue() : null;
   }
 
+  public ImmutableList<String> getHeaders(String name) {
+    return Arrays.asList(response.getHeaders(name))
+        .stream()
+        .map(Header::getValue)
+        .collect(toImmutableList());
+  }
+
   public boolean hasContent() {
     Preconditions.checkNotNull(response, "Response is not initialized.");
     return response.getEntity() != null;
diff --git a/java/com/google/gerrit/httpd/restapi/RestApiServlet.java b/java/com/google/gerrit/httpd/restapi/RestApiServlet.java
index 38fe0e2..10f2638 100644
--- a/java/com/google/gerrit/httpd/restapi/RestApiServlet.java
+++ b/java/com/google/gerrit/httpd/restapi/RestApiServlet.java
@@ -110,6 +110,7 @@
 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;
@@ -1322,11 +1323,42 @@
   }
 
   private TraceContext enableTracing(HttpServletRequest req, HttpServletResponse res) {
-    String traceValue = req.getParameter(ParameterParser.TRACE_PARAMETER);
-    return TraceContext.newTrace(
-        traceValue != null,
-        traceValue,
-        (tagName, traceId) -> res.setHeader(X_GERRIT_TRACE, traceId.toString()));
+    // There are 2 ways to enable tracing for REST calls:
+    // 1. by using the 'trace' or 'trace=<trace-id>' request parameter
+    // 2. by setting the 'X-Gerrit-Trace:' or 'X-Gerrit-Trace:<trace-id>' header
+    String traceValueFromHeader = req.getHeader(X_GERRIT_TRACE);
+    String traceValueFromRequestParam = req.getParameter(ParameterParser.TRACE_PARAMETER);
+    boolean doTrace = traceValueFromHeader != null || traceValueFromRequestParam != null;
+
+    // Check whether no trace ID, one trace ID or 2 different trace IDs have been specified.
+    String traceId1;
+    String traceId2;
+    if (!Strings.isNullOrEmpty(traceValueFromHeader)) {
+      traceId1 = traceValueFromHeader;
+      if (!Strings.isNullOrEmpty(traceValueFromRequestParam)
+          && !traceValueFromHeader.equals(traceValueFromRequestParam)) {
+        traceId2 = traceValueFromRequestParam;
+      } else {
+        traceId2 = null;
+      }
+    } else {
+      traceId1 = Strings.emptyToNull(traceValueFromRequestParam);
+      traceId2 = null;
+    }
+
+    // Use the first trace ID to start tracing. If this trace ID is null, a trace ID will be
+    // generated.
+    TraceContext traceContext =
+        TraceContext.newTrace(
+            doTrace,
+            traceId1,
+            (tagName, traceId) -> res.setHeader(X_GERRIT_TRACE, traceId.toString()));
+    // If a second trace ID was specified, add a tag for it as well.
+    if (traceId2 != null) {
+      traceContext.addTag(RequestId.Type.TRACE_ID, traceId2);
+      res.addHeader(X_GERRIT_TRACE, traceId2);
+    }
+    return traceContext;
   }
 
   private boolean isDelete(HttpServletRequest req) {
diff --git a/javatests/com/google/gerrit/acceptance/rest/TraceIT.java b/javatests/com/google/gerrit/acceptance/rest/TraceIT.java
index c682963..4de54e3 100644
--- a/javatests/com/google/gerrit/acceptance/rest/TraceIT.java
+++ b/javatests/com/google/gerrit/acceptance/rest/TraceIT.java
@@ -18,6 +18,7 @@
 import static org.apache.http.HttpStatus.SC_CREATED;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
 import com.google.gerrit.acceptance.AbstractDaemonTest;
 import com.google.gerrit.acceptance.PushOneCommit;
@@ -36,6 +37,7 @@
 import com.google.gerrit.server.validators.ValidationException;
 import com.google.inject.Inject;
 import java.util.List;
+import org.apache.http.message.BasicHeader;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -75,7 +77,7 @@
   }
 
   @Test
-  public void restCallWithTrace() throws Exception {
+  public void restCallWithTraceRequestParam() throws Exception {
     RestResponse response =
         adminRestSession.put("/projects/new2?" + ParameterParser.TRACE_PARAMETER);
     assertThat(response.getStatusCode()).isEqualTo(SC_CREATED);
@@ -85,7 +87,7 @@
   }
 
   @Test
-  public void restCallWithTraceAndProvidedTraceId() throws Exception {
+  public void restCallWithTraceRequestParamAndProvidedTraceId() throws Exception {
     RestResponse response =
         adminRestSession.put("/projects/new3?" + ParameterParser.TRACE_PARAMETER + "=issue/123");
     assertThat(response.getStatusCode()).isEqualTo(SC_CREATED);
@@ -95,6 +97,70 @@
   }
 
   @Test
+  public void restCallWithTraceHeader() throws Exception {
+    RestResponse response =
+        adminRestSession.putWithHeader(
+            "/projects/new4", new BasicHeader(RestApiServlet.X_GERRIT_TRACE, null));
+    assertThat(response.getStatusCode()).isEqualTo(SC_CREATED);
+    assertThat(response.getHeader(RestApiServlet.X_GERRIT_TRACE)).isNotNull();
+    assertThat(projectCreationListener.traceId).isNotNull();
+    assertThat(projectCreationListener.isLoggingForced).isTrue();
+  }
+
+  @Test
+  public void restCallWithTraceHeaderAndProvidedTraceId() throws Exception {
+    RestResponse response =
+        adminRestSession.putWithHeader(
+            "/projects/new5", new BasicHeader(RestApiServlet.X_GERRIT_TRACE, "issue/123"));
+    assertThat(response.getStatusCode()).isEqualTo(SC_CREATED);
+    assertThat(response.getHeader(RestApiServlet.X_GERRIT_TRACE)).isEqualTo("issue/123");
+    assertThat(projectCreationListener.traceId).isEqualTo("issue/123");
+    assertThat(projectCreationListener.isLoggingForced).isTrue();
+  }
+
+  @Test
+  public void restCallWithTraceRequestParamAndTraceHeader() throws Exception {
+    // trace ID only specified by trace header
+    RestResponse response =
+        adminRestSession.putWithHeader(
+            "/projects/new6?trace", new BasicHeader(RestApiServlet.X_GERRIT_TRACE, "issue/123"));
+    assertThat(response.getStatusCode()).isEqualTo(SC_CREATED);
+    assertThat(response.getHeader(RestApiServlet.X_GERRIT_TRACE)).isEqualTo("issue/123");
+    assertThat(projectCreationListener.traceId).isEqualTo("issue/123");
+    assertThat(projectCreationListener.isLoggingForced).isTrue();
+
+    // trace ID only specified by trace request parameter
+    response =
+        adminRestSession.putWithHeader(
+            "/projects/new7?trace=issue/123", new BasicHeader(RestApiServlet.X_GERRIT_TRACE, null));
+    assertThat(response.getStatusCode()).isEqualTo(SC_CREATED);
+    assertThat(response.getHeader(RestApiServlet.X_GERRIT_TRACE)).isEqualTo("issue/123");
+    assertThat(projectCreationListener.traceId).isEqualTo("issue/123");
+    assertThat(projectCreationListener.isLoggingForced).isTrue();
+
+    // same trace ID specified by trace header and trace request parameter
+    response =
+        adminRestSession.putWithHeader(
+            "/projects/new8?trace=issue/123",
+            new BasicHeader(RestApiServlet.X_GERRIT_TRACE, "issue/123"));
+    assertThat(response.getStatusCode()).isEqualTo(SC_CREATED);
+    assertThat(response.getHeader(RestApiServlet.X_GERRIT_TRACE)).isEqualTo("issue/123");
+    assertThat(projectCreationListener.traceId).isEqualTo("issue/123");
+    assertThat(projectCreationListener.isLoggingForced).isTrue();
+
+    // different trace IDs specified by trace header and trace request parameter
+    response =
+        adminRestSession.putWithHeader(
+            "/projects/new9?trace=issue/123",
+            new BasicHeader(RestApiServlet.X_GERRIT_TRACE, "issue/456"));
+    assertThat(response.getStatusCode()).isEqualTo(SC_CREATED);
+    assertThat(response.getHeaders(RestApiServlet.X_GERRIT_TRACE))
+        .containsExactly("issue/123", "issue/456");
+    assertThat(projectCreationListener.traceIds).containsExactly("issue/123", "issue/456");
+    assertThat(projectCreationListener.isLoggingForced).isTrue();
+  }
+
+  @Test
   public void pushWithoutTrace() throws Exception {
     PushOneCommit push = pushFactory.create(db, admin.getIdent(), testRepo);
     PushOneCommit.Result r = push.to("refs/heads/master");
@@ -155,12 +221,14 @@
   private static class TraceValidatingProjectCreationValidationListener
       implements ProjectCreationValidationListener {
     String traceId;
+    ImmutableSet<String> traceIds;
     Boolean isLoggingForced;
 
     @Override
     public void validateNewProject(CreateProjectArgs args) throws ValidationException {
       this.traceId =
           Iterables.getFirst(LoggingContext.getInstance().getTagsAsMap().get("TRACE_ID"), null);
+      this.traceIds = LoggingContext.getInstance().getTagsAsMap().get("TRACE_ID");
       this.isLoggingForced = LoggingContext.getInstance().shouldForceLogging(null, null, false);
     }
   }