Configure Buffer pools for Jetty server

This circumvents jetty/jetty.project#12670

Author: wendigo

http-server/src/main/java/io/airlift/http/server/HttpServer.java

@@ -22,6 +22,7 @@
 import io.airlift.http.server.jetty.MonitoredQueuedThreadPoolMBean;
 import io.airlift.log.Logger;
 import io.airlift.node.NodeInfo;
+import io.airlift.units.DataSize;
 import jakarta.annotation.PostConstruct;
 import jakarta.annotation.PreDestroy;
 import jakarta.servlet.Filter;
@@ -34,6 +35,8 @@
 import org.eclipse.jetty.http2.server.AuthorityCustomizer;
 import org.eclipse.jetty.http2.server.HTTP2CServerConnectionFactory;
 import org.eclipse.jetty.http2.server.HTTP2ServerConnectionFactory;
+import org.eclipse.jetty.io.ArrayByteBufferPool;
+import org.eclipse.jetty.io.ByteBufferPool;
 import org.eclipse.jetty.io.ConnectionStatistics;
 import org.eclipse.jetty.jmx.MBeanContainer;
 import org.eclipse.jetty.server.ConnectionFactory;
@@ -65,6 +68,7 @@
 import java.util.EnumSet;
 import java.util.List;
 import java.util.Optional;
+import java.util.OptionalLong;
 import java.util.Set;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
@@ -75,6 +79,7 @@
 import static com.google.common.base.Verify.verify;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static io.airlift.concurrent.Threads.daemonThreadsNamed;
+import static java.lang.Math.max;
 import static java.lang.Math.toIntExact;
 import static java.time.temporal.ChronoUnit.DAYS;
 import static java.util.Collections.list;
@@ -140,7 +145,12 @@ public HttpServer(
             log.info("Virtual threads support is enabled");
             threadPool.setVirtualThreadsExecutor(executor);
         }
-        server = new Server(threadPool);
+
+        int maxBufferSize = toIntExact(max(
+                toSafeBytes(config.getMaxRequestHeaderSize()).orElse(65536),
+                toSafeBytes(config.getMaxResponseHeaderSize()).orElse(65536)));
+
+        server = new Server(threadPool, null, createByteBufferPool(maxBufferSize, config));
         this.monitoredQueuedThreadPoolMBean = new MonitoredQueuedThreadPoolMBean(threadPool);
 
         boolean showStackTrace = config.isShowStackTrace();
@@ -274,14 +284,25 @@ public HttpServer(
         }
 
         server.setHandler(statsHandler);
-
         ErrorHandler errorHandler = new ErrorHandler();
         errorHandler.setShowMessageInTitle(showStackTrace);
         errorHandler.setShowStacks(showStackTrace);
         errorHandler.setDefaultResponseMimeType(TEXT_PLAIN.asString());
         server.setErrorHandler(errorHandler);
     }
 
+    private ByteBufferPool createByteBufferPool(int maxBufferSize, HttpServerConfig config)
+    {
+        return new ArrayByteBufferPool.Quadratic(
+                0,
+                maxBufferSize,
+                Integer.MAX_VALUE,
+                config.getMaxHeapMemory().map(DataSize::toBytes)
+                        .orElse(0L), // Use default heuristics for max heap memory
+                config.getMaxDirectMemory().map(DataSize::toBytes)
+                        .orElse(0L)); // Use default heuristics for max direct memory
+    }
+
     private ConnectionFactory[] insecureFactories(HttpServerConfig config, HttpConfiguration httpConfiguration)
     {
         HttpConnectionFactory http1 = new HttpConnectionFactory(httpConfiguration);
@@ -473,4 +494,13 @@ private static ServerConnector createServerConnector(
         connector.open(channel);
         return connector;
     }
+
+    private static OptionalLong toSafeBytes(DataSize dataSize)
+    {
+        if (dataSize == null) {
+            return OptionalLong.empty();
+        }
+
+        return OptionalLong.of(dataSize.toBytes());
+    }
 }

http-server/src/main/java/io/airlift/http/server/HttpServerConfig.java

@@ -17,14 +17,18 @@
 
 import io.airlift.configuration.Config;
 import io.airlift.configuration.ConfigDescription;
