Merge pull request #31156 from vespa-engine/vekterli/explicitly-tracked-received-doc-entries-in-v1-handler

Explicitly track number of entries received through visiting

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());