Explicitly track number of locally entries received through visiting

Using the underlying session's `VisitorStatistics` may not be 1-1
with the actual number of entries the data handler has been invoked
with. This causes issues if anyone tries to cross-check the document
count emitted as part of the results vs. the number of entries
actually present in the result array.

The session updates its statistics based on the what is returned from
the backend as part of _successful_ `CreateVisitorReply` messages.
If a CreateVisitor returns with a transient error, the statistics will
not be updated, but it's unspecified how many (if any) document entries
that particular visitor may already have pushed to the client directly
from the content nodes.

Note that the locally tracked number is only reported if the session
itself is receiving the document data; if a remote data handler is
in use we have to report what the `VisitorStatistics` give us.

2 files changed, 16 insertions, 6 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 b483d6977d6..6e07661235e 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
@@ -1397,6 +1397,7 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
             Phaser phaser = new Phaser(2); // Synchronize this thread (dispatch) with the visitor callback thread.
             AtomicReference<String> error = new AtomicReference<>(); // Set if error occurs during processing of visited documents.
             callback.onStart(response, fullyApplied);
+            final AtomicLong locallyReceivedDocCount = new AtomicLong(0);
             VisitorControlHandler controller = new VisitorControlHandler() {
                 final ScheduledFuture<?> abort = streaming ? visitDispatcher.schedule(this::abort, visitTimeout(request), MILLISECONDS) : null;
                 final AtomicReference<VisitorSession> session = new AtomicReference<>();
@@ -1410,7 +1411,10 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
                         try (response) {
                             callback.onEnd(response);
 
-                            response.writeDocumentCount(getVisitorStatistics() == null ? 0 : getVisitorStatistics().getDocumentsVisited());
+                            // Locally tracked document count is only correct if we have a local data handler.
+                            // Otherwise, we have to report the statistics received transitively from the content nodes.
+                            long statsDocCount = (getVisitorStatistics() != null ? getVisitorStatistics().getDocumentsVisited() : 0);
+                            response.writeDocumentCount(parameters.getLocalDataHandler() != null ? locallyReceivedDocCount.get() : statsDocCount);
 
                             if (session.get() != null)
                                 response.writeTrace(session.get().getTrace());
@@ -1456,6 +1460,7 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
                         if (m instanceof PutDocumentMessage put) document = put.getDocumentPut().getDocument();
                         else if (parameters.visitRemoves() && m instanceof RemoveDocumentMessage remove) removeId = remove.getDocumentId();
                         else throw new UnsupportedOperationException("Got unsupported message type: " + m.getClass().getName());
+                        locallyReceivedDocCount.getAndAdd(1);
                         callback.onDocument(response,
                                             document,
                                             removeId,
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 58cf34712aa..b2c0b1b2ce8 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
@@ -290,7 +290,7 @@ public class DocumentV1ApiTest {
             parameters.getLocalDataHandler().onMessage(new RemoveDocumentMessage(new DocumentId("id:space:music::t-square-truth")), tokens.get(3));
             VisitorStatistics statistics = new VisitorStatistics();
             statistics.setBucketsVisited(1);
-            statistics.setDocumentsVisited(3);
+            statistics.setDocumentsVisited(123); // Ignored in favor of tracking actually emitted entries
             parameters.getControlHandler().onVisitorStatistics(statistics);
             parameters.getControlHandler().onDone(VisitorControlHandler.CompletionCode.TIMEOUT, "timeout is OK");
         });
@@ -323,7 +323,7 @@ public class DocumentV1ApiTest {
                             "remove": "id:space:music::t-square-truth"
                            }
                          ],
-                         "documentCount": 3,
+                         "documentCount": 4,
                          "trace": [
                            { "message": "Tracy Chapman" },
                            {
@@ -441,13 +441,18 @@ public class DocumentV1ApiTest {
             assertEquals("[Content:cluster=content]", parameters.getRemoteDataHandler());
             assertEquals("[document]", parameters.fieldSet());
             assertEquals(60_000L, parameters.getSessionTimeoutMs());
+            VisitorStatistics statistics = new VisitorStatistics();
+            statistics.setBucketsVisited(1);
+            statistics.setDocumentsVisited(2);
+            // Visiting with remote data handlers should report the remotely aggregated statistics
+            parameters.getControlHandler().onVisitorStatistics(statistics);
             parameters.getControlHandler().onDone(VisitorControlHandler.CompletionCode.SUCCESS, "We made it!");
         });
         response = driver.sendRequest("http://localhost/document/v1/space/music/docid?destinationCluster=content&selection=true&cluster=content&timeout=60", POST);
         assertSameJson("""
                        {
                          "pathId": "/document/v1/space/music/docid",
-                         "documentCount": 0
+                         "documentCount": 2
                        }""",
                        response.readAll());
         assertEquals(200, response.getStatus());
@@ -488,7 +493,7 @@ public class DocumentV1ApiTest {
         assertSameJson("""
                        {
                          "pathId": "/document/v1/space/music/docid",
-                         "documentCount": 0
+                         "documentCount": 1
                        }""",
                        response.readAll());
         assertEquals(200, response.getStatus());
@@ -542,7 +547,7 @@ public class DocumentV1ApiTest {
         assertSameJson("""
                        {
                          "pathId": "/document/v1/space/music/docid",
-                         "documentCount": 0,
+                         "documentCount": 1,
                          "message": "boom"
                        }""",
                        response.readAll());