MLE-19222 Eval/invoke now stream results

Finally makes use of the OkHttp MultipartReader to stream body parts instead of reading them all into memory. There's no test to be added here, as the only way to verify this is to try with a sufficient amount of data to cause an OutOfMemoryError. That will be verified manually instead. The existing regression tests will suffice to ensure that eval/invoke still work properly.

Author: rjrudin

marklogic-client-api/src/main/java/com/marklogic/client/impl/IoUtil.java

@@ -0,0 +1,26 @@
+/*
+ * Copyright © 2024 MarkLogic Corporation. All Rights Reserved.
+ */
+package com.marklogic.client.impl;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+
+public interface IoUtil {
+
+	/**
+	 * Tossing this commonly used logic here so that it can be reused. Can be removed when we drop Java 8 support, as
+	 * Java 9+ has a "readAllBytes" method.
+	 */
+	static byte[] streamToBytes(InputStream stream) throws IOException {
+		ByteArrayOutputStream buffer = new ByteArrayOutputStream();
+		byte[] b = new byte[8192];
+		int len = 0;
+		while ((len = stream.read(b)) != -1) {
+			buffer.write(b, 0, len);
+		}
+		buffer.flush();
+		return buffer.toByteArray();
+	}
+}

marklogic-client-api/src/main/java/com/marklogic/client/impl/OkHttpServices.java

@@ -19,6 +19,7 @@
 import com.marklogic.client.eval.EvalResultIterator;
 import com.marklogic.client.impl.okhttp.HttpUrlBuilder;
 import com.marklogic.client.impl.okhttp.OkHttpUtil;
+import com.marklogic.client.impl.okhttp.PartIterator;
 import com.marklogic.client.io.*;
 import com.marklogic.client.io.marker.*;
 import com.marklogic.client.query.*;
@@ -3836,7 +3837,29 @@ private <U extends OkHttpResultIterator> U postIteratedResourceImpl(
 
 		Response response = sendRequestWithRetry(requestBldr, (transaction == null), doPostFunction, resendableConsumer);
 		checkStatus(response, response.code(), "apply", "resource", path, ResponseStatus.OK_OR_CREATED_OR_NO_CONTENT);
-		return makeResults(constructor, reqlog, "apply", "resource", response);
+
+		boolean shouldStreamResults = "eval".equalsIgnoreCase(path) || "invoke".equalsIgnoreCase(path);
+		boolean hasDataToStream = response.body().contentLength() != 0;
+		// If body is empty, we can use the "old" way of reading results as there's nothing to stream.
+		return shouldStreamResults && hasDataToStream ?
+			evalAndStreamResults(reqlog, response) :
+			makeResults(constructor, reqlog, "apply", "resource", response);
+	}
+
+	/**
+	 * Added to resolve MLE-19222, where the eval/invoke response was read into memory, leading to OutOfMemoryErrors.
+	 * The one thing we are not able to do here though is check for errors in the trailers, as trailers cannot be
+	 * read until the entire body has been read. But we don't want to read the entire body right away.
+	 */
+	private <U extends OkHttpResultIterator> U evalAndStreamResults(RequestLogger reqlog, Response response) {
+		if (response == null) return null;
+		try {
+			MultipartReader reader = new MultipartReader(response.body());
+			PartIterator partIterator = new PartIterator(reader);
+			return (U) new DefaultOkHttpResultIterator(reqlog, partIterator, response);
+		} catch (IOException e) {
+			throw new MarkLogicIOException(e);
+		}
 	}
 
 	@Override
