NIFI-16152 Switched from Jersey to Jetty for gzip compression (#11489)

- Added drain and close for replicated request response body to handle gzip encoded responses

Author: exceptionfactory

nifi-framework-bundle/nifi-framework/nifi-framework-cluster/src/main/java/org/apache/nifi/cluster/coordination/http/replication/StandardAsyncClusterResponse.java

@@ -17,6 +17,7 @@
 
 package org.apache.nifi.cluster.coordination.http.replication;
 
+import jakarta.ws.rs.core.Response;
 import org.apache.nifi.cluster.coordination.http.HttpResponseMapper;
 import org.apache.nifi.cluster.manager.NodeResponse;
 import org.apache.nifi.cluster.protocol.NodeIdentifier;
@@ -216,6 +217,7 @@ public synchronized NodeResponse getMergedResponse(final boolean triggerCallback
 
         final long start = System.nanoTime();
         mergedResponse = responseMapper.mapResponses(uri, method, nodeResponses, merge);
+        closeUnusedResponses(nodeResponses);
         final long nanos = System.nanoTime() - start;
         addTiming("Map/Merge Responses", "All Nodes", nanos);
 
@@ -329,6 +331,23 @@ public String toString() {
             + ", responses=" + getCompletedNodeIdentifiers().size() + "/" + responseMap.size() + "]";
     }
 
+    private void closeUnusedResponses(final Set<NodeResponse> nodeResponses) {
+        final Response clientMergedResponse = mergedResponse.getClientResponse();
+        if (mergedResponse.getUpdatedEntity() != null) {
+            // Close merged response since Updated Entity used in place of stream
+            clientMergedResponse.close();
+        }
+
+        for (final NodeResponse nodeResponse : nodeResponses) {
+            final Response clientNodeResponse = nodeResponse.getClientResponse();
+            if (clientMergedResponse == clientNodeResponse) {
+                continue;
+            }
+
+            nodeResponse.close();
+        }
+    }
+
     private static class ResponseHolder {
         private final long nanoStart;
         private long requestNanos;

nifi-framework-bundle/nifi-framework/nifi-framework-cluster/src/main/java/org/apache/nifi/cluster/coordination/http/replication/ThreadPoolRequestReplicator.java

@@ -523,8 +523,13 @@ public void onCompletion(final NodeResponse nodeResponse) {
                         // If all nodes responded with 202-Accepted, then we can replicate the original request
                         // to all nodes and we are finished.
                         if (dissentingCount == 0) {
-                            logger.debug("Received verification from all {} nodes that mutable request {} {} can be made", numNodes, method, uri.getPath());
-                            replicate(nodeIds, method, uri, entity, headers, false, clusterResponse, true, merge, monitor);
+                            try {
+                                logger.debug("Received verification from all {} nodes that mutable request {} {} can be made", numNodes, method, uri.getPath());
+                                replicate(nodeIds, method, uri, entity, headers, false, clusterResponse, true, merge, monitor);
+                            } finally {
+                                // Close HTTP Responses after replication completed
+                                nodeResponses.forEach(NodeResponse::close);
+                            }
                             return;
                         }
 
@@ -582,6 +587,9 @@ public void onCompletion(final NodeResponse nodeResponse) {
                                     clusterResponse.setFailure(failure, response.getNodeId());
                                 }
                             }
+
+                            // Close verification responses for cancelled transactions
+                            nodeResponses.forEach(NodeResponse::close);
                         } finally {
                             if (monitor != null) {
                                 synchronized (monitor) {

nifi-framework-bundle/nifi-framework/nifi-framework-cluster/src/main/java/org/apache/nifi/cluster/coordination/http/replication/client/StandardHttpReplicationClient.java

@@ -17,6 +17,7 @@
 package org.apache.nifi.cluster.coordination.http.replication.client;
 
 import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.core.JsonParser;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.module.jakarta.xmlbind.JakartaXmlBindAnnotationIntrospector;
 import jakarta.ws.rs.core.MultivaluedHashMap;
@@ -41,6 +42,7 @@
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.OutputStream;
 import java.io.UncheckedIOException;
 import java.net.URI;
 import java.util.LinkedHashMap;
@@ -113,6 +115,8 @@ public StandardHttpReplicationClient(final WebClientService webClientService, fi
 
         objectMapper.setDefaultPropertyInclusion(JsonInclude.Value.construct(JsonInclude.Include.NON_NULL, JsonInclude.Include.ALWAYS));
         objectMapper.setAnnotationIntrospector(new JakartaXmlBindAnnotationIntrospector(objectMapper.getTypeFactory()));
+        // Disable closing source streams to allow draining and subsequent closing
+        objectMapper.disable(JsonParser.Feature.AUTO_CLOSE_SOURCE);
 
         jsonSerializer = new JsonEntitySerializer(objectMapper);
         xmlSerializer = new XmlEntitySerializer();
@@ -208,11 +212,25 @@ private Response replicate(final HttpRequestBodySpec httpRequestBodySpec, final
 
         final InputStream responseBody = getResponseBody(responseEntity.body(), headers);
         final Runnable closeCallback = () -> {
+            try {
+                // Drain raw response stream before closing the Response Entity
+                responseEntity.body().transferTo(OutputStream.nullOutputStream());
+            } catch (final IOException e) {
+                logger.debug("Drain failed for Replicated {} {} HTTP {}", method, location, statusCode, e);
+            }
+
             try {
                 responseEntity.close();
             } catch (final IOException e) {
                 logger.warn("Close failed for Replicated {} {} HTTP {}", method, location, statusCode, e);
             }
+
+            try {
+                // Release resources for gzip wrapped streams
+                responseBody.close();
+            } catch (final IOException e) {
+                logger.warn("Close failed for Replicated Response Body {} {} HTTP {}", method, location, statusCode, e);
+            }
         };
 
         final long elapsed = System.currentTimeMillis() - started;

nifi-framework-bundle/nifi-framework/nifi-framework-cluster/src/main/java/org/apache/nifi/cluster/coordination/http/replication/io/ReplicatedResponse.java

@@ -17,8 +17,6 @@
 
 package org.apache.nifi.cluster.coordination.http.replication.io;
 
-import com.fasterxml.jackson.core.JsonFactory;
-import com.fasterxml.jackson.core.JsonParser;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import jakarta.ws.rs.core.EntityTag;
 import jakarta.ws.rs.core.GenericType;
@@ -51,28 +49,28 @@ public class ReplicatedResponse extends Response {
     private static final int MAXIMUM_BUFFER_SIZE = 1048576;
     private static final int CONTENT_LENGTH_UNKNOWN = -1;
 
-    private final ObjectMapper codec;
+    private final ObjectMapper objectMapper;
     private final InputStream responseBody;
     private final MultivaluedMap<String, String> responseHeaders;
     private final URI location;
     private final int statusCode;
     private final Runnable closeCallback;
     private final int contentLength;
 
-    private final JsonFactory jsonFactory = new JsonFactory();
-
     private final byte[] bufferedResponseBody;
 
+    private Object bufferedEntity;
+
     public ReplicatedResponse(
-            final ObjectMapper codec,
+            final ObjectMapper objectMapper,
             final InputStream responseBody,
             final MultivaluedMap<String, String> responseHeaders,
             final URI location,
             final int statusCode,
             final int contentLength,
             final Runnable closeCallback
     ) {
-        this.codec = codec;
+        this.objectMapper = objectMapper;
         this.responseBody = responseBody;
         this.responseHeaders = responseHeaders;
         this.location = location;
@@ -101,42 +99,30 @@ public StatusType getStatusInfo() {
 
     @Override
     public Object getEntity() {
-        final InputStream responseBodyStream = getResponseBodyStream();
-
-        try {
-            final JsonParser parser = jsonFactory.createParser(responseBodyStream);
-            parser.setCodec(codec);
-            return parser.readValueAs(Object.class);
-        } catch (final Exception e) {
-            throw new RuntimeException("Failed to parse response", e);
+        if (bufferedEntity == null) {
+            // Read response entity to buffered entity to support multiple invocations
+            bufferedEntity = readResponseEntity(Object.class);
         }
+
+        return bufferedEntity;
     }
 
     @Override
     @SuppressWarnings("unchecked")
-    public <T> T readEntity(Class<T> entityType) {
-        final InputStream responseBodyStream = getResponseBodyStream();
+    public <T> T readEntity(final Class<T> entityType) {
+        final T entity;
 
+        // Return raw response body stream when requested without buffering
         if (InputStream.class.equals(entityType)) {
-            return (T) responseBodyStream;
-        }
-
-        if (String.class.equals(entityType)) {
-            try {
-                final byte[] responseBytes = responseBodyStream.readAllBytes();
-                return (T) new String(responseBytes, StandardCharsets.UTF_8);
-            } catch (final IOException e) {
-                throw new UncheckedIOException("Read Replicated Response Body to String failed for %s".formatted(location), e);
-            }
+            return (T) getResponseBodyStream();
         }
 
-        try {
-            final JsonParser parser = jsonFactory.createParser(responseBodyStream);
-            parser.setCodec(codec);
-            return parser.readValueAs(entityType);
-        } catch (final Exception e) {
-            throw new RuntimeException("Failed to parse response as entity of type " + entityType, e);
+        if (bufferedEntity == null) {
+            // Read response entity to buffered entity to support multiple invocations
+            bufferedEntity = readResponseEntity(entityType);
         }
+        entity = (T) bufferedEntity;
+        return entity;
     }
 
     @Override
@@ -287,6 +273,26 @@ private InputStream getResponseBodyStream() {
         return responseBodyStream;
     }
 
+    @SuppressWarnings("unchecked")
+    private <T> T readResponseEntity(final Class<T> entityType) {
+        final InputStream responseBodyStream = getResponseBodyStream();
+
+        if (String.class.equals(entityType)) {
+            try {
+                final byte[] responseBytes = responseBodyStream.readAllBytes();
+                return (T) new String(responseBytes, StandardCharsets.UTF_8);
+            } catch (final IOException e) {
+                throw new UncheckedIOException("Read Replicated Response Body to String failed for %s".formatted(location), e);
+            }
+        }
+
+        try {
+            return objectMapper.readValue(responseBodyStream, entityType);
+        } catch (final Exception e) {
+            throw new RuntimeException("Failed to parse response as Entity [%s] for %s".formatted(entityType, location), e);
+        }
+    }
+
     private static byte[] readResponseBody(final InputStream inputStream, final URI location, final int statusCode) {
         try {
             return inputStream.readAllBytes();

nifi-framework-bundle/nifi-framework/nifi-framework-cluster/src/main/java/org/apache/nifi/cluster/manager/NodeResponse.java

@@ -29,6 +29,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.io.Closeable;
 import java.io.InputStream;
 import java.net.URI;
 import java.util.List;
@@ -47,7 +48,7 @@
  * This class overrides hashCode and equals and considers two instances to be equal if they have the equal NodeIdentifiers.
  *
  */
-public class NodeResponse {
+public class NodeResponse implements Closeable {
 
     private static final Logger logger = LoggerFactory.getLogger(NodeResponse.class);
     private final String httpMethod;
@@ -269,4 +270,11 @@ public String toString() {
             .append(",Duration=").append(TimeUnit.MILLISECONDS.convert(requestDurationNanos, TimeUnit.NANOSECONDS)).append(" ms]");
         return sb.toString();
     }
+
+    @Override
+    public void close() {
+        if (response != null) {
+            response.close();
+        }
+    }
 }

nifi-framework-bundle/nifi-framework/nifi-web/nifi-jetty/pom.xml

@@ -85,6 +85,14 @@
             <artifactId>nifi-web-servlet-shared</artifactId>
             <version>2.11.0-SNAPSHOT</version>
         </dependency>
+        <dependency>
+            <groupId>org.eclipse.jetty.compression</groupId>
+            <artifactId>jetty-compression-server</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.eclipse.jetty.compression</groupId>
+            <artifactId>jetty-compression-gzip</artifactId>
+        </dependency>
         <dependency>
             <groupId>org.eclipse.jetty</groupId>
             <artifactId>jetty-deploy</artifactId>

nifi-framework-bundle/nifi-framework/nifi-web/nifi-jetty/src/main/java/org/apache/nifi/web/server/JettyServer.java

@@ -69,6 +69,7 @@
 import org.apache.nifi.web.server.filter.RequestFilterProvider;
 import org.apache.nifi.web.server.filter.RestApiRequestFilterProvider;
 import org.apache.nifi.web.server.filter.StandardRequestFilterProvider;
+import org.eclipse.jetty.compression.server.CompressionHandler;
 import org.eclipse.jetty.deploy.StandardDeployer;
 import org.eclipse.jetty.ee.webapp.WebAppClassLoader;
 import org.eclipse.jetty.ee11.servlet.ErrorPageErrorHandler;
@@ -222,7 +223,13 @@ public void init() {
                 deployer = new StandardDeployer(contextHandlerCollection);
                 server.addBean(deployer);
 
-                serverHandlerCollection.addHandler(contextHandlerCollection);
+                // Deploy the web applications beneath the CompressionHandler
+                final CompressionHandler compressionHandler = serverHandlerCollection.getDescendant(CompressionHandler.class);
+                if (compressionHandler == null) {
+                    throw new IllegalStateException("Compression Handler not configured: Server Provider configuration failed");
+                } else {
+                    compressionHandler.setHandler(contextHandlerCollection);
+                }
             } else {
                 throw new IllegalStateException("Server Handler not Handler.Collection: Server Provider configuration failed");
             }

nifi-framework-bundle/nifi-framework/nifi-web/nifi-jetty/src/main/java/org/apache/nifi/web/server/StandardServerProvider.java

@@ -21,8 +21,12 @@
 import org.apache.nifi.web.server.connector.FrameworkServerConnectorFactory;
 import org.apache.nifi.web.server.handler.ContextPathRedirectPatternRule;
 import org.apache.nifi.web.server.handler.HeaderWriterHandler;
+import org.apache.nifi.web.server.handler.UnsupportedContentEncodingHandler;
 import org.apache.nifi.web.server.log.RequestLogProvider;
 import org.apache.nifi.web.server.log.StandardRequestLogProvider;
+import org.eclipse.jetty.compression.gzip.GzipCompression;
+import org.eclipse.jetty.compression.server.CompressionConfig;
+import org.eclipse.jetty.compression.server.CompressionHandler;
 import org.eclipse.jetty.rewrite.handler.RedirectPatternRule;
 import org.eclipse.jetty.rewrite.handler.RewriteHandler;
 import org.eclipse.jetty.server.Handler;
@@ -47,6 +51,8 @@
  * Standard implementation of Server Provider with default Handlers
  */
 class StandardServerProvider implements ServerProvider {
+    private static final String ROOT_PATH = "/";
+
     private static final String ALL_PATHS_PATTERN = "/*";
 
     private static final String FRONTEND_CONTEXT_PATH = "/nifi";
@@ -139,6 +145,27 @@ private Handler getStandardHandler() {
         // Set Handler for standard response headers
         standardHandler.addHandler(new HeaderWriterHandler());
 
+        // Reject requests that declare a Content-Encoding because request bodies are not decompressed
+        standardHandler.addHandler(new UnsupportedContentEncodingHandler());
+
+        // Set Handler for response compression
+        standardHandler.addHandler(getCompressionHandler());
+
         return standardHandler;
     }
+
+    private CompressionHandler getCompressionHandler() {
+        final GzipCompression gzipCompression = new GzipCompression();
+        final CompressionHandler compressionHandler = new CompressionHandler();
+        compressionHandler.putCompression(gzipCompression);
+
+        final CompressionConfig compressionConfig = CompressionConfig.builder()
+                .defaults()
+                // Disable decompression of requests
+                .decompressExcludeEncoding(gzipCompression.getEncodingName())
+                .build();
+        compressionHandler.putConfiguration(ROOT_PATH, compressionConfig);
+
+        return compressionHandler;
+    }
 }