Preserve connection keep-alive for JMeter compatibility

src/main/java/com/blazemeter/jmeter/http2/core/HTTP2JettyClient.java

@@ -1138,6 +1138,7 @@ public HTTPSampleResult retryWithHTTP11Only(HTTP2Sampler sampler, HTTPSampleResu
 
     // Copy body if present
     setBody(http11Request, sampler, result);
+    JmeterRequestHeadersSupport.prepareFromSampler(http11Request, sampler.getUseKeepAlive());
 
     // Send request
     lowLevelDebug("Sending HTTP/1.1 fallback request");
@@ -1199,6 +1200,7 @@ private ContentResponse sendWithHTTP11Only(Request originalRequest,
         // will be lost, but this is an edge case.
       }
       ensureHostHeader(http11Request, uri);
+      JmeterRequestHeadersSupport.copySamplerHeaderState(originalRequest, http11Request);
 
       configureContentDecodersAndCapture(httpClientHttp1Only, http11Request);
 
@@ -1324,6 +1326,7 @@ private void samplePrepareRequest(Request request,
     result.sampleStart();
 
     setBody(request, sampler, result);
+    JmeterRequestHeadersSupport.prepareFromSampler(request, sampler.getUseKeepAlive());
     initializeSentBytes(result, request);
 
   }
@@ -3452,7 +3455,8 @@ private String getSerializedRequestHeaders(Request request, boolean refresh) {
     if (!refresh && cached instanceof String) {
       return (String) cached;
     }
-    String serialized = buildHeadersString(request.getHeaders());
+    String serialized = buildHeadersString(
+        JmeterRequestHeadersSupport.headersForSampleResult(request));
     request.attribute(ATTR_REQUEST_HEADERS_SERIALIZED, serialized);
     return serialized;
   }

src/main/java/com/blazemeter/jmeter/http2/core/JmeterRequestHeadersSupport.java

@@ -0,0 +1,104 @@
+package com.blazemeter.jmeter.http2.core;
+
+import org.apache.jmeter.protocol.http.util.HTTPConstants;
+import org.eclipse.jetty.client.Request;
+import org.eclipse.jetty.http.HttpFields;
+import org.eclipse.jetty.http.HttpHeader;
+import org.eclipse.jetty.http.HttpVersion;
+
+/**
+ * Preserves JMeter-visible request headers in {@code HTTPSampleResult#getRequestHeaders()}
+ * when Jetty omits or strips them on the wire.
+ *
+ * <p>Counterpart of {@link JmeterCompressionHeadersSupport} for outbound request headers.
+ */
+final class JmeterRequestHeadersSupport {
+
+  static final String ATTR_USE_KEEPALIVE = "bzm.useKeepAlive";
+
+  private JmeterRequestHeadersSupport() {
+  }
+
+  /**
+   * Applies sampler-driven request headers before send, matching {@code HTTPHC4Impl#setupRequest}.
+   */
+  static void prepareFromSampler(Request request, boolean useKeepAlive) {
+    if (request == null) {
+      return;
+    }
+    request.attribute(ATTR_USE_KEEPALIVE, useKeepAlive);
+    applyConnectionHeader(request, useKeepAlive);
+    // TODO: HC4 sends explicit empty User-Agent when none is configured (disableDefaultUserAgent).
+    // prepareEmptyUserAgentHeader(request);
+  }
+
+  /**
+   * Copies sampler header settings to a cloned request (for example HTTP/1.1 fallback).
+   */
+  static void copySamplerHeaderState(Request from, Request to) {
+    if (from == null || to == null) {
+      return;
+    }
+    Object useKeepAlive = from.getAttributes().get(ATTR_USE_KEEPALIVE);
+    if (useKeepAlive != null) {
+      to.attribute(ATTR_USE_KEEPALIVE, useKeepAlive);
+    }
+  }
+
+  /**
+   * Request headers to record in the sample after Jetty may have adjusted the live fields.
+   */
+  static HttpFields headersForSampleResult(Request request) {
+    if (request == null) {
+      return HttpFields.EMPTY;
+    }
+    HttpFields.Mutable merged = HttpFields.build(request.getHeaders());
+    restoreConnectionHeaderForSample(request, merged);
+    // restoreEmptyUserAgentForSample(request, merged);
+    return merged;
+  }
+
+  private static void applyConnectionHeader(Request request, boolean useKeepAlive) {
+    if (!shouldSendConnectionHeader(request)) {
+      return;
+    }
+    HttpFields.Mutable mutableHeaders = mutableHeaders(request);
+    if (mutableHeaders == null || mutableHeaders.contains(HttpHeader.CONNECTION)) {
+      return;
+    }
+    if (useKeepAlive) {
+      mutableHeaders.put(HTTPConstants.HEADER_CONNECTION, HTTPConstants.KEEP_ALIVE);
+    } else {
+      mutableHeaders.put(HTTPConstants.HEADER_CONNECTION, HTTPConstants.CONNECTION_CLOSE);
+    }
+  }
+
+  private static void restoreConnectionHeaderForSample(
+      Request request, HttpFields.Mutable headers) {
+    Object useKeepAlive = request.getAttributes().get(ATTR_USE_KEEPALIVE);
+    if (useKeepAlive == null || !shouldSendConnectionHeader(request)) {
+      return;
+    }
+    if (Boolean.TRUE.equals(useKeepAlive)) {
+      headers.put(HTTPConstants.HEADER_CONNECTION, HTTPConstants.KEEP_ALIVE);
+    } else {
+      headers.put(HTTPConstants.HEADER_CONNECTION, HTTPConstants.CONNECTION_CLOSE);
+    }
+  }
+
+  private static boolean shouldSendConnectionHeader(Request request) {
+    if (request != null && request.getVersion() == HttpVersion.HTTP_2) {
+      return false;
+    }
+    HttpFields headers = request.getHeaders();
+    return headers == null || !headers.contains(HttpHeader.CONNECTION);
+  }
+
+  private static HttpFields.Mutable mutableHeaders(Request request) {
+    HttpFields headers = request.getHeaders();
+    if (headers instanceof HttpFields.Mutable) {
+      return (HttpFields.Mutable) headers;
+    }
+    return null;
+  }
+}

src/test/java/com/blazemeter/jmeter/http2/core/HttpKeepAliveSampleHeadersTest.java

@@ -0,0 +1,132 @@
+package com.blazemeter.jmeter.http2.core;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import com.blazemeter.jmeter.http2.sampler.HTTP2Sampler;
+import com.blazemeter.jmeter.http2.sampler.JMeterTestUtils;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.net.ServerSocket;
+import java.net.Socket;
+import java.net.URL;
+import java.nio.charset.StandardCharsets;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import org.apache.jmeter.protocol.http.sampler.HTTPSampleResult;
+import org.apache.jmeter.protocol.http.util.HTTPConstants;
+import org.apache.jmeter.util.JMeterUtils;
+import org.junit.After;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+public class HttpKeepAliveSampleHeadersTest {
+
+  private ExecutorService executor;
+
+  @BeforeClass
+  public static void setupJmeter() {
+    JMeterTestUtils.setupJmeterEnv();
+  }
+
+  @After
+  public void tearDown() throws Exception {
+    if (executor != null) {
+      executor.shutdownNow();
+      executor.awaitTermination(5, TimeUnit.SECONDS);
+      executor = null;
+    }
+  }
+
+  @Test
+  public void sampleRequestHeadersIncludeKeepAliveWhenEnabled() throws Exception {
+    configureLegacyHttp1Only();
+
+    try (ServerSocket serverSocket = new ServerSocket(0)) {
+      int port = serverSocket.getLocalPort();
+      Future<String> received = startWireCapture(serverSocket);
+
+      HTTP2JettyClient jettyClient = new HTTP2JettyClient(false, "keepalive-sample-test");
+      jettyClient.start();
+      try {
+        HTTPSampleResult sampleResult = samplePlainHttp(jettyClient, port, true);
+        String wireRequest = received.get(10, TimeUnit.SECONDS);
+
+        assertThat(sampleResult.getRequestHeaders())
+            .contains(HTTPConstants.HEADER_CONNECTION + ": " + HTTPConstants.KEEP_ALIVE);
+        assertThat(wireRequest).contains("GET /test HTTP/1.1");
+        assertThat(wireRequest.toLowerCase()).doesNotContain("connection: keep-alive");
+      } finally {
+        jettyClient.stop();
+      }
+    }
+  }
+
+  @Test
+  public void sampleRequestHeadersIncludeCloseWhenKeepAliveDisabled() throws Exception {
+    configureLegacyHttp1Only();
+
+    try (ServerSocket serverSocket = new ServerSocket(0)) {
+      int port = serverSocket.getLocalPort();
+      Future<String> received = startWireCapture(serverSocket);
+
+      HTTP2JettyClient jettyClient = new HTTP2JettyClient(false, "keepalive-close-test");
+      jettyClient.start();
+      try {
+        HTTPSampleResult sampleResult = samplePlainHttp(jettyClient, port, false);
+        String wireRequest = received.get(10, TimeUnit.SECONDS);
+
+        assertThat(sampleResult.getRequestHeaders())
+            .contains(HTTPConstants.HEADER_CONNECTION + ": "
+                + HTTPConstants.CONNECTION_CLOSE);
+        assertThat(wireRequest.toLowerCase()).contains("connection: close");
+      } finally {
+        jettyClient.stop();
+      }
+    }
+  }
+
+  private static void configureLegacyHttp1Only() {
+    JMeterUtils.setProperty("blazemeter.http.enableHttp2", "false");
+    JMeterUtils.setProperty("blazemeter.http.enableHttp3", "false");
+    JMeterUtils.setProperty("blazemeter.http.profile", "legacy");
+  }
+
+  private Future<String> startWireCapture(ServerSocket serverSocket) {
+    executor = Executors.newSingleThreadExecutor();
+    return executor.submit(() -> readRequest(serverSocket));
+  }
+
+  private static HTTPSampleResult samplePlainHttp(HTTP2JettyClient jettyClient, int port,
+                                                  boolean useKeepAlive) throws Exception {
+    HTTP2Sampler sampler = new HTTP2Sampler();
+    sampler.setMethod("GET");
+    sampler.setDomain("localhost");
+    sampler.setPort(port);
+    sampler.setPath("/test");
+    sampler.setProtocol("http");
+    sampler.setUseKeepAlive(useKeepAlive);
+
+    URL url = new URL("http", "localhost", port, "/test");
+    HTTPSampleResult result = new HTTPSampleResult();
+    result.setSampleLabel("GET_NO_PARAMETERS");
+    result.setHTTPMethod("GET");
+    result.setURL(url);
+
+    return jettyClient.sample(sampler, result, false, 0);
+  }
+
+  private static String readRequest(ServerSocket serverSocket) throws Exception {
+    try (Socket socket = serverSocket.accept();
+         InputStream in = socket.getInputStream();
+         OutputStream out = socket.getOutputStream()) {
+      byte[] buffer = new byte[4096];
+      int read = in.read(buffer);
+      out.write("HTTP/1.0 200 OK\r\nContent-Type: text/plain\r\n\r\n"
+          .getBytes(StandardCharsets.US_ASCII));
+      out.flush();
+      return new String(buffer, 0, read, StandardCharsets.US_ASCII);
+    }
+  }
+}

src/test/java/com/blazemeter/jmeter/http2/core/JmeterRequestHeadersSupportTest.java

@@ -0,0 +1,77 @@
+package com.blazemeter.jmeter.http2.core;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyString;
+
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.jmeter.protocol.http.util.HTTPConstants;
+import org.eclipse.jetty.client.Request;
+import org.eclipse.jetty.http.HttpFields;
+import org.eclipse.jetty.http.HttpHeader;
+import org.eclipse.jetty.http.HttpVersion;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+public class JmeterRequestHeadersSupportTest {
+
+  private static Request mockRequest(HttpFields.Mutable headers) {
+    Map<String, Object> attributes = new HashMap<>();
+    Request request = Mockito.mock(Request.class);
+    Mockito.when(request.getHeaders()).thenReturn(headers);
+    Mockito.when(request.getVersion()).thenReturn(HttpVersion.HTTP_1_1);
+    Mockito.when(request.getAttributes()).thenReturn(attributes);
+    Mockito.doAnswer(invocation -> {
+      attributes.put(invocation.getArgument(0), invocation.getArgument(1));
+      return invocation.getMock();
+    }).when(request).attribute(anyString(), any());
+    return request;
+  }
+
+  @Test
+  public void prepareFromSamplerAddsConnectionKeepAliveWhenEnabled() {
+    HttpFields.Mutable headers = HttpFields.build();
+    Request request = mockRequest(headers);
+
+    JmeterRequestHeadersSupport.prepareFromSampler(request, true);
+
+    assertThat(headers.get(HttpHeader.CONNECTION)).isEqualTo(HTTPConstants.KEEP_ALIVE);
+  }
+
+  @Test
+  public void headersForSampleResultRestoresKeepAliveAfterHeadersCleared() {
+    HttpFields.Mutable headers = HttpFields.build();
+    Request request = mockRequest(headers);
+    JmeterRequestHeadersSupport.prepareFromSampler(request, true);
+    headers.remove(HttpHeader.CONNECTION);
+
+    HttpFields sampleHeaders = JmeterRequestHeadersSupport.headersForSampleResult(request);
+
+    assertThat(sampleHeaders.get(HttpHeader.CONNECTION)).isEqualTo(HTTPConstants.KEEP_ALIVE);
+  }
+
+  @Test
+  public void headersForSampleResultRestoresConnectionCloseWhenDisabled() {
+    HttpFields.Mutable headers = HttpFields.build();
+    Request request = mockRequest(headers);
+    JmeterRequestHeadersSupport.prepareFromSampler(request, false);
+    headers.remove(HttpHeader.CONNECTION);
+
+    HttpFields sampleHeaders = JmeterRequestHeadersSupport.headersForSampleResult(request);
+
+    assertThat(sampleHeaders.get(HttpHeader.CONNECTION))
+        .isEqualTo(HTTPConstants.CONNECTION_CLOSE);
+  }
+
+  @Test
+  public void prepareFromSamplerDoesNotOverrideExplicitConnectionHeader() {
+    HttpFields.Mutable headers = HttpFields.build()
+        .add(HttpHeader.CONNECTION, "Upgrade");
+    Request request = mockRequest(headers);
+
+    JmeterRequestHeadersSupport.prepareFromSampler(request, true);
+
+    assertThat(headers.get(HttpHeader.CONNECTION)).isEqualTo("Upgrade");
+  }
+}