Stop writing documents once document array is closed

2 files changed, 14 insertions, 10 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 1e9cdfbbbea..29b34ff4468 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
@@ -86,7 +86,6 @@ import java.util.Optional;
 import java.util.StringJoiner;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedDeque;
-import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Phaser;
 import java.util.concurrent.ScheduledExecutorService;
@@ -583,6 +582,7 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
         private final OutputStream out = new ContentChannelOutputStream(buffer);
         private final JsonGenerator json;
         private final ResponseHandler handler;
+        private boolean documentsDone = false;
         private ContentChannel channel;
 
         private JsonResponse(ResponseHandler handler) throws IOException {
@@ -636,6 +636,7 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
         /** Closes the JSON and the output content channel of this. */
         @Override
         public synchronized void close() throws IOException {
+            documentsDone = true; // In case we were closed without explicitly closing the documents array.
             try {
                 if (channel == null) {
                     log.log(WARNING, "Close called before response was committed, in " + getClass().getName());
@@ -695,12 +696,18 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
         }
 
         synchronized void writeDocumentValue(Document document, CompletionHandler completionHandler) {
-            new JsonWriter(json).write(document);
+            if ( ! documentsDone)
+                new JsonWriter(json).write(document);
+
             if (completionHandler != null)
-                buffer.write(emptyBuffer, completionHandler);
+                if ( ! documentsDone)
+                    buffer.write(emptyBuffer, completionHandler);
+                else
+                    completionHandler.completed();
         }
 
         synchronized void writeArrayEnd() throws IOException {
+            documentsDone = true;
             json.writeEndArray();
         }
 
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 0e62b620828..412558f9a87 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
@@ -263,12 +263,13 @@ public class DocumentV1ApiTest {
             // 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));
-            parameters.getLocalDataHandler().onMessage(new PutDocumentMessage(new DocumentPut(doc3)), tokens.get(2));
             VisitorStatistics statistics = new VisitorStatistics();
             statistics.setBucketsVisited(1);
-            statistics.setDocumentsVisited(3);
+            statistics.setDocumentsVisited(2);
             parameters.getControlHandler().onVisitorStatistics(statistics);
             parameters.getControlHandler().onDone(VisitorControlHandler.CompletionCode.TIMEOUT, "timeout is OK");
+            // Extra documents are ignored.
+            parameters.getLocalDataHandler().onMessage(new PutDocumentMessage(new DocumentPut(doc3)), tokens.get(2));
         });
         response = driver.sendRequest("http://localhost/document/v1?cluster=content&bucketSpace=default&wantedDocumentCount=1025&concurrency=123" +
                                       "&selection=all%20the%20things&fieldSet=[id]&timeout=6&stream=true&slices=4&sliceId=1");
@@ -286,13 +287,9 @@ public class DocumentV1ApiTest {
                        "      \"fields\": {" +
                        "        \"artist\": \"Asa-Chan & Jun-Ray\"" +
                        "      }" +
-                       "    }," +
-                       "    {" +
-                       "     \"id\": \"id:space:music:g=a:three\"," +
-                       "     \"fields\": {}" +
                        "    }" +
                        "  ]," +
-                       "  \"documentCount\": 3" +
+                       "  \"documentCount\": 2" +
                        "}", response.readAll());
         assertEquals(200, response.getStatus());