Explicitly track number of 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.

2 files changed, 7 insertions, 5 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..f1829a1c42e 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 receivedDocsCount = 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,7 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
                         try (response) {
                             callback.onEnd(response);
 
-                            response.writeDocumentCount(getVisitorStatistics() == null ? 0 : getVisitorStatistics().getDocumentsVisited());
+                            response.writeDocumentCount(receivedDocsCount.get());
 
                             if (session.get() != null)
                                 response.writeTrace(session.get().getTrace());
@@ -1456,6 +1457,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());
+                        receivedDocsCount.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..3a8456d213d 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" },
                            {
@@ -488,7 +488,7 @@ public class DocumentV1ApiTest {
         assertSameJson("""
                        {
                          "pathId": "/document/v1/space/music/docid",
-                         "documentCount": 0
+                         "documentCount": 1
                        }""",
                        response.readAll());
         assertEquals(200, response.getStatus());
@@ -542,7 +542,7 @@ public class DocumentV1ApiTest {
         assertSameJson("""
                        {
                          "pathId": "/document/v1/space/music/docid",
-                         "documentCount": 0,
+                         "documentCount": 1,
                          "message": "boom"
                        }""",
                        response.readAll());