@@ -4587,6 +4610,12 @@ static abstract class OkHttpResultIterator<T extends OkHttpResult> {
 		private long totalSize = -1;
 		private Closeable closeable;
 
+		OkHttpResultIterator(RequestLogger reqlog, Iterator<BodyPart> partIterator, Closeable closeable) {
+			this.reqlog = reqlog;
+			this.partQueue = partIterator;
+			this.closeable = closeable;
+		}
+
 		OkHttpResultIterator(RequestLogger reqlog, List<BodyPart> partList, Closeable closeable) {
 			this.reqlog = reqlog;
 			if (partList != null && partList.size() > 0) {
@@ -4685,14 +4714,15 @@ OkHttpServiceResult constructNext(RequestLogger logger, BodyPart part) {
 		}
 	}
 
-	static class DefaultOkHttpResultIterator
-		extends OkHttpResultIterator<OkHttpResult>
-		implements Iterator<OkHttpResult> {
-		DefaultOkHttpResultIterator(RequestLogger reqlog,
-									List<BodyPart> partList, Closeable closeable) {
+	static class DefaultOkHttpResultIterator extends OkHttpResultIterator<OkHttpResult> implements Iterator<OkHttpResult> {
+		DefaultOkHttpResultIterator(RequestLogger reqlog, List<BodyPart> partList, Closeable closeable) {
 			super(reqlog, partList, closeable);
 		}
 
+		DefaultOkHttpResultIterator(RequestLogger reqlog, Iterator<BodyPart> partIterator, Closeable closeable) {
+			super(reqlog, partIterator, closeable);
+		}
+
 		OkHttpResult constructNext(RequestLogger logger, BodyPart part) {
 			return new OkHttpResult(logger, part);
 		}

marklogic-client-api/src/main/java/com/marklogic/client/impl/okhttp/PartIterator.java

@@ -0,0 +1,86 @@
+/*
+ * Copyright © 2024 MarkLogic Corporation. All Rights Reserved.
+ */
+package com.marklogic.client.impl.okhttp;
+
+import com.marklogic.client.MarkLogicIOException;
+import com.marklogic.client.impl.IoUtil;
+import jakarta.activation.DataHandler;
+import jakarta.mail.BodyPart;
+import jakarta.mail.MessagingException;
+import jakarta.mail.internet.MimeBodyPart;
+import jakarta.mail.util.ByteArrayDataSource;
+import okhttp3.Headers;
+import okhttp3.MultipartReader;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Iterator;
+
+/**
+ * Adapts the iterator over the OkHttp MultipartReader to conform to the Iterator that is required by
+ * OkHttpEvalResultIterator. By converting each MultipartReader.Part into a jakarta.mail.BodyPart, we can reuse
+ * all the existing plumbing that depends on jakarta.mail.BodyPart.
+ * <p>
+ * Added to resolve MLE-19222, where eval/invoke results are not being streamed but rather were all being read into
+ * memory, leading to OutOfMemoryErrors.
+ */
+public class PartIterator implements Iterator<BodyPart> {
+
+	private final MultipartReader reader;
+	private BodyPart nextBodyPart;
+
+	public PartIterator(MultipartReader reader) {
+		this.reader = reader;
+		readNextPart();
+	}
+
+	@Override
+	public boolean hasNext() {
+		return nextBodyPart != null;
+	}
+
+	@Override
+	public BodyPart next() {
+		BodyPart partToReturn = nextBodyPart;
+		readNextPart();
+		return partToReturn;
+	}
+
+	private void readNextPart() {
+		try {
+			// See http://okhttp.foofun.cn/4.x/okhttp/okhttp3/-multipart-reader/ for more info on the OkHttp
+			// MultipartReader. This was actually requested many moons ago by one of the original Java Client
+			// developers - https://github.com/square/okhttp/issues/3394.
+			MultipartReader.Part nextPart = reader.nextPart();
+			this.nextBodyPart = nextPart != null ? convertPartToBodyPart(nextPart) : null;
+		} catch (Exception e) {
+			throw new MarkLogicIOException(e);
+		}
+	}
+
+	private static BodyPart convertPartToBodyPart(MultipartReader.Part part) throws IOException, MessagingException {
+		MimeBodyPart bodyPart = new MimeBodyPart();
+
+		try {
+			try (InputStream inputStream = part.body().inputStream()) {
+				byte[] bytes = IoUtil.streamToBytes(inputStream);
+				bodyPart.setDataHandler(new DataHandler(new ByteArrayDataSource(bytes, part.headers().get("Content-Type"))));
+			}
+
+			// part.headers.toMultimap() is lowercasing header names, which causes later issues.
+			Headers headers = part.headers();
+			for (String headerName : headers.names()) {
+				for (String headerValue : headers.values(headerName)) {
+					bodyPart.addHeader(headerName, headerValue);
+				}
+			}
+			return bodyPart;
+		} finally {
+			// Looking at the OkHttp source code, this does not appear necessary, as closing the InputStream above should
+			// achieve the same effect. But there is no downside to doing this, as it may be required by a future version
+			// of OkHttp.
+			part.close();
+		}
+	}
+}

marklogic-client-api/src/main/java/com/marklogic/client/io/InputStreamHandle.java

@@ -10,6 +10,7 @@
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
 
+import com.marklogic.client.impl.IoUtil;
 import com.marklogic.client.io.marker.*;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -45,8 +46,6 @@ public class InputStreamHandle
   private byte[] contentBytes;
   private InputStream content;
 
-  final static private int BUFFER_SIZE = 8192;
-
   /**
    * Creates a factory to create an InputStreamHandle instance for an input stream.
    * @return  the factory
@@ -186,17 +185,9 @@ public InputStream bytesToContent(byte[] buffer) {
   public byte[] contentToBytes(InputStream content) {
     try {
       if (content == null) return null;
-
-      ByteArrayOutputStream buffer = new ByteArrayOutputStream();
-
-      byte[] b = new byte[BUFFER_SIZE];
-      int len = 0;
-      while ((len = content.read(b)) != -1) {
-        buffer.write(b, 0, len);
-      }
+	  byte[] bytes = IoUtil.streamToBytes(content);
       content.close();
-
-      return buffer.toByteArray();
+	  return bytes;
     } catch (IOException e) {
       throw new MarkLogicIOException(e);
     }