Support tracing for /document/v1/ visit requests

2 files changed, 56 insertions, 27 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 2d97c33741e..06cdccd17fc 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
@@ -38,6 +38,7 @@ import com.yahoo.documentapi.ProgressToken;
 import com.yahoo.documentapi.Response.Outcome;
 import com.yahoo.documentapi.Result;
 import com.yahoo.documentapi.VisitorControlHandler;
+import com.yahoo.documentapi.VisitorControlSession;
 import com.yahoo.documentapi.VisitorDataHandler;
 import com.yahoo.documentapi.VisitorParameters;
 import com.yahoo.documentapi.VisitorSession;
@@ -1175,6 +1176,8 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
                                                                            StringJoiner::merge)
                                                                    .toString());
 
+        getProperty(request, TRACELEVEL, integerParser).ifPresent(parameters::setTraceLevel);
+
         getProperty(request, CONTINUATION).map(ProgressToken::fromSerializedString).ifPresent(parameters::setResumeToken);
         parameters.setPriority(DocumentProtocol.Priority.NORMAL_4);
 
@@ -1310,6 +1313,11 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
             callback.onStart(response, fullyApplied);
             VisitorControlHandler controller = new VisitorControlHandler() {
                 final ScheduledFuture<?> abort = streaming ? visitDispatcher.schedule(this::abort, request.getTimeout(MILLISECONDS), MILLISECONDS) : null;
+                final AtomicReference<VisitorSession> session = new AtomicReference<>();
+                @Override public void setSession(VisitorControlSession session) { // Workaround for broken session API ಠ_ಠ
+                    super.setSession(session);
+                    if (session instanceof VisitorSession visitorSession) this.session.set(visitorSession);
+                }
                 @Override public void onDone(CompletionCode code, String message) {
                     super.onDone(code, message);
                     loggingException(() -> {
@@ -1319,6 +1327,9 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
                             if (getVisitorStatistics() != null)
                                 response.writeDocumentCount(getVisitorStatistics().getDocumentsVisited());
 
+                            if (session.get() != null)
+                                response.writeTrace(session.get().getTrace());
+
                             int status = Response.Status.BAD_GATEWAY;
                             switch (code) {
                                 case TIMEOUT:
@@ -1350,7 +1361,6 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
                         phaser.arriveAndAwaitAdvance(); // We may get here while dispatching thread is still putting us in the map.
                         visits.remove(this).destroy();
                     });
-
                 }
             };
             if (parameters.getRemoteDataHandler() == null) {
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 cd57818e74e..973d0a98b24 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
@@ -208,6 +208,12 @@ public class DocumentV1ApiTest {
 
         // GET at root is a visit. Numeric parameters have an upper bound.
         access.expect(tokens);
+        Trace visitorTrace = new Trace(9);
+        visitorTrace.trace(7, "Tracy Chapman", false);
+        visitorTrace.getRoot().addChild(new TraceNode().setStrict(false)
+                                                .addChild("Fast Car")
+                                                .addChild("Baby Can I Hold You"));
+        access.visitorTrace = visitorTrace;
         access.expect(parameters -> {
             assertEquals("content", parameters.getRoute().toString());
             assertEquals("default", parameters.getBucketSpace());
@@ -216,6 +222,7 @@ public class DocumentV1ApiTest {
             assertEquals("[id]", parameters.getFieldSet());
             assertEquals("(all the things)", parameters.getDocumentSelection());
             assertEquals(6000, parameters.getSessionTimeoutMs());
+            assertEquals(9, parameters.getTraceLevel());
             // Put some documents in the response
             parameters.getLocalDataHandler().onMessage(new PutDocumentMessage(new DocumentPut(doc1)), tokens.get(0));
             parameters.getLocalDataHandler().onMessage(new PutDocumentMessage(new DocumentPut(doc2)), tokens.get(1));
@@ -227,32 +234,43 @@ public class DocumentV1ApiTest {
             parameters.getControlHandler().onDone(VisitorControlHandler.CompletionCode.TIMEOUT, "timeout is OK");
         });
         response = driver.sendRequest("http://localhost/document/v1?cluster=content&bucketSpace=default&wantedDocumentCount=1025&concurrency=123" +
-                                      "&selection=all%20the%20things&fieldSet=[id]&timeout=6");
-        assertSameJson("{" +
-                       "  \"pathId\": \"/document/v1\"," +
-                       "  \"documents\": [" +
-                       "    {" +
-                       "      \"id\": \"id:space:music::one\"," +
-                       "      \"fields\": {" +
-                       "        \"artist\": \"Tom Waits\", " +
-                       "        \"embedding\": { \"type\": \"tensor(x[3])\", \"values\": [1.0,2.0,3.0] } " +
-                       "      }" +
-                       "    }," +
-                       "    {" +
-                       "      \"id\": \"id:space:music:n=1:two\"," +
-                       "      \"fields\": {" +
-                       "        \"artist\": \"Asa-Chan & Jun-Ray\", " +
-                       "        \"embedding\": { \"type\": \"tensor(x[3])\", \"values\": [4.0,5.0,6.0] } " +
-                       "      }" +
-                       "    }," +
-                       "    {" +
-                       "     \"id\": \"id:space:music:g=a:three\"," +
-                       "     \"fields\": {}" +
-                       "    }" +
-                       "  ]," +
-                       "  \"documentCount\": 3" +
-                       "}", response.readAll());
+                                      "&selection=all%20the%20things&fieldSet=[id]&timeout=6&tracelevel=9");
+        assertSameJson("""
+                       {
+                         "pathId": "/document/v1",
+                         "documents": [
+                           {
+                             "id": "id:space:music::one",
+                             "fields": {
+                               "artist": "Tom Waits",\s
+                               "embedding": { "type": "tensor(x[3])", "values": [1.0,2.0,3.0] }\s
+                             }
+                           },
+                           {
+                             "id": "id:space:music:n=1:two",
+                             "fields": {
+                               "artist": "Asa-Chan & Jun-Ray",\s
+                               "embedding": { "type": "tensor(x[3])", "values": [4.0,5.0,6.0] }\s
+                             }
+                           },
+                           {
+                            "id": "id:space:music:g=a:three",
+                            "fields": {}
+                           }
+                         ],
+                         "documentCount": 3,
+                         "trace": [
+                           { "message": "Tracy Chapman" },
+                           {
+                             "fork": [
+                               { "message": "Fast Car" },
+                               { "message": "Baby Can I Hold You" }
+                             ]
+                           }
+                         ]
+                       }""", response.readAll());
         assertEquals(200, response.getStatus());
+        access.visitorTrace = null;
 
         // GET at root is a visit. Streaming mode can be specified with &stream=true
         access.expect(tokens);
@@ -911,6 +929,7 @@ public class DocumentV1ApiTest {
         private final AtomicReference<Consumer<VisitorParameters>> expectations = new AtomicReference<>();
         private final Set<AckToken> outstanding = new CopyOnWriteArraySet<>();
         private final MockAsyncSession session = new MockAsyncSession();
+        private Trace visitorTrace;
 
         MockDocumentAccess(DocumentmanagerConfig config) {
             super(new DocumentAccessParams().setDocumentmanagerConfig(config));
@@ -936,7 +955,7 @@ public class DocumentV1ApiTest {
                 }
                 @Override public boolean isDone() { return false; }
                 @Override public ProgressToken getProgress() { return null; }
-                @Override public Trace getTrace() { return null; }
+                @Override public Trace getTrace() { return visitorTrace; }
                 @Override public boolean waitUntilDone(long timeoutMs) { return false; }
                 @Override public void ack(AckToken token) { assertTrue(outstanding.remove(token)); }
                 @Override public void abort() { }