Call completion handlers earlier when receiving data/close in doc/v1/handler

1 files changed, 18 insertions, 4 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 bb1dd8f20c4..c056fe46e01 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
@@ -552,16 +552,30 @@ public class DocumentV1ApiHandler extends AbstractRequestHandler {
             this.reader = reader;
         }
 
+        /** Write is complete when we have stored the buffer — call completion handler. */
         @Override
         public void write(ByteBuffer buf, CompletionHandler handler) {
-            delegate.write(buf, handler);
+            try {
+                delegate.write(buf, logException);
+                handler.completed();
+            }
+            catch (Exception e) {
+                handler.failed(e);
+            }
         }
 
+        /** Close is complete when we have close the buffer. */
         @Override
         public void close(CompletionHandler handler) {
-            delegate.close(handler);
-            try (UnsafeContentInputStream in = new UnsafeContentInputStream(delegate)) {
-                reader.accept(in);
+            try {
+                delegate.close(logException);
+                try (UnsafeContentInputStream in = new UnsafeContentInputStream(delegate)) {
+                    reader.accept(in);
+                }
+                handler.completed();
+            }
+            catch (Exception e) {
+                handler.failed(e);
             }
         }