Add traces to document/v1

2 files changed, 85 insertions, 21 deletions

diff --git a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/DocumentV1ApiHandler.java b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/DocumentV1ApiHandler.java
index 43358aa8699..caf1403bdf8 100644
--- a/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/DocumentV1ApiHandler.java
+++ b/vespaclient-container-plugin/src/main/java/com/yahoo/document/restapi/resource/DocumentV1ApiHandler.java
@@ -52,6 +52,8 @@ import com.yahoo.jdisc.handler.UnsafeContentInputStream;
 import com.yahoo.jdisc.http.HttpRequest;
 import com.yahoo.jdisc.http.HttpRequest.Method;
 import com.yahoo.messagebus.StaticThrottlePolicy;
+import com.yahoo.messagebus.Trace;
+import com.yahoo.messagebus.TraceNode;
 import com.yahoo.metrics.simple.MetricReceiver;
 import com.yahoo.restapi.Path;
 import com.yahoo.text.Text;
@@ -309,13 +311,8 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
 
     private ContentChannel getDocument(HttpRequest request, DocumentPath path, ResponseHandler handler) {
         enqueueAndDispatch(request, handler, () -> {
-            DocumentOperationParameters rawParameters = parameters();
-            rawParameters = getProperty(request, CLUSTER).map(cluster -> resolveCluster(Optional.of(cluster), clusters).route())
-                                                         .map(rawParameters::withRoute)
-                                                         .orElse(rawParameters);
-            rawParameters = getProperty(request, FIELD_SET).map(rawParameters::withFieldSet)
-                                                           .orElse(rawParameters);
-            DocumentOperationParameters parameters = rawParameters.withResponseHandler(response -> {
+            DocumentOperationParameters parameters = parametersFromRequest(request, CLUSTER, FIELD_SET)
+                    .withResponseHandler(response -> {
                 handle(path, handler, response, (document, jsonResponse) -> {
                     if (document != null) {
                         jsonResponse.writeSingleDocument(document);
@@ -336,9 +333,8 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
             enqueueAndDispatch(request, handler, () -> {
                 DocumentPut put = parser.parsePut(in, path.id().toString());
                 getProperty(request, CONDITION).map(TestAndSetCondition::new).ifPresent(put::setCondition);
-                DocumentOperationParameters rawParameters = getProperty(request, ROUTE).map(parameters()::withRoute)
-                                                                                       .orElse(parameters());
-                DocumentOperationParameters parameters = rawParameters.withResponseHandler(response -> handle(path, handler, response));
+                DocumentOperationParameters parameters = parametersFromRequest(request, ROUTE)
+                        .withResponseHandler(response -> handle(path, handler, response));
                 return () -> dispatchOperation(() -> asyncSession.put(put, parameters));
             });
         });
@@ -350,10 +346,9 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
             enqueueAndDispatch(request, handler, () -> {
                 DocumentUpdate update = parser.parseUpdate(in, path.id().toString());
                 getProperty(request, CONDITION).map(TestAndSetCondition::new).ifPresent(update::setCondition);
-                getProperty(request, CREATE).map(booleanParser::parse).ifPresent(update::setCreateIfNonExistent);
-                DocumentOperationParameters rawParameters = getProperty(request, ROUTE).map(parameters()::withRoute)
-                                                                                       .orElse(parameters());
-                DocumentOperationParameters parameters = rawParameters.withResponseHandler(response -> handle(path, handler, response));
+                getProperty(request, CREATE, booleanParser).ifPresent(update::setCreateIfNonExistent);
+                DocumentOperationParameters parameters = parametersFromRequest(request, ROUTE)
+                        .withResponseHandler(response -> handle(path, handler, response));
                 return () -> dispatchOperation(() -> asyncSession.update(update, parameters));
             });
         });
@@ -364,14 +359,40 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
         enqueueAndDispatch(request, handler, () -> {
             DocumentRemove remove = new DocumentRemove(path.id());
             getProperty(request, CONDITION).map(TestAndSetCondition::new).ifPresent(remove::setCondition);
-            DocumentOperationParameters rawParameters = getProperty(request, ROUTE).map(parameters()::withRoute)
-                                                                                   .orElse(parameters());
-            DocumentOperationParameters parameters = rawParameters.withResponseHandler(response -> handle(path, handler, response));
+            DocumentOperationParameters parameters = parametersFromRequest(request, ROUTE)
+                    .withResponseHandler(response -> handle(path, handler, response));
             return () -> dispatchOperation(() -> asyncSession.remove(remove, parameters));
         });
         return ignoredContent;
     }
 
+    private DocumentOperationParameters parametersFromRequest(HttpRequest request, String... names) {
+        DocumentOperationParameters parameters = getProperty(request, TRACE_LEVEL, integerParser).map(parameters()::withTraceLevel)
+                                                                                                 .orElse(parameters());
+        for (String name : names) switch (name) {
+            case CLUSTER:
+                parameters = getProperty(request, CLUSTER).map(cluster -> resolveCluster(Optional.of(cluster), clusters).route())
+                                                          .map(parameters::withRoute)
+                                                          .orElse(parameters);
+                break;
+            case CONDITION:
+                parameters = getProperty(request, ROUTE).map(parameters::withRoute)
+                                                        .orElse(parameters);
+                break;
+            case FIELD_SET:
+                parameters = getProperty(request, FIELD_SET).map(parameters::withFieldSet)
+                                                            .orElse(parameters);
+                break;
+            case ROUTE:
+                parameters = getProperty(request, ROUTE).map(parameters::withRoute)
+                                                        .orElse(parameters);
+                break;
+            default:
+                throw new IllegalArgumentException("Unrecognized document operation parameter name '" + name + "'");
+        }
+        return parameters;
+    }
+
     /** Dispatches enqueued requests until one is blocked. */
     void dispatchEnqueued() {
         try {
@@ -501,6 +522,26 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
             json.writeStringField("id", id.toString());
         }
 
+        synchronized void writeTrace(Trace trace) throws IOException {
+            if (trace != null && ! trace.getRoot().isEmpty()) {
+                writeTrace(trace.getRoot());
+            }
+        }
+
+        private void writeTrace(TraceNode node) throws IOException {
+            if (node.hasNote())
+                json.writeStringField("message", node.getNote());
+            if ( ! node.isLeaf()) {
+                json.writeArrayFieldStart(node.isStrict() ? "trace" : "fork");
+                for (int i = 0; i < node.getNumChildren(); i++) {
+                    json.writeStartObject();
+                    writeTrace(node.getChild(i));
+                    json.writeEndObject();
+                }
+                json.writeEndArray();
+            }
+        }
+
         synchronized void writeSingleDocument(Document document) throws IOException {
             new JsonWriter(json).writeFields(document);
         }
@@ -714,6 +755,7 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
 
     private static void handle(DocumentPath path, ResponseHandler handler, com.yahoo.documentapi.Response response, SuccessCallback callback) {
         try (JsonResponse jsonResponse = JsonResponse.create(path, handler)) {
+            jsonResponse.writeTrace(response.getTrace());
             if (response.isSuccess())
                 callback.onSuccess((response instanceof DocumentResponse) ? ((DocumentResponse) response).getDocument() : null, jsonResponse);
             else {
diff --git a/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/DocumentV1ApiTest.java b/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/DocumentV1ApiTest.java
index 4aeb0c5f8cc..723c84935f1 100644
--- a/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/DocumentV1ApiTest.java
+++ b/vespaclient-container-plugin/src/test/java/com/yahoo/document/restapi/resource/DocumentV1ApiTest.java
@@ -43,6 +43,7 @@ import com.yahoo.documentapi.VisitorSession;
 import com.yahoo.jdisc.Metric;
 import com.yahoo.messagebus.StaticThrottlePolicy;
 import com.yahoo.messagebus.Trace;
+import com.yahoo.messagebus.TraceNode;
 import com.yahoo.metrics.simple.MetricReceiver;
 import com.yahoo.searchdefinition.derived.Deriver;
 import com.yahoo.slime.Inspector;
@@ -61,6 +62,7 @@ import java.io.UncheckedIOException;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.OptionalInt;
 import java.util.TreeMap;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
@@ -320,11 +322,16 @@ public class DocumentV1ApiTest {
             DocumentPut expectedPut = new DocumentPut(doc2);
             expectedPut.setCondition(new TestAndSetCondition("test it"));
             assertEquals(expectedPut, put);
-            assertEquals(parameters(), parameters);
-            parameters.responseHandler().get().handleResponse(new DocumentResponse(0, doc2));
+            assertEquals(parameters().withTraceLevel(9), parameters);
+            Trace trace = new Trace(9);
+            trace.trace(7, "Tracy Chapman", false);
+            trace.getRoot().addChild(new TraceNode().setStrict(false)
+                                                    .addChild("Fast Car")
+                                                    .addChild("Baby Can I Hold You"));
+            parameters.responseHandler().get().handleResponse(new DocumentResponse(0, doc2, trace));
             return new Result(Result.ResultType.SUCCESS, null);
         });
-        response = driver.sendRequest("http://localhost/document/v1/space/music/number/1/two?condition=test%20it", POST,
+        response = driver.sendRequest("http://localhost/document/v1/space/music/number/1/two?condition=test%20it&traceLevel=9", POST,
                                       "{" +
                                       "  \"fields\": {" +
                                       "    \"artist\": \"Asa-Chan & Jun-Ray\"" +
@@ -332,7 +339,22 @@ public class DocumentV1ApiTest {
                                       "}");
         assertSameJson("{" +
                        "  \"pathId\": \"/document/v1/space/music/number/1/two\"," +
-                       "  \"id\": \"id:space:music:n=1:two\"" +
+                       "  \"id\": \"id:space:music:n=1:two\"," +
+                       "  \"trace\": [" +
+                       "    {" +
+                       "      \"message\": \"Tracy Chapman\"" +
+                       "    }," +
+                       "    {" +
+                       "      \"fork\": [" +
+                       "        {" +
+                       "          \"message\": \"Fast Car\"" +
+                       "        }," +
+                       "        {" +
+                       "          \"message\": \"Baby Can I Hold You\"" +
+                       "        }" +
+                       "      ]" +
+                       "    }" +
+                       "  ]" +
                        "}", response.readAll());
         assertEquals(200, response.getStatus());