HSEARCH-5464 Move around gson entity to keep the common code in the backend

Author: marko-bekhta

backend/elasticsearch-client/elasticsearch-rest5-client/src/main/java/org/hibernate/search/backend/elasticsearch/client/rest5/impl/ByteBufferDataStreamChannel.java

@@ -97,7 +97,7 @@ private CompletableFuture<Response> send(ElasticsearchRequest elasticsearchReque
 
 		HttpEntity entity;
 		try {
-			entity = GsonHttpEntity.toEntity( gson, elasticsearchRequest );
+			entity = ClientRest5GsonHttpEntity.toEntity( gson, elasticsearchRequest );
 		}
 		catch (IOException | RuntimeException e) {
 			completableFuture.completeExceptionally( e );

backend/elasticsearch-client/elasticsearch-rest5-client/src/main/java/org/hibernate/search/backend/elasticsearch/client/rest5/impl/ClientRest5ElasticsearchClient.java

@@ -0,0 +1,151 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.search.backend.elasticsearch.client.rest5.impl;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.Set;
+
+import org.hibernate.search.backend.elasticsearch.client.common.gson.entity.spi.ContentEncoder;
+import org.hibernate.search.backend.elasticsearch.client.common.gson.entity.spi.GsonHttpEntityContentProvider;
+import org.hibernate.search.backend.elasticsearch.client.common.spi.ElasticsearchRequest;
+
+import com.google.gson.Gson;
+import com.google.gson.JsonObject;
+
+import org.apache.hc.core5.function.Supplier;
+import org.apache.hc.core5.http.ContentType;
+import org.apache.hc.core5.http.Header;
+import org.apache.hc.core5.http.HttpEntity;
+import org.apache.hc.core5.http.nio.AsyncEntityProducer;
+import org.apache.hc.core5.http.nio.DataStreamChannel;
+
+/**
+ * Optimised adapter to encode GSON objects into HttpEntity instances.
+ * The naive approach was using various StringBuilders; the objects we
+ * need to serialise into JSON might get large and this was causing the
+ * internal StringBuilder buffers to need frequent resizing and cause
+ * problems with excessive allocations.
+ *
+ * Rather than trying to guess reasonable default sizes for these buffers,
+ * we can defer the serialisation to write directly into the ByteBuffer
+ * of the HTTP client, this has the additional benefit of making the
+ * intermediary buffers short lived.
+ *
+ * The one complexity to watch for is flow control: when writing into
+ * the output buffer chances are that not all bytes are accepted; in
+ * this case we have to hold on the remaining portion of data to
+ * be written when the flow control is re-enabled.
+ *
+ * A side effect of this strategy is that the total content length which
+ * is being produced is not known in advance. Not reporting the length
+ * in advance to the Apache Http client causes it to use chunked-encoding,
+ * which is great for large blocks but not optimal for small messages.
+ * For this reason we attempt to start encoding into a small buffer
+ * upfront: if all data we need to produce fits into that then we can
+ * report the content length; if not the encoding completion will be deferred
+ * but not resetting so to avoid repeating encoding work.
+ *
+ * @author Sanne Grinovero (C) 2017 Red Hat Inc.
+ */
+final class ClientRest5GsonHttpEntity extends GsonHttpEntityContentProvider implements HttpEntity, AsyncEntityProducer {
+
+	private static final String CONTENT_TYPE = ContentType.APPLICATION_JSON.toString();
+
+
+	public static HttpEntity toEntity(Gson gson, ElasticsearchRequest request) throws IOException {
+		final List<JsonObject> bodyParts = request.bodyParts();
+		if ( bodyParts.isEmpty() ) {
+			return null;
+		}
+		return new ClientRest5GsonHttpEntity( gson, bodyParts );
+	}
+
+	public ClientRest5GsonHttpEntity(Gson gson, List<JsonObject> bodyParts) throws IOException {
+		super( gson, bodyParts );
+	}
+
+	@Override
+	public boolean isRepeatable() {
+		return true;
+	}
+
+	@Override
+	public void failed(Exception cause) {
+
+	}
+
+	@Override
+	public boolean isChunked() {
+		return false;
+	}
+
+	@Override
+	public Set<String> getTrailerNames() {
+		return Set.of();
+	}
+
+	@Override
+	public String getContentType() {
+		return CONTENT_TYPE;
+	}
+
+	@Override
+	public String getContentEncoding() {
+		//Apparently this is the correct value:
+		return null;
+	}
+
+	@Override
+	public boolean isStreaming() {
+		return false;
+	}
+
+	@Override
+	public Supplier<List<? extends Header>> getTrailers() {
+		return null;
+	}
+
+	@Override
+	public int available() {
+		return 0;
+	}
+
+	@Override
+	public void produce(DataStreamChannel channel) throws IOException {
+		produceContent( new ClientRest5ContentEncoder( channel ) );
+	}
+
+	@Override
+	public void releaseResources() {
+		close();
+	}
+
+
+	private static class ClientRest5ContentEncoder implements ContentEncoder {
+
+		private final DataStreamChannel channel;
+
+		public ClientRest5ContentEncoder(DataStreamChannel channel) {
+			this.channel = channel;
+		}
+
+		@Override
+		public int write(ByteBuffer src) throws IOException {
+			return channel.write( src );
+		}
+
+		@Override
+		public void complete() throws IOException {
+			channel.endStream();
+		}
+
+		@Override
+		public boolean isCompleted() {
+			return false;
+		}
+	}
+}

backend/elasticsearch-client/elasticsearch-rest5-client/src/main/java/org/hibernate/search/backend/elasticsearch/client/rest5/impl/ClientRest5GsonHttpEntity.java

@@ -7,22 +7,28 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URISyntaxException;
+import java.nio.ByteBuffer;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
+import org.hibernate.search.backend.elasticsearch.client.common.gson.entity.spi.ContentEncoder;
+import org.hibernate.search.backend.elasticsearch.client.common.gson.entity.spi.ContentProducer;
+import org.hibernate.search.backend.elasticsearch.client.common.gson.entity.spi.HttpAsyncContentProducerInputStream;
 import org.hibernate.search.backend.elasticsearch.client.common.spi.ElasticsearchRequestInterceptor;
 import org.hibernate.search.backend.elasticsearch.client.common.spi.ElasticsearchRequestInterceptorContext;
 import org.hibernate.search.util.common.AssertionFailure;
 
 import org.apache.hc.client5.http.protocol.HttpClientContext;
 import org.apache.hc.core5.http.EntityDetails;
+import org.apache.hc.core5.http.Header;
 import org.apache.hc.core5.http.HttpEntity;
 import org.apache.hc.core5.http.HttpEntityContainer;
 import org.apache.hc.core5.http.HttpRequest;
 import org.apache.hc.core5.http.HttpRequestInterceptor;
 import org.apache.hc.core5.http.NameValuePair;
 import org.apache.hc.core5.http.nio.AsyncEntityProducer;
+import org.apache.hc.core5.http.nio.DataStreamChannel;
 import org.apache.hc.core5.http.protocol.HttpContext;
 import org.apache.hc.core5.net.URIBuilder;
 
@@ -68,7 +74,7 @@ else if ( request instanceof HttpEntityContainer entityContainer ) {
 				if ( !producer.isRepeatable() ) {
 					throw new AssertionFailure( "Cannot sign AWS requests with non-repeatable entities" );
 				}
-				return new HttpAsyncEntityProducerInputStream( producer, 1024 );
+				return new HttpAsyncContentProducerInputStream( new ClientRest5GsonContentProducer( producer ), 1024 );
 			}
 			return null;
 		}
@@ -140,4 +146,50 @@ public String toString() {
 			return request.toString();
 		}
 	}
+
+	private static class GsonEncoderDataStreamChannel implements DataStreamChannel {
+		private final ContentEncoder encoder;
+
+		public GsonEncoderDataStreamChannel(ContentEncoder encoder) {
+			this.encoder = encoder;
+		}
+
+		@Override
+		public void requestOutput() {
+
+		}
+
+		@Override
+		public int write(ByteBuffer src) throws IOException {
+			return encoder.write( src );
+		}
+
+		@Override
+		public void endStream(List<? extends Header> trailers) throws IOException {
+			encoder.complete();
+		}
+
+		@Override
+		public void endStream() throws IOException {
+			encoder.complete();
+		}
+	}
+
+	private static class ClientRest5GsonContentProducer implements ContentProducer {
+		private final AsyncEntityProducer producer;
+
+		public ClientRest5GsonContentProducer(AsyncEntityProducer producer) {
+			this.producer = producer;
+		}
+
+		@Override
+		public void produceContent(ContentEncoder encoder) throws IOException {
+			producer.produce( new GsonEncoderDataStreamChannel( encoder ) );
+		}
+
+		@Override
+		public void close() throws IOException {
+			producer.releaseResources();
+		}
+	}
 }