Merge pull request #1732 from marklogic/feature/19222-oom

MLE-19222 Eval/invoke now stream results

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