Skip to content

Commit

Permalink
NIFI-13807: Ensure that we always close InputStream / OutputSTrema fo…
Browse files Browse the repository at this point in the history 
…r HttpExchange. Verified fix using InvokeHTTP to hit the health endpoint hundreds of thousands of times and verified NiFi heap usage was still at 200 MB after a young GC event
  • Loading branch information
@markap14
markap14 committed Sep 26, 2024
1 parent 1548a21 commit 4ee052a
Show file tree
Hide file tree
Showing 5 changed files with 86 additions and 51 deletions.
43 changes: 23 additions & 20 deletions ...ramework/nifi-runtime/src/main/java/org/apache/nifi/runtime/HealthClusterHttpHandler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import org.apache.nifi.cluster.ClusterDetailsFactory;
import org.apache.nifi.cluster.ConnectionState;
import org.apache.nifi.controller.DecommissionTask;
import org.apache.nifi.util.HttpExchangeUtils;

import java.io.IOException;
import java.io.OutputStream;
Expand Down Expand Up @@ -57,28 +58,30 @@ class HealthClusterHttpHandler implements HttpHandler {

@Override
public void handle(final HttpExchange exchange) throws IOException {
HttpExchangeUtils.drainRequestBody(exchange);

final String requestMethod = exchange.getRequestMethod();

final OutputStream responseBody = exchange.getResponseBody();

if (GET_METHOD.contentEquals(requestMethod)) {
exchange.getResponseHeaders().set(CONTENT_TYPE_HEADER, TEXT_PLAIN);
final ConnectionState connectionState = getConnectionState();
final String status = STATUS.formatted(connectionState);
final byte[] response = status.getBytes(StandardCharsets.UTF_8);
final int responseCode = getResponseCode(connectionState);
exchange.sendResponseHeaders(responseCode, response.length);
responseBody.write(response);
} else if (DELETE_METHOD.contentEquals(requestMethod)) {
startDecommission();

exchange.getResponseHeaders().set(CONTENT_TYPE_HEADER, TEXT_PLAIN);
final String status = STATUS.formatted(ConnectionState.OFFLOADING);
final byte[] response = status.getBytes(StandardCharsets.UTF_8);
exchange.sendResponseHeaders(HTTP_ACCEPTED, response.length);
responseBody.write(response);
} else {
exchange.sendResponseHeaders(HTTP_BAD_METHOD, NO_RESPONSE_BODY);
try (final OutputStream responseBody = exchange.getResponseBody()) {
if (GET_METHOD.contentEquals(requestMethod)) {
exchange.getResponseHeaders().set(CONTENT_TYPE_HEADER, TEXT_PLAIN);
final ConnectionState connectionState = getConnectionState();
final String status = STATUS.formatted(connectionState);
final byte[] response = status.getBytes(StandardCharsets.UTF_8);
final int responseCode = getResponseCode(connectionState);
exchange.sendResponseHeaders(responseCode, response.length);
responseBody.write(response);
} else if (DELETE_METHOD.contentEquals(requestMethod)) {
startDecommission();

exchange.getResponseHeaders().set(CONTENT_TYPE_HEADER, TEXT_PLAIN);
final String status = STATUS.formatted(ConnectionState.OFFLOADING);
final byte[] response = status.getBytes(StandardCharsets.UTF_8);
exchange.sendResponseHeaders(HTTP_ACCEPTED, response.length);
responseBody.write(response);
} else {
exchange.sendResponseHeaders(HTTP_BAD_METHOD, NO_RESPONSE_BODY);
}
}
}

Expand Down
24 changes: 13 additions & 11 deletions ...work/nifi-runtime/src/main/java/org/apache/nifi/runtime/HealthDiagnosticsHttpHandler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import org.apache.nifi.NiFiServer;
import org.apache.nifi.diagnostics.DiagnosticsDump;
import org.apache.nifi.diagnostics.DiagnosticsFactory;
import org.apache.nifi.util.HttpExchangeUtils;

import java.io.IOException;
import java.io.OutputStream;
Expand Down Expand Up @@ -54,22 +55,23 @@ class HealthDiagnosticsHttpHandler implements HttpHandler {

@Override
public void handle(final HttpExchange exchange) throws IOException {
final String requestMethod = exchange.getRequestMethod();
HttpExchangeUtils.drainRequestBody(exchange);

if (GET_METHOD.contentEquals(requestMethod)) {
exchange.getResponseHeaders().set(CONTENT_TYPE_HEADER, TEXT_PLAIN);
exchange.sendResponseHeaders(HTTP_OK, STREAM_RESPONSE_BODY);
final String requestMethod = exchange.getRequestMethod();
try (final OutputStream responseBody = exchange.getResponseBody()) {
if (GET_METHOD.contentEquals(requestMethod)) {
exchange.getResponseHeaders().set(CONTENT_TYPE_HEADER, TEXT_PLAIN);
exchange.sendResponseHeaders(HTTP_OK, STREAM_RESPONSE_BODY);

final URI requestUri = exchange.getRequestURI();
final boolean verboseRequested = getVerboseRequested(requestUri);
final URI requestUri = exchange.getRequestURI();
final boolean verboseRequested = getVerboseRequested(requestUri);

final DiagnosticsFactory diagnosticsFactory = server.getDiagnosticsFactory();
final DiagnosticsDump diagnosticsDump = diagnosticsFactory.create(verboseRequested);
try (OutputStream responseBody = exchange.getResponseBody()) {
final DiagnosticsFactory diagnosticsFactory = server.getDiagnosticsFactory();
final DiagnosticsDump diagnosticsDump = diagnosticsFactory.create(verboseRequested);
diagnosticsDump.writeTo(responseBody);
} else {
exchange.sendResponseHeaders(HTTP_BAD_METHOD, NO_RESPONSE_BODY);
}
} else {
exchange.sendResponseHeaders(HTTP_BAD_METHOD, NO_RESPONSE_BODY);
}
}

Expand Down
21 changes: 12 additions & 9 deletions .../nifi-framework/nifi-runtime/src/main/java/org/apache/nifi/runtime/HealthHttpHandler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

import com.sun.net.httpserver.HttpExchange;
import com.sun.net.httpserver.HttpHandler;
import org.apache.nifi.util.HttpExchangeUtils;

import java.io.IOException;
import java.io.OutputStream;
Expand All @@ -42,17 +43,19 @@ class HealthHttpHandler implements HttpHandler {

@Override
public void handle(final HttpExchange exchange) throws IOException {
final String requestMethod = exchange.getRequestMethod();
HttpExchangeUtils.drainRequestBody(exchange);

final OutputStream responseBody = exchange.getResponseBody();
final String requestMethod = exchange.getRequestMethod();

if (GET_METHOD.contentEquals(requestMethod)) {
exchange.getResponseHeaders().set(CONTENT_TYPE_HEADER, TEXT_PLAIN);
final byte[] response = STATUS_UP.getBytes(StandardCharsets.UTF_8);
exchange.sendResponseHeaders(HTTP_OK, response.length);
responseBody.write(response);
} else {
exchange.sendResponseHeaders(HTTP_BAD_METHOD, NO_RESPONSE_BODY);
try (final OutputStream responseBody = exchange.getResponseBody()) {
if (GET_METHOD.contentEquals(requestMethod)) {
exchange.getResponseHeaders().set(CONTENT_TYPE_HEADER, TEXT_PLAIN);
final byte[] response = STATUS_UP.getBytes(StandardCharsets.UTF_8);
exchange.sendResponseHeaders(HTTP_OK, response.length);
responseBody.write(response);
} else {
exchange.sendResponseHeaders(HTTP_BAD_METHOD, NO_RESPONSE_BODY);
}
}
}
}
24 changes: 13 additions & 11 deletions ...rk/nifi-runtime/src/main/java/org/apache/nifi/runtime/HealthStatusHistoryHttpHandler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import org.apache.nifi.NiFiServer;
import org.apache.nifi.controller.status.history.StatusHistoryDump;
import org.apache.nifi.controller.status.history.StatusHistoryDumpFactory;
import org.apache.nifi.util.HttpExchangeUtils;

import java.io.IOException;
import java.io.OutputStream;
Expand Down Expand Up @@ -60,23 +61,24 @@ class HealthStatusHistoryHttpHandler implements HttpHandler {

@Override
public void handle(final HttpExchange exchange) throws IOException {
final String requestMethod = exchange.getRequestMethod();
HttpExchangeUtils.drainRequestBody(exchange);

if (GET_METHOD.contentEquals(requestMethod)) {
exchange.getResponseHeaders().set(CONTENT_TYPE_HEADER, APPLICATION_JSON);
exchange.sendResponseHeaders(HTTP_OK, STREAM_RESPONSE_BODY);
final String requestMethod = exchange.getRequestMethod();

final URI requestUri = exchange.getRequestURI();
final int daysRequested = getDaysRequested(requestUri);
try (final OutputStream responseBody = exchange.getResponseBody()) {
if (GET_METHOD.contentEquals(requestMethod)) {
exchange.getResponseHeaders().set(CONTENT_TYPE_HEADER, APPLICATION_JSON);
exchange.sendResponseHeaders(HTTP_OK, STREAM_RESPONSE_BODY);

final StatusHistoryDumpFactory statusHistoryDumpFactory = server.getStatusHistoryDumpFactory();
final StatusHistoryDump statusHistoryDump = statusHistoryDumpFactory.create(daysRequested);
final URI requestUri = exchange.getRequestURI();
final int daysRequested = getDaysRequested(requestUri);

try (OutputStream responseBody = exchange.getResponseBody()) {
final StatusHistoryDumpFactory statusHistoryDumpFactory = server.getStatusHistoryDumpFactory();
final StatusHistoryDump statusHistoryDump = statusHistoryDumpFactory.create(daysRequested);
statusHistoryDump.writeTo(responseBody);
} else {
exchange.sendResponseHeaders(HTTP_BAD_METHOD, NO_RESPONSE_BODY);
}
} else {
exchange.sendResponseHeaders(HTTP_BAD_METHOD, NO_RESPONSE_BODY);
}
}

Expand Down
25 changes: 25 additions & 0 deletions ...dle/nifi-framework/nifi-runtime/src/main/java/org/apache/nifi/util/HttpExchangeUtils.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
package org.apache.nifi.util;

import com.sun.net.httpserver.HttpExchange;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.InputStream;

public class HttpExchangeUtils {
private static final Logger logger = LoggerFactory.getLogger(HttpExchangeUtils.class);

public static void drainRequestBody(final HttpExchange exchange) throws IOException {
final byte[] buffer = new byte[4096];
try (final InputStream in = exchange.getRequestBody()) {
while ((in.read(buffer)) != -1) {
// Ignore the data read, just drain the input stream
}
} catch (final IOException ioe) {
// Since we don't actually care about the contents of the input, we will ignore any Exceptions when reading from it.
logger.debug("Failed to fully drain HttpExchange InputStream", ioe);
}
}

}

0 comments on commit 4ee052a

Please sign in to comment.