+import io.airlift.configuration.ConfigHidden;
 import io.airlift.configuration.DefunctConfig;
 import io.airlift.units.DataSize;
 import io.airlift.units.Duration;
 import io.airlift.units.MaxDataSize;
 import io.airlift.units.MinDataSize;
+import jakarta.validation.constraints.AssertTrue;
 import jakarta.validation.constraints.Min;
 import jakarta.validation.constraints.NotNull;
 
+import java.util.Optional;
+
 import static io.airlift.http.server.HttpServerConfig.ProcessForwardedMode.ACCEPT;
 import static io.airlift.http.server.HttpServerConfig.ProcessForwardedMode.REJECT;
 import static io.airlift.units.DataSize.Unit.KILOBYTE;
@@ -94,6 +98,9 @@ public class HttpServerConfig
 
     private boolean compressionEnabled = true;
 
+    private Optional<DataSize> maxHeapMemory = Optional.empty();
+    private Optional<DataSize> maxDirectMemory = Optional.empty();
+
     public boolean isHttpEnabled()
     {
         return httpEnabled;
@@ -473,6 +480,38 @@ public HttpServerConfig setCompressionEnabled(boolean compressionEnabled)
         return this;
     }
 
+    public Optional<@MinDataSize("16MB") DataSize> getMaxHeapMemory()
+    {
+        return maxHeapMemory;
+    }
+
+    @Config("http-server.max-heap-memory")
+    @ConfigHidden
+    public HttpServerConfig setMaxHeapMemory(DataSize maxHeapMemory)
+    {
+        this.maxHeapMemory = Optional.ofNullable(maxHeapMemory);
+        return this;
+    }
+
+    public Optional<@MinDataSize("16MB") DataSize> getMaxDirectMemory()
+    {
+        return maxDirectMemory;
+    }
+
+    @Config("http-server.max-direct-memory")
+    @ConfigHidden
+    public HttpServerConfig setMaxDirectMemory(DataSize maxDirectMemory)
+    {
+        this.maxDirectMemory = Optional.ofNullable(maxDirectMemory);
+        return this;
+    }
+
+    @AssertTrue(message = "either both http-server.max-heap-memory and http-server.max-direct-memory are set or none of them")
+    public boolean eitherBothMemorySettingsAreSetOrNone()
+    {
+        return maxHeapMemory.isPresent() == maxDirectMemory.isPresent();
+    }
+
     public enum ProcessForwardedMode
     {
         ACCEPT,

http-server/src/test/java/io/airlift/http/server/TestHttpServerConfig.java

@@ -67,7 +67,9 @@ public void testDefaults()
                 .setHttp2InputBufferSize(DataSize.of(8, KILOBYTE))
                 .setHttp2InitialStreamReceiveWindowSize(DataSize.of(16, MEGABYTE))
                 .setHttp2StreamIdleTimeout(new Duration(15, SECONDS))
-                .setCompressionEnabled(true));
+                .setCompressionEnabled(true)
+                .setMaxHeapMemory(null)
+                .setMaxDirectMemory(null));
     }
 
     @Test
@@ -103,6 +105,8 @@ public void testExplicitPropertyMappings()
                 .put("http-server.http2.input-buffer-size", "4MB")
                 .put("http-server.http2.stream-idle-timeout", "23s")
                 .put("http-server.compression.enabled", "false")
+                .put("http-server.max-heap-memory", "127GB")
+                .put("http-server.max-direct-memory", "129GB")
                 .build();
 
         HttpServerConfig expected = new HttpServerConfig()
@@ -134,7 +138,9 @@ public void testExplicitPropertyMappings()
                 .setHttp2InitialStreamReceiveWindowSize(DataSize.of(4, MEGABYTE))
                 .setHttp2InputBufferSize(DataSize.of(4, MEGABYTE))
                 .setHttp2StreamIdleTimeout(new Duration(23, SECONDS))
-                .setCompressionEnabled(false);
+                .setCompressionEnabled(false)
+                .setMaxHeapMemory(DataSize.of(127, GIGABYTE))
+                .setMaxDirectMemory(DataSize.of(129, GIGABYTE));
 
         assertFullMapping(properties, expected);
     }